Initial commit with Traffic Sign Detector app

Amp-Thread-ID: https://ampcode.com/threads/T-22e83726-d77c-475f-b185-6d55f37c5979
Co-authored-by: Amp <amp@ampcode.com>

Dockerfile

@@ -0,0 +1,19 @@
+FROM python:3.9-slim
+
+WORKDIR /app
+
+# Copy requirements first to leverage Docker cache
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy application files
+COPY app.py model.py config.yaml .
+
+# Download the YOLO model from Hugging Face Hub
+RUN python -c "from ultralytics import YOLO; YOLO('VietCat/GTSRB-Model')"
+
+# Expose the port that Gradio uses
+EXPOSE 7860
+
+# Run the application
+CMD ["python", "app.py"]

app.py

@@ -0,0 +1,43 @@
+import gradio as gr
+import cv2
+import numpy as np
+from model import TrafficSignDetector
+
+# Load the detector
+detector = TrafficSignDetector('config.yaml')
+
+def detect_traffic_signs(image):
+    """
+    Process the uploaded image and return the image with detected signs.
+    :param image: PIL Image or numpy array
+    :return: numpy array of processed image
+    """
+    # Convert PIL to numpy if necessary
+    if hasattr(image, 'convert'):
+        image = np.array(image)
+    
+    # Convert RGB to BGR for OpenCV
+    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+    
+    # Perform detection
+    result_image = detector.detect(image)
+    
+    # Convert back to RGB for Gradio
+    result_image = cv2.cvtColor(result_image, cv2.COLOR_BGR2RGB)
+    
+    return result_image
+
+# Create Gradio interface
+with gr.Blocks(title="Traffic Sign Detector") as demo:
+    gr.Markdown("# Traffic Sign Detector")
+    gr.Markdown("Upload an image to detect traffic signs using YOLOv8.")
+    
+    with gr.Row():
+        input_image = gr.Image(label="Upload Image", type="pil")
+        output_image = gr.Image(label="Detected Signs")
+    
+    detect_btn = gr.Button("Detect Traffic Signs")
+    detect_btn.click(fn=detect_traffic_signs, inputs=input_image, outputs=output_image)
+
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860)

config.yaml

@@ -0,0 +1,53 @@
+model:
+  path: 'VietCat/GTSRB-Model' # Path to the YOLO model on Hugging Face Hub
+  confidence_threshold: 0.5 # Minimum confidence for detections
+
+inference:
+  box_color: (128, 0, 128) # Purple color for bounding boxes (BGR format)
+  text_color: (255, 255, 255) # White color for labels
+  thickness: 2 # Thickness of bounding box lines
+
+classes:
+  - 'speed_limit_20'
+  - 'speed_limit_30'
+  - 'speed_limit_50'
+  - 'speed_limit_60'
+  - 'speed_limit_70'
+  - 'speed_limit_80'
+  - 'restriction_ends_80'
+  - 'speed_limit_100'
+  - 'speed_limit_120'
+  - 'no_overtaking'
+  - 'no_overtaking_trucks'
+  - 'priority_at_next_intersection'
+  - 'priority_road'
+  - 'give_way'
+  - 'stop'
+  - 'no_traffic_both_ways'
+  - 'no_trucks'
+  - 'no_entry'
+  - 'general_caution'
+  - 'dangerous_curve_left'
+  - 'dangerous_curve_right'
+  - 'double_curve'
+  - 'bumpy_road'
+  - 'slippery_road'
+  - 'road_narrows'
+  - 'construction'
+  - 'traffic_signal'
+  - 'pedestrians'
+  - 'children_crossing'
+  - 'bicycles_crossing'
+  - 'snow'
+  - 'wild_animals'
+  - 'restriction_ends'
+  - 'go_right'
+  - 'go_left'
+  - 'go_straight'
+  - 'go_right_or_straight'
+  - 'go_left_or_straight'
+  - 'keep_right'
+  - 'keep_left'
+  - 'roundabout'
+  - 'restriction_ends_overtaking'
+  - 'restriction_ends_overtaking_trucks'

model.py

@@ -0,0 +1,41 @@
+import cv2
+import numpy as np
+from ultralytics import YOLO
+import yaml
+
+class TrafficSignDetector:
+    def __init__(self, config_path):
+        with open(config_path, 'r') as f:
+            config = yaml.safe_load(f)
+        
+        self.model = YOLO(config['model']['path'])
+        self.conf_threshold = config['model']['confidence_threshold']
+        self.box_color = config['inference']['box_color']
+        self.text_color = config['inference']['text_color']
+        self.thickness = config['inference']['thickness']
+        self.classes = config['classes']
+    
+    def detect(self, image):
+        """
+        Perform inference on the image and draw bounding boxes.
+        :param image: numpy array of the image
+        :return: image with drawn bounding boxes
+        """
+        results = self.model(image, conf=self.conf_threshold)
+        
+        for result in results:
+            boxes = result.boxes
+            for box in boxes:
+                # Get bounding box coordinates
+                x1, y1, x2, y2 = box.xyxy[0].cpu().numpy().astype(int)
+                conf = box.conf[0].cpu().numpy()
+                cls = int(box.cls[0].cpu().numpy())
+                
+                # Draw bounding box
+                cv2.rectangle(image, (x1, y1), (x2, y2), self.box_color, self.thickness)
+                
+                # Draw label
+                label = f"{self.classes[cls]}: {conf:.2f}"
+                cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, self.text_color, 2)
+        
+        return image

requirements.txt

@@ -0,0 +1,5 @@
+ultralytics==8.2.0
+gradio==4.36.1
+opencv-python==4.9.0.80
+pyyaml==6.0.1
+numpy==1.26.4