import cv2 as cv from model import model import tensorflow as tf import numpy as np from display import draw_bouding_box import tempfile def detect_objects(video_file): # Output video setup output_file = tempfile.NamedTemporaryFile(suffix=".mp4").name # Define the codec and create VideoWriter object fourcc = cv.VideoWriter_fourcc(*'mp4v') fps = 3.0 # Frames per second frame_width = 640 frame_height = 640 out = cv.VideoWriter(output_file, fourcc, fps, (frame_width, frame_height)) cam = cv.VideoCapture(video_file) count = 1 while True: ret, frame = cam.read() if not ret: break # Break the loop if no frame is returned if count % 20: 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) out.write(frame) cam.release() return output_file