Improve OCR and set python 3.10

README.md

@@ -5,7 +5,7 @@ colorFrom: gray
 colorTo: gray
 sdk: gradio
 sdk_version: 6.14.0
-python_version: '3.13'
+python_version: '3.10'
 app_file: app.py
 pinned: false
 license: mit

app.py

@@ -25,7 +25,9 @@ def run_ocr(image_bgr):
         return ocr.ocr(image_bgr, cls=True)
     except TypeError:
         return ocr.ocr(image_bgr)
-plate_regex = re.compile(r"[A-Z]{2}[0-9]{2}[A-Z]{1,2}[0-9]{4}")
+
+
+plate_regex = re.compile(r"[A-Z]{2}\d{1,2}[A-Z]{1,3}\d{3,4}")
 
 
 def preprocess_plate(crop_rgb):
@@ -40,6 +42,29 @@ def preprocess_plate(crop_rgb):
     return cv2.cvtColor(filtered, cv2.COLOR_GRAY2BGR)
 
 
+def deskew(image_bgr):
+    gray = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2GRAY)
+    coords = np.column_stack(np.where(gray > 0))
+    if len(coords) < 10:
+        return image_bgr
+
+    angle = cv2.minAreaRect(coords)[-1]
+    if angle < -45:
+        angle = -(90 + angle)
+    else:
+        angle = -angle
+
+    h, w = image_bgr.shape[:2]
+    matrix = cv2.getRotationMatrix2D((w // 2, h // 2), angle, 1.0)
+    return cv2.warpAffine(
+        image_bgr,
+        matrix,
+        (w, h),
+        flags=cv2.INTER_CUBIC,
+        borderMode=cv2.BORDER_REPLICATE,
+    )
+
+
 def clean_text(text):
     return re.sub(r"[^A-Z0-9]", "", text.upper())
 
@@ -54,6 +79,33 @@ def fix_common(text):
     )
 
 
+def smart_fix(text):
+    if len(text) >= 10:
+        chars = list(text)
+        chars[2] = chars[2].replace("O", "0")
+        chars[3] = chars[3].replace("O", "0")
+        return "".join(chars)
+    return text
+
+
+def build_crops(crop_bgr):
+    scaled = cv2.resize(crop_bgr, None, fx=1.2, fy=1.2, interpolation=cv2.INTER_CUBIC)
+    blurred = cv2.GaussianBlur(crop_bgr, (3, 3), 0)
+    return [crop_bgr, scaled, blurred]
+
+
+def combine_ocr_text(ocr_out):
+    if not ocr_out or not ocr_out[0]:
+        return "", 0.0
+    texts = [item[1][0] for item in ocr_out[0] if item and item[1]]
+    confidences = [float(item[1][1]) for item in ocr_out[0] if item and item[1]]
+    if not texts or not confidences:
+        return "", 0.0
+    full_text = "".join(texts)
+    avg_conf = sum(confidences) / len(confidences)
+    return full_text, avg_conf
+
+
 def select_best(plates):
     if not plates:
         return "", None
@@ -78,7 +130,12 @@ def detect(image):
 
     height, width = image.shape[:2]
     plates = []
-    for idx, (x1, y1, x2, y2) in enumerate(boxes.xyxy.cpu().numpy(), start=1):
+    xyxy = boxes.xyxy.cpu().numpy()
+    confs = boxes.conf.cpu().numpy()
+
+    for idx, (x1, y1, x2, y2) in enumerate(xyxy, start=1):
+        if float(confs[idx - 1]) < 0.5:
+            continue
         pad = int(0.03 * max(x2 - x1, y2 - y1))
         left = max(int(x1) - pad, 0)
         top = max(int(y1) - pad, 0)
@@ -88,10 +145,28 @@ def detect(image):
             continue
 
         crop = image[top:bottom, left:right]
-        prepped = preprocess_plate(crop)
-        ocr_out = run_ocr(prepped)
+        best_candidate = None
+
+        for variant in build_crops(crop):
+            prepped = preprocess_plate(variant)
+            prepped = deskew(prepped)
+            ocr_out = run_ocr(prepped)
+            raw_text, confidence = combine_ocr_text(ocr_out)
+            if not raw_text:
+                continue
+
+            normalized = smart_fix(fix_common(clean_text(raw_text)))
+            if len(normalized) < 8 or confidence < 0.5:
+                continue
+
+            if best_candidate is None or confidence > best_candidate["confidence"]:
+                best_candidate = {
+                    "raw": raw_text,
+                    "normalized": normalized,
+                    "confidence": confidence,
+                }
 
-        if not ocr_out or not ocr_out[0]:
+        if best_candidate is None:
             plates.append(
                 {
                     "box": idx,
@@ -103,17 +178,12 @@ def detect(image):
             )
             continue
 
-        best = max(ocr_out[0], key=lambda item: item[1][1])
-        raw_text = best[1][0]
-        confidence = float(best[1][1])
-        normalized = fix_common(clean_text(raw_text))
-
         plates.append(
             {
                 "box": idx,
-                "raw": raw_text,
-                "normalized": normalized,
-                "confidence": confidence,
+                "raw": best_candidate["raw"],
+                "normalized": best_candidate["normalized"],
+                "confidence": best_candidate["confidence"],
                 "bbox": [left, top, right, bottom],
             }
         )