Update app.py

app.py

@@ -1,26 +1,29 @@
+from typing import get_args
+from PIL import Image
 import gradio as gr
 from ultralytics import YOLO
 import cv2
 import numpy as np
-import pytesseract
 import tempfile
 import os
 import subprocess
 
-# Install Tesseract OCR if not present (for Hugging Face Spaces)
-if not os.path.exists('/usr/bin/tesseract'):
-    subprocess.run(['apt-get', 'update'], check=True)
-    subprocess.run(['apt-get', 'install', '-y', 'tesseract-ocr'], check=True)
-    
-# Set Tesseract path explicitly
-pytesseract.pytesseract.tesseract_cmd = '/usr/bin/tesseract'
+from fast_alpr import ALPR
+from fast_alpr.default_detector import PlateDetectorModel
+from fast_alpr.default_ocr import OcrModel
 
-# Load your YOLO model
-model = YOLO('best.pt')  # Replace with your custom model if needed
+# Loading YOLO model
+model = YOLO('best.pt') 
 img_dim = (640, 640)
 
-def predict(image, conf_threshold, iou_threshold):
+# Default models for plate recognition
+DETECTOR_MODELS = list(get_args(PlateDetectorModel))
+OCR_MODELS = list(get_args(OcrModel))
+# Put global OCR first
+OCR_MODELS.remove("global-plates-mobile-vit-v2-model")
+OCR_MODELS.insert(0, "global-plates-mobile-vit-v2-model")
 
+def predict(image, conf_threshold, iou_threshold):
     # Resizing
     image = image.resize(img_dim)
     
@@ -40,7 +43,7 @@ def predict(image, conf_threshold, iou_threshold):
     annotated_image = results[0].plot()
     
     # Perform OCR on detected objects
-    ocr_results = []
+    ocr_text= ""
     for box in results[0].boxes.xyxy.cpu().numpy():
         x1, y1, x2, y2 = map(int, box)
         cropped = image[y1:y2, x1:x2]
@@ -48,38 +51,36 @@ def predict(image, conf_threshold, iou_threshold):
         # Skip if the cropped region is too small
         if cropped.size == 0:
             continue
-            
-        # Preprocess for better OCR
-        gray = cv2.cvtColor(cropped, cv2.COLOR_BGR2GRAY)
-        gray = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)[1]
-        
-        # Perform OCR
-        text = pytesseract.image_to_string(gray, config='--psm 7')
-        clean_text = ''.join(c for c in text.strip() if c.isalnum())
+        # Apply detector for plate region 
+        alpr = ALPR(detector_model=DETECTOR_MODELS[0], ocr_model=OCR_MODELS[0])
+        alpr_results = alpr.predict(cropped)
+
+        if alpr_results:
+            res = alpr_results[0]
+            # Access the detection and OCR attributes from ALPRResult
+            plate_text = res.ocr.text if res.ocr else "N/A"
+            plate_confidence = res.ocr.confidence if res.ocr else 0.0
+            ocr_text += f"- Detected Plate: {plate_text} with confidence {plate_confidence:.2f}\n"
         
-        if clean_text:
-            ocr_results.append(clean_text)
             # Add text to annotated image
-            cv2.putText(annotated_image, clean_text, (x1, y1-10), 
+            cv2.putText(annotated_image, plate_text, (x1, y1-10), 
                        cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
     
     # Convert back to RGB for Gradio display
-    annotated_image = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)
-    
-    # Format OCR results
-    ocr_text = "Detected text:\n" + "\n".join(ocr_results) if ocr_results else "No text detected"
+    annotated_image = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)     
     
     return annotated_image, ocr_text
+    
 
 # Create Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("# YOLO Object Detection with OCR")
+    gr.Markdown("# MIA-Yolov8 for peruvian plate recognition")
     with gr.Row():
         with gr.Column():
             input_image = gr.Image(label="Input Image", type="pil")
             conf_slider = gr.Slider(0, 1, value=0.25, label="Confidence Threshold")
             iou_slider = gr.Slider(0, 1, value=0.45, label="IOU Threshold")
-            submit_btn = gr.Button("Detect Objects")
+            submit_btn = gr.Button("Run model")
         with gr.Column():
             output_image = gr.Image(label="Detected Objects")
             ocr_output = gr.Textbox(label="OCR Results")