predict.py

import os
import base64
import cv2
import numpy as np

MODEL_DIR = os.path.join(os.path.dirname(__file__), "models")
MODEL_PATH = os.path.join(MODEL_DIR, "best.onnx")
_model = None

CLASS_NAMES = {0: "comedone", 1: "nodules", 2: "papules", 3: "pustules"}
CLASS_ID_TO_NAME = {0: "comedone", 1: "nodules", 2: "papules", 3: "pustules"}

CLASS_COLORS = {
    "comedone": (0, 200, 255),
    "nodules":  (0, 0, 255),
    "papules":  (0, 165, 255),
    "pustules": (255, 0, 0),
}


def get_model():
    from ultralytics import YOLO
    global _model
    if _model is None:
        _model = YOLO(MODEL_PATH)
    return _model


def draw_annotations(image_bgr: np.ndarray, results) -> str:
    """Gambar bbox + label pada gambar, return sebagai base64 JPEG string."""
    annotated = image_bgr.copy()

    if results.boxes is not None and len(results.boxes) > 0:
        boxes   = results.boxes.xyxy.cpu().numpy()
        confs   = results.boxes.conf.cpu().numpy()
        cls_ids = results.boxes.cls.cpu().numpy().astype(int)

        for box, conf, cls_id in zip(boxes, confs, cls_ids):
            cls_name = CLASS_ID_TO_NAME.get(cls_id, "unknown")
            color = CLASS_COLORS.get(cls_name, (0, 255, 0))
            x1, y1, x2, y2 = map(int, box)

            cv2.rectangle(annotated, (x1, y1), (x2, y2), color, 2)

            label = f"{cls_name} {conf:.2f}"
            (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
            cv2.rectangle(annotated, (x1, y1 - th - 6), (x1 + tw + 4, y1), color, -1)
            cv2.putText(annotated, label, (x1 + 2, y1 - 3),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)

    _, buffer = cv2.imencode('.jpg', annotated, [cv2.IMWRITE_JPEG_QUALITY, 85])
    return base64.b64encode(buffer).decode('utf-8')


def detect_acne(image_bgr: np.ndarray, conf_threshold: float = 0.05, iou_threshold: float = 0.35) -> dict:
    """
    Deteksi jerawat pada gambar (numpy array BGR).
    Return dict berisi list deteksi, ringkasan, dan gambar anotasi (base64).
    """
    model = get_model()

    if model is None:
        return {
            "detections": [],
            "detected_classes": [],
            "summary": {},
            "total_detections": 0,
            "model_loaded": False,
            "annotated_image": None,
        }

    results = model.predict(
        source=image_bgr,
        conf=conf_threshold,
        iou=iou_threshold,
        verbose=False,
        imgsz=640,
    )[0]

    detections = []
    summary: dict[str, int] = {}

    if results.boxes is not None and len(results.boxes) > 0:
        for box, conf, cls_id in zip(
            results.boxes.xyxy.tolist(),
            results.boxes.conf.tolist(),
            results.boxes.cls.tolist(),
        ):
            cls_name = CLASS_ID_TO_NAME.get(int(cls_id), "unknown")
            summary[cls_name] = summary.get(cls_name, 0) + 1
            detections.append({
                "class": cls_name,
                "confidence": round(float(conf), 4),
                "bbox_xyxy": [round(v, 4) for v in box],
            })

    annotated_b64 = draw_annotations(image_bgr, results)

    return {
        "detections": detections,
        "detected_classes": list(summary.keys()),
        "summary": summary,
        "total_detections": len(detections),
        "model_loaded": True,
        "annotated_image": annotated_b64,
        "conf_threshold": conf_threshold,
        "iou_threshold": iou_threshold,
    }