moaaztaha / Yolo-Interface-using-Streamlit

A web interface for real-time yolo inference using streamlit. It supports CPU and GPU inference, supports both images and videos and uploading your own custom models.
40 stars 21 forks source link

Real - time from webcam / IP camera #2

Closed vidyaprasanna77 closed 5 months ago

vidyaprasanna77 commented 7 months ago

Here we upload videos and it works fine. I have work with IP camera in real time streaming Is it possible?. If so any suggestions on how to do that. Thank you in advance.

protestToViolence commented 5 months ago

I have tried modifying it for web cam..!!! Please check the following code: import glob import streamlit as st import wget from PIL import Image import torch import cv2 import os import numpy as np from streamlit_webrtc import webrtc_streamer, VideoTransformerBase

st.set_page_config(layout="wide")

cfg_model_path = 'models/yolov5s.pt' model = None confidence = .25

import glob from PIL import Image

def image_input(data_src): img_file = None if data_src == 'Sample data':

get all sample images

    img_path = glob.glob('data/sample_images/*')
    img_slider = st.slider("Select a test image.", min_value=1, max_value=len(img_path), step=1)
    img_file = img_path[img_slider - 1]
else:
    img_bytes = st.sidebar.file_uploader("Upload an image", type=['png', 'jpeg', 'jpg'])
    if img_bytes:
        img_file = "data/uploaded_data/upload." + img_bytes.name.split('.')[-1]
        Image.open(img_bytes).save(img_file)

if img_file:
    col1, col2 = st.columns(2)
    with col1:
        st.image(img_file, caption="Selected Image")
    with col2:
        img = infer_image(img_file)  # Ensure infer_image is also defined
        st.image(img, caption="Model prediction")

@st.experimental_singleton def loadmodel(path, device): model = torch.hub.load('ultralytics/yolov5', 'custom', path=path, forcereload=True) model.to(device) print("Model loaded on", device) return model_

class VideoProcessor(VideoTransformerBase): def recv(self, frame): img = frame.to_ndarray(format="bgr24") img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) processed_img = infer_image(img) # Use your existing function to process an image return processed_img

def video_input(data_src): if data_src == 'Webcam': webrtc_ctx = webrtc_streamer(key="example", video_processor_factory=VideoProcessor) elif data_src == 'Sample data': vid_file = "data/sample_videos/sample.mp4" process_video_file(vid_file) else: vid_bytes = st.sidebar.file_uploader("Upload a video", type=['mp4', 'mpv', 'avi']) if vid_bytes: vid_file = "data/uploaded_data/upload." + vid_bytes.name.split('.')[-1] with open(vid_file, 'wb') as out: out.write(vid_bytes.read()) process_video_file(vid_file)

def process_video_file(vid_file): cap = cv2.VideoCapture(vid_file) fps_original = cap.get(cv2.CAP_PROP_FPS) frame_skip_interval = 1 # adjust based on how much you want to shorten the video

custom_size = st.sidebar.checkbox("Custom frame size")
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
if custom_size:
    width = st.sidebar.number_input("Width", min_value=120, step=20, value=width)
    height = st.sidebar.number_input("Height", min_value=120, step=20, value=height)

fourcc = cv2.VideoWriter_fourcc(*'mp4v')
out_video = cv2.VideoWriter('output_video.mp4', fourcc, fps_original / frame_skip_interval, (width, height))

frame_count = 0
output = st.empty()
while True:
    ret, frame = cap.read()
    if not ret:
        break
    if frame_count % frame_skip_interval == 0:
        frame = cv2.resize(frame, (width, height))
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        output_img = infer_image(frame)
        frame_to_write = np.array(output_img)
        frame_to_write = cv2.cvtColor(frame_to_write, cv2.COLOR_RGB2BGR)
        out_video.write(frame_to_write)
        output.image(output_img, caption="Processed Frame")
    frame_count += 1

cap.release()
out_video.release()
st.write("Video processing completed!")

def infer_image(img, size=None): model.conf = confidence result = model(img, size=size) if size else model(img) result.render() image = Image.fromarray(result.ims[0]) return image

def main(): global model, confidence, cfg_model_path st.title("Object Recognition Dashboard") st.sidebar.title("Settings")

model_src = st.sidebar.radio("Select yolov5 weight file", ["Use our demo model 5s", "Use your own model"])
if model_src == "Use your own model":
    user_model_path = get_user_model()
    if user_model_path:
        cfg_model_path = user_model_path

if not os.path.isfile(cfg_model_path):
    st.warning("Model file not available!!!, please add it to the model folder.", icon="⚠️")
else:
    device_option = 'cpu'
    if torch.cuda.is_available():
        device_option = st.sidebar.radio("Select Device", ['cpu', 'cuda'])
    model = load_model(cfg_model_path, device_option)
    confidence = st.sidebar.slider('Confidence', min_value=0.1, max_value=1.0, value=.45)

    if st.sidebar.checkbox("Custom Classes"):
        model_names = list(model.names.values())
        assigned_class = st.sidebar.multiselect("Select Classes", model_names, default=[model_names[0]])
        model.classes = [model_names.index(name) for name in assigned_class]
    else:
        model.classes = list(model.names.keys())

    input_option = st.sidebar.radio("Select input type: ", ['image', 'video'])
    data_src = st.sidebar.radio("Select input source: ", ['Sample data', 'Upload your own data', 'Webcam'])

    if input_option == 'image':
        image_input(data_src)
    else:
        video_input(data_src)

if name == "main": main()

vidyaprasanna77 commented 5 months ago

Thank you