weldvision/ensemble.py

import cv2
import numpy as np
from pathlib import Path
from ultralytics import YOLO
from huggingface_hub import hf_hub_download


class WeldVision:
    """
    Local four-model WeldVision ensemble.

    Usage:
        model = WeldVision.from_pretrained(
            "bhavibhatt/weldvision-ensemble"
        )
        result = model.predict("weld.jpg")
    """

    CLASS_NAMES = {
        0: "Bad Welding",
        1: "Crack",
        2: "Excess Reinforcement",
        3: "Good Welding",
        4: "Porosity",
        5: "Spatters",
    }

    PENALTIES = {
        "Crack": 40,
        "Porosity": 15,
        "Spatters": 5,
        "Excess Reinforcement": 20,
        "Bad Welding": 50,
        "Good Welding": 0,
    }

    def __init__(self, model_paths, conf=0.25, ensemble_iou=0.50, imgsz=640):
        self.conf = conf
        self.ensemble_iou = ensemble_iou
        self.imgsz = imgsz

        self.base_1 = YOLO(str(model_paths["best.pt"]))
        self.base_2 = YOLO(str(model_paths["best_v0.pt"]))
        self.crack = YOLO(str(model_paths["crack_specialist.pt"]))
        self.spatters = YOLO(str(model_paths["spatters_specialist.pt"]))

    @classmethod
    def from_pretrained(
        cls,
        repo_id,
        revision=None,
        cache_dir=None,
        conf=0.25,
        ensemble_iou=0.50,
        imgsz=640,
    ):
        """
        Download the four weights from Hugging Face Hub and load them locally.
        """
        names = [
            "best.pt",
            "best_v0.pt",
            "crack_specialist.pt",
            "spatters_specialist.pt",
        ]

        paths = {}
        for name in names:
            paths[name] = hf_hub_download(
                repo_id=repo_id,
                filename=f"weights/{name}",
                revision=revision,
                cache_dir=cache_dir,
            )

        return cls(
            paths,
            conf=conf,
            ensemble_iou=ensemble_iou,
            imgsz=imgsz,
        )

    @staticmethod
    def _load_image(image):
        if isinstance(image, (str, Path)):
            image = cv2.imread(str(image))
            if image is None:
                raise ValueError(f"Could not read image: {image}")
            return image

        if isinstance(image, np.ndarray):
            if image.ndim != 3 or image.shape[2] != 3:
                raise ValueError("Image must have shape H x W x 3")
            return image

        raise TypeError("image must be a file path or HxWx3 numpy array")

    @staticmethod
    def _mask_iou(a, b):
        a = a.astype(bool)
        b = b.astype(bool)
        inter = np.logical_and(a, b).sum()
        union = np.logical_or(a, b).sum()
        return float(inter / union) if union else 0.0

    def _extract(self, result, source):
        if result.boxes is None or result.masks is None:
            return []

        boxes = result.boxes.data.cpu().numpy()
        masks = result.masks.data.cpu().numpy()
        out = []

        for box, mask in zip(boxes, masks):
            x1, y1, x2, y2, conf, cls_id = box
            cls_id = int(cls_id)

            if source == "crack_specialist":
                class_name = "Crack"
            elif source == "spatters_specialist":
                class_name = "Spatters"
            else:
                class_name = self.CLASS_NAMES.get(cls_id, str(cls_id))

            out.append({
                "box": np.array([x1, y1, x2, y2], dtype=np.float32),
                "conf": float(conf),
                "class_name": class_name,
                "mask": mask.astype(np.float32),
                "source": source,
            })

        return out

    def _run(self, model, image, source):
        result = model.predict(
            image,
            conf=self.conf,
            imgsz=self.imgsz,
            verbose=False,
        )[0]
        return self._extract(result, source)

    def _merge(self, predictions):
        predictions = sorted(
            predictions,
            key=lambda p: p["conf"],
            reverse=True,
        )

        selected = []
        for candidate in predictions:
            duplicate = False

            for existing in selected:
                if candidate["class_name"] != existing["class_name"]:
                    continue

                if self._mask_iou(
                    candidate["mask"],
                    existing["mask"],
                ) >= self.ensemble_iou:
                    duplicate = True
                    break

            if not duplicate:
                selected.append(candidate)

        return selected

    def _severity(self, name):
        penalty = self.PENALTIES.get(name, 0)
        if penalty >= 30:
            return "HIGH"
        if penalty >= 15:
            return "MEDIUM"
        if penalty > 0:
            return "LOW"
        return "NONE"

    def predict(self, image):
        """
        Run the four-model ensemble.

        Returns a JSON-serializable dictionary.
        """
        image_bgr = self._load_image(image)
        h, w = image_bgr.shape[:2]

        p1 = self._run(self.base_1, image_bgr, "best.pt")
        p2 = self._run(self.base_2, image_bgr, "best_v0.pt")
        pc = self._run(self.crack, image_bgr, "crack_specialist")
        ps = self._run(self.spatters, image_bgr, "spatters_specialist")

        merged = self._merge(p1 + p2 + pc + ps)

        detections = []
        score = 100
        highest = "NONE"
        rank = {"NONE": 0, "LOW": 1, "MEDIUM": 2, "HIGH": 3}

        for p in merged:
            name = p["class_name"]
            if name == "Good Welding":
                continue

            severity = self._severity(name)
            score -= self.PENALTIES.get(name, 0)
            if rank[severity] > rank[highest]:
                highest = severity

            x1, y1, x2, y2 = p["box"]
            detections.append({
                "class": name,
                "confidence": round(float(p["conf"]), 4),
                "severity": severity,
                "box": [
                    round(float(max(0, min(w, x1))), 2),
                    round(float(max(0, min(h, y1))), 2),
                    round(float(max(0, min(w, x2))), 2),
                    round(float(max(0, min(h, y2))), 2),
                ],
                "source": p["source"],
            })

        score = max(0, score)

        if score < 70 or highest == "HIGH":
            decision = "FAIL"
        elif score < 85 or highest == "MEDIUM":
            decision = "REVIEW"
        else:
            decision = "PASS"

        return {
            "model": "WeldVision-Ensemble",
            "version": "1.0",
            "decision": decision,
            "score": score,
            "highest_severity": highest,
            "model_counts": {
                "best.pt": len(p1),
                "best_v0.pt": len(p2),
                "crack_specialist.pt": len(pc),
                "spatters_specialist.pt": len(ps),
                "ensemble": len(merged),
            },
            "detections": detections,
        }