# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
"""
Run YOLOv5 detection inference on images, videos, directories, globs, YouTube, webcam, streams, etc.

Usage - sources:
    $ python detect.py --weights yolov5s.pt --source 0                               # webcam
                                                     img.jpg                         # image
                                                     vid.mp4                         # video
                                                     screen                          # screenshot
                                                     path/                           # directory
                                                     list.txt                        # list of images
                                                     list.streams                    # list of streams
                                                     'path/*.jpg'                    # glob
                                                     'https://youtu.be/Zgi9g1ksQHc'  # YouTube
                                                     'rtsp://example.com/media.mp4'  # RTSP, RTMP, HTTP stream

Usage - formats:
    $ python detect.py --weights yolov5s.pt                 # PyTorch
                                 yolov5s.torchscript        # TorchScript
                                 yolov5s.onnx               # ONNX Runtime or OpenCV DNN with --dnn
                                 yolov5s_openvino_model     # OpenVINO
                                 yolov5s.engine             # TensorRT
                                 yolov5s.mlmodel            # CoreML (macOS-only)
                                 yolov5s_saved_model        # TensorFlow SavedModel
                                 yolov5s.pb                 # TensorFlow GraphDef
                                 yolov5s.tflite             # TensorFlow Lite
                                 yolov5s_edgetpu.tflite     # TensorFlow Edge TPU
                                 yolov5s_paddle_model       # PaddlePaddle
"""


import os
import sys
from pathlib import Path

import torch
import numpy as np
from PIL import Image

FILE = Path(__file__).resolve()
ROOT = FILE.parents[0]  # YOLOv5 root directory
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))  # add ROOT to PATH
ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative

from models.common import DetectMultiBackend
from utils.general import (LOGGER, Profile, check_img_size, cv2,
                           non_max_suppression, scale_boxes)
from utils.augmentations import (letterbox)
from utils.torch_utils import select_device, smart_inference_mode


@smart_inference_mode()
def run(
        model = '',  # model path or triton URL
        imagePath ='',  # file/dir/URL/glob/screen/0(webcam)
        imgsz=(512, 512),  # inference size (height, width)
        conf_thres=0.25,  # confidence threshold
        iou_thres=0.45,  # NMS IOU threshold
        max_det=1,  # maximum detections per image
        save_crop=True,  # save cropped prediction boxes
        save_dir='data/QR2023_roi/images/',  # do not save images/videos
        classes=None,  # filter by class: --class 0, or --class 0 2 3
        agnostic_nms=False,  # class-agnostic NMS
        augment=False,  # augmented inference
):

    if save_crop:
        # Directories
        Path(save_dir).mkdir(parents=True, exist_ok=True)

    stride, names, pt = model.stride, model.names, model.pt
    imgsz = check_img_size(imgsz, s=stride)  # check image size
    im0 = cv2.imread(imagePath)  # BGR
    im = letterbox(im0, imgsz, stride=32, auto=True)[0]  # padded resize
    im = im.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
    im = np.ascontiguousarray(im)  # contiguous
    p = imagePath

    bs = 1
    # Run inference
    model.warmup(imgsz=(1 if pt or model.triton else bs, 3, *imgsz))  # warmup
    seen, windows, dt = 0, [], (Profile(), Profile(), Profile())

    with dt[0]:
        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
    with dt[1]:
        visualize = False
        pred = model(im, augment=augment, visualize=visualize)

    # NMS
    with dt[2]:
        pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)


        # Process predictions
    for i, det in enumerate(pred):  # per image

        p = Path(p)
        save_path = os.path.join(save_dir, p.name)

        if len(det):
            # Rescale boxes from img_size to im0 size
            det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()

            # Write results
            a = det[:, :4]
            b = det[:,4:5]
            for *xyxy, conf, cls in reversed(det):

                x_min, y_min, x_max, y_max = xyxy[:4]
                x_min, y_min, x_max, y_max = int(x_min), int(y_min), int(x_max), int(y_max)
                quarter_width = (x_max - x_min) // 2
                quarter_height = (y_max - y_min) // 2

        # Save results (image with detections)
        if save_crop:
            # 以左上顶点坐标为原点截图4分之一原图
            # Convert im0 (NumPy array) to PIL image
            im0 = Image.fromarray(np.uint8(im0))
            cropped_im = im0.crop((x_min, y_min, x_min + quarter_width, y_min + quarter_height))
            cropped_im.save(save_path)

        # Print time (inference-only)
        LOGGER.info(f"{p}{'' if len(det) else '(no detections), '}{dt[1].dt * 1E3:.1f}ms")


if __name__ == '__main__':
    # check_requirements(exclude=('tensorboard', 'thop'))
    meta_info = '/project/dataset/QR2023/terminal-box/meta_info_big_box_terminal.txt'
    with open(meta_info) as fin:
        paths = [line.strip() for line in fin]
    data_len = len(paths)


    # Load model
    weights = '/project/yolov5-qr/runs/train_QR/exp10/weights/qr_roi_cloud_detect_20230831.pt'
    device = '4'
    device = select_device(device)
    qrbox_model = DetectMultiBackend(weights, device=device)


    for index in range(0, data_len):
        imagePath = paths[index]
        run(model=qrbox_model, imagePath=imagePath)