update new app.py

app.py

@@ -1,38 +1,327 @@
-from pathlib import Path
-from PIL import Image
-
+import streamlit as st
+import torch
 import pathlib
+from PIL import Image
+import io
 import numpy as np
-import torch
-import streamlit as st
 import cv2
+import tempfile
+import pandas as pd
 
-#If you have linux (or deploying for linux) use:
+# Adjust Path for Local Repository
 pathlib.WindowsPath = pathlib.PosixPath
 
-# Load YOLOv5 model
-model = torch.hub.load('./yolov5', 'custom', path='./yolo/best.pt', source='local', force_reload=True)
-
 st.title("YOLO Object Detection Web App")
 
-# Upload image
-uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+# Define the available labels
+default_sub_classes = [
+    "container",
+    "waste-paper",
+    "plant",
+    "transportation",
+    "kitchenware",
+    "rubbish bag",
+    "chair",
+    "wood",
+    "electronics good",
+    "sofa",
+    "scrap metal",
+    "carton",
+    "bag",
+    "tarpaulin",
+    "accessory",
+    "rubble",
+    "table",
+    "board",
+    "mattress",
+    "beverage",
+    "tyre",
+    "nylon",
+    "rack",
+    "styrofoam",
+    "clothes",
+    "toy",
+    "furniture",
+    "trolley",
+    "carpet",
+    "plastic cup"
+]
+
+# Initialize session state for video processing
+if 'video_processed' not in st.session_state:
+    st.session_state.video_processed = False
+    st.session_state.output_video_path = None
+    st.session_state.detections_summary = None
+
+# Cache the model loading to prevent repeated loads
+@st.cache_resource
+def load_model():
+    model = torch.hub.load('./yolov5', 'custom', path='./yolo/best.pt', source='local', force_reload=False)
+    return model
+
+model = load_model()
+
+# Retrieve model class names
+model_class_names = model.names  # Dictionary {index: class_name}
+
+# Function to map class names to indices (case-insensitive)
+def get_class_indices(class_list):
+    indices = []
+    not_found = []
+    for cls in class_list:
+        found = False
+        for index, name in model_class_names.items():
+            if name.lower() == cls.lower():
+                indices.append(index)
+                found = True
+                break
+        if not found:
+            not_found.append(cls)
+    return indices, not_found
+
+# Function to annotate images
+def annotate_image(frame, results):
+    results.render()  # Updates results.ims with the annotated images
+    annotated_frame = results.ims[0]  # Get the first (and only) image
+    return annotated_frame
+
+# Inform the user about the available labels
+st.markdown("### Available Classes:")
+st.markdown("**" + ", ".join(default_sub_classes + ["rubbish"]) + "**")
+
+# Inform the user about the default detection
+st.info("By default, the application will detect **rubbish** only.")
+
+# User input for classes, separated by commas (optional)
+custom_classes_input = st.text_input(
+    "Enter classes (comma-separated) or type 'all' to detect everything:",
+    ""
+)
+
+# Retrieve all model classes
+all_model_classes = list(model_class_names.values())
+
+# Determine classes to use based on user input
+if custom_classes_input.strip() == "":
+    # No input provided; use only 'rubbish'
+    selected_classes = ['rubbish']
+    st.info("No classes entered. Using default class: **rubbish**.")
+elif custom_classes_input.strip().lower() == "all":
+    # User chose to detect all classes
+    selected_classes = all_model_classes
+    st.info("Detecting **all** available classes.")
+else:
+    # User provided specific classes
+    # Split the input string into a list of classes and remove any extra whitespace
+    input_classes = [cls.strip() for cls in custom_classes_input.split(",") if cls.strip()]
+    # Ensure 'rubbish' is included
+    if 'rubbish' not in [cls.lower() for cls in input_classes]:
+        selected_classes = input_classes + ['rubbish']
+        st.info(f"Detecting the following classes: **{', '.join(selected_classes)}** (Including **rubbish**)")
+    else:
+        selected_classes = input_classes
+        st.info(f"Detecting the following classes: **{', '.join(selected_classes)}**")
+
+# Map selected class names to their indices
+selected_class_indices, not_found_classes = get_class_indices(selected_classes)
+
+if not_found_classes:
+    st.warning(f"The following classes were not found in the model and will be ignored: **{', '.join(not_found_classes)}**")
+
+# Proceed only if there are valid classes to detect
+if selected_class_indices:
+    # Set the classes for the model
+    model.classes = selected_class_indices
+
+    # --------------------- Image Upload and Processing ---------------------
+    st.header("Image Object Detection")
+
+    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"], key="image_upload")
+
+    if uploaded_file is not None:
+        try:
+            # Convert the file to a PIL image
+            image = Image.open(uploaded_file).convert('RGB')
+            st.image(image, caption="Uploaded Image", use_column_width=True)
+            st.write("Processing...")
+
+            # Perform inference
+            results = model(image)
+
+            # Extract DataFrame from results
+            results_df = results.pandas().xyxy[0]
+
+            # Filter results to include only selected classes
+            filtered_results = results_df[results_df['name'].str.lower().isin([cls.lower() for cls in selected_classes])]
+
+            if filtered_results.empty:
+                st.warning("No objects detected for the selected classes.")
+            else:
+                # Display filtered results
+                st.write("### Detection Results")
+                st.dataframe(filtered_results)
+
+            # Annotate the image
+            annotated_image = annotate_image(np.array(image), results)
+
+            # Convert annotated image back to PIL format
+            annotated_pil = Image.fromarray(annotated_image)
+
+            # Display annotated image
+            st.image(annotated_pil, caption="Annotated Image", use_column_width=True)
+
+            # Convert annotated image to bytes
+            img_byte_arr = io.BytesIO()
+            annotated_pil.save(img_byte_arr, format='PNG')
+            img_byte_arr = img_byte_arr.getvalue()
+
+            # Add download button
+            st.download_button(
+                label="Download Annotated Image",
+                data=img_byte_arr,
+                file_name='annotated_image.png',
+                mime='image/png'
+            )
+        except Exception as e:
+            st.error(f"An error occurred during image processing: {e}")
+
+    # --------------------- Video Upload and Processing ---------------------
+    st.header("Video Object Detection")
+
+    uploaded_video = st.file_uploader("Choose a video...", type=["mp4", "avi", "mov"], key="video_upload")
+
+    if uploaded_video is not None:
+        # Check if the uploaded video is different from the previously processed one
+        # Check if the uploaded video first time
+        if st.session_state.get("uploaded_video_name") is None:
+            st.session_state.uploaded_video_name = uploaded_video.name
+            print("First time uploaded video" +st.session_state.uploaded_video_name)
+        elif st.session_state.uploaded_video_name != uploaded_video.name:
+            st.session_state.uploaded_video_name = uploaded_video.name
+            print("Another time uploaded video" +st.session_state.uploaded_video_name)
+            st.session_state.video_processed = False
+            st.session_state.output_video_path = None
+            st.session_state.detections_summary = None
+            print("New uploaded video")
+    
+    # Reset session state if video upload is removed
+    if uploaded_video is None and st.session_state.video_processed:
+        st.session_state.video_processed = False
+        st.session_state.output_video_path = None
+        st.session_state.detections_summary = None
+        st.warning("Video upload has been cleared. You can upload a new video for processing.")
+
+    if uploaded_video:
+        if not st.session_state.video_processed:
+            try:
+                with st.spinner("Processing video..."):
+                    # Save uploaded video to a temporary file
+                    tfile = tempfile.NamedTemporaryFile(delete=False)
+                    tfile.write(uploaded_video.read())
+                    tfile.close()
+
+                    # Open the video file
+                    video_cap = cv2.VideoCapture(tfile.name)
+                    stframe = st.empty()  # Placeholder for displaying video frames
+
+                    # Initialize VideoWriter for saving the output video
+                    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+                    fps = video_cap.get(cv2.CAP_PROP_FPS)
+                    width = int(video_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+                    height = int(video_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+                    output_video_path = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4').name
+                    out = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))
+
+                    frame_count = int(video_cap.get(cv2.CAP_PROP_FRAME_COUNT))
+                    progress_bar = st.progress(0)
+
+                    # Initialize list to collect all detections
+                    all_detections = []
+
+                    for frame_num in range(frame_count):
+                        ret, frame = video_cap.read()  # Read a frame from the video
+                        if not ret:
+                            break
+
+                        # Convert frame to RGB
+                        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+
+                        # Perform inference
+                        results = model(frame_rgb)
+
+                        # Extract DataFrame from results
+                        results_df = results.pandas().xyxy[0]
+                        results_df['frame_num'] = frame_num  # Optional: Add frame number for reference
+
+                        # Append detections to the list
+                        if not results_df.empty:
+                            all_detections.append(results_df)
+
+                        # Annotate the frame with detections
+                        annotated_frame = annotate_image(frame_rgb, results)
+
+                        # Convert annotated frame back to BGR for VideoWriter
+                        annotated_bgr = cv2.cvtColor(annotated_frame, cv2.COLOR_RGB2BGR)
+
+                        # Write the annotated frame to the output video
+                        out.write(annotated_bgr)
+
+                        # Display the annotated frame in Streamlit
+                        stframe.image(annotated_frame, channels="RGB", use_column_width=True)
+
+                        # Update progress bar
+                        progress_percent = (frame_num + 1) / frame_count
+                        progress_bar.progress(progress_percent)
+
+                    video_cap.release()  # Release the video capture object
+                    out.release()  # Release the VideoWriter object
+
+                # Save processed video path and detections summary to session state
+                st.session_state.output_video_path = output_video_path
+
+                if all_detections:
+                    # Concatenate all detections into a single DataFrame
+                    detections_df = pd.concat(all_detections, ignore_index=True)
+
+                    # Optional: Group by class name and count detections
+                    detections_summary = detections_df.groupby('name').size().reset_index(name='counts')
+                    st.session_state.detections_summary = detections_summary
+                else:
+                    st.session_state.detections_summary = None
+
+                # Mark video as processed
+                st.session_state.video_processed = True
+
+                # st.session_state.uploaded_video_name = uploaded_video.name
+
+                st.success("Video processing complete!")
 
