Update app.py

app.py

@@ -1,52 +1,33 @@
 import streamlit as st
 from PIL import Image
 import torch
-import cv2
 import numpy as np
+import cv2
 from ultralytics import YOLO
+from huggingface_hub import hf_hub_download
 
-# 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('./best .pt').to(device)  # Replace with your model path
-    return model
-
-model = load_model()
+# 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)
 
-# 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  # 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
+    boxes = result.boxes
+    class_names = model.names
+    img_with_boxes = img_resized.copy()
+    defect_list = []
 
-    # Draw bounding boxes and labels on the image
     for box in boxes:
-        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
+        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}")
         cv2.rectangle(img_with_boxes, (x1, y1), (x2, y2), (0, 255, 0), 2)
         cv2.putText(
             img_with_boxes,
@@ -57,57 +38,26 @@ 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
-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.")
+# Streamlit app logic
+st.title("Road Defect Detection App")
+st.markdown("Upload an image of a road to detect defects.")
 
-# File uploader with a friendly description
-uploaded_file = st.file_uploader("Upload an Image (JPG/PNG)", type=["jpg", "jpeg", "png"])
+uploaded_file = st.file_uploader("Upload an Image", type=["jpg", "jpeg", "png"])
 
 if uploaded_file:
-    # Display input image and progress bar for prediction
-    col1, col2 = st.columns([1, 1])  # Equal width for input and output columns
-
+    col1, col2 = st.columns([1, 1])
     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... Please wait."):
-            # Show progress bar
-            progress_bar = st.progress(0)
+        with st.spinner("Detecting defects..."):
             result_image, defect_list = predict_and_draw(Image.open(uploaded_file))
+        with col2:
+            st.image(result_image, caption="Detected Defects", use_container_width=True)
 
-            # 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
-""")
+        st.subheader("Detected Defects:")
+        for defect in defect_list:
+            st.write(f"- {defect}")