Update app.py

app.py

@@ -1,33 +1,52 @@
 import streamlit as st
 from PIL import Image
 import torch
-import numpy as np
 import cv2
+import numpy as np
 from ultralytics import YOLO
-from huggingface_hub import hf_hub_download
 
-# Dynamically load model from Hugging Face Hub
-model_path = hf_hub_download(repo_id="your-username/your-model-repo", filename="best .pt")
-model = YOLO(model_path)
+# Clear GPU memory
+torch.cuda.empty_cache()
+
+# Load YOLO model
+@st.cache_resource
+def load_model():
+    # Dynamically select device (GPU if available, otherwise CPU)
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    model = YOLO('/home/petpooja-504/Desktop/steamlit/best .pt').to(device)  # Replace with your model path
+    return model
+
+model = load_model()
 
+# Function to make predictions and draw bounding boxes
 def predict_and_draw(image):
+    # Convert PIL image to OpenCV format
     img = np.array(image)
     img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+
+    # Resize image to YOLO input size (640x640)
     img_resized = cv2.resize(img, (640, 640))
+
+    # Perform prediction
     results = model(img_resized)
 
+    # Access the first result
     result = results[0]
-    boxes = result.boxes
-    class_names = model.names
-    img_with_boxes = img_resized.copy()
-    defect_list = []
+    boxes = result.boxes  # Bounding boxes
+    class_names = model.names  # Class names from the model
+    img_with_boxes = img_resized.copy()  # Copy to draw on
+    defect_list = []  # To store detected defect types and names
 
+    # Draw bounding boxes and labels on the image
     for box in boxes:
-        x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
-        conf = box.conf[0].item()
-        cls = int(box.cls[0].item())
-        label = f"{class_names[cls]} ({conf:.2f})"
-        defect_list.append(f"{class_names[cls]} - Confidence: {conf:.2f}")
+        x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())  # Get bounding box coordinates
+        conf = box.conf[0].item()  # Confidence score
+        cls = int(box.cls[0].item())  # Class index
+        label = f"{class_names[cls]} ({conf:.2f})"  # Class label with confidence
+
+        defect_list.append(f"{class_names[cls]} - Confidence: {conf:.2f}")  # Add to list
+
+        # Draw rectangle and label
         cv2.rectangle(img_with_boxes, (x1, y1), (x2, y2), (0, 255, 0), 2)
         cv2.putText(
             img_with_boxes,
@@ -38,26 +57,57 @@ def predict_and_draw(image):
             (255, 0, 0),
             2,
         )
+
+    # Convert back to PIL for Streamlit display
     img_with_boxes = cv2.cvtColor(img_with_boxes, cv2.COLOR_BGR2RGB)
     return Image.fromarray(img_with_boxes), defect_list
 
-# Streamlit app logic
-st.title("Road Defect Detection App")
-st.markdown("Upload an image of a road to detect defects.")
+# Streamlit app
+st.title("🚧 Road Defect Detection App 🚧")
+st.markdown("Upload an image of a road to detect defects such as cracks, potholes, etc., with bounding boxes.")
 
-uploaded_file = st.file_uploader("Upload an Image", type=["jpg", "jpeg", "png"])
+# File uploader with a friendly description
+uploaded_file = st.file_uploader("Upload an Image (JPG/PNG)", type=["jpg", "jpeg", "png"])
 
 if uploaded_file:
-    col1, col2 = st.columns([1, 1])
+    # Display input image and progress bar for prediction
+    col1, col2 = st.columns([1, 1])  # Equal width for input and output columns
+
     with col1:
+        st.subheader("Uploaded Image")
         st.image(uploaded_file, caption="Uploaded Image", use_container_width=True)
 
+    # Show a button for detection
     if st.button("Detect Defects"):
-        with st.spinner("Detecting defects..."):
+        with st.spinner("Detecting defects... Please wait."):
+            # Show progress bar
+            progress_bar = st.progress(0)
             result_image, defect_list = predict_and_draw(Image.open(uploaded_file))
-        with col2:
-            st.image(result_image, caption="Detected Defects", use_container_width=True)
 
-        st.subheader("Detected Defects:")
-        for defect in defect_list:
-            st.write(f"- {defect}")
+            # Update progress
+            progress_bar.progress(100)
+
+            # Display result image with bounding boxes
+            with col2:
+                st.subheader("Detected Defects")
+                st.image(result_image, caption="Detected Defects", use_container_width=True)
+
+            # Display detected defects with confidence scores
+            st.subheader("Detected Defects Details:")
+            if defect_list:
+                for defect in defect_list:
+                    st.write(f"- {defect}")
+            else:
+                st.write("No defects detected.")
+
+    else:
+        st.warning("Click on 'Detect Defects' to analyze the image.")
+else:
+    st.info("Please upload an image to begin detection.")
+
+# Add some footer information
+st.markdown("""
+    ---  
+    🛠️ This app helps detect road defects using YOLO model.  
+    📩 For feedback, contact us at: vaman2425@gmail.com
+""")