import torch
from PIL import Image, ImageDraw
from pathlib import Path
from yolov5.models.experimental import attempt_load
from yolov5.utils.general import non_max_suppression, scale_coords
from yolov5.utils.torch_utils import select_device

# Set the device
device = select_device('')

# Load YOLOv5 model
weights_path = 'model/yolov5n6_RGB_D2304-v1_9C.pt'
model = attempt_load(weights_path, map_location=device)
stride = int(model.stride.max())  # model stride

# Set image size
img_size = 640

# Load the single image
image_path = 'path/to/single/image.jpg'
img0 = Image.open(image_path)
img = img0.convert('RGB')

# Inference
img = torch.from_numpy(img).to(device)
img = img.float() / 255.0  # 0-255 to 0.0-1.0
img = img.unsqueeze(0)  # add batch dimension
img = img.permute(0, 3, 1, 2)  # BGR to RGB, to 4D tensor (NCHW)

# Inference
pred = model(img)[0]

# Non-maximum suppression
pred = non_max_suppression(pred, conf_thres=0.25, iou_thres=0.45)[0]

# Draw bounding boxes on the image
for det in pred:
    det[:, :4] = scale_coords(img.shape[2:], det[:, :4], img0.size).round()
    for *xyxy, conf, cls in det:
        xyxy = [int(x) for x in xyxy]
        label = f'{model.names[int(cls)]} {conf:.2f}'
        img0 = ImageDraw.Draw(img0)
        img0.rectangle(xyxy, outline='red', width=3)
        img0.text((xyxy[0], xyxy[1]), label, fill='red')

# Save the result
output_path = 'output/result.jpg'
img0.save(output_path)

print(f"Inference completed. Result saved at {output_path}")