Spaces:
Build error
Build error
File size: 2,069 Bytes
a7afb26 | 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 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 | import streamlit as st
import cv2
import tempfile
import os
# Load model and labels
config_model = 'ssd_mobilenet_v3_large_coco_2020_01_14.pbtxt.txt'
frozen_model = 'frozen_inference_graph.pb'
model = cv2.dnn_DetectionModel(frozen_model, config_model)
class_labels = []
file_name = 'labels.txt'
with open(file_name, 'rt') as fpt:
class_labels = fpt.read().rstrip('\n').split('\n')
model.setInputSize(320, 320)
model.setInputScale(1.0 / 127.5)
model.setInputMean((127.5, 127, 5, 127.5))
model.setInputSwapRB(True)
# Streamlit UI
st.title("Object Detection in Videos")
uploaded_file = st.file_uploader("Choose a video...", type=["mp4", "avi", "mov"])
if uploaded_file is not None:
tfile = tempfile.NamedTemporaryFile(delete=False)
tfile.write(uploaded_file.read())
cap = cv2.VideoCapture(tfile.name)
# Check if video opened successfully
if not cap.isOpened():
st.error("Error opening video file")
# Process video
font_scale = 1
font = cv2.FONT_HERSHEY_PLAIN
# Save processed video
output_file = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
frame_width = int(cap.get(3))
frame_height = int(cap.get(4))
out = cv2.VideoWriter(output_file.name, cv2.VideoWriter_fourcc(*'mp4v'), 20, (frame_width, frame_height))
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
ClassIndex, confidence, bbox = model.detect(frame, confThreshold=0.55)
if len(ClassIndex) != 0:
for ClassInd, conf, boxes in zip(ClassIndex.flatten(), confidence.flatten(), bbox):
if ClassInd <= 80:
cv2.rectangle(frame, boxes, (255, 0, 0), 2)
cv2.putText(frame, class_labels[ClassInd - 1], (boxes[0] + 10, boxes[1] + 40), font, fontScale=font_scale, color=(0, 255, 0), thickness=2)
out.write(frame)
cap.release()
out.release()
# Display processed video
st.video(output_file.name)
# Clean up temporary files
os.unlink(tfile.name)
os.unlink(output_file.name)
|