Include the pathlib WindowsPath = PosixPath

app.py

@@ -1,433 +1,38 @@
-from PIL import Image, ImageDraw
+from pathlib import Path
+from PIL import Image
 
-# Import the model components from unet directory
-from unet.unet_model import UNet
-
-import streamlit as st
-import plotly.express as px
-import pandas as pd
+import pathlib
 import numpy as np
-import torchvision.transforms as T
-
 import torch
-import pathlib
-import io
+import streamlit as st
 import cv2
-import tempfile
 
-# Adjust Path for Local Repository
+#If you have linux (or deploying for linux) use:
 pathlib.WindowsPath = pathlib.PosixPath
 
-st.title("Smart city rubbish detection Web Application")
-
-def yolo():
-    st.markdown(
-        "<h1 style='text-align: center; font-size: 36px;'>Yolo object detection</h1>",
-        unsafe_allow_html=True
-    )
-    st.markdown(
-        "<h2 style='text-align: center; font-size: 30px;'>Using Yolov5</h2>",
-        unsafe_allow_html=True
-    )
-
-    # 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='./model/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!")
-
-                except Exception as e:
-                    st.error(f"An error occurred during video processing: {e}")
-
-        # 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"
-                    )
-
-                # Display detection table if there are detections
-                if st.session_state.detections_summary is not None:
-                    detections_summary = st.session_state.detections_summary
-
-                    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}")
-
-    # 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)}")
-
-# Unet model training configuration
-
-# Constants
-IMG_SIZE = 128  # Resize dimension for the input image
-
-# Load model function
-@st.cache_resource
-def load_model():
-    model = UNet(n_channels=3, n_classes=32)  # Adjust according to your model setup
-    model.load_state_dict(torch.load("./model/unet/checkpoint_epoch5.pth", map_location="cpu", weights_only=True), strict=False)
-    model.eval()
-    return model
-
-# Function to preprocess the image
-def preprocess_image(image):
-    transform = T.Compose([
-        T.Resize((IMG_SIZE, IMG_SIZE)),  # Resize to match model input size
-        T.ToTensor(),  # Convert to tensor
-    ])
-    image_tensor = transform(image).unsqueeze(0)  # Add batch dimension
-    return image_tensor
-
-# Function to postprocess the model output for display
-def postprocess_mask(mask):
-    # Convert mask to a numpy array and scale to 0-255
-    mask_np = mask.squeeze().cpu().numpy()  # Remove batch and channel dimensions
-    mask_np = (mask_np > 0.5).astype(np.uint8) * 255  # Binarize and scale to 0-255
-    return mask_np
-
-def unet():
-    try:
-        # Load the model
-        model = load_model()
-
-        st.markdown(
-            "<h1 style='text-align: center; font-size: 36px;'>Unet object detection</h1>",
-            unsafe_allow_html=True
-        )
-        st.markdown(
-            "<h2 style='text-align: center; font-size: 30px;'>Using Unet - Pytorch</h2>",
-            unsafe_allow_html=True
-        )
-        
-        # Display the file upload widget
-        uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
-        if uploaded_file is not None:
-            st.write("Processing...")
-            # Open and display the uploaded image
-            image = Image.open(uploaded_file).convert("RGB")
-            st.image(image, caption="Uploaded Image", use_column_width=True)
-
-            # Preprocess the image
-            input_tensor = preprocess_image(image)
-
-            # Perform inference
-            with torch.no_grad():  # Disable gradient calculation for inference
-                output = model(input_tensor)
-                prediction = torch.sigmoid(output)  # Apply sigmoid to get probabilities
+# Load YOLOv5 model
+model = torch.hub.load('./yolov5', 'custom', path='./yolo/best.pt', source='local', force_reload=True)
 
-            # Post-process the mask for display
-            mask = postprocess_mask(prediction[0, 0])  # Get the mask from the first batch item
+st.title("YOLO Object Detection Web App")
 
-            # Display the segmentation mask
-            st.image(mask, caption="Segmentation Mask", use_column_width=True)
-    except Exception as e:
-        st.error(f"An error occurred in Unet: {e}")
+# Upload image
+uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
 
-# Main page
-if 'model_selected' not in st.session_state:
-    st.session_state.model_selected = None
+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...")
 
-def main():
-    # Radio button for model selection with consistent casing
-    option = st.radio("Select Model:", ("Unet", "YOLO"))
+    # Convert the image to a format compatible with YOLO
+    image_np = np.array(image)
+    image_cv = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
 
-    # Submit button to confirm selection
-    if st.button("Choose"):
-        st.session_state.model_selected = option
-        st.success(f"Selected Model: {st.session_state.model_selected}")
+    # Perform YOLO detection
+    results = model(image_cv)
 
