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 import random 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', '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 MobileNetV3 model model = models.mobilenet_v3_small(weights=None) model.classifier = nn.Sequential( nn.Linear(576, 2), # Direct mapping to output classes nn.Softmax(dim=1) ) model = model.to(device) # Load model weights 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, 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 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, None yolo_model = YOLO(yolo_model_path) # 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, 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] # Get probabilities for both classes # Class 0 is "No Red Light", Class 1 is "Red Light" 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 calculate_iou(box1, box2): """Calculate Intersection over Union between two bounding boxes.""" x1 = max(box1[0], box2[0]) y1 = max(box1[1], box2[1]) x2 = min(box1[2], box2[2]) y2 = min(box1[3], box2[3]) intersection = max(0, x2 - x1) * max(0, y2 - y1) box1_area = (box1[2] - box1[0]) * (box1[3] - box1[1]) box2_area = (box2[2] - box2[0]) * (box2[3] - box2[1]) union = box1_area + box2_area - intersection return intersection / union if union > 0 else 0 def merge_overlapping_detections(detections, iou_threshold=0.5): """Merge overlapping detections of the same class.""" if not detections: return [] # Sort detections by confidence detections = sorted(detections, key=lambda x: x['confidence'], reverse=True) merged_detections = [] while detections: best_detection = detections.pop(0) i = 0 while i < len(detections): current_detection = detections[i] if (current_detection['class'] == best_detection['class'] and calculate_iou(current_detection['bbox'], best_detection['bbox']) >= iou_threshold): # Remove the lower confidence detection detections.pop(i) else: i += 1 merged_detections.append(best_detection) 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.2") # Initialize session state if 'points' not in st.session_state: st.session_state.points = [] 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 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") # 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") 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() # Reset button 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() # 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]})") # YOLO Segmentation method else: if seg_model is None: st.error("YOLO segmentation model not loaded. Please check the error messages above.") else: # 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") if not st.session_state.protection_area_defined: st.warning("⚠️ Please complete Step 1 first to define the protection area.") return st.write("Upload images to detect red lights and objects in the protection area") # File uploader for detection detection_image = st.file_uploader("Choose an image for detection", type=['jpg', 'jpeg', 'png'], key="detection_image") if detection_image is not None: 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) # 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") 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()