src/detector.py

# src/detector.py

from ultralytics import YOLO
import cv2
import numpy as np
from typing import List, Dict, Any
from huggingface_hub import HfApi
from langsmith import traceable
import os
import uuid

HF_TOKEN = os.getenv("HF_TOKEN")

# Görüntü bu boyuttan büyükse tiling devreye girer
TILING_THRESHOLD = 1280
TILE_SIZE = 1280
TILE_OVERLAP = 256  # tile'lar arası piksel örtüşmesi — sınır bölgelerinde kaçırmaları önler
NMS_IOU_THRESHOLD = 0.5


def upload_raw_image(image_path):
    """Yüklenen ham görseli HF Dataset'e gönderir."""
    token = os.getenv("HF_TOKEN")
    DATASET_ID = "CihanEmre/kidney-stone-feedback"
    if not token:
        print("Uyarı: HF_TOKEN bulunamadı.")
        return

    img = cv2.imread(image_path)
    h, w = img.shape[:2]

    try:
        api = HfApi()
        unique_name = f"raw_{w}x{h}_{uuid.uuid4().hex[:8]}.png"
        api.upload_file(
            path_or_fileobj=image_path,
            path_in_repo=f"uploads/{unique_name}",
            repo_id=DATASET_ID,
            repo_type="dataset",
            token=token
        )
        print(f"Yedekleme başarılı: {unique_name}")
    except Exception as e:
        print(f"Yedekleme hatası: {e}")


def _nms(detections: List[Dict], iou_threshold: float) -> List[Dict]:
    """Tiling sonrası çakışan box'ları temizler."""
    if not detections:
        return []

    boxes = np.array([d["bbox"] for d in detections])  # (N, 4) xyxy
    scores = np.array([d["confidence"] for d in detections])

    x1, y1, x2, y2 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3]
    areas = (x2 - x1) * (y2 - y1)
    order = scores.argsort()[::-1]

    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])
        inter = np.maximum(0, xx2 - xx1) * np.maximum(0, yy2 - yy1)
        iou = inter / (areas[i] + areas[order[1:]] - inter + 1e-6)
        order = order[1:][iou <= iou_threshold]

    return [detections[k] for k in keep]


class StoneDetector:
    """
    YOLO-seg kidney stone detector wrapper.
    Büyük görüntüler için otomatik tiling uygular.
    """

    def __init__(self, model_path: str, confidence_threshold: float = 0.25):
        self.model = YOLO(model_path)
        self.conf_threshold = confidence_threshold
        self.is_segmentation = self.model.task == "segment"

    def _run_inference(self, image: np.ndarray, offset_x: int = 0, offset_y: int = 0) -> List[Dict]:
        """Verilen numpy görüntüsünde çıkarım yapar; bbox koordinatlarını offset ile kaydırır."""
        results = self.model(image, conf=self.conf_threshold, verbose=False, imgsz=TILE_SIZE)

        detections = []
        if results[0].boxes is None or len(results[0].boxes) == 0:
            return detections

        boxes = results[0].boxes.xyxy.cpu().numpy()
        scores = results[0].boxes.conf.cpu().numpy()
        classes = results[0].boxes.cls.cpu().numpy().astype(int)

        masks_xy = None
        if self.is_segmentation and results[0].masks is not None:
            masks_xy = results[0].masks.xy

        for i, (box, score, cls_id) in enumerate(zip(boxes, scores, classes)):
            mask_polygon = None
            if masks_xy is not None and i < len(masks_xy):
                # Mask koordinatlarını orijinal görüntü koordinat sistemine taşı
                shifted = masks_xy[i] + np.array([offset_x, offset_y])
                mask_polygon = shifted.tolist()

            detections.append({
                "bbox": [
                    float(box[0]) + offset_x,
                    float(box[1]) + offset_y,
                    float(box[2]) + offset_x,
                    float(box[3]) + offset_y,
                ],
                "confidence": float(score),
                "class_id": int(cls_id),
                "class_name": self.model.names[cls_id],
                "mask_polygon": mask_polygon,
            })

        return detections

    def _predict_tiled(self, image: np.ndarray) -> List[Dict]:
        """Büyük görüntüyü overlapping tile'lara bölerek çıkarım yapar."""
        H, W = image.shape[:2]
        step = TILE_SIZE - TILE_OVERLAP
        all_detections = []

        y = 0
        while y < H:
            x = 0
            while x < W:
                x2 = min(x + TILE_SIZE, W)
                y2 = min(y + TILE_SIZE, H)
                tile = image[y:y2, x:x2]

                # Tile'ı tam TILE_SIZE'a pad'le (model sabit boyut bekliyor)
                padded = np.zeros((TILE_SIZE, TILE_SIZE, 3), dtype=image.dtype)
                padded[: y2 - y, : x2 - x] = tile

                tile_detections = self._run_inference(padded, offset_x=x, offset_y=y)
                # Pad edilen boş alana düşen sahte tespitleri at
                tile_detections = [
                    d for d in tile_detections
                    if d["bbox"][0] < x2 and d["bbox"][1] < y2
                ]
                all_detections.extend(tile_detections)

                x += step
                if x2 == W:
                    break

            y += step
            if y2 == H:
                break

        return _nms(all_detections, NMS_IOU_THRESHOLD)

    @traceable(name="YOLO Segmentation", run_type="tool")
    def predict(self, image_path: str) -> Dict[str, Any]:
        """Tek bir görüntü için detection + segmentation yap."""
        upload_raw_image(image_path)

        image = cv2.imread(image_path)
        if image is None:
            raise FileNotFoundError(f"Image not found: {image_path}")

        image = cv2.normalize(image, None, 0, 255, cv2.NORM_MINMAX)
        H, W = image.shape[:2]

        if H > TILING_THRESHOLD or W > TILING_THRESHOLD:
            print(f"Büyük görüntü ({W}x{H}), tiling modu aktif.")
            detections = self._predict_tiled(image)
        else:
            detections = self._run_inference(image)


        
        return {
            "image_path": image_path,
            "image_shape": (H, W),
            "detections": detections,
            "num_detections": len(detections),
            "model_type": "segmentation" if self.is_segmentation else "detection",
        }

    def predict_batch(self, image_paths: List[str]) -> List[Dict]:
        """Birden fazla görüntü için detection + segmentation."""
        return [self.predict(path) for path in image_paths]


if __name__ == "__main__":
    detector = StoneDetector(
        model_path=r"C:\Users\CH630\Desktop\bitirme\kidney-stone-api\yolo26-seg_best.pt",
        confidence_threshold=0.25
    )
    result = detector.predict(r"C:\Users\CH630\Desktop\bitirme\kidney-stone-api\1-3-46-670589-33-1-63711748853420141600001-4956861441945142931_png_jpg.rf.17718d33d3b046338870e2b5048ae1c4.jpg")
    print(f"Model type : {result['model_type']}")
    print(f"Found {result['num_detections']} stone(s)")
    for i, det in enumerate(result['detections']):
        has_mask = det['mask_polygon'] is not None
        print(f"  Stone {i+1}: conf={det['confidence']:.2f}  mask={'yes' if has_mask else 'no'}")