from __future__ import annotations from dataclasses import dataclass from pathlib import Path from typing import Any import cv2 import numpy as np from core.config import settings from core.logger import logger try: from ultralytics import YOLO except Exception: # pragma: no cover - optional import guard YOLO = None @dataclass(frozen=True) class SurfaceGateResult: is_steel: bool confidence: float mode: str label: str reason: str metrics: dict[str, float] roi_bbox: tuple[int, int, int, int] | None = None roi_area_ratio: float | None = None def to_metadata(self) -> dict[str, Any]: return { "passed": self.is_steel, "confidence": round(self.confidence, 3), "mode": self.mode, "label": self.label, "reason": self.reason, "metrics": self.metrics, "roi_bbox": list(self.roi_bbox) if self.roi_bbox else None, "roi_area_ratio": round(self.roi_area_ratio, 4) if self.roi_area_ratio is not None else None, } class SteelSurfaceGate: def __init__(self) -> None: self.detector_model = self._load_model(settings.SURFACE_DETECTOR_PATH, "detector") self.classifier_model = self._load_model(settings.SURFACE_CLASSIFIER_PATH, "classifier") def _load_model(self, model_path: str | None, kind: str): if not model_path: return None resolved_path = Path(model_path) if not resolved_path.exists(): logger.warning("Surface %s path does not exist: %s", kind, resolved_path) return None if YOLO is None: logger.warning("Ultralytics is unavailable for the surface %s model", kind) return None try: model = YOLO(str(resolved_path)) logger.info("Surface %s model loaded from: %s", kind, resolved_path) return model except Exception as exc: # pragma: no cover - defensive runtime fallback logger.warning("Surface %s model failed to load: %s", kind, exc) return None def evaluate(self, image: np.ndarray) -> SurfaceGateResult: if self.detector_model is not None: return self._evaluate_with_detector(image) if self.classifier_model is not None: return self._evaluate_with_classifier(image) return self._evaluate_with_heuristic(image) def _evaluate_with_detector(self, image: np.ndarray) -> SurfaceGateResult: results = self.detector_model( image, conf=settings.SURFACE_DETECTOR_CONFIDENCE, imgsz=settings.SURFACE_DETECTOR_IMAGE_SIZE, verbose=False, ) result = results[0] boxes = getattr(result, "boxes", None) names = getattr(result, "names", None) or getattr(self.detector_model, "names", {}) if boxes is None or len(boxes) == 0: return SurfaceGateResult( is_steel=False, confidence=0.99, mode="detector", label="non_steel", reason=( "Frame skipped because no steel-surface region was localized. Aim the camera closer to the sheet or coil." ), metrics={"boxes_detected": 0.0}, ) height, width = image.shape[:2] image_area = max(height * width, 1) candidates: list[dict[str, float | tuple[int, int, int, int] | str]] = [] for raw_box in boxes: xyxy = raw_box.xyxy[0].tolist() x1, y1, x2, y2 = [int(round(value)) for value in xyxy] x1 = max(0, min(x1, width - 1)) y1 = max(0, min(y1, height - 1)) x2 = max(x1 + 1, min(x2, width)) y2 = max(y1 + 1, min(y2, height)) box_width = x2 - x1 box_height = y2 - y1 area_ratio = (box_width * box_height) / image_area cls_tensor = getattr(raw_box, "cls", None) cls_index = int(cls_tensor[0].item()) if cls_tensor is not None else 0 label = self._resolve_label(names, cls_index) confidence = float(raw_box.conf[0].item()) if not self._detector_label_matches(label, cls_index, names): continue expand_ratio = settings.SURFACE_DETECTOR_EXPAND_RATIO pad_x = int(box_width * expand_ratio) pad_y = int(box_height * expand_ratio) expanded_x1 = max(0, x1 - pad_x) expanded_y1 = max(0, y1 - pad_y) expanded_x2 = min(width, x2 + pad_x) expanded_y2 = min(height, y2 + pad_y) expanded_bbox = ( expanded_x1, expanded_y1, expanded_x2 - expanded_x1, expanded_y2 - expanded_y1, ) expanded_area_ratio = (expanded_bbox[2] * expanded_bbox[3]) / image_area selector_score = (confidence * 0.8) + (min(expanded_area_ratio / 0.55, 1.0) * 0.2) candidates.append( { "confidence": confidence, "label": label, "bbox": expanded_bbox, "area_ratio": expanded_area_ratio, "selector_score": selector_score, } ) if not candidates: return SurfaceGateResult( is_steel=False, confidence=0.99, mode="detector", label="non_steel", reason=( "Frame skipped because the localized objects did not match the expected steel-surface class." ), metrics={"boxes_detected": float(len(boxes))}, ) best_candidate = max( candidates, key=lambda candidate: ( float(candidate["selector_score"]), float(candidate["confidence"]), float(candidate["area_ratio"]), ), ) roi_area_ratio = float(best_candidate["area_ratio"]) if roi_area_ratio < settings.SURFACE_DETECTOR_MIN_AREA_RATIO: return SurfaceGateResult( is_steel=False, confidence=float(best_candidate["confidence"]), mode="detector", label=str(best_candidate["label"]), reason=( "Frame skipped because the localized steel region is too small for reliable defect inspection. Move closer to the material." ), metrics={ "boxes_detected": float(len(candidates)), "roi_area_ratio": round(roi_area_ratio, 4), }, roi_bbox=best_candidate["bbox"], roi_area_ratio=roi_area_ratio, ) roi_x, roi_y, roi_width, roi_height = best_candidate["bbox"] roi_image = image[roi_y:roi_y + roi_height, roi_x:roi_x + roi_width] heuristic_result = self._evaluate_with_heuristic(roi_image) if not heuristic_result.is_steel: return SurfaceGateResult( is_steel=False, confidence=float(best_candidate["confidence"]), mode="detector+heuristic", label=str(best_candidate["label"]), reason=( "Frame skipped because the localized ROI did not pass the steel-surface texture validation step. " "Reduce background content and center the actual material." ), metrics={ "boxes_detected": float(len(candidates)), "roi_area_ratio": round(roi_area_ratio, 4), "selector_score": round(float(best_candidate["selector_score"]), 4), **heuristic_result.metrics, }, roi_bbox=best_candidate["bbox"], roi_area_ratio=roi_area_ratio, ) return SurfaceGateResult( is_steel=True, confidence=float(best_candidate["confidence"]), mode="detector+heuristic", label=str(best_candidate["label"]), reason=( f"Steel-surface detector localized an inspection ROI with {float(best_candidate['confidence']):.0%} confidence, " "and the ROI passed surface-texture validation." ), metrics={ "boxes_detected": float(len(candidates)), "roi_area_ratio": round(roi_area_ratio, 4), "selector_score": round(float(best_candidate["selector_score"]), 4), **heuristic_result.metrics, }, roi_bbox=best_candidate["bbox"], roi_area_ratio=roi_area_ratio, ) def _evaluate_with_classifier(self, image: np.ndarray) -> SurfaceGateResult: results = self.classifier_model(image, verbose=False) result = results[0] probs = getattr(result, "probs", None) names = getattr(result, "names", None) or getattr(self.classifier_model, "names", {}) if probs is None: raise RuntimeError("Classification model returned no probabilities") top_index = int(getattr(probs, "top1", 0)) top_confidence_raw = getattr(probs, "top1conf") top_confidence = float( top_confidence_raw.item() if hasattr(top_confidence_raw, "item") else top_confidence_raw ) label = self._resolve_label(names, top_index) steel_label = settings.SURFACE_CLASSIFIER_STEEL_LABEL.strip().lower() normalized_label = label.strip().lower() is_steel = ( normalized_label == steel_label or steel_label in normalized_label or normalized_label in steel_label ) meets_threshold = top_confidence >= settings.SURFACE_MIN_STEEL_CONFIDENCE if is_steel and meets_threshold: reason = ( f"Steel surface classifier accepted the frame with {top_confidence:.0%} confidence." ) elif is_steel: reason = ( "Frame resembles steel, but the classifier confidence is too low for reliable defect analysis." ) else: reason = ( f'Classifier labeled the frame as "{label}" instead of steel, so the defect model was skipped.' ) return SurfaceGateResult( is_steel=bool(is_steel and meets_threshold), confidence=top_confidence, mode="classifier", label=label, reason=reason, metrics={}, roi_bbox=(0, 0, image.shape[1], image.shape[0]) if is_steel and meets_threshold else None, roi_area_ratio=1.0 if is_steel and meets_threshold else None, ) def _evaluate_with_heuristic(self, image: np.ndarray) -> SurfaceGateResult: features = self._extract_features(image) checks = { "gray_ratio": features["gray_ratio"] >= settings.SURFACE_MIN_GRAY_RATIO, "low_sat_ratio": features["low_saturation_ratio"] >= settings.SURFACE_MIN_LOW_SAT_RATIO, "mean_saturation": features["mean_saturation"] <= settings.SURFACE_MAX_MEAN_SATURATION, "colorfulness": features["colorfulness"] <= settings.SURFACE_MAX_COLORFULNESS, "skin_ratio": features["skin_ratio"] <= settings.SURFACE_MAX_SKIN_RATIO, "texture_variance": features["texture_variance"] >= settings.SURFACE_MIN_TEXTURE_VARIANCE, } weights = { "gray_ratio": 0.24, "low_sat_ratio": 0.22, "mean_saturation": 0.16, "colorfulness": 0.16, "skin_ratio": 0.12, "texture_variance": 0.10, } score = sum(weights[name] for name, passed in checks.items() if passed) chroma_gate = checks["gray_ratio"] and checks["low_sat_ratio"] is_steel = score >= 0.72 and chroma_gate and checks["skin_ratio"] failed_checks = [name for name, passed in checks.items() if not passed] confidence = score if is_steel else min(0.99, max(0.55, 1.0 - score + 0.08 * len(failed_checks))) if is_steel: reason = ( f"Frame passed the steel-surface gate with {confidence:.0%} confidence and proceeded to defect segmentation." ) else: reason = self._build_failure_reason(features, failed_checks) return SurfaceGateResult( is_steel=is_steel, confidence=confidence, mode="heuristic", label="steel" if is_steel else "non_steel", reason=reason, metrics={key: round(value, 4) for key, value in features.items()}, roi_bbox=(0, 0, image.shape[1], image.shape[0]) if is_steel else None, roi_area_ratio=1.0 if is_steel else None, ) def _resolve_label(self, names: Any, index: int) -> str: if isinstance(names, dict): return str(names.get(index, index)) if isinstance(names, list) and 0 <= index < len(names): return str(names[index]) return str(index) def _detector_label_matches(self, label: str, cls_index: int, names: Any) -> bool: target = settings.SURFACE_DETECTOR_CLASS_NAME.strip().lower() normalized_label = label.strip().lower() if not target: return True if normalized_label == target or target in normalized_label or normalized_label in target: return True if isinstance(names, dict) and len(names) == 1 and cls_index == 0: return True if isinstance(names, list) and len(names) == 1 and cls_index == 0: return True return False def _extract_features(self, image: np.ndarray) -> dict[str, float]: height, width = image.shape[:2] target_width = min(320, max(96, width)) target_height = max(96, int(height * target_width / max(width, 1))) resized = cv2.resize(image, (target_width, target_height)) hsv = cv2.cvtColor(resized, cv2.COLOR_BGR2HSV) gray = cv2.cvtColor(resized, cv2.COLOR_BGR2GRAY) ycrcb = cv2.cvtColor(resized, cv2.COLOR_BGR2YCrCb) b_channel, g_channel, r_channel = [channel.astype(np.float32) for channel in cv2.split(resized)] saturation = hsv[:, :, 1].astype(np.float32) gray_delta = settings.SURFACE_GRAY_DELTA gray_mask = ( (np.abs(r_channel - g_channel) <= gray_delta) & (np.abs(r_channel - b_channel) <= gray_delta) & (np.abs(g_channel - b_channel) <= gray_delta) ) rg = np.abs(r_channel - g_channel) yb = np.abs(0.5 * (r_channel + g_channel) - b_channel) colorfulness = ( np.sqrt(float(rg.std()) ** 2 + float(yb.std()) ** 2) + 0.3 * np.sqrt(float(rg.mean()) ** 2 + float(yb.mean()) ** 2) ) luminance, cr_channel, cb_channel = cv2.split(ycrcb) skin_mask = ( (cr_channel > 135) & (cr_channel < 180) & (cb_channel > 85) & (cb_channel < 135) & (luminance > 60) ) return { "gray_ratio": float(gray_mask.mean()), "low_saturation_ratio": float( (saturation <= settings.SURFACE_LOW_SAT_PIXEL_THRESHOLD).mean() ), "mean_saturation": float(saturation.mean()), "colorfulness": float(colorfulness), "skin_ratio": float(skin_mask.mean()), "texture_variance": float(cv2.Laplacian(gray, cv2.CV_32F).var()), } def _build_failure_reason(self, features: dict[str, float], failed_checks: list[str]) -> str: if features["skin_ratio"] > settings.SURFACE_MAX_SKIN_RATIO: return ( "Frame skipped because prominent skin-tone regions were detected. Aim the camera only at the steel surface." ) if features["gray_ratio"] < settings.SURFACE_MIN_GRAY_RATIO: return ( "Frame skipped because it contains too much color variation to match the expected steel surface appearance." ) if features["low_saturation_ratio"] < settings.SURFACE_MIN_LOW_SAT_RATIO: return ( "Frame skipped because the image is too saturated. Move closer to the metal surface and reduce background content." ) if features["texture_variance"] < settings.SURFACE_MIN_TEXTURE_VARIANCE: return ( "Frame skipped because the visible area is too flat or out of focus for reliable steel-surface validation." ) if features["colorfulness"] > settings.SURFACE_MAX_COLORFULNESS: return ( "Frame skipped because the scene looks like a general object view instead of a steel inspection close-up." ) failed_text = ", ".join(failed_checks) if failed_checks else "multiple surface validation checks" return ( f"Frame skipped because it did not pass the steel-surface gate ({failed_text}). Reposition the camera toward the material and retry." )