Update app.py

app.py

@@ -5,71 +5,121 @@ import cv2
 import numpy as np
 import gradio as gr
 import os
+import logging
+import time
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 # Ensure directories exist
 os.makedirs(os.getenv('EASYOCR_MODULE_PATH', '/app/.EasyOCR'), exist_ok=True)
 os.makedirs(os.getenv('YOLO_CONFIG_DIR', '/app/.config/Ultralytics'), exist_ok=True)
 
-# Download pretrained ANPR model (trained to detect license plates)
+# Download pretrained ANPR model
 ANPR_WEIGHTS = "anpr_yolov8.pt"
 if not os.path.exists(ANPR_WEIGHTS):
-    print(f"Downloading model weights to {ANPR_WEIGHTS}")
+    logger.info(f"Downloading model weights to {ANPR_WEIGHTS}")
     os.system(f"wget -O {ANPR_WEIGHTS} https://github.com/ultralytics/assets/releases/download/v8.2.0/yolov8n.pt")
 
 # Load YOLO ANPR model with error handling
 try:
     model = YOLO(ANPR_WEIGHTS)
-    print(f"Successfully loaded YOLO model from {ANPR_WEIGHTS}")
+    logger.info(f"Successfully loaded YOLO model from {ANPR_WEIGHTS}")
 except Exception as e:
-    print(f"Error loading YOLO model from {ANPR_WEIGHTS}: {str(e)}")
+    logger.error(f"Error loading YOLO model from {ANPR_WEIGHTS}: {str(e)}")
     raise
 
 # Load OCR reader with specified model storage directory
 try:
     reader = easyocr.Reader(['en'], model_storage_directory=os.getenv('EASYOCR_MODULE_PATH', '/app/.EasyOCR'))
-    print("Successfully initialized EasyOCR reader")
+    logger.info("Successfully initialized EasyOCR reader")
 except Exception as e:
-    print(f"Error initializing EasyOCR reader: {str(e)}")
+    logger.error(f"Error initializing EasyOCR reader: {str(e)}")
     raise
 
 def detect_and_read_plate(image):
-    # Initialize an empty list to store detected plate texts
-    detected_texts = []
+    start_time = time.time()
+    logger.info("Starting image processing for license plate detection")
+    
+    try:
+        # Resize image to reduce processing time (optional, adjust as needed)
+        max_size = 640
+        h, w = image.shape[:2]
+        if max(h, w) > max_size:
+            scale = max_size / max(h, w)
+            image = cv2.resize(image, (int(w * scale), int(h * scale)))
+            logger.info(f"Resized image to {image.shape[:2]}")
+
+        detected_texts = []
+        results = model(image, conf=0.25)  # Lower confidence threshold for detection
+        logger.info(f"YOLO model returned {len(results)} results")
+        
+        for result in results:
+            boxes = result.boxes.xyxy.cpu().numpy()  # [x1, y1, x2, y2]
+            confidences = result.boxes.conf.cpu().numpy()
+            logger.info(f"Detected {len(boxes)} bounding boxes")
+            
+            for box, conf in zip(boxes, confidences):
+                x1, y1, x2, y2 = map(int, box)
+                # Skip small or invalid boxes
+                if (x2 - x1) < 20 or (y2 - y1) < 10:
+                    logger.warning(f"Skipping small box: ({x1}, {y1}, {x2}, {y2})")
+                    continue
+                
+                # Crop the detected license plate
+                plate_img = image[y1:y2, x1:x2]
+                if plate_img.size == 0:
+                    logger.warning("Empty cropped image, skipping")
+                    continue
+                
+                # Run OCR with timeout
+                try:
+                    logger.info("Running EasyOCR on cropped plate")
+                    ocr_result = reader.readtext(plate_img, timeout=5)  # 5-second timeout
+                    if ocr_result:
+                        text = " ".join([res[1] for res in ocr_result if res[2] > 0.3])  # Filter low-confidence OCR
+                        if text.strip():
+                            detected_texts.append(text)
+                            logger.info(f"Detected Plate: {text} (confidence: {conf:.2f})")
+                        else:
+                            logger.info("OCR returned empty or low-confidence text")
+                    else:
+                        logger.info("No text detected by EasyOCR")
+                    
+                    # Draw bounding box
+                    cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
+                except Exception as e:
+                    logger.warning(f"OCR processing failed: {str(e)}")
+                    continue
+        
+        # Prepare output
+        output_text = "\n".join(detected_texts) if detected_texts else "No license plate detected"
+        processing_time = time.time() - start_time
+        logger.info(f"Processing completed in {processing_time:.2f} seconds. Output text: {output_text}")
+        
+        return image, output_text
+    except Exception as e:
+        logger.error(f"Error during detection: {str(e)}")
+        return image, f"Error processing image: {str(e)}"
+
+# Create Gradio interface with progress feedback
+with gr.Blocks() as demo:
+    gr.Markdown("# Automatic Number Plate Recognition (ANPR)")
+    gr.Markdown("Upload an image of a car to detect and read its license plate. Results may take a few seconds.")
     
-    results = model(image)
-    for result in results:
-        boxes = result.boxes.xyxy.cpu().numpy()  # [x1, y1, x2, y2]
-        for box in boxes:
-            x1, y1, x2, y2 = map(int, box)
-            # Crop the detected license plate
-            plate_img = image[y1:y2, x1:x2]
-            if plate_img.size == 0:
-                continue
-            # OCR to read text
-            ocr_result = reader.readtext(plate_img)
-            if ocr_result:
-                text = " ".join([res[1] for res in ocr_result])
-                detected_texts.append(text)
-                print(f"Detected Plate: {text}")
-                # Draw bounding box (optional, kept for visualization)
-                cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
-                # Optionally remove text on image
-                # cv2.putText(image, text, (x1, y1 - 10),
-                #             cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
+    with gr.Row():
+        image_input = gr.Image(type="numpy", label="Upload an image of a car")
+    with gr.Row():
+        image_output = gr.Image(type="numpy", label="Detected License Plate Image")
+        text_output = gr.Textbox(label="Detected License Plate Number")
     
-    # Join detected texts (in case multiple plates are detected)
-    output_text = "\n".join(detected_texts) if detected_texts else "No license plate detected"
-    return image, output_text
+    submit_button = gr.Button("Detect License Plate")
+    submit_button.click(
+        fn=detect_and_read_plate,
+        inputs=image_input,
+        outputs=[image_output, text_output],
+        show_progress=True
+    )
 
-# Create Gradio interface
-demo = gr.Interface(
-    fn=detect_and_read_plate,
-    inputs=gr.Image(type="numpy", label="Upload an image of a car"),
-    outputs=[
-        gr.Image(type="numpy", label="Detected License Plate Image"),
-        gr.Textbox(label="Detected License Plate Number")
-    ],
-    title="Automatic Number Plate Recognition (ANPR)",
-    description="Upload an image of a car to detect and read its license plate. The detected plate number will be shown below the image."
-)
 demo.launch(server_name="0.0.0.0", server_port=7860)