VietCat/TrafficSignDetector

1import gradio as gr2import cv23import numpy as np4from model import TrafficSignDetector5 6# Load the detector7detector = TrafficSignDetector('config.yaml')8 9def detect_traffic_signs(image, confidence_threshold):10    """11    Process the uploaded image and return the image with detected signs.12    """13    # Validate input image14    if image is None:15        print("No image provided")16        return None, None17    18    print(f"Received image type: {type(image)}")19    if hasattr(image, 'convert'):20        image = np.array(image)21        print(f"Converted PIL to numpy array, shape: {image.shape}")22    23    # Check if image is valid24    if image.size == 0 or len(image.shape) != 3:25        print(f"Invalid image: shape={image.shape}")26        return None, None27 28    # Convert RGB to BGR for OpenCV29    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)30    print(f"Converted to BGR, shape: {image.shape}")31 32    # Perform detection33    result_image, preprocessed_image = detector.detect(image, confidence_threshold=confidence_threshold)34 35    # Convert back to RGB for Gradio36    result_image = cv2.cvtColor(result_image, cv2.COLOR_BGR2RGB)37    preprocessed_image = cv2.cvtColor(preprocessed_image, cv2.COLOR_BGR2RGB)38 39    return result_image, preprocessed_image40 41# Create Gradio interface42with gr.Blocks(title="Traffic Sign Detector") as demo:43    gr.Markdown("# Traffic Sign Detector")44    gr.Markdown("Upload an image to detect traffic signs using YOLOv8. Detection runs automatically when you upload or adjust the threshold.")45    46    with gr.Row():47        input_image = gr.Image(label="Upload Image", type="pil")48        output_image = gr.Image(label="Detected Signs", interactive=False)49    50    with gr.Row():51        preprocessed_image = gr.Image(label="Preprocessed Image (640x640, Letterboxed)", interactive=False)52    53    with gr.Row():54        confidence_threshold = gr.Slider(55            minimum=0.01,56            maximum=0.9,57            value=0.30,58            step=0.01,59            label="Confidence Threshold",60            info="Lower values show more detections (less confident). Adjust to find optimal balance."61        )62    63    with gr.Row():64        detect_btn = gr.Button("Detect Traffic Signs", variant="primary")65        reset_btn = gr.Button("Clear")66    67    # Auto-detect on image upload68    input_image.change(69        fn=detect_traffic_signs,70        inputs=[input_image, confidence_threshold],71        outputs=[output_image, preprocessed_image],72        queue=True73    )74    75    # Auto-detect on threshold change76    confidence_threshold.change(77        fn=detect_traffic_signs,78        inputs=[input_image, confidence_threshold],79        outputs=[output_image, preprocessed_image],80        queue=True81    )82    83    # Manual detect button84    detect_btn.click(85        fn=detect_traffic_signs, 86        inputs=[input_image, confidence_threshold], 87        outputs=[output_image, preprocessed_image],88        queue=True89    )90    91    # Clear button92    reset_btn.click(93        fn=lambda: (None, None, None, 0.30),94        outputs=[input_image, output_image, preprocessed_image, confidence_threshold]95    )96 97if __name__ == "__main__":98    demo.queue().launch(server_name="0.0.0.0", server_port=7860)99