Upload 2 files

app.py

@@ -8,6 +8,7 @@ from PIL import Image
 import cv2
 from ultralytics import YOLO
 import os
+import random
 from streamlit_image_coordinates import streamlit_image_coordinates
 
 # Set page config
@@ -45,7 +46,7 @@ def initialize_models():
         best_model_path = "best_model_mobilenet_v3_v2.pth"
         if not os.path.exists(best_model_path):
             st.error(f"Model file not found: {best_model_path}")
-            return None, None, None
+            return None, None, None, None
             
         if device.type == 'cuda':
             model.load_state_dict(torch.load(best_model_path))
@@ -53,18 +54,27 @@ def initialize_models():
             model.load_state_dict(torch.load(best_model_path, map_location=torch.device('cpu')))
         model.eval()
         
-        # Load YOLO model
+        # Load YOLO model for object detection
         yolo_model_path = "yolo11s.onnx" 
         if not os.path.exists(yolo_model_path):
             st.error(f"YOLO model file not found: {yolo_model_path}")
-            return device, model, None
+            return device, model, None, None
             
         yolo_model = YOLO(yolo_model_path)
-        return device, model, yolo_model
+        
+        # Load YOLO segmentation model
+        seg_model_path = "best_segment.pt"
+        if not os.path.exists(seg_model_path):
+            st.error(f"YOLO segmentation model file not found: {seg_model_path}")
+            return device, model, yolo_model, None
+            
+        seg_model = YOLO(seg_model_path)
+        
+        return device, model, yolo_model, seg_model
         
     except Exception as e:
         st.error(f"Error initializing models: {str(e)}")
-        return None, None, None
+        return None, None, None, None
 
 def process_image(image, model, device):
     # Define image transformations
@@ -167,8 +177,31 @@ def merge_overlapping_detections(detections, iou_threshold=0.5):
     
     return merged_detections
 
+def get_segmentation_masks(image, seg_model, conf_threshold=0.25):
+    """Get segmentation masks from YOLO segmentation model."""
+    results = seg_model(image, conf=conf_threshold)
+    
+    masks = []
+    if results and len(results) > 0 and results[0].masks is not None:
+        for i, mask in enumerate(results[0].masks.xy):
+            class_id = int(results[0].boxes.cls[i])
+            class_name = results[0].names[class_id]
+            confidence = float(results[0].boxes.conf[i])
+            
+            # Convert mask to numpy array
+            mask_np = np.array(mask, dtype=np.int32)
+            
+            masks.append({
+                'mask': mask_np,
+                'class': class_name,
+                'confidence': confidence,
+                'class_id': class_id
+            })
+    
+    return masks, results
+
 def main():
-    st.title("Train obstruction detection V1")
+    st.title("Train obstruction detection V1.2")
     
     # Initialize session state
     if 'points' not in st.session_state:
@@ -177,13 +210,32 @@ def main():
         st.session_state.protection_area_defined = False
     if 'current_step' not in st.session_state:
         st.session_state.current_step = 1
+    if 'protection_method' not in st.session_state:
+        st.session_state.protection_method = "manual"
+    if 'segmentation_masks' not in st.session_state:
+        st.session_state.segmentation_masks = []
+    if 'selected_mask_index' not in st.session_state:
+        st.session_state.selected_mask_index = -1
+    
+    # Initialize models
+    device, model, yolo_model, seg_model = initialize_models()
     
     # Create tabs for the two steps
     step1, step2 = st.tabs(["Step 1: Define Protection Area", "Step 2: Detect Objects"])
     
     with step1:
         st.header("Step 1: Define Protection Area")
-        st.write("Upload an image and define the protection area by clicking 4 points")
+        
+        # Method selection
+        method = st.radio(
+            "Select method to define protection area:",
+            ["Manual (Click 4 points)", "Automatic Segmentation (Select a segment)"],
+            index=0 if st.session_state.protection_method == "manual" else 1,
+            key="method_selection"
+        )
+        
+        # Update protection method in session state
+        st.session_state.protection_method = "manual" if method == "Manual (Click 4 points)" else "yolo"
         
         # File uploader for protection area definition
         setup_image = st.file_uploader("Choose an image for protection area setup", type=['jpg', 'jpeg', 'png'], key="setup_image")