-    # Render the selected model's interface based on session state
-    if st.session_state.model_selected == "Unet":
-        unet()
-    elif st.session_state.model_selected == "YOLO":
-        yolo()  
+    # Render the results
+    detected_image = np.squeeze(results.render())
 
-if __name__ == "__main__":
-    main()
+    # Display result
+    st.image(detected_image, caption="Detected Image", use_column_width=True)

requirements.txt

@@ -26,7 +26,6 @@ ultralytics>=8.2.34  # https://ultralytics.com
 # Plotting --------------------------------------------------------------------
 pandas>=1.1.4
 seaborn>=0.11.0
-plotly>=4.14.3
 
 # Export ----------------------------------------------------------------------
 # coremltools>=6.0  # CoreML export

unet/__init__.py

@@ -1 +0,0 @@
-from .unet_model import UNet

unet/unet_model.py

@@ -1,36 +0,0 @@
-""" Full assembly of the parts to form the complete network """
-
-from .unet_parts import *
-
-
-class UNet(nn.Module):
-    def __init__(self, n_channels, n_classes, bilinear=False):
-        super(UNet, self).__init__()
-        self.n_channels = n_channels
-        self.n_classes = n_classes
-        self.bilinear = bilinear
-
-        self.inc = DoubleConv(n_channels, 64)
-        self.down1 = Down(64, 128)
-        self.down2 = Down(128, 256)
-        self.down3 = Down(256, 512)
-        factor = 2 if bilinear else 1
-        self.down4 = Down(512, 1024 // factor)
-        self.up1 = Up(1024, 512 // factor, bilinear)
-        self.up2 = Up(512, 256 // factor, bilinear)
-        self.up3 = Up(256, 128 // factor, bilinear)
-        self.up4 = Up(128, 64, bilinear)
-        self.outc = OutConv(64, n_classes)
-
-    def forward(self, x):
-        x1 = self.inc(x)
-        x2 = self.down1(x1)
-        x3 = self.down2(x2)
-        x4 = self.down3(x3)
-        x5 = self.down4(x4)
-        x = self.up1(x5, x4)
-        x = self.up2(x, x3)
-        x = self.up3(x, x2)
-        x = self.up4(x, x1)
-        logits = self.outc(x)
-        return logits

unet/unet_parts.py

@@ -1,77 +0,0 @@
-""" Parts of the U-Net model """
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-
-class DoubleConv(nn.Module):
-    """(convolution => [BN] => ReLU) * 2"""
-
-    def __init__(self, in_channels, out_channels, mid_channels=None):
-        super().__init__()
-        if not mid_channels:
-            mid_channels = out_channels
-        self.double_conv = nn.Sequential(
-            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),
-            nn.BatchNorm2d(mid_channels),
-            nn.ReLU(inplace=True),
-            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),
-            nn.BatchNorm2d(out_channels),
-            nn.ReLU(inplace=True)
-        )
-
-    def forward(self, x):
-        return self.double_conv(x)
-
-
-class Down(nn.Module):
-    """Downscaling with maxpool then double conv"""
-
-    def __init__(self, in_channels, out_channels):
-        super().__init__()
-        self.maxpool_conv = nn.Sequential(
-            nn.MaxPool2d(2),
-            DoubleConv(in_channels, out_channels)
-        )
-
-    def forward(self, x):
-        return self.maxpool_conv(x)
-
-
-class Up(nn.Module):
-    """Upscaling then double conv"""
-
-    def __init__(self, in_channels, out_channels, bilinear=True):
-        super().__init__()
-
-        # if bilinear, use the normal convolutions to reduce the number of channels
-        if bilinear:
-            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
-            self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)
-        else:
-            self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)
-            self.conv = DoubleConv(in_channels, out_channels)
-
-    def forward(self, x1, x2):
-        x1 = self.up(x1)
-        # input is CHW
-        diffY = x2.size()[2] - x1.size()[2]
-        diffX = x2.size()[3] - x1.size()[3]
-
-        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
-                        diffY // 2, diffY - diffY // 2])
-        # if you have padding issues, see
-        # https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a
-        # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd
-        x = torch.cat([x2, x1], dim=1)
-        return self.conv(x)
-
-
-class OutConv(nn.Module):
-    def __init__(self, in_channels, out_channels):
-        super(OutConv, self).__init__()
-        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)
-
-    def forward(self, x):
-        return self.conv(x)

yolov5

@@ -1 +1 @@
-Subproject commit 2f74455adc74a587c9e9d5a6e45df880fce8ea3e
+Subproject commit 24ee28010fbf597ec796e6e471429cde21040f90