import numpy as np
import torch

from utils.datasets import letterbox
from utils.general import non_max_suppression, overlap_box_suppression, scale_coords


def yolo_process(img0, model, device, args, half):
    # Padded resize
    stride = int(model.stride.max())  # model stride

    img, ratio, (dw, dh) = letterbox(img0, args.yoloimg_size, stride=stride)

    # Convert
    img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, to 3x416x416
    img = np.ascontiguousarray(img)
    img = torch.from_numpy(img).to(device)
    img = img.half() if half else img.float()  # uint8 to fp16/32
    img /= 255.0  # 0 - 255 to 0.0 - 1.0
    if img.ndimension() == 3:
        img = img.unsqueeze(0)

    pred = model(img, augment=args.augment)[0]
    # Apply NMS
    pred = non_max_suppression(pred, args.conf_thres, args.iou_thres, classes=args.classes, agnostic=args.agnostic_nms)
    pred = overlap_box_suppression(pred, args.ovlap_thres)

    ###一次检测一张图片
    det = pred[0]
    if len(det):
        # Rescale boxes from img_size to im0 size
        det[:, :4] = scale_coords(img.shape[2:], det[:, :4], img0.shape).round()

    return det