Upload folder using huggingface_hub

4a2938a verified about 9 hours ago

7.41 kB

"""
Utility functions for preprocessing, TTA inference, postprocessing,
visualization, and Kaggle RLE formatting.
"""

import cv2
import numpy as np
import torch
from PIL import Image

IMG_H = 128
IMG_W = 800
NORM_MEAN = 0.449
NORM_STD = 0.226

CLASS_NAMES = ["defect_a", "defect_b", "defect_c", "defect_d"]
CLASS_LABELS = {
    "defect_a": "Defect A (Pitted / Scratch)",
    "defect_b": "Defect B (Minor defect)",
    "defect_c": "Defect C (Inclusion / Slag)",
    "defect_d": "Defect D (Surface patch / Scuff)"
}

# Color palette for classes (R, G, B)
CLASS_COLORS = [
    (255, 60, 60),    # Defect A: Red
    (50, 220, 50),    # Defect B: Green
    (30, 150, 255),   # Defect C: Cyan-Blue
    (255, 180, 20)    # Defect D: Orange-Gold
]

# Calibrated optimal thresholds & min component areas from notebook
DEFAULT_THRESHOLDS = [0.35, 0.22, 0.42, 0.42]
DEFAULT_MIN_AREAS = [25, 25, 220, 350]


def rle_encode(mask):
    """
    1-indexed, column-major (Fortran) Run-Length Encoding.
    Matches the official Kaggle competition format.
    """
    mask = np.asarray(mask)
    if mask.ndim != 2:
        raise ValueError(f"rle_encode expects a 2D mask, got shape {mask.shape}")
    pixels = np.concatenate([[0], (mask > 0).astype(np.uint8).flatten(order="F"), [0]])
    runs = np.where(pixels[1:] != pixels[:-1])[0] + 1
    runs[1::2] -= runs[::2]
    return " ".join(map(str, runs))


def rle_decode(rle, shape=(IMG_H, IMG_W)):
    """Decodes 1-indexed column-major RLE string back to 2D numpy mask."""
    total = int(shape[0] * shape[1])
    flat = np.zeros(total, dtype=np.uint8)
    if not isinstance(rle, str) or not rle.strip():
        return flat.reshape(shape, order="F")
    nums = np.asarray(rle.split(), dtype=np.int64)
    if nums.size % 2:
        return flat.reshape(shape, order="F")
    starts = nums[0::2] - 1
    lengths = nums[1::2]
    for start, length in zip(starts, lengths):
        end = min(int(start + length), total)
        if 0 <= start < total:
            flat[int(start):end] = 1
    return flat.reshape(shape, order="F")


def filter_small_components(mask, min_area):
    """Removes 8-connected components smaller than min_area pixels."""
    if min_area <= 1 or not mask.any():
        return mask
    mask_u8 = mask.astype(np.uint8)
    count, labels, stats, _ = cv2.connectedComponentsWithStats(mask_u8, connectivity=8)
    if count <= 1:
        return np.zeros_like(mask, dtype=bool)
    keep = np.flatnonzero(stats[:, cv2.CC_STAT_AREA] >= min_area)
    keep = keep[keep != 0]  # exclude background (0)
    return np.isin(labels, keep)


def preprocess_image(input_img):
    """
    Prepares input image for model:
    - Converts to Grayscale
    - Resizes to (IMG_H=128, IMG_W=800)
    - Normalizes with dataset mean & std
    Returns:
        tensor: shape (1, 1, 128, 800)
        display_rgb: numpy uint8 (128, 800, 3) for overlay rendering
    """
    if isinstance(input_img, str):
        gray = cv2.imread(input_img, cv2.IMREAD_GRAYSCALE)
    elif isinstance(input_img, np.ndarray):
        if input_img.ndim == 3:
            gray = cv2.cvtColor(input_img, cv2.COLOR_RGB2GRAY)
        else:
            gray = input_img.copy()
    elif isinstance(input_img, Image.Image):
        gray = np.array(input_img.convert("L"))
    else:
        raise ValueError("Unsupported image input format.")

    if gray.shape != (IMG_H, IMG_W):
        gray = cv2.resize(gray, (IMG_W, IMG_H), interpolation=cv2.INTER_AREA)

    norm_img = gray.astype(np.float32) / 255.0
    norm_img = (norm_img - NORM_MEAN) / NORM_STD
    tensor = torch.from_numpy(norm_img).unsqueeze(0).unsqueeze(0)  # (1, 1, 128, 800)

    display_rgb = cv2.cvtColor(gray, cv2.COLOR_GRAY2RGB)
    return tensor, display_rgb


