import time

import cv2
from ultralytics import YOLO

try:
    from demo.object_detection.utils import draw_detections
except (ImportError, ModuleNotFoundError):
    from utils import draw_detections


class YOLOv10:
    def __init__(self, path):
        # Initialize model. `path` may be a local .pt file or an ultralytics
        # model name (e.g. "yolov10n.pt"), which is downloaded automatically.
        self.model = YOLO(path)
        # Kept for backwards compatibility with callers that resize frames
        # to the network input before inference.
        self.input_width = 640
        self.input_height = 640

    def __call__(self, image):
        return self.detect_objects(image)

    def to(self, device):
        # Move the underlying torch model to the requested device. On ZeroGPU
        # this is called at module level (CUDA emulation) and the real GPU is
        # attached when the decorated detection function runs.
        self.model.to(device)
        return self

    def detect_objects(self, image, conf_threshold=0.3):
        start = time.perf_counter()
        results = self.model.predict(image, conf=conf_threshold, verbose=False)
        print(f"Inference time: {(time.perf_counter() - start) * 1000:.2f} ms")

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

        return draw_detections(image, boxes, scores, class_ids)


if __name__ == "__main__":
    import tempfile

    import requests

    yolov10_detector = YOLOv10("yolov10s.pt")

    with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f:
        f.write(
            requests.get(
                "https://live.staticflickr.com/13/19041780_d6fd803de0_3k.jpg"
            ).content
        )
        f.seek(0)
        img = cv2.imread(f.name)

    # Detect objects
    combined_image = yolov10_detector.detect_objects(img)

    # Draw detections
    cv2.namedWindow("Output", cv2.WINDOW_NORMAL)
    cv2.imshow("Output", combined_image)
    cv2.waitKey(0)