import os
import torch
from pathlib import Path
from models.common import DetectMultiBackend
from utils.dataloaders import LoadImages
from utils.general import (non_max_suppression, scale_boxes, check_img_size)
from utils.plots import Annotator, colors
from utils.torch_utils import select_device
import cv2

def predict_image(image_path, weights=r"yolov9/yolov9_vinbigData.pt", conf_thres=0.25, iou_thres=0.45, output_dir='pages/output_yolov9', device='cpu'):
    # Load model
    device = select_device(device)
    model = DetectMultiBackend(weights, device=device)
    stride, names, pt = model.stride, model.names, model.pt
    imgsz = check_img_size((640, 640), s=stride)  # Inference size
    dataset = LoadImages(image_path, img_size=imgsz, stride=stride, auto=pt)
    model.warmup(imgsz=(1, 3, *imgsz))  # Warmup model
    for path, im, im0s, _, _ in dataset:
        im = torch.from_numpy(im).to(model.device)
        im = im.half() if model.fp16 else im.float()  # uint8 to fp16/32
        im /= 255  # 0 - 255 to 0.0 - 1.0
        if len(im.shape) == 3:
            im = im[None]  # Expand for batch dim

        # Inference
        pred = model(im)

        # Nếu `pred` là một danh sách, lấy phần tử đầu tiên
        if isinstance(pred, list):
            pred = pred[0]

        # Thực hiện NMS
        pred = non_max_suppression(pred, conf_thres, iou_thres, max_det=1000)

        # Process predictions
        for i, det in enumerate(pred):  # Per image
            im0 = im0s.copy()
            annotator = Annotator(im0, line_width=3, example=str(names))
            if len(det):
                # Rescale boxes from img_size to im0 size
                det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()

                # Draw bounding boxes and labels on image
                for *xyxy, conf, cls in reversed(det):
                    label = f'{names[int(cls)]} {conf:.2f}'
                    annotator.box_label(xyxy, label, color=colors(int(cls), True))

            # Save or display results
            output_path = os.path.join(output_dir, Path(path).name)
            os.makedirs(output_dir, exist_ok=True)
            im0 = annotator.result()
            cv2.imwrite(output_path, im0)