165 lines
4.9 KiB
Python
165 lines
4.9 KiB
Python
import argparse
|
|
import shutil
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
|
|
import matplotlib.patches as mpatches
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
from matplotlib.colors import LinearSegmentedColormap, ListedColormap
|
|
from PIL import Image
|
|
|
|
# --- Setup output folders ---
|
|
output_path = Path("/home/fedex/mt/plots/data_2d_projections_training")
|
|
datetime_folder_name = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
|
|
output_datetime_path = output_path / datetime_folder_name
|
|
latest_folder_path = output_path / "latest"
|
|
archive_folder_path = output_path / "archive"
|
|
|
|
for folder in (
|
|
output_path,
|
|
output_datetime_path,
|
|
latest_folder_path,
|
|
archive_folder_path,
|
|
):
|
|
folder.mkdir(exist_ok=True, parents=True)
|
|
|
|
# --- Parse command-line arguments ---
|
|
parser = argparse.ArgumentParser(
|
|
description="Plot two 2D projections as used in training (unstretched, grayscale)"
|
|
)
|
|
parser.add_argument(
|
|
"--input1",
|
|
type=Path,
|
|
default=Path(
|
|
"/home/fedex/mt/data/subter/new_projection/1_loop_closure_illuminated_2023-01-23.npy"
|
|
),
|
|
help="Path to first .npy file containing 2D projection data",
|
|
)
|
|
parser.add_argument(
|
|
"--input2",
|
|
type=Path,
|
|
default=Path(
|
|
"/home/fedex/mt/data/subter/new_projection/3_smoke_human_walking_2023-01-23.npy"
|
|
),
|
|
help="Path to second .npy file containing 2D projection data",
|
|
)
|
|
parser.add_argument(
|
|
"--frame1",
|
|
type=int,
|
|
default=955,
|
|
help="Frame index to plot from first file (0-indexed)",
|
|
)
|
|
parser.add_argument(
|
|
"--frame2",
|
|
type=int,
|
|
default=242,
|
|
help="Frame index to plot from second file (0-indexed)",
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
# --- Load the numpy projection data ---
|
|
proj_data1 = np.load(args.input1)
|
|
proj_data2 = np.load(args.input2)
|
|
|
|
# Choose the desired frames
|
|
try:
|
|
frame1 = proj_data1[args.frame1]
|
|
frame2 = proj_data2[args.frame2]
|
|
except IndexError as e:
|
|
raise ValueError(f"Frame index out of range: {e}")
|
|
|
|
# Debug info: Print the percentage of missing data in each frame
|
|
print(f"Frame 1 missing data percentage: {np.isnan(frame1).mean() * 100:.2f}%")
|
|
print(f"Frame 2 missing data percentage: {np.isnan(frame2).mean() * 100:.2f}%")
|
|
|
|
# --- Create a figure with 2 vertical subplots ---
|
|
fig, (ax1, ax2) = plt.subplots(nrows=2, ncols=1, figsize=(10, 5))
|
|
|
|
# Create custom colormap for missing data visualization
|
|
missing_color = [1, 0, 0, 1] # Red with full alpha
|
|
cmap_missing = ListedColormap([missing_color])
|
|
|
|
# Replace the plotting section
|
|
for ax, frame, title in zip(
|
|
(ax1, ax2),
|
|
(frame1, frame2),
|
|
(
|
|
"Normal LiDAR Frame",
|
|
"Degraded LiDAR Frame (Smoke)",
|
|
),
|
|
):
|
|
# Create mask for missing data (directly from NaN values)
|
|
missing_mask = np.isnan(frame)
|
|
|
|
# Plot the valid data in grayscale
|
|
frame_valid = np.copy(frame)
|
|
frame_valid[missing_mask] = 0 # Set missing values to black in base image
|
|
im = ax.imshow(frame_valid, cmap="gray", aspect="equal", vmin=0, vmax=0.8)
|
|
|
|
# Overlay missing data in red with reduced alpha
|
|
ax.imshow(
|
|
missing_mask,
|
|
cmap=ListedColormap([[1, 0, 0, 1]]), # Pure red
|
|
alpha=0.3, # Reduced alpha for better visibility
|
|
)
|
|
|
|
ax.set_title(title)
|
|
ax.axis("off")
|
|
|
|
# Adjust layout
|
|
plt.tight_layout()
|
|
|
|
# Create a more informative legend
|
|
legend_elements = [
|
|
mpatches.Patch(facecolor="red", alpha=0.7, label="Missing Data"),
|
|
mpatches.Patch(facecolor="white", label="Close Distance (0m)"),
|
|
mpatches.Patch(facecolor="gray", label="Mid Distance"),
|
|
mpatches.Patch(facecolor="black", label="Far Distance (70m)"),
|
|
]
|
|
|
|
# Add legend with better positioning and formatting
|
|
fig.legend(
|
|
handles=legend_elements,
|
|
loc="center right",
|
|
bbox_to_anchor=(0.98, 0.5),
|
|
title="Distance Information",
|
|
framealpha=1.0,
|
|
)
|
|
|
|
# Save the plot
|
|
output_file = output_datetime_path / "data_2d_projections_training.png"
|
|
plt.savefig(output_file, dpi=300, bbox_inches="tight", pad_inches=0.1)
|
|
plt.close()
|
|
|
|
print(f"Plot saved to: {output_file}")
|
|
|
|
# --- Create grayscale training images ---
|
|
for degradation_status, frame_number, frame in (
|
|
("normal", args.frame1, frame1),
|
|
("smoke", args.frame2, frame2),
|
|
):
|
|
frame_gray = np.nan_to_num(frame, nan=0).astype(np.float32)
|
|
gray_image = Image.fromarray(frame_gray, mode="F")
|
|
gray_output_file = (
|
|
output_datetime_path
|
|
/ f"frame_{frame_number}_training_{degradation_status}.tiff"
|
|
)
|
|
gray_image.save(gray_output_file)
|
|
print(f"Training image saved to: {gray_output_file}")
|
|
|
|
# --- Handle folder structure ---
|
|
shutil.rmtree(latest_folder_path, ignore_errors=True)
|
|
latest_folder_path.mkdir(exist_ok=True, parents=True)
|
|
|
|
for file in output_datetime_path.iterdir():
|
|
shutil.copy2(file, latest_folder_path)
|
|
|
|
script_path = Path(__file__)
|
|
shutil.copy2(script_path, output_datetime_path)
|
|
shutil.copy2(script_path, latest_folder_path)
|
|
|
|
shutil.move(output_datetime_path, archive_folder_path)
|
|
print(f"Output archived to: {archive_folder_path}")
|