Upload 2 files

weldvision/__init__.py

@@ -0,0 +1,3 @@
+from .ensemble import WeldVision
+
+__all__ = ["WeldVision"]

weldvision/ensemble.py

@@ -0,0 +1,249 @@
+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,
+        }