Closed vidyaprasanna77 closed 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':
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()
Thank you
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.