| """ |
| Instance Processor — DefectDiffu Edition |
| |
| Stripped-down version: DefectDiffu does NOT use 16x16 patch mappings, |
| so all patch-based artifact logic has been removed. |
| |
| Retained utilities: |
| - bbox ↔ mask helpers (for verification cropping) |
| - IoU calculation |
| - Visualization helpers |
| """ |
|
|
| import random |
| import torch |
| import numpy as np |
| from PIL import Image |
| from typing import List, Dict, Tuple, Optional, Union |
| import matplotlib.pyplot as plt |
| import matplotlib.patches as patches |
|
|
|
|
| class InstanceProcessor: |
| """Utility class for detection post-processing and mask operations.""" |
|
|
| @staticmethod |
| def calculate_iou(box1: Union[List, np.ndarray], box2: Union[List, np.ndarray]) -> float: |
| x1 = max(box1[0], box2[0]) |
| y1 = max(box1[1], box2[1]) |
| x2 = min(box1[2], box2[2]) |
| y2 = min(box1[3], box2[3]) |
| if x2 <= x1 or y2 <= y1: |
| return 0.0 |
| intersection = (x2 - x1) * (y2 - y1) |
| area1 = (box1[2] - box1[0]) * (box1[3] - box1[1]) |
| area2 = (box2[2] - box2[0]) * (box2[3] - box2[1]) |
| union = area1 + area2 - intersection |
| return intersection / union if union > 0 else 0.0 |
|
|
| @staticmethod |
| def mask_from_bbox(bbox: Tuple[int, int, int, int], img_shape: Tuple[int, ...]) -> np.ndarray: |
| """Create a binary mask from a bounding box.""" |
| h, w = img_shape[:2] |
| mask = np.zeros((h, w), dtype=np.uint8) |
| x1, y1, x2, y2 = bbox |
| x1, y1 = max(0, x1), max(0, y1) |
| x2, y2 = min(w, x2), min(h, y2) |
| if x2 > x1 and y2 > y1: |
| mask[y1:y2, x1:x2] = 1 |
| return mask |
|
|
| @staticmethod |
| def get_bbox_from_mask(mask: np.ndarray, margin: int = 0) -> Tuple[int, int, int, int]: |
| """Compute tight bounding box from binary mask, with optional margin.""" |
| ys, xs = np.where(mask > 0) |
| if len(ys) == 0: |
| return (0, 0, 0, 0) |
| y1, y2 = ys.min(), ys.max() |
| x1, x2 = xs.min(), xs.max() |
| h, w = mask.shape |
| x1 = max(0, x1 - margin) |
| y1 = max(0, y1 - margin) |
| x2 = min(w, x2 + margin) |
| y2 = min(h, y2 + margin) |
| return (x1, y1, x2, y2) |
|
|
| @staticmethod |
| def visualize_generation_result( |
| original_image: np.ndarray, |
| generated_image: np.ndarray, |
| defect_mask: np.ndarray, |
| output_path: str, |
| title: str = "DefectDiffu Generation Result" |
| ): |
| """Create a 3-panel visualization: original, generated, mask overlay.""" |
| fig, axes = plt.subplots(1, 3, figsize=(18, 6)) |
|
|
| axes[0].imshow(original_image) |
| axes[0].set_title("Original (Planning Reference)") |
| axes[0].axis("off") |
|
|
| axes[1].imshow(generated_image) |
| axes[1].set_title("Generated Defect Image") |
| axes[1].axis("off") |
|
|
| axes[2].imshow(generated_image) |
| axes[2].imshow(defect_mask, alpha=0.5, cmap="Reds") |
| axes[2].set_title("Defect Mask Overlay") |
| axes[2].axis("off") |
|
|
| fig.suptitle(title, fontsize=14) |
| plt.tight_layout() |
| plt.savefig(output_path, dpi=150, bbox_inches="tight") |
| plt.close(fig) |
| print(f"[Viz] Saved result visualization to {output_path}") |
|
|
| @staticmethod |
| def resize_to_square(image: np.ndarray, size: int = 512) -> np.ndarray: |
| """Resize image to square (DefectDiffu expects 512x512).""" |
| pil_img = Image.fromarray(image) if isinstance(image, np.ndarray) else image |
| pil_img = pil_img.resize((size, size), Image.LANCZOS) |
| return np.array(pil_img) |
|
|
