Upload 5 files

README.md

@@ -0,0 +1,19 @@
+---
+license: apache-2.0
+tags:
+- computer-vision
+- image-segmentation
+- welding
+- weld-inspection
+- yolo
+pipeline_tag: image-segmentation
+---
+# WeldVision Ensemble
+
+Four-model YOLO segmentation ensemble for prototype weld-surface defect inspection.
+
+Models: `best.pt`, `best_v0.pt`, `crack_specialist.pt`, `spatters_specialist.pt`.
+
+The custom `handler.py` loads all four weights and performs class-aware mask-IoU merging.
+
+**Prototype only:** the score and PASS/REVIEW/FAIL rules are not welding-code acceptance criteria and do not replace qualified human inspection.

config.json

@@ -0,0 +1 @@
+{"name":"WeldVision-Ensemble","version":"1.0","models":["best.pt","best_v0.pt","crack_specialist.pt","spatters_specialist.pt"],"confidence":0.25,"ensemble_mask_iou":0.50,"image_size":640}

ensemble.py

@@ -0,0 +1,43 @@
+import cv2
+import numpy as np
+from ultralytics import YOLO
+
+class WeldVisionEnsemble:
+    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_dir, conf=.25, ensemble_iou=.50, imgsz=640):
+        self.conf=conf; self.ensemble_iou=ensemble_iou; self.imgsz=imgsz
+        self.base1=YOLO(f'{model_dir}/best.pt'); self.base2=YOLO(f'{model_dir}/best_v0.pt')
+        self.crack=YOLO(f'{model_dir}/crack_specialist.pt'); self.spatters=YOLO(f'{model_dir}/spatters_specialist.pt')
+    @staticmethod
+    def mask_iou(a,b):
+        a=a.astype(bool); b=b.astype(bool); u=np.logical_or(a,b).sum(); return float(np.logical_and(a,b).sum()/u) if u else 0.0
+    def extract(self,r,source):
+        if r.boxes is None or r.masks is None: return []
+        out=[]
+        for box,mask in zip(r.boxes.data.cpu().numpy(),r.masks.data.cpu().numpy()):
+            x1,y1,x2,y2,conf,cls=box; cls=int(cls)
+            name='Crack' if source=='crack' else 'Spatters' if source=='spatters' else self.CLASS_NAMES.get(cls,str(cls))
+            out.append({'box':np.array([x1,y1,x2,y2],dtype=np.float32),'conf':float(conf),'class_name':name,'mask':mask.astype(np.float32),'source':source})
+        return out
+    def run(self,model,img,source):
+        return self.extract(model.predict(img,conf=self.conf,imgsz=self.imgsz,verbose=False)[0],source)
+    def merge(self,preds):
+        selected=[]
+        for p in sorted(preds,key=lambda x:x['conf'],reverse=True):
+            if not any(p['class_name']==q['class_name'] and self.mask_iou(p['mask'],q['mask'])>=self.ensemble_iou for q in selected): selected.append(p)
+        return selected
+    def predict(self,image):
+        image=np.asarray(image)
+        if image.ndim!=3 or image.shape[2]!=3: raise ValueError('image must be HxWx3 RGB')
+        bgr=cv2.cvtColor(image,cv2.COLOR_RGB2BGR)
+        p1=self.run(self.base1,bgr,'base_1'); p2=self.run(self.base2,bgr,'base_2'); pc=self.run(self.crack,bgr,'crack'); ps=self.run(self.spatters,bgr,'spatters')
+        merged=self.merge(p1+p2+pc+ps); detections=[]; h,w=bgr.shape[:2]
+        for p in merged:
+            c=p['class_name']
+            if c=='Good Welding': continue
+            sev='HIGH' if self.PENALTIES.get(c,0)>=30 else 'MEDIUM' if self.PENALTIES.get(c,0)>=15 else 'LOW' if self.PENALTIES.get(c,0)>0 else 'NONE'
+            x1,y1,x2,y2=p['box']; detections.append({'class':c,'confidence':round(p['conf'],4),'severity':sev,'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,100-sum(self.PENALTIES.get(d['class'],0) for d in detections)); high=any(d['severity']=='HIGH' for d in detections); med=any(d['severity']=='MEDIUM' for d in detections)
+        decision='FAIL' if score<70 or high else 'REVIEW' if score<85 or med else 'PASS'
+        return {'model':'WeldVision-Ensemble','version':'1.0','decision':decision,'score':score,'detections':detections,'model_counts':{'best.pt':len(p1),'best_v0.pt':len(p2),'crack_specialist.pt':len(pc),'spatters_specialist.pt':len(ps),'ensemble':len(merged)},'recommendation':'Significant visible defect detected. Human inspection required.' if decision=='FAIL' else 'Moderate defect or quality concern detected. Human review required.' if decision=='REVIEW' else 'No significant visible defect detected.'}

handler.py

@@ -0,0 +1,13 @@
+import io, os
+import numpy as np
+from PIL import Image
+from ensemble import WeldVisionEnsemble
+MODEL=WeldVisionEnsemble(os.path.join(os.path.dirname(__file__),'models'))
+def _image(data):
+    if isinstance(data,dict): data=data.get('image',data.get('inputs'))
+    if isinstance(data,str):
+        import base64; data=base64.b64decode(data)
+    if isinstance(data,(bytes,bytearray)): return np.array(Image.open(io.BytesIO(data)).convert('RGB'))
+    if isinstance(data,np.ndarray): return data
+    raise ValueError('Expected image bytes, base64, or numpy array')
+def handler(data): return MODEL.predict(_image(data))

requirements.txt

@@ -0,0 +1,6 @@
+ultralytics==8.4.126
+opencv-python-headless
+numpy
+Pillow
+torch
+torchvision