@torch.no_grad()
def predict_with_tta(model, tensor_x, tta=True, device="cpu"):
    """
    Inference with optional 4-view Test-Time Augmentation (flips).
    Returns numpy probability array of shape (4, 128, 800) with values in [0, 1].
    """
    model.eval()
    tensor_x = tensor_x.to(device)
    probs = torch.sigmoid(model(tensor_x).float())
    if tta:
        # 4-view flip TTA: original, hflip, vflip, hvflip
        for dims in ([3], [2], [2, 3]):
            flipped_in = torch.flip(tensor_x, dims)
            flipped_pred = torch.sigmoid(model(flipped_in).float())
            probs = probs + torch.flip(flipped_pred, dims)
        probs = probs / 4.0
    return probs.squeeze(0).cpu().numpy()


def apply_postprocess(probs_np, thresholds=None, min_areas=None):
    """
    Applies calibrated probability thresholds and minimum connected component areas.
    """
    if thresholds is None:
        thresholds = DEFAULT_THRESHOLDS
    if min_areas is None:
        min_areas = DEFAULT_MIN_AREAS

    n_cls = probs_np.shape[0]
    masks = np.zeros_like(probs_np, dtype=bool)

    for c in range(n_cls):
        th = float(thresholds[c])
        area = int(min_areas[c])
        raw_mask = probs_np[c] >= th
        masks[c] = filter_small_components(raw_mask, area)

    return masks


