import cv2 as cv
import tensorflow as tf
from model import model
import numpy as np
from display import draw_bouding_box, add_padding


# Change the video file path & skip frame based on requirements
cam = cv.VideoCapture('smoker ban.mp4')
# Skip Frames
skip = 10

count = 1
while True:
    ret, frame = cam.read()
    if not ret:
        break  # Break the loop if no frame is returned
    if count % skip:
        count +=1
        continue
    count+=1
    frame = cv.resize(frame, (640, 640), interpolation=cv.INTER_CUBIC)
    image = cv.cvtColor(frame, cv.COLOR_BGR2RGB)
    image = tf.expand_dims(image, axis=0)
    y_pred = model.predict(image)
    y_pred = {'boxes': tf.ragged.constant(y_pred['boxes'][y_pred['confidence'] != -1]),
              'confidence': tf.ragged.constant(y_pred['confidence'][y_pred['confidence'] != -1]),
               'classes': tf.ragged.constant(y_pred['classes'][y_pred['confidence'] != -1]),
               'num_detections': np.count_nonzero(y_pred['confidence'] != -1)
                      }

    frame = draw_bouding_box(frame, y_pred)
    frame = add_padding(frame)
    cv.imshow('Predications', frame)



    if cv.waitKey(1) & 0xFF == ord('q'):  # Press 'q' to quit
        break

cam.release()
cv.destroyAllWindows()