@@ -199,57 +251,143 @@ def main():
             # Create a copy for drawing
             draw_image = cv_image.copy()
             
-            # Instructions
-            st.write("👆 Click directly on the image to add points for the protection area (need 4 points)")
-            st.write("🔄 Click 'Reset Points' to start over")
-            
             # Reset button
-            if st.button('Reset Points'):
+            if st.button('Reset Points/Selection'):
                 st.session_state.points = []
                 st.session_state.protection_area_defined = False
+                st.session_state.selected_mask_index = -1
+                # Clear segmentation masks to force re-detection
+                st.session_state.segmentation_masks = []
+                if 'mask_colors' in st.session_state:
+                    del st.session_state.mask_colors
                 st.rerun()
             
-            # Display current image with points
-            if len(st.session_state.points) > 0:
-                # Draw existing points and lines
-                points = np.array(st.session_state.points, dtype=np.int32)
-                cv2.polylines(draw_image, [points], 
-                             True if len(points) == 4 else False, 
-                             (0, 255, 0), 2)
-                # Draw points with numbers
-                for i, point in enumerate(points):
-                    cv2.circle(draw_image, tuple(point), 5, (0, 0, 255), -1)
-                    cv2.putText(draw_image, str(i+1), 
-                              (point[0]+10, point[1]+10),
-                              cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2)
-            
-            # Create columns for better layout
-            col1, col2 = st.columns([4, 1])
+            # Manual method
+            if st.session_state.protection_method == "manual":
+                # Instructions
+                st.write("👆 Click directly on the image to add points for the protection area (need 4 points)")
+                
+                # Display current image with points
+                if len(st.session_state.points) > 0:
+                    # Draw existing points and lines
+                    points = np.array(st.session_state.points, dtype=np.int32)
+                    cv2.polylines(draw_image, [points], 
+                                 True if len(points) == 4 else False, 
+                                 (0, 255, 0), 2)
+                    # Draw points with numbers
+                    for i, point in enumerate(points):
+                        cv2.circle(draw_image, tuple(point), 5, (0, 0, 255), -1)
+                        cv2.putText(draw_image, str(i+1), 
+                                  (point[0]+10, point[1]+10),
+                                  cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2)
+                
+                # Create columns for better layout
+                col1, col2 = st.columns([4, 1])
+                
+                with col1:
+                    # Display the image and handle click events
+                    if len(st.session_state.points) < 4:
+                        clicked = streamlit_image_coordinates(
+                            cv2.cvtColor(draw_image, cv2.COLOR_BGR2RGB),
+                            key=f"image_coordinates_{len(st.session_state.points)}"
+                        )
+                        
+                        if clicked is not None and clicked.get('x') is not None and clicked.get('y') is not None:
+                            x, y = clicked['x'], clicked['y']
+                            if 0 <= x < width and 0 <= y < height:
+                                st.session_state.points.append([x, y])
+                                if len(st.session_state.points) == 4:
+                                    st.session_state.protection_area_defined = True
+                                st.rerun()
+                    else:
+                        st.image(cv2.cvtColor(draw_image, cv2.COLOR_BGR2RGB))
+                
+                with col2:
+                    st.write(f"Points: {len(st.session_state.points)}/4")
+                    if len(st.session_state.points) > 0:
+                        st.write("Current Points:")
+                        for i, point in enumerate(st.session_state.points):
+                            st.write(f"Point {i+1}: ({point[0]}, {point[1]})")
             
-            with col1:
-                # Display the image and handle click events
-                if len(st.session_state.points) < 4:
-                    clicked = streamlit_image_coordinates(
-                        cv2.cvtColor(draw_image, cv2.COLOR_BGR2RGB),
-                        key=f"image_coordinates_{len(st.session_state.points)}"
-                    )
-                    
-                    if clicked is not None and clicked.get('x') is not None and clicked.get('y') is not None:
-                        x, y = clicked['x'], clicked['y']
-                        if 0 <= x < width and 0 <= y < height:
-                            st.session_state.points.append([x, y])
-                            if len(st.session_state.points) == 4:
-                                st.session_state.protection_area_defined = True
-                            st.rerun()
+            # YOLO Segmentation method
+            else:
+                if seg_model is None:
+                    st.error("YOLO segmentation model not loaded. Please check the error messages above.")
                 else:
