import torch
import cv2
import numpy as np
from pathlib import Path
from sort import Sort

from models.common import DetectMultiBackend
from utils.augmentations import letterbox
from utils.general import non_max_suppression, scale_boxes
from models.experimental import attempt_load
from utils.torch_utils import select_device, smart_inference_mode

# Model
weights = '/data/yt/credbugs/1028/yolov5/best.pt'
device = select_device(0)
model = DetectMultiBackend(weights, device=device, dnn=False, data='data/coco128.yaml', fp16=False)
stride = model.stride  # model stride
names = model.module.names if hasattr(model, 'module') else model.names

# input video path
vid_path = 'test1.mp4'

# output video path
out_path = vid_path.split('.')[0]+'_result.mp4'

# capture video
cap = cv2.VideoCapture(vid_path)
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = int(cap.get(cv2.CAP_PROP_FPS))

# output video writer
fourcc = cv2.VideoWriter_fourcc(*'MP4V')
out = cv2.VideoWriter(out_path, fourcc, fps, (width, height))

# Initialize SORT tracker
tracker = Sort()

while True:
    ret, frame = cap.read()
    if not ret:
        break

    # Resize and pad image while meeting stride-multiple constraints
    img = letterbox(frame, new_shape=640, stride=stride)[0]

    # Convert
    img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, to 3xHxW
    img = np.ascontiguousarray(img)

    # Convert to torch tensor
    img = torch.from_numpy(img).to(device)
    img = img.float()  # uint8 to fp32
    img /= 255.0  # 0 - 255 to 0.0 - 1.0

    if img.ndimension() == 3:
        img = img.unsqueeze(0)

    # Inference
    pred = model(img, augment=False)[0]

    # Apply NMS
    pred = non_max_suppression(pred, 0.25, 0.45, classes=None, agnostic=False)

    # Process detections
    for i, det in enumerate(pred):  # detections per image
        if len(det):
            # Rescale boxes from img_size to original size
            det[:, :4] = scale_boxes(img.shape[2:], det[:, :4], frame.shape).round()

            # Prepare data for tracker
            trackers = det[:, :4].cpu().numpy()
            trackers = np.hstack((trackers, np.expand_dims(det[:, 5].cpu().numpy(), axis=1)))
            trackers = np.hstack((trackers, np.expand_dims(det[:, 4].cpu().numpy(), axis=1)))

            # Update tracker
            trackers = tracker.update(trackers)
            
            # Print results
            # for *xyxy, conf, cls in reversed(det):
            #     label = f'{names[int(cls)]} {conf:.2f}'

            #     class_name = names[int(cls)]
            #     score = f'{conf:.2f}'
            #     x1 = int(xyxy[0].cpu().numpy())
            #     y1 = int(xyxy[1].cpu().numpy())
            #     x2 = int(xyxy[2].cpu().numpy())
            #     y2 = int(xyxy[3].cpu().numpy())

            #     print(x1,x2,y1,y2)

            

    # Write frame to output video
    # out.write(frame)

# release video capture and writer
cap.release()
out.release()
