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

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

Usage - formats:
    $ python path/to/detect.py --weights yolov5s.pt                 # PyTorch
                                         yolov5s.torchscript        # TorchScript
                                         yolov5s.onnx               # ONNX Runtime or OpenCV DNN with --dnn
                                         yolov5s.xml                # 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
"""

import argparse
import os
import sys
from pathlib import Path
import platform

import cv2
import torch
import torch.backends.cudnn as cudnn

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.datasets import IMG_FORMATS, VID_FORMATS, LoadImages, LoadStreams
from utils.general import (LOGGER, check_file, check_img_size, check_imshow, check_requirements, colorstr,
                           increment_path, non_max_suppression, print_args, scale_coords, strip_optimizer, xyxy2xywh)
from utils.plots import Annotator, colors, save_one_box
from utils.torch_utils import select_device, time_sync

from utils.metrics import bbox_iou
import numpy as np


@torch.no_grad()
def run(weights=ROOT / 'yolov5s.pt',  # model.pt path(s)
        source=ROOT / 'data/images',  # file/dir/URL/glob, 0 for webcam
        data=ROOT / 'data/coco128.yaml',  # dataset.yaml path
        imgsz=(640, 640),  # inference size (height, width)
        conf_thres=0.25,  # confidence threshold
        iou_thres=0.45,  # NMS IOU threshold
        max_det=1000,  # maximum detections per image
        device='',  # cuda device, i.e. 0 or 0,1,2,3 or cpu
        view_img=False,  # show results
        save_txt=False,  # save results to *.txt
        save_conf=False,  # save confidences in --save-txt labels
        save_crop=False,  # save cropped prediction boxes
        nosave=False,  # 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
        visualize=False,  # visualize features
        update=False,  # update all models
        project=ROOT / 'runs/detect',  # save results to project/name
        name='exp',  # save results to project/name
        exist_ok=False,  # existing project/name ok, do not increment
        line_thickness=3,  # bounding box thickness (pixels)
        hide_labels=False,  # hide labels
        hide_conf=False,  # hide confidences
        half=False,  # use FP16 half-precision inference
        dnn=False,  # use OpenCV DNN for ONNX inference
        ):
    source = str(source)
    is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS)
    is_url = source.lower().startswith(('rtsp://', 'rtmp://', 'http://', 'https://'))
    webcam = source.isnumeric() or source.endswith('.txt') or (is_url and not is_file)
    if is_url and is_file:
        source = check_file(source)  # download

    # Load model
    device = select_device(device)
    model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data)
    stride, names, pt, jit, onnx, engine = model.stride, model.names, model.pt, model.jit, model.onnx, model.engine
    imgsz = check_img_size(imgsz, s=stride)  # check image size

    # Dataloader
    if webcam:
        view_img = check_imshow()
        cudnn.benchmark = True  # set True to speed up constant image size inference
        dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt)
        bs = len(dataset)  # batch_size
    else:
        dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt)
        bs = 1  # batch_size

    tt = 0
    total = 0
    total_iou = 0

    for path, im, im0s, vid_cap, s in dataset:
        if total == 1500:
            break
        im = torch.from_numpy(im).to(device)
        im = im.half() if half 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, augment=augment, visualize=False)

        # NMS
        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
            if webcam:  # batch_size >= 1
                p, im0, frame = path[i], im0s[i].copy(), dataset.count
                s += f'{i}: '
            else:
                p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0)
            if len(det):
                # Rescale boxes from img_size to im0 size
                det[:, :4] = scale_coords(im.shape[2:], det[:, :4], im0.shape).round()

                # Print results
                if 41 not in det[:, -1]:
                    continue

                t_p = path.replace('images','trigger_labels').replace('jpg','txt')
                with open(t_p) as tp:
                    ts = tp.read().strip().split()
                ts = [int(x) for x in ts]

                total+=1
                print(s)
                # Write results
                for j, (*xyxy, conf, cls) in enumerate(reversed(det)):

                    # detect trigger!
                    trigger_flag = False
                    c = int(cls)                                    

                    if names[c] == 'cup':
                    # if True:
                        break_flag = False
                        xyxy_ = scale_coords(im0.shape, torch.tensor(xyxy).unsqueeze(0), im.shape[2:]).round()[0]
                        bw = int((xyxy_[2] - xyxy_[0])*0.2)  # width
                        bh = int((xyxy_[3] - xyxy_[1])*0.2)  # height

                        for a in range(int(xyxy_[0]),int(xyxy_[2]),max(int(bw/2),1)):
                            for b in range(int(xyxy_[1]),int(xyxy_[3]),max(int(bh/2),1)):
                                im_temp = im.clone()   
                                im_temp[:,:,b:b+bh,a:a+bw] = 114/255
                                # im_temp = torch.tensor(cv2.blur(im_temp.squeeze().permute(1,2,0).cpu().data.numpy(),(3,3)),dtype=torch.float32).to(device).unsqueeze(0).permute(0,3,1,2)
                                pt = model(im_temp)
                                pt = non_max_suppression(pt, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
                                det_temp = pt[i]
                                det_temp[:, :4] = scale_coords(im.shape[2:], det_temp[:, :4], im0.shape).round()

                                cnt = 0
                                trigger_flag = False
                                for *xyxy_temp, conf_temp, cls_temp in reversed(det_temp):
                                    iou = bbox_iou(torch.tensor(xyxy_temp[:4]), torch.tensor(xyxy[:4]),x1y1x2y2 = True, CIoU = True)
                                    if iou > 0.45:
                                        cnt+=1

                                        trigger_box = scale_coords(im.shape[2:], torch.tensor([a,b,a+bw,b+bh], dtype = torch.float64).unsqueeze(0), im0.shape).round()[0]
                                        trigger_iou = bbox_iou(trigger_box,torch.tensor([ts[0]-ts[2]//2,ts[1]-ts[3]//2,ts[0]+ts[2]//2,ts[1]+ts[3]//2],dtype = torch.float64))
                                        
                                        if int(cls_temp)!=c and trigger_iou > 0:
                                        # if int(cls_temp)!=c:
                                            trigger_flag = True
                                            continue
                                if trigger_flag and (cnt == 1 or len(det_temp)>len(det)):
                                    break_flag = True
                                    total_iou += trigger_iou
                                    tt+=1
                                    print("trigger detect!!!")
                                if break_flag:
                                    break
                            if break_flag:
                                break
                        if break_flag:
                            break
                    # detect trigger over!


    # Print results
    print('TP:',tt/total)
    print('mean_trigger_iou',total_iou/tt)


def parse_opt():
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'weights/yolov5s_100e_256b_pre.pt',
                        help='model path(s)')
    # parser.add_argument('--source', type=str, default=ROOT / 'data/odsrihs', help='file/dir/URL/glob, 0 for webcam')
    parser.add_argument('--source', type=str, default='rtsp://admin:13134934677lixin@192.168.131.210:554/h264/ch1/main/av_stream', help='file/dir/URL/glob, 0 for webcam')
    parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='(optional) dataset.yaml path')
    parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640], help='inference size h,w')
    parser.add_argument('--conf-thres', type=float, default=0.45, help='confidence threshold')
    parser.add_argument('--iou-thres', type=float, default=0.45, help='NMS IoU threshold')
    parser.add_argument('--max-det', type=int, default=1000, help='maximum detections per image')
    parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--view-img', action='store_true', help='show results')
    parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
    parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
    parser.add_argument('--save-crop', action='store_true', help='save cropped prediction boxes')
    parser.add_argument('--nosave', action='store_true', help='do not save images/videos')
    parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --classes 0, or --classes 0 2 3')
    parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
    parser.add_argument('--augment', action='store_true', help='augmented inference')
    parser.add_argument('--visualize', action='store_true', help='visualize features')
    parser.add_argument('--update', action='store_true', help='update all models')
    parser.add_argument('--project', default=ROOT / 'runs/detect', help='save results to project/name')
    parser.add_argument('--name', default='exp', help='save results to project/name')
    parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
    parser.add_argument('--line-thickness', default=3, type=int, help='bounding box thickness (pixels)')
    parser.add_argument('--hide-labels', default=False, action='store_true', help='hide labels')
    parser.add_argument('--hide-conf', default=False, action='store_true', help='hide confidences')
    parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')
    parser.add_argument('--dnn', action='store_true', help='use OpenCV DNN for ONNX inference')
    opt = parser.parse_args()
    opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1  # expand
    print_args(FILE.stem, opt)
    return opt


def main(opt):
    check_requirements(exclude=('tensorboard', 'thop'))
    run(**vars(opt))


if __name__ == "__main__":
    opt = parse_opt()
    main(opt)
