File size: 3,534 Bytes
c8c00f0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | """
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)
|