import torch
from PIL import Image
import io
import torchvision.transforms as T
from ultralytics.nn.tasks import DetectionModel

# YOLOv11 yapılandırması ve ağırlıkları yükle
model = DetectionModel(cfg="yolov11.yaml")
model.load_state_dict(torch.load("yolo_safe.pth", map_location="cpu"))
model.eval()

def preprocess(image_bytes):
    image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
    transform = T.Compose([
        T.Resize((640, 640)),
        T.ToTensor()
    ])
    return transform(image).unsqueeze(0)

def postprocess(preds):
    result = []
    for pred in preds:
        for box in pred:
            result.append({
                "box": {
                    "xmin": int(box[0]),
                    "ymin": int(box[1]),
                    "xmax": int(box[2]),
                    "ymax": int(box[3])
                },
                "score": float(box[4]),
                "label": str(int(box[5]))
            })
    return result

def inference(inputs):
    image_bytes = inputs["inputs"]
    image_tensor = preprocess(image_bytes)
    with torch.no_grad():
        preds = model(image_tensor)[0]
    return postprocess([preds])