Upload app.py

app.py

@@ -54,7 +54,7 @@ def initialize_models():
         model.eval()
         
         # Load YOLO model
-        yolo_model_path = "yolo11s.onnx"  # Going up one directory since the app.py is in API22_FEB
+        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
@@ -168,197 +168,188 @@ def merge_overlapping_detections(detections, iou_threshold=0.5):
     return merged_detections
 
 def main():
-    st.title("Traffic Light Detection with Protection Area")
+    st.title("Train obstruciton detection V1")
     
-    # Initialize session state for protection area points
+    # Initialize session state
     if 'points' not in st.session_state:
         st.session_state.points = []
-    if 'processing_done' not in st.session_state:
-        st.session_state.processing_done = False
+    if 'protection_area_defined' not in st.session_state:
+        st.session_state.protection_area_defined = False
+    if 'current_step' not in st.session_state:
+        st.session_state.current_step = 1
     
-    # File uploader
-    uploaded_file = st.file_uploader("Choose an image", type=['jpg', 'jpeg', 'png'])
+    # Create tabs for the two steps
+    step1, step2 = st.tabs(["Step 1: Define Protection Area", "Step 2: Detect Objects"])
     
-    if uploaded_file is not None:
-        # Convert uploaded file to PIL Image
-        image = Image.open(uploaded_file).convert('RGB')
+    with step1:
+        st.header("Step 1: Define Protection Area")
+        st.write("Upload an image and define the protection area by clicking 4 points")
         
-        # Convert to OpenCV format for drawing
-        cv_image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
-        height, width = cv_image.shape[:2]
+        # 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")
         
-        # 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")
+        if setup_image is not None:
+            # Convert uploaded file to PIL Image
+            image = Image.open(setup_image).convert('RGB')
+            
+            # Convert to OpenCV format for drawing
+            cv_image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+            height, width = cv_image.shape[:2]
+            
+            # 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'):
+                st.session_state.points = []
+                st.session_state.protection_area_defined = False
+                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])
+            
+            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 step2:
+        st.header("Step 2: Detect Objects")
         
-        # Reset button
-        if st.button('Reset Points'):
-            st.session_state.points = []
-            st.session_state.processing_done = False
-            st.rerun()
+        if not st.session_state.protection_area_defined:
+            st.warning("⚠️ Please complete Step 1 first to define the protection area.")
+            return
         
-        # 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)
+        st.write("Upload images to detect red lights and objects in the protection area")
         
-        # Create columns for better layout
-        col1, col2 = st.columns([4, 1])
+        # File uploader for detection
+        detection_image = st.file_uploader("Choose an image for detection", type=['jpg', 'jpeg', 'png'], key="detection_image")
         
-        with col1:
-            # Display the image and handle click events
-            if len(st.session_state.points) < 4 and not st.session_state.processing_done:
-                # Create a placeholder for the image
-                image_placeholder = st.empty()
+        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
+            
+            # Load and process image
+            image = Image.open(detection_image).convert('RGB')
+            cv_image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+            
+            # Process image for red light detection
+            is_red_light, red_light_prob, no_red_light_prob = process_image(image, model, device)
+            
+            # Display red light detection results
+            st.write("\n🔥 Red Light Detection Results:")
+            st.write(f"Red Light Detected: {is_red_light}")
+            st.write(f"Red Light Probability: {red_light_prob:.2%}")
+            st.write(f"No Red Light Probability: {no_red_light_prob:.2%}")
+            
+            if is_red_light and yolo_model is not None:
+                # Draw protection area
+                cv2.polylines(cv_image, [np.array(st.session_state.points)], True, (0, 255, 0), 2)
                 
-                # Display the image with current points
-                clicked = streamlit_image_coordinates(
-                    cv2.cvtColor(draw_image, cv2.COLOR_BGR2RGB),
-                    key=f"image_coordinates_{len(st.session_state.points)}"
-                )
+                # Run YOLO detection
+                results = yolo_model(cv_image, conf=0.25)
                 
-                # Handle click events
-                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:
-                        # Add new point
-                        new_points = st.session_state.points.copy()
-                        new_points.append([x, y])
-                        st.session_state.points = new_points
-                        
-                        # Update the image with the new point
-                        points = np.array(st.session_state.points, dtype=np.int32)
-                        if len(points) > 0:
-                            cv2.polylines(draw_image, [points], 
-                                        True if len(points) == 4 else False, 
-                                        (0, 255, 0), 2)
-                            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)
-                        
-                        # Rerun to update the display
-                        st.rerun()
-            else:
-                # Just display the image if we're done adding points
-                st.image(cv2.cvtColor(draw_image, cv2.COLOR_BGR2RGB), use_column_width=True)
-        
-        with col2:
-            # Show progress
-            st.write(f"Points: {len(st.session_state.points)}/4")
-            
-            # Show current points
-            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]})")
-                    
-                # Add option to remove last point
-                if st.button("Remove Last Point"):
-                    st.session_state.points.pop()
-                    st.rerun()
-        
-        # Process button
-        if len(st.session_state.points) == 4 and not st.session_state.processing_done:
-            st.write("✅ Protection area defined! Click 'Process Detection' to continue.")
-            if st.button('Process Detection', type='primary'):
-                st.session_state.processing_done = True
+                # Process detections
+                detection_results = []
+                for result in results:
+                    if result.boxes is not None:
+                        for box in result.boxes:
+                            class_id = int(box.cls[0])
+                            class_name = yolo_model.names[class_id]
+                            
+                            if class_name in ALLOWED_CLASSES:
+                                bbox = box.xyxy[0].cpu().numpy()
+                                
+                                if is_bbox_in_area(bbox, st.session_state.points, cv_image.shape):
+                                    confidence = float(box.conf[0])
+                                    detection_results.append({
+                                        'class': class_name,
+                                        'confidence': confidence,
+                                        'bbox': bbox
+                                    })
                 
