program

Signed-off-by: 葛峻恺 <202115006@mail.sdu.edu.cn>

utils/plot.py

@@ -0,0 +1,135 @@
+import os
+import sys
+import torch
+from torchsummary import summary
+from torchvision.utils import make_grid, save_image
+import numpy as np
+import matplotlib.pyplot as plt
+import scipy.ndimage
+# Add parent directory to path for config import
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+from config import Network_train_Config as cfg
+
+
+def plot_BSCAN_data(data, path, line_length=100, time_length=200, ratio=1):
+    """
+    Plot the inverted permittivity constant map and adjust the colormap range based on the ratio parameter.
+    If ratio < 1, values exceeding ratio * max(abs(data)) will be saturated at the colormap maximum.
+    
+    :param data: 2D NumPy array (N, M), where N is time/depth and M is survey line direction
+    :param path: Path to save the output image
+    :param line_length: Survey line length in meters (default: 400m)
+    :param time_length: Time range in nanoseconds (default: 200ns)
+    :param ratio: Scaling factor for colormap range (default: 1)
+    """
+    num_points, num_lines = data.shape
+
+    # Compute the maximum absolute value for normalization
+    max_abs = np.max(np.abs(data))
+    vmin = -ratio * max_abs
+    vmax = ratio * max_abs
+
+    # Set font style
+    plt.rcParams.update({'font.family': 'Times New Roman', 'font.size': 20})
+
+    # Plot the permittivity image
+    plt.figure(figsize=(10, 4))
+    im = plt.imshow(data, aspect='auto', cmap='gray',
+                    extent=[0, line_length, time_length, 0],
+                    vmin=vmin, vmax=vmax)
+    
+    # Configure axis labels and ticks
+    plt.xlabel('Distance (m)', fontsize=20)
+    plt.xticks([0, 20, 40, 60, 80, 100, line_length])
+    plt.ylabel('Time (ns)', fontsize=20)
+    plt.yticks([0, 50, 100, 150, 200, time_length])
+
+    # Customize tick and border appearance
+    plt.tick_params(axis='both', direction='in', width=1)
+    for spine in plt.gca().spines.values():
+        spine.set_linewidth(1)
+    
+    # Remove grid lines, adjust layout, and save the figure
+    plt.grid(False)
+    plt.tight_layout()
+    plt.savefig(path, dpi=300)
+    plt.show()
+
+
+
+def plot_permittivity_constant(data, path, line_length=100, time_length=200):
+    """
+    Plot the inverted permittivity constant.
+    
+    :param data: 2D NumPy array, shape (N, M), where N is depth (or time) and M is distance along the survey line.
+    :param path: Path to save the output image.
+    :param line_length: Survey line length (meters), default is 100m.
+    :param time_length: Time range (nanoseconds), default is 200ns.
+    """
+    num_points, num_lines = data.shape
+    
+    # Configure font settings
+    plt.rcParams.update({'font.family': 'Times New Roman', 'font.size': 20})
+    
+    # Plot the permittivity map
+    plt.figure(figsize=(12, 4))
+    im = plt.imshow(data, aspect='auto', cmap='rainbow_r', extent=[0, 100, 200, 0], vmin=9, vmax=26)
+    
+    # Set axis labels and ticks
+    plt.xlabel('Distance (m)', fontsize=20)
+    plt.xticks([0, 20, 40, 60, 80, 100])
+    plt.ylabel('Time (ns)', fontsize=20)
+    plt.yticks([0, 50, 100, 150, 200])
+    
+    # Add colorbar
+    cbar = plt.colorbar(im)
+    cbar.set_label('Permittivity', fontsize=20)
+    
+    # Adjust axis formatting
+    plt.tick_params(axis='both', direction='in', width=1)
+    for spine in plt.gca().spines.values():
+        spine.set_linewidth(1)
+    
+    plt.grid(False)
+    plt.tight_layout()
+    plt.savefig(path, dpi=300)
+    plt.show()
+
+
+def plot_depth_permittivity_constant(data, path, line_length=100, time_length=200):
+    """
+    Plot the inverted permittivity constant as a 2D colormap.
+    
+    :param data: 2D NumPy array (depth/time, distance), representing permittivity values.
+    :param path: File path to save the output image.
+    :param line_length: Length of the survey line in meters (default: 100m).
+    :param time_length: Time range in nanoseconds (default: 200ns).
+    """
+    num_points, num_lines = data.shape
+    
+    # Set font properties
+    plt.rcParams.update({'font.family': 'Times New Roman', 'font.size': 20})
+    
+    # Create the plot
+    plt.figure(figsize=(12, 4))
+    im = plt.imshow(data, aspect='auto', cmap='rainbow_r', extent=[0, line_length, time_length, 0], vmin=8, vmax=30)
+    
+    # Label axes and set ticks
+    plt.xlabel('Distance (m)', fontsize=20)
+    plt.xticks([0, 20, 40, 60, 80, 100])
+    plt.ylabel('Time (ns)', fontsize=20)
+    plt.yticks([0, 2, 4, 6, 8])
+    
+    # Add colorbar
+    cbar = plt.colorbar(im)
+    cbar.set_label('Permittivity', fontsize=20)
+    
+    # Format axis appearance
+    plt.tick_params(axis='both', direction='in', width=1)
+    for spine in plt.gca().spines.values():
+        spine.set_linewidth(1)
+    
+    plt.grid(False)
+    plt.tight_layout()
+    plt.savefig(path, dpi=300)
+    plt.show()