Spaces:
Runtime error
Runtime error
File size: 1,509 Bytes
43c8d64 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 |
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
|