johnwesley756/instance-segmentation

1import cv22from ultralytics import YOLO3 4MODEL_PATH = "best.pt"5 6# Load YOLO once7model = YOLO(MODEL_PATH)8 9 10def generate_summary(severity: str) -> str:11    summaries = {12        "Cavity": "Advanced tooth decay detected. Immediate dental treatment is recommended.",13        "Caries": "Early-stage caries detected. Preventive treatment can stop progression.",14        "Tooth": "Healthy tooth detected. Maintain good oral hygiene.",15        "No Detection": "No dental issues detected or image unclear."16    }17    return summaries.get(severity, "Unknown condition detected.")18 19 20def run_inference(image):21    results = model.predict(image, conf=0.25)22    annotated = results[0].plot()23 24    detections = []25    class_names = []26 27    if results[0].boxes is not None:28        for box in results[0].boxes:29            cid = int(box.cls.item())30            cname = results[0].names[cid].lower()31            conf = float(box.conf.item())32            bbox = box.xyxy[0].tolist()33 34            detections.append({35                "class": cname,36                "confidence": round(conf, 3),37                "bbox": [round(x, 2) for x in bbox]38            })39            class_names.append(cname)40 41    if "cavity" in class_names:42        severity = "Cavity"43    elif "caries" in class_names:44        severity = "Caries"45    elif class_names:46        severity = "Tooth"47    else:48        severity = "No Detection"49 50    return severity, generate_summary(severity), detections, annotated51