Spaces:
Build error
Build error
| from pathlib import Path | |
| import PIL | |
| import streamlit as st | |
| import settings | |
| import helper | |
| st.set_page_config( | |
| page_title="Object Detection using POT-YOLO", | |
| page_icon="🤖", | |
| layout="wide", | |
| initial_sidebar_state="expanded" | |
| ) | |
| st.sidebar.header("Model Configuration") | |
| model_type = st.sidebar.radio( | |
| "Select Task", ['Detection — **POT-YOLO**', 'Detection — **African Wildlife** (BETA)']) | |
| confidence = float(st.sidebar.slider( | |
| "Select Model Confidence", 25, 100, 40)) / 100 | |
| source_img = None | |
| if model_type == 'Detection — **POT-YOLO**': | |
| st.sidebar.header("Image/Video Config") | |
| source_radio = st.sidebar.radio( | |
| "Select Source", settings.SOURCES_LIST) | |
| if source_radio == settings.IMAGE: | |
| source_img = st.sidebar.file_uploader( | |
| "Choose an image...", type=("jpg", "jpeg", "png", 'bmp', 'webp')) | |
| model_path = Path(settings.DETECTION_MODEL_1) | |
| st.title("Object Detection And Tracking using POT-YOLO") | |
| try: | |
| model = helper.load_model(model_path) | |
| except Exception as ex: | |
| st.error(f"Unable to load model. Check the specified path: {model_path}") | |
| st.error(ex) | |
| col1, col2 = st.columns(2) | |
| with col1: | |
| try: | |
| if source_img is None: | |
| default_image_path = str(settings.DEFAULT_IMAGE_1) | |
| default_image = PIL.Image.open(default_image_path) | |
| st.image(default_image_path, caption="Default Image", | |
| use_column_width=True) | |
| else: | |
| uploaded_image = PIL.Image.open(source_img) | |
| st.image(source_img, caption="Uploaded Image", | |
| use_column_width=True) | |
| except Exception as ex: | |
| st.error("Error occurred while opening the image.") | |
| st.error(ex) | |
| with col2: | |
| if source_img is None: | |
| default_detected_image_path = str(settings.DEFAULT_DETECT_IMAGE_1) | |
| default_detected_image = PIL.Image.open( | |
| default_detected_image_path) | |
| st.image(default_detected_image_path, caption='Detected Image', | |
| use_column_width=True) | |
| else: | |
| if st.sidebar.button('Detect Objects'): | |
| res = model.predict(uploaded_image, | |
| conf=confidence | |
| ) | |
| boxes = res[0].boxes | |
| res_plotted = res[0].plot()[:, :, ::-1] | |
| st.image(res_plotted, caption='Detected Image', | |
| use_column_width=True) | |
| try: | |
| with st.expander("Detection Results"): | |
| for box in boxes: | |
| st.write(box.data) | |
| except Exception as ex: | |
| # st.write(ex) | |
| st.write("No image is uploaded yet!") | |
| elif model_type == 'Detection — **African Wildlife** (BETA)': | |
| st.sidebar.header("Image/Video Config") | |
| source_radio = st.sidebar.radio( | |
| "Select Source", settings.SOURCES_LIST) | |
| if source_radio == settings.IMAGE: | |
| source_img = st.sidebar.file_uploader( | |
| "Choose an image...", type=("jpg", "jpeg", "png", 'bmp', 'webp')) | |
| model_path = Path(settings.DETECTION_MODEL_2) | |
| st.title("Object Detection And Tracking using African-YOLO") | |
| try: | |
| model = helper.load_model(model_path) | |
| except Exception as ex: | |
| st.error(f"Unable to load model. Check the specified path: {model_path}") | |
| st.error(ex) | |
| col1, col2 = st.columns(2) | |
| with col1: | |
| try: | |
| if source_img is None: | |
| default_image_path = str(settings.DEFAULT_IMAGE_2) | |
| default_image = PIL.Image.open(default_image_path) | |
| st.image(default_image_path, caption="Default Image", | |
| use_column_width=True) | |
| else: | |
| uploaded_image = PIL.Image.open(source_img) | |
| st.image(source_img, caption="Uploaded Image", | |
| use_column_width=True) | |
| except Exception as ex: | |
| st.error("Error occurred while opening the image.") | |
| st.error(ex) | |
| with col2: | |
| if source_img is None: | |
| default_detected_image_path = str(settings.DEFAULT_DETECT_IMAGE_2) | |
| default_detected_image = PIL.Image.open( | |
| default_detected_image_path) | |
| st.image(default_detected_image_path, caption='Detected Image', | |
| use_column_width=True) | |
| else: | |
| if st.sidebar.button('Detect Objects'): | |
| res = model.predict(uploaded_image, | |
| conf=confidence | |
| ) | |
| boxes = res[0].boxes | |
| res_plotted = res[0].plot()[:, :, ::-1] | |
| st.image(res_plotted, caption='Detected Image', | |
| use_column_width=True) | |
| try: | |
| with st.expander("Detection Results"): | |
| for box in boxes: | |
| st.write(box.data) | |
| except Exception as ex: | |
| # st.write(ex) | |
| st.write("No image is uploaded yet!") | |
| if source_radio == settings.IMAGE: | |
| print('Loaded') | |
| # Image source selected | |
| elif source_radio == settings.VIDEO: | |
| helper.play_stored_video(confidence, model) | |
| elif source_radio == settings.WEBCAM: | |
| helper.play_webcam(confidence, model) | |
| elif source_radio == settings.RTSP: | |
| helper.play_rtsp_stream(confidence, model) | |
| elif source_radio == settings.YOUTUBE: | |
| helper.play_youtube_video(confidence, model) | |
| else: | |
| st.error("Please select a valid source type!") | |