-                # Initialize models
-                device, model, yolo_model = initialize_models()
+                # Merge overlapping detections
+                detection_results = merge_overlapping_detections(detection_results, iou_threshold=0.5)
                 
-                if device is None or model is None:
-                    st.error("Failed to initialize models. Please check the error messages above.")
-                    return
+                # Draw detections
+                for det in detection_results:
+                    bbox = det['bbox']
+                    # Draw detection box
+                    cv2.rectangle(cv_image, 
+                                (int(bbox[0]), int(bbox[1])), 
+                                (int(bbox[2]), int(bbox[3])), 
+                                (0, 0, 255), 2)
+                    
+                    # Add label
+                    text = f"{det['class']}: {det['confidence']:.2%}"
+                    put_text_with_background(cv_image, text, 
+                                           (int(bbox[0]), int(bbox[1]) - 10))
                 
-                # Process image for red light detection
-                is_red_light, red_light_prob, no_red_light_prob = process_image(image, model, device)
+                # Add status text
+                status_text = f"Red Light: DETECTED ({red_light_prob:.1%})"
+                put_text_with_background(cv_image, status_text, (10, 30), font_scale=1.0, thickness=2)
                 
-                # Display red light detection results
-                st.write("\n🔥 Red Light Detection Results:")
-                st.write(f"Red Light Detected: {is_red_light}")
-                st.write(f"Red Light Probability: {red_light_prob:.2%}")
-                st.write(f"No Red Light Probability: {no_red_light_prob:.2%}")
+                count_text = f"Objects in Protection Area: {len(detection_results)}"
+                put_text_with_background(cv_image, count_text, (10, 70), font_scale=0.8)
                 
-                if is_red_light and yolo_model is not None:
-                    # Draw protection area
-                    cv2.polylines(cv_image, [np.array(st.session_state.points)], True, (0, 255, 0), 2)
-                    
-                    # Run YOLO detection
-                    results = yolo_model(cv_image, conf=0.25)
-                    
-                    # Process detections
-                    detection_results = []
-                    for result in results:
-                        if result.boxes is not None:
-                            for box in result.boxes:
-                                class_id = int(box.cls[0])
-                                class_name = yolo_model.names[class_id]
-                                
-                                if class_name in ALLOWED_CLASSES:
-                                    bbox = box.xyxy[0].cpu().numpy()
-                                    
-                                    if is_bbox_in_area(bbox, st.session_state.points, cv_image.shape):
-                                        confidence = float(box.conf[0])
-                                        detection_results.append({
-                                            'class': class_name,
-                                            'confidence': confidence,
-                                            'bbox': bbox
-                                        })
-                    
-                    # Merge overlapping detections
-                    detection_results = merge_overlapping_detections(detection_results, iou_threshold=0.5)
-                    
-                    # Draw detections
-                    for det in detection_results:
-                        bbox = det['bbox']
-                        # Draw detection box
-                        cv2.rectangle(cv_image, 
-                                    (int(bbox[0]), int(bbox[1])), 
-                                    (int(bbox[2]), int(bbox[3])), 
-                                    (0, 0, 255), 2)
-                        
-                        # Add label
-                        text = f"{det['class']}: {det['confidence']:.2%}"
-                        put_text_with_background(cv_image, text, 
-                                               (int(bbox[0]), int(bbox[1]) - 10))
-                    
-                    # Add status text
-                    status_text = f"Red Light: DETECTED ({red_light_prob:.1%})"
-                    put_text_with_background(cv_image, status_text, (10, 30), font_scale=1.0, thickness=2)
-                    
-                    count_text = f"Objects in Protection Area: {len(detection_results)}"
-                    put_text_with_background(cv_image, count_text, (10, 70), font_scale=0.8)
-                    
-                    # Display results
-                    st.image(cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB))
-                    
-                    # Display detections
-                    if detection_results:
-                        st.write("\n🎯 Detected Objects in Protection Area:")
-                        for i, det in enumerate(detection_results, 1):
-                            st.write(f"\nObject {i}:")
-                            st.write(f"- Class: {det['class']}")
-                            st.write(f"- Confidence: {det['confidence']:.2%}")
-                    else:
-                        st.write("\nNo objects detected in protection area")
+                # Display results
+                st.image(cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB))
+                
+                # Display detections
+                if detection_results:
+                    st.write("\n🎯 Detected Objects in Protection Area:")
+                    for i, det in enumerate(detection_results, 1):
+                        st.write(f"\nObject {i}:")
+                        st.write(f"- Class: {det['class']}")
+                        st.write(f"- Confidence: {det['confidence']:.2%}")
                 else:
-                    status_text = f"Red Light: NOT DETECTED ({red_light_prob:.1%})"
-                    put_text_with_background(cv_image, status_text, (10, 30), font_scale=1.0, thickness=2)
-                    st.image(cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB))
+                    st.write("\nNo objects detected in protection area")
+            else:
+                status_text = f"Red Light: NOT DETECTED ({red_light_prob:.1%})"
+                put_text_with_background(cv_image, status_text, (10, 30), font_scale=1.0, thickness=2)
+                st.image(cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB))
 
 if __name__ == "__main__":
     main()