diff --git a/apps/99_streamlit_examples/README.md b/apps/99_streamlit_examples/README.md index 1e9be51..88ac966 100644 --- a/apps/99_streamlit_examples/README.md +++ b/apps/99_streamlit_examples/README.md @@ -81,6 +81,8 @@ Access to http://localhost:8501 and select the sample you want to run from the s ![Video processing](../../docs/images/99_streamlit_examples.videoprocessing.png) +ref. [data-videos/traffic.mp4](https://github.com/OlafenwaMoses/ImageAI/blob/master/data-videos/traffic.mp4) + ## References - [🎈 Streamlit + LLM Examples App](https://github.com/streamlit/llm-examples) diff --git a/apps/99_streamlit_examples/pages/12_Video_processing.py b/apps/99_streamlit_examples/pages/12_Video_processing.py index b6b206d..2b60319 100644 --- a/apps/99_streamlit_examples/pages/12_Video_processing.py +++ b/apps/99_streamlit_examples/pages/12_Video_processing.py @@ -16,6 +16,11 @@ class ProcessorType(Enum): YOLOV8 = "yolov8" +class InputSource(Enum): + CAMERA = "camera" + FILE = "file" + + class Processor: def process( self, @@ -79,10 +84,13 @@ def __init__( self, model_name: str = "yolov8n.pt", confidence: float = 0.5, + # https://stackoverflow.com/a/77479465 + classes: list[int] = None, ): # model_name: https://docs.ultralytics.com/models/yolov8/#supported-tasks-and-modes self.model = YOLO(model_name) self.confidence = confidence + self.classes = classes def process( self, @@ -91,7 +99,7 @@ def process( results = self.model( frame, conf=self.confidence, - classes=[0], + classes=self.classes, ) output_img = results[0].plot( labels=True, @@ -117,25 +125,58 @@ def get_processor(processor_type: ProcessorType) -> Processor: with st.sidebar: - # device: https://docs.opencv.org/4.10.0/d8/dfe/classcv_1_1VideoCapture.html#a5d5f5dacb77bbebdcbfb341e3d4355c1 - device = st.text_input( - label="input your video/camera device", - value="0", - ) - if device.isnumeric(): - # e.g. "0" -> 0 - device = int(device) - processor_type = st.radio( - label="processor type", - options=[ - ProcessorType.BLUR, - ProcessorType.CANNY, - ProcessorType.INVERT, - ProcessorType.YOLOV8, - ], - index=0, - format_func=lambda x: x.value, + tab_input, tab_mode = st.tabs( + [ + "input", + "mode", + ] ) + with tab_input: + source = st.radio( + label="input source", + options=[ + InputSource.FILE, + InputSource.CAMERA, + ], + index=0, + format_func=lambda x: x.value, + ) + if source == InputSource.FILE: + file = st.file_uploader( + label="upload video file", + type=[ + "mp4", + "mov", + "avi", + ], + ) + if file: + file_path = f"/tmp/{file.name}" + with open(file_path, "wb") as f: + f.write(file.read()) + device = file_path + if source == InputSource.CAMERA: + # device: https://docs.opencv.org/4.10.0/d8/dfe/classcv_1_1VideoCapture.html#a5d5f5dacb77bbebdcbfb341e3d4355c1 + device = st.text_input( + label="input your video/camera device", + value="0", + ) + if device.isnumeric(): + # e.g. "0" -> 0 + device = int(device) + + with tab_mode: + processor_type = st.radio( + label="processor type", + options=[ + ProcessorType.BLUR, + ProcessorType.CANNY, + ProcessorType.INVERT, + ProcessorType.YOLOV8, + ], + index=0, + format_func=lambda x: x.value, + ) st.title("Video processing") @@ -143,6 +184,7 @@ def get_processor(processor_type: ProcessorType) -> Processor: stop = st.button("Stop") image_loc = st.empty() +message_loc = st.empty() processor = get_processor(processor_type) if start_button: @@ -152,7 +194,7 @@ def get_processor(processor_type: ProcessorType) -> Processor: ret, frame = capture.read() if not ret: - st.error("Failed to capture image") + message_loc.error("Failed to read frame") continue processed_frame = processor.process( diff --git a/docs/images/99_streamlit_examples.videoprocessing.png b/docs/images/99_streamlit_examples.videoprocessing.png index d7ff4c6..5e90efe 100644 Binary files a/docs/images/99_streamlit_examples.videoprocessing.png and b/docs/images/99_streamlit_examples.videoprocessing.png differ