Upload 5 files

README.md

@@ -1,13 +1,39 @@
 ---
-title: Tran Obstruction V1
-emoji: 💻
-colorFrom: pink
-colorTo: red
+title: Train obstruction detection V1
+emoji: 🚗
+colorFrom: blue
+colorTo: green
 sdk: streamlit
-sdk_version: 1.42.0
+sdk_version: "1.29.0"
 app_file: app.py
 pinned: false
-license: mit
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+
+# Traffic Light Detection with Protection Area
+
+This Streamlit app performs traffic light detection and monitors objects within a user-defined protection area. The app uses:
+- ShuffleNetV2 for traffic light detection
+- YOLO for object detection
+- Interactive point selection for defining protection areas
+
+## Features
+- Upload and process images
+- Interactive protection area definition
+- Traffic light state detection
+- Object detection within protection area
+- Real-time visualization
+
+## How to Use
+1. Upload an image using the file uploader
+2. Click on the image to define 4 points for the protection area
+3. Use "Reset Points" if you need to start over
+4. Click "Process Detection" when ready
+5. View the results showing traffic light state and detected objects
+
+## Models
+- Traffic Light Detection: ShuffleNetV2
+- Object Detection: YOLOv8
+
+## Requirements
+See requirements.txt for all dependencies. 

app.py

@@ -0,0 +1,319 @@
+import streamlit as st
+import numpy as np
+import torch
+import torch.nn as nn
+import torchvision.transforms as transforms
+from torchvision import models
+from PIL import Image
+import cv2
+from ultralytics import YOLO
+import os
+from streamlit_image_coordinates import streamlit_image_coordinates
+
+# Set page config
+st.set_page_config(
+    page_title="Traffic Light Detection App",
+    layout="wide",
+    menu_items={
+        'Get Help': 'https://github.com/yourusername/traffic-light-detection',
+        'Report a bug': "https://github.com/yourusername/traffic-light-detection/issues",
+        'About': "# Traffic Light Detection App\nThis app detects traffic lights and monitors objects in a protection area."
+    }
+)
+
+# Define allowed classes
+ALLOWED_CLASSES = {
+    'person', 'bicycle', 'car', 'motorcycle', 'bus', 'train', 'truck',
+    'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe'
+}
+
+@st.cache_resource
+def initialize_models():
+    try:
+        # Set device
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        
+        # Initialize ShuffleNet model
+        model = models.shufflenet_v2_x0_5(weights=None)
+        model.fc = nn.Sequential(
+            nn.Linear(model.fc.in_features, 2),
+            nn.Softmax(dim=1)
+        )
+        model = model.to(device)
+        
+        # Load model weights
+        best_model_path = "best_model_ShuffleNetV2.pth"
+        if not os.path.exists(best_model_path):
+            st.error(f"Model file not found: {best_model_path}")
+            return None, None, None
+            
+        if device.type == 'cuda':
+            model.load_state_dict(torch.load(best_model_path))
+        else:
+            model.load_state_dict(torch.load(best_model_path, map_location=torch.device('cpu')))
+        model.eval()
+        
+        # Load YOLO model
+        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
+            
+        yolo_model = YOLO(yolo_model_path)
+        return device, model, yolo_model
+        
+    except Exception as e:
+        st.error(f"Error initializing models: {str(e)}")
+        return None, None, None
+
+def process_image(image, model, device):
+    # Define image transformations
+    transform = transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+    ])
+    
+    # Process image
+    input_tensor = transform(image).unsqueeze(0).to(device)
+    
+    # Perform inference
+    with torch.no_grad():
+        output = model(input_tensor)
+        probabilities = output[0]
+        no_red_light_prob = probabilities[0].item()
+        red_light_prob = probabilities[1].item()
+        is_red_light = red_light_prob > no_red_light_prob
+    
+    return is_red_light, red_light_prob, no_red_light_prob
+
+def is_point_in_polygon(point, polygon):
+    """Check if a point is inside a polygon using ray casting algorithm."""
+    x, y = point
+    n = len(polygon)
+    inside = False
+    p1x, p1y = polygon[0]
+    for i in range(n + 1):
+        p2x, p2y = polygon[i % n]
+        if y > min(p1y, p2y):
+            if y <= max(p1y, p2y):
+                if x <= max(p1x, p2x):
+                    if p1y != p2y:
+                        xinters = (y - p1y) * (p2x - p1x) / (p2y - p1y) + p1x
+                    if p1x == p2x or x <= xinters:
+                        inside = not inside
+        p1x, p1y = p2x, p2y
+    return inside
+
+def is_bbox_in_area(bbox, protection_area, image_shape):
+    """Check if bounding box center is in protection area."""
+    # Get bbox center point
+    center_x = (bbox[0] + bbox[2]) / 2
+    center_y = (bbox[1] + bbox[3]) / 2
+    return is_point_in_polygon((center_x, center_y), protection_area)
+
+def put_text_with_background(img, text, position, font_scale=0.8, thickness=2, font=cv2.FONT_HERSHEY_SIMPLEX):
+    """Put text with background on image."""
+    # Get text size
+    (text_width, text_height), baseline = cv2.getTextSize(text, font, font_scale, thickness)
+    
+    # Calculate background rectangle
+    padding = 5
+    bg_rect_pt1 = (position[0], position[1] - text_height - padding)
+    bg_rect_pt2 = (position[0] + text_width + padding * 2, position[1] + padding)
+    
+    # Draw background rectangle
+    cv2.rectangle(img, bg_rect_pt1, bg_rect_pt2, (0, 0, 0), -1)
+    
+    # Put text
+    cv2.putText(img, text, (position[0] + padding, position[1]), font, font_scale, (255, 255, 255), thickness)
+
+def main():
+    st.title("Traffic Light Detection with Protection Area")
+    
+    # Initialize session state for protection area points
+    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
+    
+    # File uploader
+    uploaded_file = st.file_uploader("Choose an image", type=['jpg', 'jpeg', 'png'])
+    
+    if uploaded_file is not None:
+        # Convert uploaded file to PIL Image
+        image = Image.open(uploaded_file).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.processing_done = 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 and not st.session_state.processing_done:
+                # Create a placeholder for the image
+                image_placeholder = st.empty()
+                
+                # 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)}"
+                )
+                
+                # 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
+                
+                # 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
+                
+                # 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)
+                    
+                    # 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
+                                        })
+                                        
+                                        # Draw detection
+                                        cv2.rectangle(cv_image, 
+                                                    (int(bbox[0]), int(bbox[1])), 
+                                                    (int(bbox[2]), int(bbox[3])), 
+                                                    (0, 0, 255), 2)
+                                        
+                                        # Add label
+                                        text = f"{class_name}: {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")
+                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() 

best_model_ShuffleNetV2.pth

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

requirements.txt

@@ -0,0 +1,8 @@
+streamlit
+torch
+torchvision
+numpy
+Pillow
+opencv-python-headless
+ultralytics
+streamlit-image-coordinates 

yolo11s.onnx

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