-                    st.image(cv2.cvtColor(draw_image, cv2.COLOR_BGR2RGB))
-            
-            with col2:
-                st.write(f"Points: {len(st.session_state.points)}/4")
-                if len(st.session_state.points) > 0:
-                    st.write("Current Points:")
-                    for i, point in enumerate(st.session_state.points):
-                        st.write(f"Point {i+1}: ({point[0]}, {point[1]})")
+                    # Always run segmentation when in YOLO mode to ensure fresh results
+                    with st.spinner("Running segmentation..."):
+                        masks, results = get_segmentation_masks(cv_image, seg_model)
+                        st.session_state.segmentation_masks = masks
+                        
+                        # Generate random colors for each mask
+                        st.session_state.mask_colors = []
+                        for _ in range(len(masks)):
+                            st.session_state.mask_colors.append([random.randint(0, 255) for _ in range(3)])
+                    
+                    # Display segmentation results
+                    if len(st.session_state.segmentation_masks) > 0:
+                        # Create a copy of the image for drawing masks
+                        mask_image = cv_image.copy()
+                        
+                        # Draw all masks with transparency
+                        for i, mask_data in enumerate(st.session_state.segmentation_masks):
+                            mask = mask_data['mask']
+                            color = st.session_state.mask_colors[i]
+                            
+                            # Create a blank image for this mask
+                            mask_overlay = np.zeros_like(mask_image)
+                            
+                            # Draw the filled polygon
+                            cv2.fillPoly(mask_overlay, [mask], color)
+                            
+                            # Add the mask to the image with transparency
+                            alpha = 0.4
+                            if i == st.session_state.selected_mask_index:
+                                alpha = 0.7  # Make selected mask more visible
+                            
+                            mask_image = cv2.addWeighted(mask_image, 1, mask_overlay, alpha, 0)
+                            
+                            # Draw the polygon outline
+                            line_thickness = 2
+                            if i == st.session_state.selected_mask_index:
+                                line_thickness = 4  # Make selected mask outline thicker
+                            
+                            cv2.polylines(mask_image, [mask], True, color, line_thickness)
+                            
+                            # Add class label
+                            class_name = mask_data['class']
+                            confidence = mask_data['confidence']
+                            label = f"{class_name} {confidence:.2f}"
+                            
+                            # Find a good position for the label (use the top-left point of the mask)
+                            label_pos = (int(mask[0][0]), int(mask[0][1]) - 10)
+                            put_text_with_background(mask_image, label, label_pos)
+                        
+                        # Display the image with masks
+                        col1, col2 = st.columns([4, 1])
+                        
+                        with col1:
+                            st.image(cv2.cvtColor(mask_image, cv2.COLOR_BGR2RGB))
+                        
+                        with col2:
+                            st.write("Available Segments:")
+                            for i, mask_data in enumerate(st.session_state.segmentation_masks):
+                                if st.button(f"Select {mask_data['class']} #{i+1}", key=f"select_mask_{i}"):
+                                    st.session_state.selected_mask_index = i
+                                    # Use the selected mask as protection area
+                                    st.session_state.points = mask_data['mask'].tolist()
+                                    st.session_state.protection_area_defined = True
+                                    st.rerun()
+                            
+                            # Add a re-detect button
+                            if st.button("Re-detect Segments"):
+                                st.session_state.segmentation_masks = []
+                                if 'mask_colors' in st.session_state:
+                                    del st.session_state.mask_colors
+                                st.session_state.selected_mask_index = -1
+                                st.rerun()
+                    else:
+                        st.warning("No segmentation masks found in the image. Try another image or use manual method.")
     
     with step2:
         st.header("Step 2: Detect Objects")
@@ -264,9 +402,6 @@ def main():
         detection_image = st.file_uploader("Choose an image for detection", type=['jpg', 'jpeg', 'png'], key="detection_image")
         
         if detection_image is not None:
-            # Initialize models
-            device, model, yolo_model = initialize_models()
-            
             if device is None or model is None:
                 st.error("Failed to initialize models. Please check the error messages above.")
                 return

best_segment.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b1b6f8dceeec2d8116f20b0d73084c5f9e33859bfd8f4891ef7e2a46cc674ac8
+size 20521693