from __future__ import annotations import argparse import csv import io import json import shutil import time from collections import defaultdict from pathlib import Path import matplotlib.pyplot as plt import matplotlib.patches as patches import numpy as np from PIL import Image from mpl_toolkits.mplot3d import Axes3D # noqa: F401 from scipy import ndimage import torch import torch.nn.functional as F from train import build_model, read_labels CLASS_ORDER = ["DAUTEL EVO nano", "DJI MAVIC3 PRO", "DJI MINI3", "DJI MINI4 PRO"] PLOT_3D_BG = "#050505" PLOT_3D_FG = "#e8e8e8" PLOT_3D_PANE = (0.46, 0.46, 0.46, 1) PLOT_3D_GRID = (0.72, 0.72, 0.72, 0.65) def load_manifest(processed_dir: Path) -> list[dict[str, str]]: manifest = processed_dir / "manifest.csv" if not manifest.exists(): raise FileNotFoundError(f"Missing processed manifest: {manifest}") with manifest.open() as f: return list(csv.DictReader(f)) def plot_one_spectrogram(npz_path: Path, title: str, out_path: Path) -> None: data = np.load(npz_path) x = data["x"].astype(np.float32) fig, ax = plt.subplots(figsize=(7, 4.5)) im = ax.imshow(x, aspect="auto", origin="lower", cmap="viridis") ax.set_title(title) ax.set_xlabel("Time frame") ax.set_ylabel("Frequency bin") fig.colorbar(im, ax=ax, label="Log magnitude") fig.tight_layout() fig.savefig(out_path, dpi=160) plt.close(fig) def normalized_surface( x: np.ndarray, max_freq_bins: int = 160, max_time_frames: int = 120, front_crop_fraction: float = 0.12, ): front_crop = int(x.shape[0] * front_crop_fraction) if front_crop > 0 and x.shape[0] > front_crop: x = x[front_crop:, :] freq_step = max(1, int(np.ceil(x.shape[0] / max_freq_bins))) time_step = max(1, int(np.ceil(x.shape[1] / max_time_frames))) z = x[::freq_step, ::time_step].astype(np.float32) lo, hi = np.percentile(z, [2, 98]) z = np.clip(z, lo, hi) z = (z - z.min()) / (z.max() - z.min() + 1e-6) t = np.arange(z.shape[1]) f = np.arange(z.shape[0]) tt, ff = np.meshgrid(t, f) return tt, ff, z def plot_one_spectrogram_3d(npz_path: Path, title: str, out_path: Path, cmap: str = "turbo") -> None: data = np.load(npz_path) x = data["x"].astype(np.float32) tt, ff, z = normalized_surface(x) fig = plt.figure(figsize=(8, 5), facecolor=PLOT_3D_BG) ax = fig.add_subplot(111, projection="3d", facecolor=PLOT_3D_BG) ax.plot_surface( tt, ff, z, cmap=cmap, linewidth=0, antialiased=True, shade=True, rstride=1, cstride=1, ) ax.view_init(elev=32, azim=-58) ax.set_box_aspect((1.65, 1.0, 0.45)) ax.set_title(title, color=PLOT_3D_FG, pad=12) ax.set_xlabel("Time", color=PLOT_3D_FG) ax.set_ylabel("Frequency", color=PLOT_3D_FG) ax.set_zlabel("Power", color=PLOT_3D_FG) ax.tick_params(colors=PLOT_3D_FG) ax.xaxis.pane.set_facecolor(PLOT_3D_PANE) ax.yaxis.pane.set_facecolor(PLOT_3D_PANE) ax.zaxis.pane.set_facecolor(PLOT_3D_PANE) for axis in (ax.xaxis, ax.yaxis, ax.zaxis): axis._axinfo["grid"]["color"] = PLOT_3D_GRID axis._axinfo["grid"]["linewidth"] = 0.8 ax.grid(True) fig.tight_layout() fig.savefig(out_path, dpi=180, facecolor=fig.get_facecolor()) plt.close(fig) def style_3d_axis(ax, title: str) -> None: ax.set_facecolor(PLOT_3D_BG) ax.view_init(elev=32, azim=-58) ax.set_box_aspect((1.65, 1.0, 0.45)) ax.set_title(title, color=PLOT_3D_FG, pad=10, fontsize=12) ax.set_xlabel("Time", color=PLOT_3D_FG, labelpad=6) ax.set_ylabel("Frequency", color=PLOT_3D_FG, labelpad=6) ax.set_zlabel("Power", color=PLOT_3D_FG, labelpad=6) ax.tick_params(colors=PLOT_3D_FG, labelsize=7) ax.xaxis.pane.set_facecolor(PLOT_3D_PANE) ax.yaxis.pane.set_facecolor(PLOT_3D_PANE) ax.zaxis.pane.set_facecolor(PLOT_3D_PANE) for axis in (ax.xaxis, ax.yaxis, ax.zaxis): axis._axinfo["grid"]["color"] = PLOT_3D_GRID axis._axinfo["grid"]["linewidth"] = 0.8 ax.grid(True) def render_3d_spectrogram_frame(x: np.ndarray, title: str) -> Image.Image: tt, ff, z = normalized_surface(x) fig = plt.figure(figsize=(8, 5), facecolor=PLOT_3D_BG) ax = fig.add_subplot(111, projection="3d") ax.plot_surface( tt, ff, z, cmap="turbo", linewidth=0, antialiased=True, shade=True, rstride=1, cstride=1, ) style_3d_axis(ax, title) fig.tight_layout() buffer = io.BytesIO() fig.savefig(buffer, format="png", dpi=120, facecolor=fig.get_facecolor()) plt.close(fig) buffer.seek(0) return Image.open(buffer).convert("RGB") def plot_spectrogram_previews( processed_dir: Path, plots_dir: Path, max_per_class: int, max_3d_per_class: int, ) -> list[dict[str, str]]: rows = load_manifest(processed_dir) sample_dir = processed_dir / "samples" out_dir = plots_dir / "spectrograms" out_3d_dir = plots_dir / "spectrograms_3d" out_dir.mkdir(parents=True, exist_ok=True) out_3d_dir.mkdir(parents=True, exist_ok=True) by_label: dict[str, list[dict[str, str]]] = defaultdict(list) for row in rows: by_label[row["label"]].append(row) summary = [] grid_items = [] for label, label_rows in sorted(by_label.items()): selected = label_rows[:max_per_class] safe_label = label.replace(" ", "_").replace("/", "_") for idx, row in enumerate(selected): npz_path = sample_dir / row["path"] out_path = out_dir / f"{safe_label}_{idx}.png" plot_one_spectrogram(npz_path, f"{label} ({row['split']})", out_path) item = {"label": label, "split": row["split"], "plot": str(out_path)} if idx < max_3d_per_class: out_3d_path = out_3d_dir / f"{safe_label}_{idx}_3d.png" plot_one_spectrogram_3d(npz_path, f"{label} ({row['split']})", out_3d_path) item["plot_3d"] = str(out_3d_path) summary.append(item) if idx == 0: grid_items.append((label, np.load(npz_path)["x"].astype(np.float32))) if grid_items: fig, axes = plt.subplots(1, len(grid_items), figsize=(5 * len(grid_items), 4), squeeze=False) for ax, (label, x) in zip(axes[0], grid_items): ax.imshow(x, aspect="auto", origin="lower", cmap="viridis") ax.set_title(label) ax.set_xlabel("Time") ax.set_ylabel("Frequency") fig.tight_layout() fig.savefig(out_dir / "spectrogram_grid.png", dpi=160) plt.close(fig) return summary def plot_resnet18_spectrogram_previews( processed_dir: Path, plots_dir: Path, max_per_class: int, max_3d_per_class: int, ) -> list[dict[str, str]]: rows = load_manifest(processed_dir) sample_dir = processed_dir / "samples" out_root = plots_dir / "spectrograms_resnet18" out_2d_dir = out_root / "2d" out_3d_dir = out_root / "3d" out_2d_dir.mkdir(parents=True, exist_ok=True) out_3d_dir.mkdir(parents=True, exist_ok=True) for stale_plot in out_3d_dir.glob("*.png"): stale_plot.unlink() by_label: dict[str, list[dict[str, str]]] = defaultdict(list) for row in rows: by_label[row["label"]].append(row) summary = [] for label, label_rows in sorted(by_label.items()): selected = label_rows[:max_per_class] safe_label = label.replace(" ", "_").replace("/", "_") for idx, row in enumerate(selected): npz_path = sample_dir / row["path"] out_path = out_2d_dir / f"{safe_label}_{idx}.png" title = f"ResNet18 input: {label} ({row['split']})" plot_one_spectrogram(npz_path, title, out_path) item = {"label": label, "split": row["split"], "plot": str(out_path)} if idx < max_3d_per_class: out_3d_path = out_3d_dir / f"{safe_label}_{idx}_3d.png" plot_one_spectrogram_3d(npz_path, title, out_3d_path) item["plot_3d"] = str(out_3d_path) summary.append(item) return summary def plot_resnet18_3d_grid(processed_dir: Path, plots_dir: Path) -> dict[str, object]: rows = load_manifest(processed_dir) sample_dir = processed_dir / "samples" out_root = plots_dir / "spectrograms_resnet18" out_grid_dir = out_root / "grid" out_grid_dir.mkdir(parents=True, exist_ok=True) by_label: dict[str, list[dict[str, str]]] = defaultdict(list) for row in rows: by_label[row["label"]].append(row) selected = [] for label in CLASS_ORDER: label_rows = sorted(by_label.get(label, []), key=lambda row: (row["split"], row["path"])) if label_rows: selected.append((label, label_rows[0])) if len(selected) < 4: return {"created": False, "reason": "Need one sample for each selected class"} out_path = out_grid_dir / "class_comparison_3d_grid.png" fig = plt.figure(figsize=(14, 9), facecolor=PLOT_3D_BG) fig.suptitle("RFUAV 3D Spectrogram Class Comparison", color=PLOT_3D_FG, fontsize=18) for idx, (label, row) in enumerate(selected[:4], start=1): npz_path = sample_dir / row["path"] data = np.load(npz_path) x = data["x"].astype(np.float32) tt, ff, z = normalized_surface(x) ax = fig.add_subplot(2, 2, idx, projection="3d") ax.plot_surface( tt, ff, z, cmap="turbo", linewidth=0, antialiased=True, shade=True, rstride=1, cstride=1, ) style_3d_axis(ax, label) fig.tight_layout(rect=(0, 0, 1, 0.95)) fig.savefig(out_path, dpi=180, facecolor=fig.get_facecolor()) plt.close(fig) return {"created": True, "path": str(out_path), "labels": [label for label, _ in selected[:4]]} def plot_mavic3_3d_spectrogram_gif( processed_dir: Path, plots_dir: Path, frame_count: int, fps: int, ) -> dict[str, object]: rows = load_manifest(processed_dir) sample_dir = processed_dir / "samples" out_dir = plots_dir / "gifs" out_dir.mkdir(parents=True, exist_ok=True) mavic_rows = sorted( [row for row in rows if row["label"] == "DJI MAVIC3 PRO"], key=lambda row: (row["split"], row["path"]), ) if not mavic_rows: return {"created": False, "reason": "No DJI MAVIC3 PRO samples available"} selected_rows = evenly_sample_rows(mavic_rows, min(frame_count, len(mavic_rows))) gif_frames = [] for idx, row in enumerate(selected_rows, start=1): x = np.load(sample_dir / row["path"])["x"].astype(np.float32) gif_frames.append(render_3d_spectrogram_frame(x, f"DJI MAVIC3 PRO | frame {idx}/{len(selected_rows)}")) out_path = out_dir / "dji_mavic3_pro_3d_spectrogram.gif" duration_ms = int(1000 / max(1, fps)) gif_frames[0].save( out_path, save_all=True, append_images=gif_frames[1:], duration=duration_ms, loop=0, optimize=True, ) return { "created": True, "path": str(out_path), "frames": len(gif_frames), "fps": fps, "label": "DJI MAVIC3 PRO", } class GradCam: def __init__(self, model: torch.nn.Module, target_layer: torch.nn.Module): self.model = model self.activations: torch.Tensor | None = None self.gradients: torch.Tensor | None = None self.forward_handle = target_layer.register_forward_hook(self._save_activation) self.backward_handle = target_layer.register_full_backward_hook(self._save_gradient) def _save_activation(self, _module, _inputs, output) -> None: self.activations = output def _save_gradient(self, _module, _grad_input, grad_output) -> None: self.gradients = grad_output[0] def close(self) -> None: self.forward_handle.remove() self.backward_handle.remove() def __call__(self, logits: torch.Tensor, class_idx: int, out_shape: tuple[int, int]) -> np.ndarray: self.model.zero_grad(set_to_none=True) score = logits[0, class_idx] score.backward(retain_graph=True) if self.activations is None or self.gradients is None: raise RuntimeError("Grad-CAM hooks did not capture activations/gradients") weights = self.gradients.mean(dim=(2, 3), keepdim=True) cam = (weights * self.activations).sum(dim=1, keepdim=True) cam = torch.relu(cam) cam = F.interpolate(cam, size=out_shape, mode="bilinear", align_corners=False) cam = cam[0, 0].detach().cpu().numpy() cam = cam - cam.min() cam = cam / (cam.max() + 1e-8) return cam def load_resnet18_for_gradcam(processed_dir: Path, checkpoint_dir: Path, device: torch.device): labels = read_labels(processed_dir) model = build_model("resnet18", len(labels)).to(device) checkpoint_path = checkpoint_dir / "resnet18" / "best_model.pt" checkpoint = torch.load(checkpoint_path, map_location=device) model.load_state_dict(checkpoint["model"]) model.eval() return model, labels, checkpoint_path def attention_boxes(cam: np.ndarray, threshold: float = 0.72, min_area: int = 18, max_boxes: int = 5) -> list[tuple[int, int, int, int, float]]: mask = cam >= threshold labeled, count = ndimage.label(mask) regions = [] for idx, region_slice in enumerate(ndimage.find_objects(labeled), start=1): if region_slice is None: continue y_slice, x_slice = region_slice area = int((labeled[region_slice] == idx).sum()) if area < min_area: continue score = float(cam[region_slice][labeled[region_slice] == idx].mean()) regions.append((x_slice.start, y_slice.start, x_slice.stop, y_slice.stop, score, area)) if not regions and count: y, x = np.unravel_index(int(np.argmax(cam)), cam.shape) pad_y = max(4, cam.shape[0] // 18) pad_x = max(3, cam.shape[1] // 18) regions.append(( max(0, x - pad_x), max(0, y - pad_y), min(cam.shape[1], x + pad_x), min(cam.shape[0], y + pad_y), float(cam[y, x]), 1, )) regions = sorted(regions, key=lambda item: (item[4], item[5]), reverse=True)[:max_boxes] return [(x1, y1, x2, y2, score) for x1, y1, x2, y2, score, _area in regions] def render_gradcam_frame( spec: np.ndarray, cam: np.ndarray, boxes: list[tuple[int, int, int, int, float]], frame_idx: int, total_frames: int, true_label: str, pred_label: str, latency_ms: float, ) -> Image.Image: peak_y, peak_x = np.unravel_index(int(np.argmax(spec)), spec.shape) fig = plt.figure(figsize=(12, 6), facecolor="#05070a") gs = fig.add_gridspec(1, 2, width_ratios=[3.2, 1.25], wspace=0.08) ax = fig.add_subplot(gs[0, 0]) side = fig.add_subplot(gs[0, 1]) ax.set_facecolor("#05070a") vmin, vmax = np.percentile(spec, [2, 98]) height, width = spec.shape image_extent = (0, width, 0, height) ax.imshow(spec, aspect="auto", origin="lower", extent=image_extent, cmap="turbo", vmin=vmin, vmax=vmax) ax.imshow( cam, aspect="auto", origin="lower", extent=image_extent, cmap="cool", alpha=np.clip(cam, 0.0, 0.65) * 0.52, ) for box_idx, (x1, y1, x2, y2, score) in enumerate(boxes, start=1): x1 = max(0, min(width, x1)) x2 = max(0, min(width, x2)) y1 = max(0, min(height, y1)) y2 = max(0, min(height, y2)) rect = patches.Rectangle( (x1, y1), x2 - x1, y2 - y1, linewidth=2.0, edgecolor="#00f5ff" if box_idx == 1 else "#ffe66d", facecolor="none", ) ax.add_patch(rect) ax.text( x1, min(height - 2, y2 + 2), f"A{box_idx} {score:.2f}", color="#05070a", fontsize=8, fontweight="bold", bbox={"facecolor": "#00f5ff" if box_idx == 1 else "#ffe66d", "alpha": 0.9, "pad": 1.5}, ) ax.scatter([peak_x + 0.5], [peak_y + 0.5], marker="x", s=50, linewidths=1.5, color="white") ax.set_title("DJI MINI4 PRO | RF spectrogram + ResNet18 Grad-CAM", color="white", fontsize=14, pad=10) ax.set_xlabel("Time frame", color="#d7e3f4") ax.set_ylabel("Frequency bin", color="#d7e3f4") ax.set_xlim(0, width) ax.set_ylim(0, height) ax.tick_params(colors="#aeb9c8") side.set_facecolor("#0c1119") side.axis("off") panel_lines = [ ("FRAME", f"{frame_idx:02d} / {total_frames:02d}"), ("TRUE CLASS", true_label), ("PREDICTION", pred_label), ("LATENCY", f"{latency_ms:.2f} ms"), ("PEAK BIN", f"f={peak_y}, t={peak_x}"), ("PEAK POWER", f"{float(spec[peak_y, peak_x]):.1f} dB"), ("ATTN BOXES", str(len(boxes))), ] side.text(0.04, 0.94, "CLASSIFIER TELEMETRY", color="white", fontsize=12, fontweight="bold", transform=side.transAxes) y = 0.84 for label, value in panel_lines: side.text(0.05, y, label, color="#6f7f93", fontsize=8, fontweight="bold", transform=side.transAxes) side.text(0.05, y - 0.045, value, color="white", fontsize=12, transform=side.transAxes, wrap=True) y -= 0.105 buffer = io.BytesIO() fig.savefig(buffer, format="png", dpi=120, facecolor=fig.get_facecolor()) plt.close(fig) buffer.seek(0) return Image.open(buffer).convert("RGB") def plot_mini4_gradcam_gif( processed_dir: Path, results_dir: Path, checkpoint_dir: Path, plots_dir: Path, frame_count: int, fps: int, ) -> dict[str, object]: rows = load_manifest(processed_dir) sample_dir = processed_dir / "samples" out_dir = plots_dir / "gifs" out_dir.mkdir(parents=True, exist_ok=True) mini4_rows = sorted( [row for row in rows if row["label"] == "DJI MINI4 PRO"], key=lambda row: (row["split"], row["path"]), ) if not mini4_rows: return {"created": False, "reason": "No DJI MINI4 PRO samples available"} metrics_path = results_dir / "resnet18" / "metrics.json" test_accuracy = None if metrics_path.exists(): metrics = json.loads(metrics_path.read_text()) test_accuracy = metrics.get("test_accuracy") device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model, labels, checkpoint_path = load_resnet18_for_gradcam(processed_dir, checkpoint_dir, device) gradcam = GradCam(model, model.layer4[-1]) selected_rows = evenly_sample_rows(mini4_rows, min(frame_count, len(mini4_rows))) gif_frames: list[Image.Image] = [] try: for idx, row in enumerate(selected_rows, start=1): spec = np.load(sample_dir / row["path"])["x"].astype(np.float32) normalized = (spec - spec.mean()) / (spec.std() + 1e-6) x = torch.from_numpy(normalized).unsqueeze(0).unsqueeze(0).to(device) start = time.perf_counter() logits = model(x) if device.type == "cuda": torch.cuda.synchronize() latency_ms = (time.perf_counter() - start) * 1000.0 probs = torch.softmax(logits.detach(), dim=1)[0] pred_id = int(torch.argmax(probs).item()) cam = gradcam(logits, pred_id, spec.shape) boxes = attention_boxes(cam) frame = render_gradcam_frame( spec=spec, cam=cam, boxes=boxes, frame_idx=idx, total_frames=len(selected_rows), true_label=row["label"], pred_label=labels[pred_id], latency_ms=latency_ms, ) gif_frames.append(frame) finally: gradcam.close() out_path = out_dir / "dji_mini4_pro_gradcam_boxes.gif" duration_ms = int(1000 / max(1, fps)) gif_frames[0].save( out_path, save_all=True, append_images=gif_frames[1:], duration=duration_ms, loop=0, optimize=True, ) return { "created": True, "path": str(out_path), "frames": len(gif_frames), "fps": fps, "label": "DJI MINI4 PRO", "model": "resnet18", "checkpoint": str(checkpoint_path), } def evenly_sample_rows(rows: list[dict[str, str]], count: int) -> list[dict[str, str]]: rows = sorted(rows, key=lambda row: (row["split"], row["path"])) if len(rows) <= count: return rows indices = np.linspace(0, len(rows) - 1, count, dtype=int) return [rows[int(idx)] for idx in indices] def plot_time_sweep_comparison_gif( processed_dir: Path, plots_dir: Path, frame_count: int, fps: int, ) -> dict[str, object]: rows = load_manifest(processed_dir) sample_dir = processed_dir / "samples" out_dir = plots_dir / "gifs" out_dir.mkdir(parents=True, exist_ok=True) by_label: dict[str, list[dict[str, str]]] = defaultdict(list) for row in rows: by_label[row["label"]].append(row) selected_labels = [label for label in CLASS_ORDER if by_label.get(label)] if len(selected_labels) < 2: return {"created": False, "reason": "Need at least two classes for comparison GIF"} frames_available = min(len(by_label[label]) for label in selected_labels) frame_count = min(frame_count, frames_available) if frame_count <= 0: return {"created": False, "reason": "No spectrogram samples available"} selected_rows = { label: evenly_sample_rows(by_label[label], frame_count) for label in selected_labels } spectrograms: dict[str, list[np.ndarray]] = {} all_values = [] for label, label_rows in selected_rows.items(): specs = [] for row in label_rows: x = np.load(sample_dir / row["path"])["x"].astype(np.float32) specs.append(x) all_values.append(x.reshape(-1)) spectrograms[label] = specs values = np.concatenate(all_values) vmin, vmax = np.percentile(values, [2, 98]) gif_frames: list[Image.Image] = [] for frame_idx in range(frame_count): fig, axes = plt.subplots(2, 2, figsize=(10, 7), facecolor="black", squeeze=False) fig.suptitle( f"RFUAV spectrogram time-sweep comparison | frame {frame_idx + 1}/{frame_count}", color="white", fontsize=15, ) for ax in axes.ravel(): ax.set_facecolor("black") ax.axis("off") for ax, label in zip(axes.ravel(), selected_labels): row = selected_rows[label][frame_idx] ax.imshow( spectrograms[label][frame_idx], aspect="auto", origin="lower", cmap="turbo", vmin=vmin, vmax=vmax, ) ax.set_title(label, color="white", fontsize=12) ax.set_xlabel("Time frame", color="white") ax.set_ylabel("Frequency bin", color="white") ax.tick_params(colors="white", labelsize=8) ax.axis("on") fig.tight_layout(rect=(0, 0, 1, 0.93)) buffer = io.BytesIO() fig.savefig(buffer, format="png", dpi=120, facecolor=fig.get_facecolor()) plt.close(fig) buffer.seek(0) gif_frames.append(Image.open(buffer).convert("RGB")) out_path = out_dir / "time_sweep_comparison.gif" duration_ms = int(1000 / max(1, fps)) gif_frames[0].save( out_path, save_all=True, append_images=gif_frames[1:], duration=duration_ms, loop=0, optimize=True, ) return { "created": True, "path": str(out_path), "frames": frame_count, "fps": fps, "labels": selected_labels, } def find_metric_files(results_dir: Path) -> list[Path]: if not results_dir.exists(): return [] return sorted(results_dir.rglob("metrics.json")) def plot_training_curves(metric_path: Path, plots_dir: Path) -> dict[str, object]: metrics = json.loads(metric_path.read_text()) model_name = str(metrics.get("model") or metric_path.parent.name) safe_name = model_name.replace("/", "_") model_dir = plots_dir / safe_name history = metrics.get("history", []) model_dir.mkdir(parents=True, exist_ok=True) if history: epochs = [row["epoch"] for row in history] fig, axes = plt.subplots(1, 2, figsize=(11, 4)) axes[0].plot(epochs, [row["train_loss"] for row in history], label="train_loss") axes[0].plot(epochs, [row["val_loss"] for row in history], label="val_loss") axes[0].set_title(f"{model_name}: loss") axes[0].set_xlabel("Epoch") axes[0].legend() axes[1].plot(epochs, [row["val_accuracy"] for row in history], label="val_accuracy") axes[1].set_title(f"{model_name}: validation accuracy") axes[1].set_xlabel("Epoch") axes[1].set_ylim(0, 1) axes[1].legend() fig.tight_layout() fig.savefig(model_dir / "training_curves.png", dpi=160) plt.close(fig) confusion = metric_path.parent / "confusion_matrix.png" if confusion.exists(): shutil.copy2(confusion, model_dir / "confusion_matrix.png") shutil.copy2(metric_path, model_dir / "metrics.json") return { "model": model_name, "best_val_accuracy": metrics.get("best_val_accuracy"), "test_accuracy": metrics.get("test_accuracy"), "metrics_path": str(metric_path), "plot_dir": str(model_dir), } def plot_metric_artifacts(results_dir: Path, plots_dir: Path) -> list[dict[str, object]]: summaries = [plot_training_curves(path, plots_dir) for path in find_metric_files(results_dir)] if summaries: out_dir = plots_dir / "comparison" out_dir.mkdir(parents=True, exist_ok=True) names = [str(item["model"]) for item in summaries] test_acc = [float(item.get("test_accuracy") or 0.0) for item in summaries] val_acc = [float(item.get("best_val_accuracy") or 0.0) for item in summaries] x = np.arange(len(names)) width = 0.36 fig, ax = plt.subplots(figsize=(max(7, 2.5 * len(names)), 4.5)) ax.bar(x - width / 2, val_acc, width, label="best val acc") ax.bar(x + width / 2, test_acc, width, label="test acc") ax.set_xticks(x, names, rotation=20, ha="right") ax.set_ylim(0, 1) ax.set_ylabel("Accuracy") ax.set_title("Model accuracy comparison") ax.legend() fig.tight_layout() fig.savefig(out_dir / "comparison_accuracy.png", dpi=160) plt.close(fig) return summaries def main() -> None: parser = argparse.ArgumentParser(description="Generate plot artifacts for RFUAV spectrograms and CNN metrics.") parser.add_argument("--processed-dir", default="/data/RFUAV_processed") parser.add_argument("--results-dir", default="/data/results") parser.add_argument("--checkpoint-dir", default="/data/checkpoints") parser.add_argument("--plots-dir", default="/data/plots") parser.add_argument("--max-spectrograms-per-class", type=int, default=3) parser.add_argument("--max-3d-spectrograms-per-class", type=int, default=1) parser.add_argument("--gif-frames", type=int, default=30) parser.add_argument("--gif-fps", type=int, default=2) parser.add_argument("--mavic3-3d-gif-frames", type=int, default=20) parser.add_argument("--mavic3-3d-gif-fps", type=int, default=2) parser.add_argument("--mini4-gradcam-gif-frames", type=int, default=30) parser.add_argument("--mini4-gradcam-gif-fps", type=int, default=2) args = parser.parse_args() processed_dir = Path(args.processed_dir) results_dir = Path(args.results_dir) checkpoint_dir = Path(args.checkpoint_dir) plots_dir = Path(args.plots_dir) plots_dir.mkdir(parents=True, exist_ok=True) summary = { "spectrograms": [], "spectrograms_resnet18": [], "gifs": {}, "mavic3_3d_gif": {}, "mini4_gradcam_gif": {}, "metrics": [], } if processed_dir.exists() and (processed_dir / "manifest.csv").exists(): summary["spectrograms"] = plot_spectrogram_previews( processed_dir, plots_dir, args.max_spectrograms_per_class, args.max_3d_spectrograms_per_class, ) summary["spectrograms_resnet18"] = plot_resnet18_spectrogram_previews( processed_dir, plots_dir, args.max_spectrograms_per_class, args.max_3d_spectrograms_per_class, ) summary["spectrograms_resnet18_3d_grid"] = plot_resnet18_3d_grid(processed_dir, plots_dir) summary["gifs"] = plot_time_sweep_comparison_gif( processed_dir, plots_dir, args.gif_frames, args.gif_fps, ) summary["mavic3_3d_gif"] = plot_mavic3_3d_spectrogram_gif( processed_dir, plots_dir, args.mavic3_3d_gif_frames, args.mavic3_3d_gif_fps, ) summary["mini4_gradcam_gif"] = plot_mini4_gradcam_gif( processed_dir, results_dir, checkpoint_dir, plots_dir, args.mini4_gradcam_gif_frames, args.mini4_gradcam_gif_fps, ) else: print(f"Skipping spectrogram plots; missing {processed_dir / 'manifest.csv'}") summary["metrics"] = plot_metric_artifacts(results_dir, plots_dir) (plots_dir / "summary.json").write_text(json.dumps(summary, indent=2)) print(f"Generated plots under {plots_dir}") print( json.dumps( { "spectrogram_plots": len(summary["spectrograms"]), "resnet18_spectrogram_plots": len(summary["spectrograms_resnet18"]), "gif": summary["gifs"], "mavic3_3d_gif": summary["mavic3_3d_gif"], "mini4_gradcam_gif": summary["mini4_gradcam_gif"], "metric_runs": len(summary["metrics"]), }, indent=2, ) ) if __name__ == "__main__": main()