Spaces:
Paused
Paused
| from pathlib import Path | |
| import numpy as np | |
| import matplotlib.pyplot as plt | |
| from matplotlib.colors import ListedColormap | |
| from matplotlib.patches import Patch | |
| from src.operations import remap_mask, create_overlay, to_binary_clothes, morphological_operations | |
| from src.utils import denormalize_image | |
| Path("assets").mkdir(parents=True, exist_ok=True) | |
| CLASS_NAMES = [ | |
| "Background", | |
| "Upper-Clothes", | |
| "Skirt", | |
| "Pants", | |
| "Dress" | |
| ] | |
| CLASS_COLORS = [ | |
| "#000000", # Background | |
| "#FF0000", # Upper-Clothes | |
| "#00FF00", # Skirt | |
| "#0000FF", # Pants | |
| "#FFFF00", # Dress | |
| ] | |
| SEGMENTATION_CMAP = ListedColormap(CLASS_COLORS) | |
| def colorize_mask(mask): | |
| """ | |
| Convert class-index mask into RGB colored mask. | |
| """ | |
| mask = np.asarray(mask).astype(np.uint8) | |
| colored_mask = np.zeros( | |
| (*mask.shape, 3), | |
| dtype=np.uint8 | |
| ) | |
| colors = [ | |
| (0, 0, 0), # Background | |
| (255, 0, 0), # Upper-Clothes | |
| (0, 255, 0), # Skirt | |
| (0, 0, 255), # Pants | |
| (255, 255, 0), # Dress | |
| ] | |
| for class_id, color in enumerate(colors): | |
| colored_mask[mask == class_id] = color | |
| return colored_mask | |
| def add_segmentation_legend(fig): | |
| legend_elements = [ | |
| Patch( | |
| facecolor=CLASS_COLORS[i], | |
| label=CLASS_NAMES[i] | |
| ) | |
| for i in range(len(CLASS_NAMES)) | |
| ] | |
| fig.legend( | |
| handles=legend_elements, | |
| loc="lower center", | |
| ncol=5, | |
| frameon=False | |
| ) | |
| def visualize_batch(images, masks, num_samples = 4, save_path="assets/visualization_batch.png"): | |
| num_samples = min(num_samples, len(images)) | |
| fig, axes = plt.subplots(2, num_samples, figsize = (4 * num_samples, 8)) | |
| if num_samples == 1: | |
| axes = np.expand_dims(axes, axis = 1) | |
| for i in range(num_samples): | |
| image = images[i].permute(1, 2, 0).cpu().numpy() | |
| image = denormalize_image(image) | |
| axes[0, i].imshow(image) | |
| axes[0, i].set_title("Image") | |
| axes[0, i].axis("off") | |
| mask = masks[i].cpu().numpy() | |
| axes[1, i].imshow(mask, cmap = SEGMENTATION_CMAP, vmin = 0, vmax = 4) | |
| axes[1, i].set_title("Mask") | |
| axes[1, i].axis("off") | |
| plt.tight_layout() | |
| plt.savefig(save_path, dpi=150, bbox_inches="tight") | |
| plt.close() | |
| def visualize_preprocessing_sample(image, mask, label_map, alpha=0.5, save_path="assets/preprocessing.png"): | |
| if hasattr(image, "permute"): | |
| image = image.permute(1, 2, 0).cpu().numpy() | |
| if hasattr(mask, "numpy"): | |
| mask = mask.cpu().numpy() | |
| # If image is uint8 -> convert to [0, 1] | |
| if image.dtype == np.uint8: | |
| image_display = image / 255.0 | |
| else: | |
| image_display = np.clip(image, 0, 1) | |
| remapped_mask = remap_mask(mask, label_map) | |
| overlay = create_overlay((image_display * 255).astype(np.uint8), remapped_mask, alpha=alpha) | |
| fig, axes = plt.subplots(1, 4, figsize=(16, 4)) | |
| axes[0].imshow(image_display) | |
| axes[0].set_title("Original Image") | |
| axes[0].axis("off") | |
| axes[1].imshow(mask, cmap=SEGMENTATION_CMAP, vmin=0, vmax=4) | |
| axes[1].set_title("Original Mask") | |
| axes[1].axis("off") | |
| axes[2].imshow(remapped_mask, cmap=SEGMENTATION_CMAP, vmin=0, vmax=4) | |
| axes[2].set_title("Remapped Mask") | |
| axes[2].axis("off") | |
| axes[3].imshow(overlay) | |
| axes[3].set_title("Overlay") | |
| axes[3].axis("off") | |
| plt.tight_layout() | |
| plt.savefig(save_path, dpi=150, bbox_inches="tight") | |
| plt.close() | |
| def visualize_random_samples(dataset, label_map, num_samples=4, alpha=0.5, save_path="assets/random_samples.png"): | |
| num_samples = min(num_samples, len(dataset)) | |
| indices = np.random.choice(len(dataset), size=num_samples, replace=False) | |
| fig, axes = plt.subplots(num_samples, 4, figsize=(16, 4 * num_samples)) | |
| if num_samples == 1: | |
| axes = np.expand_dims(axes, axis=0) | |
| for row, idx in enumerate(indices): | |
| sample = dataset[idx] | |
| image = sample["pixel_values"].convert("RGB") | |
| image = np.array(image, dtype=np.uint8) | |
| mask = np.array(sample["label"], dtype=np.uint8) | |
| remapped_mask = remap_mask(mask, label_map) | |
| overlay = create_overlay(image, remapped_mask, alpha=alpha) | |
| axes[row, 0].imshow(image) | |
| axes[row, 0].set_title(f"Image #{idx}") | |
| axes[row, 0].axis("off") | |
| axes[row, 1].imshow(mask, cmap="tab10") | |
| axes[row, 1].set_title("Original Mask") | |
| axes[row, 1].axis("off") | |
| axes[row, 2].imshow(remapped_mask, cmap="tab10", vmin=0, vmax=4) | |
| axes[row, 2].set_title("Remapped Mask") | |
| axes[row, 2].axis("off") | |
| axes[row, 3].imshow(overlay) | |
| axes[row, 3].set_title("Overlay") | |
| axes[row, 3].axis("off") | |
| plt.tight_layout() | |
| plt.savefig(save_path, dpi=150, bbox_inches="tight") | |
| plt.close() | |
| def visualize_prediction( | |
| image, | |
| gt_mask, | |
| pred_mask, | |
| opening_kernel=3, | |
| closing_kernel=11, | |
| opening_iterations=1, | |
| closing_iterations=1, | |
| alpha=0.5, | |
| save_path="assets/prediction.png" | |
| ): | |
| # Convert image to numpy | |
| if hasattr(image, "permute"): | |
| image = image.permute(1, 2, 0).cpu().numpy() | |
| # Convert masks to numpy | |
| if hasattr(gt_mask, "numpy"): | |
| gt_mask = gt_mask.cpu().numpy() | |
| if hasattr(pred_mask, "numpy"): | |
| pred_mask = pred_mask.cpu().numpy() | |
| # If image is normalized, convert it back to display range | |
| if image.dtype != np.uint8: | |
| image = np.clip(image, 0, 1) | |
| image = (image * 255).astype(np.uint8) | |
| # Ground Truth | |
| gt_mask = gt_mask.astype(np.uint8) | |
| # Prediction -> Binary | |
| pred_binary = to_binary_clothes(pred_mask) | |
| # Prediction -> Morphological Operations | |
| final_prediction = morphological_operations( | |
| pred_binary, | |
| opening_kernel=opening_kernel, | |
| closing_kernel=closing_kernel, | |
| opening_iterations=opening_iterations, | |
| closing_iterations=closing_iterations | |
| ) | |
| # Overlay | |
| overlay = create_overlay(image, final_prediction, alpha=alpha) | |
| # Visualization | |
| fig, axes = plt.subplots(1, 4, figsize=(16, 4)) | |
| axes[0].imshow(image) | |
| axes[0].set_title("Original Image") | |
| axes[0].axis("off") | |
| axes[1].imshow(gt_mask, cmap="tab10", vmin=0, vmax=4) | |
| axes[1].set_title("Ground Truth") | |
| axes[1].axis("off") | |
| axes[2].imshow(final_prediction, cmap="gray", vmin=0, vmax=1) | |
| axes[2].set_title("Final Prediction") | |
| axes[2].axis("off") | |
| axes[3].imshow(overlay) | |
| axes[3].set_title("Overlay") | |
| axes[3].axis("off") | |
| plt.tight_layout() | |
| plt.savefig(save_path, dpi=150, bbox_inches="tight") | |
| plt.close() | |
| def visualize_test_predictions( | |
| images, | |
| gt_masks, | |
| pred_masks, | |
| num_samples=4, | |
| opening_kernel=3, | |
| closing_kernel=11, | |
| opening_iterations=1, | |
| closing_iterations=1, | |
| alpha=0.5, | |
| save_path="assets/test_predictions.png" | |
| ): | |
| num_samples = min(num_samples, len(images)) | |
| fig, axes = plt.subplots(num_samples, 4, figsize=(16, 4 * num_samples)) | |
| if num_samples == 1: | |
| axes = np.expand_dims(axes, axis=0) | |
| for row in range(num_samples): | |
| image = images[row] | |
| gt_mask = gt_masks[row] | |
| pred_mask = pred_masks[row] | |
| # Convert image | |
| if hasattr(image, "permute"): | |
| image = image.permute(1, 2, 0).cpu().numpy() | |
| if image.dtype != np.uint8: | |
| image = np.clip(image, 0, 1) | |
| image = (image * 255).astype(np.uint8) | |
| # Convert masks | |
| if hasattr(gt_mask, "numpy"): | |
| gt_mask = gt_mask.cpu().numpy() | |
| if hasattr(pred_mask, "numpy"): | |
| pred_mask = pred_mask.cpu().numpy() | |
| gt_mask = gt_mask.astype(np.uint8) | |
| pred_mask = pred_mask.astype(np.uint8) | |
| # Binary prediction | |
| pred_binary = to_binary_clothes(pred_mask) | |
| # Morphological post-processing | |
| final_prediction = morphological_operations( | |
| pred_binary, | |
| opening_kernel=opening_kernel, | |
| closing_kernel=closing_kernel, | |
| opening_iterations=opening_iterations, | |
| closing_iterations=closing_iterations | |
| ) | |
| # Overlay | |
| overlay = create_overlay(image, final_prediction, alpha=alpha) | |
| # Original image | |
| axes[row, 0].imshow(image) | |
| axes[row, 0].set_title("Original Image") | |
| axes[row, 0].axis("off") | |
| # Ground truth | |
| axes[row, 1].imshow(gt_mask, cmap=SEGMENTATION_CMAP, vmin=0, vmax=4) | |
| axes[row, 1].set_title("Ground Truth") | |
| axes[row, 1].axis("off") | |
| # Final prediction | |
| axes[row, 2].imshow(final_prediction, cmap="gray", vmin=0, vmax=1) | |
| axes[row, 2].set_title("Final Prediction") | |
| axes[row, 2].axis("off") | |
| # Overlay | |
| axes[row, 3].imshow(overlay) | |
| axes[row, 3].set_title("Overlay") | |
| axes[row, 3].axis("off") | |
| plt.tight_layout() | |
| plt.savefig(save_path, dpi=150, bbox_inches="tight") | |
| plt.close() | |
| def plot_training_curves(history): | |
| epochs = range(1, len(history["train_loss"]) + 1) | |
| # Loss Curves | |
| plt.figure(figsize=(8, 5)) | |
| plt.plot(epochs, history["train_loss"], label="Train Loss") | |
| plt.plot(epochs, history["val_loss"], label="Validation Loss") | |
| plt.xlabel("Epoch") | |
| plt.ylabel("Loss") | |
| plt.title("Training and Validation Loss") | |
| plt.legend() | |
| plt.grid(True) | |
| plt.tight_layout() | |
| plt.savefig("assets/training_loss.png", dpi=150, bbox_inches="tight") | |
| plt.close() | |
| # Loss Components | |
| plt.figure(figsize=(8, 5)) | |
| plt.plot(epochs, history["train_ce"], label="Train CE") | |
| plt.plot(epochs, history["val_ce"], label="Validation CE") | |
| plt.plot(epochs, history["train_dice_loss"], label="Train Dice Loss") | |
| plt.plot(epochs, history["val_dice_loss"], label="Validation Dice Loss") | |
| plt.xlabel("Epoch") | |
| plt.ylabel("Loss") | |
| plt.title("Loss Components") | |
| plt.legend() | |
| plt.grid(True) | |
| plt.tight_layout() | |
| plt.savefig("assets/training_components.png", dpi=150, bbox_inches="tight") | |
| plt.close() | |
| # Segmentation Metrics | |
| plt.figure(figsize=(8, 5)) | |
| plt.plot(epochs, history["val_miou"], label="mIoU") | |
| plt.plot(epochs, history["val_miou_clothes"], label="mIoU Clothes") | |
| plt.plot(epochs, history["val_dice"], label="Dice") | |
| plt.plot(epochs, history["val_dice_clothes"], label="Dice Clothes") | |
| plt.xlabel("Epoch") | |
| plt.ylabel("Score") | |
| plt.title("Validation Segmentation Metrics") | |
| plt.legend() | |
| plt.grid(True) | |
| plt.tight_layout() | |
| plt.savefig("assets/validation_metrics.png", dpi=150, bbox_inches="tight") | |
| plt.close() | |
| def visualize_inference(image, pred_mask, final_mask, alpha=0.5, save_path="assets/inference_result.png"): | |
| if hasattr(image, "permute"): | |
| image = image.permute(1, 2, 0).cpu().numpy() | |
| if image.dtype != np.uint8: | |
| image = np.clip(image, 0, 1) | |
| image = (image * 255).astype(np.uint8) | |
| if hasattr(pred_mask, "numpy"): | |
| pred_mask = pred_mask.cpu().numpy() | |
| if hasattr(final_mask, "numpy"): | |
| final_mask = final_mask.cpu().numpy() | |
| pred_mask = pred_mask.astype(np.uint8) | |
| final_mask = final_mask.astype(np.uint8) | |
| overlay = create_overlay(image, final_mask, alpha=alpha) | |
| fig, axes = plt.subplots(1, 4, figsize=(16, 4)) | |
| # Original Image | |
| axes[0].imshow(image) | |
| axes[0].set_title("Original Image") | |
| axes[0].axis("off") | |
| # Raw Prediction | |
| colored_pred = colorize_mask(pred_mask) | |
| axes[1].imshow(colored_pred) | |
| axes[1].set_title("Raw Prediction") | |
| axes[1].axis("off") | |
| # Final Prediction | |
| axes[2].imshow(final_mask, cmap="gray", vmin=0, vmax=1) | |
| axes[2].set_title("Final Prediction") | |
| axes[2].axis("off") | |
| # Overlay | |
| axes[3].imshow(overlay) | |
| axes[3].set_title("Overlay") | |
| axes[3].axis("off") | |
| add_segmentation_legend(fig) | |
| plt.tight_layout(rect = [0, 0.08, 1, 1]) | |
| plt.savefig(save_path, dpi=150, bbox_inches="tight") | |
| plt.close() |