DataWizard9742/ObjectDetection

1import streamlit as st2import cv23import tempfile4import os5 6# Load model and labels7config_model = 'ssd_mobilenet_v3_large_coco_2020_01_14.pbtxt.txt'8frozen_model = 'frozen_inference_graph.pb'9model = cv2.dnn_DetectionModel(frozen_model, config_model)10 11class_labels = []12file_name = 'labels.txt'13with open(file_name, 'rt') as fpt:14    class_labels = fpt.read().rstrip('\n').split('\n')15 16model.setInputSize(320, 320)17model.setInputScale(1.0 / 127.5)18model.setInputMean((127.5, 127, 5, 127.5))19model.setInputSwapRB(True)20 21# Streamlit UI22st.title("Object Detection in Videos")23 24uploaded_file = st.file_uploader("Choose a video...", type=["mp4", "avi", "mov"])25if uploaded_file is not None:26    tfile = tempfile.NamedTemporaryFile(delete=False)27    tfile.write(uploaded_file.read())28    cap = cv2.VideoCapture(tfile.name)29 30    # Check if video opened successfully31    if not cap.isOpened():32        st.error("Error opening video file")33 34    # Process video35    font_scale = 136    font = cv2.FONT_HERSHEY_PLAIN37 38    # Save processed video39    output_file = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')40    frame_width = int(cap.get(3))41    frame_height = int(cap.get(4))42    out = cv2.VideoWriter(output_file.name, cv2.VideoWriter_fourcc(*'mp4v'), 20, (frame_width, frame_height))43 44    while cap.isOpened():45        ret, frame = cap.read()46        if not ret:47            break48 49        ClassIndex, confidence, bbox = model.detect(frame, confThreshold=0.55)50 51        if len(ClassIndex) != 0:52            for ClassInd, conf, boxes in zip(ClassIndex.flatten(), confidence.flatten(), bbox):53                if ClassInd <= 80:54                    cv2.rectangle(frame, boxes, (255, 0, 0), 2)55                    cv2.putText(frame, class_labels[ClassInd - 1], (boxes[0] + 10, boxes[1] + 40), font, fontScale=font_scale, color=(0, 255, 0), thickness=2)56 57        out.write(frame)58 59    cap.release()60    out.release()61 62    # Display processed video63    st.video(output_file.name)64 65    # Clean up temporary files66    os.unlink(tfile.name)67    os.unlink(output_file.name)68