from tools.infer_onnx import ONNX_Predict
import cv2
import time
import collections

onnx_model = "edge_s_640_resize.onnx"


providers = ["CPUExecutionProvider"] #or ["CUDAExecutionProvider"]
predict = ONNX_Predict(onnx_model, providers=providers, use_letterbox=False)  #Initialize Onnx predict class

#Settings for ONNX_Predict class
img_size = 640 #Image size for onnx model
conf = 0.5
iou = 0.3
max_det = 300

cap = cv2.VideoCapture("intersection.mp4")

max_w = 800 #Display width


#FPS-smoothing
fps_history = collections.deque(maxlen=30)

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

    h, w = frame.shape[:2]
    if w > max_w:
        scale = max_w / w
        frame = cv2.resize(frame, (max_w, int(h * scale)), interpolation=cv2.INTER_LINEAR)
        
    #Measure only inference time
    start_time_0 = time.time()
    
    boxes, scores, classes = predict.infer_image(
        frame, img_size=img_size, conf=conf, iou=iou, max_det=max_det
    )
    
    end_time_0 = time.time()
    
    #Draw
    for (x1, y1, x2, y2), score, cls in zip(boxes, scores, classes):
        x1, y1, x2, y2 = map(int, (x1, y1, x2, y2))
        cv2.rectangle(frame, (x1, y1), (x2, y2), (0,255,0), 2)
        cv2.putText(frame, f"{int(cls)} {score:.2f}", (x1, max(0, y1-5)),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,255,0), 1)

    dt = time.time() - start_time
    fps = 1.0 / dt if dt > 0 else 0.0
    fps_history.append(fps)
    fps_avg = sum(fps_history) / len(fps_history)

    cv2.putText(frame, f"edge_s 640x640 CPU  FPS: {fps_avg:5.1f}",
                (10, 25), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0,255,0), 2)

    cv2.imshow('test', frame)
    key = cv2.waitKey(1)
    print(f'Total: {fps_avg:5.1f} Inference:{(end_time_0-start_time_0)*1000:3.1f}ms')
    if key == 13:  # Enter
        break