def create_segmentation_overlay(base_rgb, masks, alpha=0.55):
    """
    Overlays colored semi-transparent defect masks and clear contour borders onto base image.
    """
    overlay = base_rgb.copy().astype(np.float32)

    for c, mask in enumerate(masks):
        if not mask.any():
            continue
        color = np.array(CLASS_COLORS[c], dtype=np.float32)
        # Apply semi-transparent fill
        idx = mask > 0
        overlay[idx] = (1 - alpha) * overlay[idx] + alpha * color

    overlay = np.clip(overlay, 0, 255).astype(np.uint8)

    # Draw sharp contour borders on top for precision visibility
    for c, mask in enumerate(masks):
        if not mask.any():
            continue
        mask_u8 = mask.astype(np.uint8)
        contours, _ = cv2.findContours(mask_u8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        color = CLASS_COLORS[c]
        cv2.drawContours(overlay, contours, -1, color, 1)

    return overlay


def compute_defect_diagnostics(masks, display_shape=(IMG_H, IMG_W)):
    """
    Computes per-class analytics:
    - Detection boolean
    - Component count
    - Total pixel area
    - Coverage percentage
    - RLE encoding string
    """
    total_canvas_pixels = float(display_shape[0] * display_shape[1])
    diagnostics = []
    has_any_defect = False

    for c, name in enumerate(CLASS_NAMES):
        mask = masks[c]
        is_detected = bool(mask.any())
        if is_detected:
            has_any_defect = True
            count, _, stats, _ = cv2.connectedComponentsWithStats(mask.astype(np.uint8), connectivity=8)
            comp_count = max(0, count - 1)
            area_px = int(mask.sum())
            area_pct = round(100.0 * (area_px / total_canvas_pixels), 3)
            rle_str = rle_encode(mask)
        else:
            comp_count = 0
            area_px = 0
            area_pct = 0.0
            rle_str = ""

        diagnostics.append({
            "Class Name": name,
            "Label": CLASS_LABELS[name],
            "Detected": "YES" if is_detected else "NO",
            "Components": comp_count,
            "Area (px)": area_px,
            "Surface %": f"{area_pct}%",
            "RLE Encoded": (rle_str[:60] + "...") if len(rle_str) > 60 else (rle_str or "EMPTY")
        })

    overall_status = "DEFECT DETECTED" if has_any_defect else "CLEAN / NO DEFECT"
    return diagnostics, overall_status

	"""
	Utility functions for preprocessing, TTA inference, postprocessing,
	visualization, and Kaggle RLE formatting.
	"""

	import cv2
	import numpy as np
	import torch
	from PIL import Image

	IMG_H = 128
	IMG_W = 800
	NORM_MEAN = 0.449
	NORM_STD = 0.226

	CLASS_NAMES = ["defect_a", "defect_b", "defect_c", "defect_d"]
	CLASS_LABELS = {
	"defect_a": "Defect A (Pitted / Scratch)",
	"defect_b": "Defect B (Minor defect)",
	"defect_c": "Defect C (Inclusion / Slag)",
	"defect_d": "Defect D (Surface patch / Scuff)"
	}

	# Color palette for classes (R, G, B)
	CLASS_COLORS = [
	(255, 60, 60), # Defect A: Red
	(50, 220, 50), # Defect B: Green
	(30, 150, 255), # Defect C: Cyan-Blue
	(255, 180, 20) # Defect D: Orange-Gold
	]

	# Calibrated optimal thresholds & min component areas from notebook
	DEFAULT_THRESHOLDS = [0.35, 0.22, 0.42, 0.42]
	DEFAULT_MIN_AREAS = [25, 25, 220, 350]


	def rle_encode(mask):
	"""
	1-indexed, column-major (Fortran) Run-Length Encoding.
	Matches the official Kaggle competition format.
	"""
	mask = np.asarray(mask)
	if mask.ndim != 2:
	raise ValueError(f"rle_encode expects a 2D mask, got shape {mask.shape}")
	pixels = np.concatenate([[0], (mask > 0).astype(np.uint8).flatten(order="F"), [0]])
	runs = np.where(pixels[1:] != pixels[:-1])[0] + 1
	runs[1::2] -= runs[::2]
	return " ".join(map(str, runs))


	def rle_decode(rle, shape=(IMG_H, IMG_W)):
	"""Decodes 1-indexed column-major RLE string back to 2D numpy mask."""
	total = int(shape[0] * shape[1])
	flat = np.zeros(total, dtype=np.uint8)
	if not isinstance(rle, str) or not rle.strip():
	return flat.reshape(shape, order="F")
	nums = np.asarray(rle.split(), dtype=np.int64)
	if nums.size % 2:
	return flat.reshape(shape, order="F")
	starts = nums[0::2] - 1
	lengths = nums[1::2]
	for start, length in zip(starts, lengths):
	end = min(int(start + length), total)
	if 0 <= start < total:
	flat[int(start):end] = 1
	return flat.reshape(shape, order="F")


	def filter_small_components(mask, min_area):
	"""Removes 8-connected components smaller than min_area pixels."""
	if min_area <= 1 or not mask.any():
	return mask
	mask_u8 = mask.astype(np.uint8)
	count, labels, stats, _ = cv2.connectedComponentsWithStats(mask_u8, connectivity=8)
	if count <= 1:
	return np.zeros_like(mask, dtype=bool)
	keep = np.flatnonzero(stats[:, cv2.CC_STAT_AREA] >= min_area)
	keep = keep[keep != 0] # exclude background (0)
	return np.isin(labels, keep)


	def preprocess_image(input_img):
	"""
	Prepares input image for model:
	- Converts to Grayscale
	- Resizes to (IMG_H=128, IMG_W=800)
	- Normalizes with dataset mean & std
	Returns:
	tensor: shape (1, 1, 128, 800)
	display_rgb: numpy uint8 (128, 800, 3) for overlay rendering
	"""
	if isinstance(input_img, str):
	gray = cv2.imread(input_img, cv2.IMREAD_GRAYSCALE)
	elif isinstance(input_img, np.ndarray):
	if input_img.ndim == 3:
	gray = cv2.cvtColor(input_img, cv2.COLOR_RGB2GRAY)
	else:
	gray = input_img.copy()
	elif isinstance(input_img, Image.Image):
	gray = np.array(input_img.convert("L"))
	else:
	raise ValueError("Unsupported image input format.")

	if gray.shape != (IMG_H, IMG_W):
	gray = cv2.resize(gray, (IMG_W, IMG_H), interpolation=cv2.INTER_AREA)

	norm_img = gray.astype(np.float32) / 255.0
	norm_img = (norm_img - NORM_MEAN) / NORM_STD
	tensor = torch.from_numpy(norm_img).unsqueeze(0).unsqueeze(0) # (1, 1, 128, 800)

	display_rgb = cv2.cvtColor(gray, cv2.COLOR_GRAY2RGB)
	return tensor, display_rgb


	@torch.no_grad()
	def predict_with_tta(model, tensor_x, tta=True, device="cpu"):
	"""
	Inference with optional 4-view Test-Time Augmentation (flips).
	Returns numpy probability array of shape (4, 128, 800) with values in [0, 1].
	"""
	model.eval()
	tensor_x = tensor_x.to(device)
	probs = torch.sigmoid(model(tensor_x).float())
	if tta:
	# 4-view flip TTA: original, hflip, vflip, hvflip
	for dims in ([3], [2], [2, 3]):
	flipped_in = torch.flip(tensor_x, dims)
	flipped_pred = torch.sigmoid(model(flipped_in).float())
	probs = probs + torch.flip(flipped_pred, dims)
	probs = probs / 4.0
	return probs.squeeze(0).cpu().numpy()


	def apply_postprocess(probs_np, thresholds=None, min_areas=None):
	"""
	Applies calibrated probability thresholds and minimum connected component areas.
	"""
	if thresholds is None:
	thresholds = DEFAULT_THRESHOLDS
	if min_areas is None:
	min_areas = DEFAULT_MIN_AREAS

	n_cls = probs_np.shape[0]
	masks = np.zeros_like(probs_np, dtype=bool)

	for c in range(n_cls):
	th = float(thresholds[c])
	area = int(min_areas[c])
	raw_mask = probs_np[c] >= th
	masks[c] = filter_small_components(raw_mask, area)

	return masks


	def create_segmentation_overlay(base_rgb, masks, alpha=0.55):
	"""
	Overlays colored semi-transparent defect masks and clear contour borders onto base image.
	"""
	overlay = base_rgb.copy().astype(np.float32)

	for c, mask in enumerate(masks):
	if not mask.any():
	continue
	color = np.array(CLASS_COLORS[c], dtype=np.float32)
	# Apply semi-transparent fill
	idx = mask > 0
	overlay[idx] = (1 - alpha) * overlay[idx] + alpha * color

	overlay = np.clip(overlay, 0, 255).astype(np.uint8)

	# Draw sharp contour borders on top for precision visibility
	for c, mask in enumerate(masks):
	if not mask.any():
	continue
	mask_u8 = mask.astype(np.uint8)
	contours, _ = cv2.findContours(mask_u8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
	color = CLASS_COLORS[c]
	cv2.drawContours(overlay, contours, -1, color, 1)

	return overlay


	def compute_defect_diagnostics(masks, display_shape=(IMG_H, IMG_W)):
	"""
	Computes per-class analytics:
	- Detection boolean
	- Component count
	- Total pixel area
	- Coverage percentage
	- RLE encoding string
	"""
	total_canvas_pixels = float(display_shape[0] * display_shape[1])
	diagnostics = []
	has_any_defect = False

	for c, name in enumerate(CLASS_NAMES):
	mask = masks[c]
	is_detected = bool(mask.any())
	if is_detected:
	has_any_defect = True
	count, _, stats, _ = cv2.connectedComponentsWithStats(mask.astype(np.uint8), connectivity=8)
	comp_count = max(0, count - 1)
	area_px = int(mask.sum())
	area_pct = round(100.0 * (area_px / total_canvas_pixels), 3)
	rle_str = rle_encode(mask)
	else:
	comp_count = 0
	area_px = 0
	area_pct = 0.0
	rle_str = ""

	diagnostics.append({
	"Class Name": name,
	"Label": CLASS_LABELS[name],
	"Detected": "YES" if is_detected else "NO",
	"Components": comp_count,
	"Area (px)": area_px,
	"Surface %": f"{area_pct}%",
	"RLE Encoded": (rle_str[:60] + "...") if len(rle_str) > 60 else (rle_str or "EMPTY")
	})

	overall_status = "DEFECT DETECTED" if has_any_defect else "CLEAN / NO DEFECT"
	return diagnostics, overall_status