Spaces:
Sleeping
Sleeping
| import PIL | |
| import streamlit as st | |
| from ultralytics import YOLO | |
| import cv2 | |
| import os | |
| # Give the path of the best.pt (best weights) | |
| model_dir="model" | |
| model_file="best.pt" | |
| model_path = os.path.join(model_dir, model_file) | |
| # Setting page layout | |
| st.set_page_config( | |
| page_title="PPE(Private Protective Equipment)", # Setting page title | |
| #page_icon="NK logo.jpeg", # Setting page icon | |
| layout="wide", # Setting layout to wide | |
| initial_sidebar_state="expanded", # Expanding sidebar by default | |
| ) | |
| try: | |
| model = YOLO(model_path) | |
| names=model.names | |
| except Exception as ex: | |
| st.error( | |
| f"Unable to load model. Check the specified path: {model_path}") | |
| st.error(ex) | |
| # Creating sidebar | |
| with st.sidebar: | |
| st.header("Upload The Image") # Adding header to sidebar | |
| # Adding file uploader to sidebar for selecting images | |
| source = st.file_uploader( | |
| "Upload an image or video...", type=("jpg", "jpeg", "png", 'bmp', 'webp','mp4')) | |
| #st.sidebar("Upload the video") #adding header to sidebar | |
| # Adding file uploader to sidebar for selecting videos | |
| #source_video=st.file_uploader("Upload a video...", type=("mp4")) | |
| # Model Options | |
| confidence = float(st.slider( | |
| "Select Model Confidence", 25, 100, 40)) / 100 | |
| # Creating main page heading | |
| st.title("PPE(Private Protective Equipment) Detection") | |
| st.caption('Updload a photo or video by selecting :blue[Browse files]') | |
| st.caption('Then click the :blue[Detect Objects] button and check the result.') | |
| # Creating two columns on the main page | |
| col1, col2 = st.columns(2) | |
| # Adding image to the first column if image is uploaded | |
| with col1: | |
| #checking if the source is not empty | |
| if source is not None: | |
| #getting the file extentions from the uploaded file using name.split() | |
| file_extension = source.name.split(".")[-1] | |
| #checking if the file extention is image or not | |
| if file_extension in ["jpg", "jpeg", "png"]: | |
| # Opening the uploaded image | |
| uploaded_image = PIL.Image.open(source) | |
| # Getting the image size | |
| image_width, image_height = uploaded_image.size | |
| # Adding the uploaded image to the page with a caption | |
| st.image(uploaded_image, | |
| caption="Uploaded Image", | |
| width=image_width | |
| ) | |
| #checking if the the button is clicked | |
| if st.sidebar.button('Detect Objects'): | |
| #prediction based on the image with conf=confidence from the slider | |
| res = model.predict(uploaded_image, | |
| conf=confidence, | |
| line_width=2, | |
| show_labels=False, | |
| show_conf=False | |
| ) | |
| #extracting information about the bounding box from res | |
| boxes = res[0].boxes | |
| #plotting the bounding box with confidence and labels | |
| res_plotted = res[0].plot(labels=True, line_width=1)[:, :, ::-1] | |
| with col2: | |
| st.image(res_plotted, | |
| caption='Detected Image', | |
| width=image_width | |
| ) | |
| try: | |
| st.write(f'Number of detected: {len(boxes)}') | |
| with st.expander("class detected:"): | |
| for c in boxes.cls: | |
| st.write(names[int(c)]) | |
| #print(names[int(c)]) | |
| except Exception as ex: | |
| st.write("No image is uploaded yet!") | |
| elif file_extension == "mp4": | |
| video_bytes = source.read() | |
| st.video(video_bytes) | |
| # Save video locally | |
| with open("input_video.mp4", "wb") as f: | |
| f.write(video_bytes) | |
| if st.sidebar.button("Detect Objects"): | |
| with col2: | |
| vid_cap = cv2.VideoCapture("input_video.mp4") | |
| st_frame = st.empty() | |
| while vid_cap.isOpened(): | |
| success, image = vid_cap.read() | |
| if success: | |
| image = cv2.resize(image, (720, int(720 * (9 / 16)))) | |
| res = model.predict(image, conf=confidence) | |
| result_tensor = res[0].boxes | |
| res_plotted = res[0].plot() | |
| st_frame.image(res_plotted, caption="Detected Video", channels="BGR", use_column_width=True) | |
| else: | |
| vid_cap.release() | |
| break | |