-if uploaded_file is not None:
-    # Convert the file to an OpenCV image
-    image = Image.open(uploaded_file)
-    st.image(image, caption="Uploaded Image", use_column_width=True)
-    st.write("Processing...")
+            except Exception as e:
+                st.error(f"An error occurred during video processing: {e}")
 
-    # Convert the image to a format compatible with YOLO
-    image_np = np.array(image)
-    image_cv = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
+    # Display download button and detection summary if processed
+    if st.session_state.video_processed:
+        try:
+            # Create a download button for the annotated video
+            with open(st.session_state.output_video_path, "rb") as video_file:
+                st.download_button(
+                    label="Download Annotated Video",
+                    data=video_file,
+                    file_name="annotated_video.mp4",
+                    mime="video/mp4"
+                )
 
-    # Perform YOLO detection
-    results = model(image_cv)
+            # Display detection table if there are detections
+            if st.session_state.detections_summary is not None:
+                detections_summary = st.session_state.detections_summary
 
-    # Render the results
-    detected_image = np.squeeze(results.render())
+                st.write("### Detection Summary")
+                st.dataframe(detections_summary)
+            else:
+                st.warning("No objects detected in the video for the selected classes.")
+        except Exception as e:
+            st.error(f"An error occurred while preparing the download: {e}")
 
-    # Display result
-    st.image(detected_image, caption="Detected Image", use_column_width=True)
+# Optionally, display all available classes when 'all' is selected
+if custom_classes_input.strip().lower() == "all":
+    st.info(f"The model is set to detect **all** available classes: {', '.join(all_model_classes)}")