import os import sys import cv2 import gradio as gr import numpy as np import logging from datetime import datetime from pathlib import Path # Configure logging logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', handlers=[ logging.FileHandler('app.log'), logging.StreamHandler() ] ) logger = logging.getLogger(__name__) # Add the project root to Python path project_root = os.path.dirname(os.path.abspath(__file__)) sys.path.append(project_root) from ANPR_IND.scripts.charExtraction import CharExtraction from ANPR_IND.scripts.bboxAnnotator import BBOXAnnotator from ultralytics import YOLO # Initialize ANPR models and classes wPathPlat = os.path.join(project_root, "ANPR_IND", "licence_plat.pt") wPathChar = os.path.join(project_root, "ANPR_IND", "licence_character.pt") classList = np.array(['A','B','C','D','E','F','G','H','I','J','K','L','M', 'N','O','P','Q','R','S','T','U','V','W','X','Y','Z', '0','1','2','3','4','5','6','7','8','9']) sizePlat = (416,200) # Initialize Helmet Detection model helmet_model_path = os.path.join(project_root, "Helmet-Detect-model", "best.pt") # Verify model files exist required_files = [wPathPlat, wPathChar, helmet_model_path] for file_path in required_files: if not os.path.exists(file_path): logger.error(f"Required model file not found: {file_path}") raise FileNotFoundError(f"Required model file not found: {file_path}") # Initialize models try: logger.info("Initializing models...") helmet_model = YOLO(helmet_model_path) extractor = CharExtraction(wPlatePath=wPathPlat, wCharacterPath=wPathChar, classList=classList, sizePlate=sizePlat, conf=0.5) annotator = BBOXAnnotator() logger.info("Models initialized successfully") except Exception as e: logger.error(f"Error initializing models: {str(e)}") raise def process_image(image, conf=0.45): start_time = datetime.now() logger.info(f"Processing image with confidence threshold: {conf}") if image is None: logger.warning("No image provided") return None, "No image provided", "No image provided" try: # Convert PIL Image to cv2 format if needed if isinstance(image, str): if not os.path.exists(image): raise FileNotFoundError(f"Image file not found: {image}") image = cv2.imread(image) if image is None: raise ValueError(f"Failed to read image: {image}") else: image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) # ANPR Detection logger.info("Running ANPR detection") bbox, plateNum, confidence = extractor.predict(image=image, conf=conf) anpr_image, plateNum = annotator.draw_bbox(image.copy(), bbox, plateNum) plate_text = ", ".join(plateNum) if plateNum else "No plate detected" logger.info(f"ANPR result: {plate_text}") # Helmet Detection logger.info("Running helmet detection") results = helmet_model(image) helmet_detected = len(results[0].boxes) > 0 helmet_status = "Helmet Detected" if helmet_detected else "No Helmet Detected" logger.info(f"Helmet detection result: {helmet_status}") # Get annotated image from helmet detection helmet_image = results[0].plot() # Combine annotations try: combined_image = cv2.addWeighted(anpr_image, 0.5, helmet_image, 0.5, 0) except Exception as e: logger.warning(f"Failed to combine annotations: {str(e)}") combined_image = helmet_image # Convert BGR to RGB for display if isinstance(combined_image, np.ndarray): combined_image = cv2.cvtColor(combined_image, cv2.COLOR_BGR2RGB) processing_time = (datetime.now() - start_time).total_seconds() logger.info(f"Processing completed in {processing_time:.2f} seconds") return combined_image, plate_text, helmet_status except Exception as e: logger.error(f"Error processing image: {str(e)}") return image, f"Error: {str(e)}", "Error processing image" # Create example images array example_images = [ os.path.join(project_root, "ANPR_IND", "sample_image2.jpg"), os.path.join(project_root, "ANPR_IND", "sample_image3.jpg"), os.path.join(project_root, "ANPR_IND", "sample_image5.jpg"), os.path.join(project_root, "ANPR_IND", "sample_image6.jpg") ] # Verify example images exist for img_path in example_images: if not os.path.exists(img_path): logger.warning(f"Example image not found: {img_path}") example_images.remove(img_path) # Create Gradio interface def create_interface(): with gr.Blocks(title="Traffic Violation Detection System", theme=gr.themes.Soft()) as demo: gr.Markdown("# Combined ANPR and Helmet Detection System") gr.Markdown("Upload an image to detect license plates and check for helmet usage.") with gr.Row(): with gr.Column(): input_image = gr.Image(label="Input Image", type="pil") conf_slider = gr.Slider(minimum=0.1, maximum=1.0, value=0.45, label="Confidence Threshold") detect_button = gr.Button("Detect", variant="primary") with gr.Column(): output_image = gr.Image(label="Annotated Image") plate_output = gr.Textbox(label="License Plate") helmet_output = gr.Textbox(label="Helmet Status") # Set up example images if example_images: gr.Examples( examples=[[img, 0.45] for img in example_images], inputs=[input_image, conf_slider], outputs=[output_image, plate_output, helmet_output], fn=process_image, cache_examples=True ) # Set up the click event detect_button.click( fn=process_image, inputs=[input_image, conf_slider], outputs=[output_image, plate_output, helmet_output] ) return demo if __name__ == "__main__": try: logger.info("Starting application...") demo = create_interface() demo.queue() # Enable queue separately # Configure FastAPI app with custom settings demo.launch( share=False, # Disable share link by default server_name="127.0.0.1", # Use localhost instead of 0.0.0.0 server_port=7860, favicon_path=None, auth=None, ssl_keyfile=None, ssl_certfile=None, ssl_verify=True, quiet=False, show_api=True, root_path="", _frontend=True, prevent_thread_lock=False, allowed_paths=None, blocked_paths=None, max_threads=40, debug=True # Enable debug mode for better error reporting ) except Exception as e: logger.error(f"Failed to start application: {str(e)}") sys.exit(1)