src/visualizer.py

# src/visualizer.py

import cv2
import numpy as np
import os
from pathlib import Path
from typing import Dict, Optional

# Çoklu taş için BGR renk paleti
_STONE_COLORS = [
    (60,  20, 220),   # Kırmızı
    (220, 100, 20),   # Mavi
    (20,  180, 220),  # Turuncu
    (180, 20,  220),  # Mor
    (20,  220, 120),  # Yeşil-sarı
]
_BBOX_COLOR = (30, 200, 30)   # Yeşil — bbox her zaman yeşil


def _stone_color(idx: int) -> tuple:
    return _STONE_COLORS[idx % len(_STONE_COLORS)]


def draw_detections(
    image_path: str,
    detection_result: Dict,
    features: Dict,
    output_path: Optional[str] = None
) -> str:
    """
    Görüntü üzerine her taş için:
      - Yeşil bbox (detection çıktısı)
      - Renkli yarı saydam mask dolgusu + kontur (segmentation çıktısı)
      - Etiket: taş ID, boyut (mm), confidence

    Returns:
        Kaydedilen annotated görüntünün path'i
    """
    image = cv2.imread(image_path)
    if image is None:
        raise FileNotFoundError(f"Image not found: {image_path}")

    detections = detection_result.get("detections", [])
    stone_features = features.get("stones", [])

    # ── Adım 1: Mask dolguları ayrı overlay'de çiz ──────────────────────────
    overlay = image.copy()
    for idx, det in enumerate(detections):
        mask_polygon = det.get("mask_polygon")
        if mask_polygon:
            pts = np.array(mask_polygon, dtype=np.int32)
            cv2.fillPoly(overlay, [pts], _stone_color(idx))

    # Yarı saydamlık uygula (dolgu görüntüyle karışsın, çizgiler keskin kalsın)
    result = cv2.addWeighted(overlay, 0.30, image, 0.70, 0)

    # ── Adım 2: Bbox + mask kontur + etiket üst katmanda çiz ────────────────
    for idx, det in enumerate(detections):
        color = _stone_color(idx)

        # Bbox — yeşil ince dikdörtgen
        x1, y1, x2, y2 = [int(c) for c in det["bbox"]]
        cv2.rectangle(result, (x1, y1), (x2, y2), _BBOX_COLOR, 1)

        # Mask kontur — taşın gerçek sınırı
        mask_polygon = det.get("mask_polygon")
        if mask_polygon:
            pts = np.array(mask_polygon, dtype=np.int32)
            cv2.polylines(result, [pts], isClosed=True, color=color, thickness=2)

        # Etiket metni
        size_mm = 0.0
        if idx < len(stone_features):
            size_mm = stone_features[idx]["size"].get("max_dimension_mm", 0.0)
        conf = det["confidence"]
        label = f"#{idx + 1}  {size_mm}mm  conf:{conf:.2f}"

        # Etiket arka planı (okunabilirlik için)
        font = cv2.FONT_HERSHEY_SIMPLEX
        font_scale = 0.45
        thickness = 1
        (tw, th), baseline = cv2.getTextSize(label, font, font_scale, thickness)
        lx = x1
        ly = max(y1 - 6, th + 6)
        cv2.rectangle(result,
                      (lx, ly - th - 4),
                      (lx + tw + 6, ly + baseline),
                      _BBOX_COLOR, cv2.FILLED)
        cv2.putText(result, label,
                    (lx + 3, ly - 1),
                    font, font_scale, (255, 255, 255), thickness, cv2.LINE_AA)

    # Taş tespit edilmediyse küçük bir uyarı yaz
    if not detections:
        cv2.putText(result, "No stones detected",
                    (10, 30), cv2.FONT_HERSHEY_SIMPLEX,
                    0.7, (0, 0, 220), 2, cv2.LINE_AA)

    # ── Adım 3: Kaydet ──────────────────────────────────────────────────────
    if output_path is None:
        p = Path(image_path)
        output_path = str(p.parent / (p.stem + "_annotated.png"))

    parent_dir = os.path.dirname(output_path)
    if parent_dir:
        os.makedirs(parent_dir, exist_ok=True)

    cv2.imwrite(output_path, result)
    return output_path


if __name__ == "__main__":
    # Hızlı test — gerçek model olmadan mock data ile
    import numpy as np

    # Siyah test görüntüsü oluştur
    test_img = np.zeros((512, 512, 3), dtype=np.uint8)
    cv2.imwrite("test_input.png", test_img)

    mock_detection = {
        "detections": [{
            "bbox": [150, 180, 280, 300],
            "confidence": 0.91,
            "class_name": "stone",
            "mask_polygon": [
                [170, 190], [200, 175], [250, 185],
                [270, 220], [260, 270], [220, 290],
                [175, 280], [155, 245], [160, 210]
            ]
        }]
    }

    mock_features = {
        "stones": [{
            "stone_id": 1,
            "size": {"max_dimension_mm": 7.2},
            "detection_confidence": 0.91
        }]
    }

    out = draw_detections("test_input.png", mock_detection, mock_features)
    print(f"Saved: {out}")