Update app.py

app.py

@@ -4,27 +4,39 @@ import cv2
 import pandas as pd
 import gradio as gr
 import pytesseract
+import spaces
+
 from PIL import Image
 
 
+# =========================================================
+# ZERO GPU STARTUP CHECK
+# =========================================================
+
+@spaces.GPU
+def zerogpu_probe():
+    return "OK"
+
+
 # =========================================================
 # OCR CORRECTION
 # =========================================================
 
 def correct_ocr_mistakes(text):
+
     if not text:
         return ""
 
     text = text.upper()
 
     replacements = {
-        'O': '0',
-        'I': '1',
-        'L': '1',
-        'S': '5',
-        'B': '8',
-        'G': '6',
-        'Z': '2'
+        "O": "0",
+        "I": "1",
+        "L": "1",
+        "S": "5",
+        "B": "8",
+        "G": "6",
+        "Z": "2"
     }
 
     for k, v in replacements.items():
@@ -34,7 +46,7 @@ def correct_ocr_mistakes(text):
 
 
 # =========================================================
-# GET LAST 3 DIGITS
+# EXTRACT LAST 3 DIGITS
 # =========================================================
 
 def extract_last_3_digits(val):
@@ -44,7 +56,7 @@ def extract_last_3_digits(val):
 
     val = correct_ocr_mistakes(val)
 
-    digits = re.findall(r'\d', val)
+    digits = re.findall(r"\d", val)
 
     if len(digits) >= 3:
         return "".join(digits[-3:])
@@ -78,12 +90,12 @@ def parse_ocr_text(text):
         # -------------------------------------------------
 
         if (
-            re.search(r'AW[B8]\s*[:\-]', line, re.IGNORECASE)
+            re.search(r"AW[B8]\s*[:\-]", line, re.IGNORECASE)
             or "AWB" in line.upper()
         ):
 
             m = re.search(
-                r'AW[B8]\s*[:\-]?\s*([A-Z0-9_]+)',
+                r"AW[B8]\s*[:\-]?\s*([A-Z0-9_]+)",
                 line,
                 re.IGNORECASE
             )
@@ -100,13 +112,13 @@ def parse_ocr_text(text):
         # -------------------------------------------------
 
         if re.search(
-            r'Packet\s*[I1][D0]',
+            r"Packet\s*[I1][D0]",
             line,
             re.IGNORECASE
         ):
 
             m = re.search(
-                r'Packet\s*[I1][D0]\s*[:\-]?\s*([A-Z0-9_]+)',
+                r"Packet\s*[I1][D0]\s*[:\-]?\s*([A-Z0-9_]+)",
                 line,
                 re.IGNORECASE
             )
@@ -123,13 +135,13 @@ def parse_ocr_text(text):
         # -------------------------------------------------
 
         if re.search(
-            r'Ship\s*T[o0]',
+            r"Ship\s*T[o0]",
             line,
             re.IGNORECASE
         ):
 
             m = re.search(
-                r'Ship\s*T[o0]\s*[:\-]?\s*(.+)',
+                r"Ship\s*T[o0]\s*[:\-]?\s*(.+)",
                 line,
                 re.IGNORECASE
             )
@@ -159,11 +171,13 @@ def preprocess_image(image_path):
     if img is None:
         return []
 
+    # Grayscale
     gray = cv2.cvtColor(
         img,
         cv2.COLOR_BGR2GRAY
     )
 
+    # Resize
     gray = cv2.resize(
         gray,
         None,
@@ -172,6 +186,7 @@ def preprocess_image(image_path):
         interpolation=cv2.INTER_CUBIC
     )
 
+    # Denoise
     denoised = cv2.fastNlMeansDenoising(
         gray,
         None,
@@ -180,6 +195,7 @@ def preprocess_image(image_path):
         searchWindowSize=21
     )
 
+    # Adaptive threshold
     adaptive_thresh = cv2.adaptiveThreshold(
         denoised,
         255,
@@ -202,19 +218,36 @@ def preprocess_image(image_path):
 def process_images(files, progress=gr.Progress()):
 
     if not files:
-        return pd.DataFrame(
-            columns=["AWB", "Packet ID", "Ship To"]
-        ), None
+
+        return (
+            pd.DataFrame(
+                columns=[
+                    "AWB",
+                    "Packet ID",
+                    "Ship To"
+                ]
+            ),
+            None
+        )
+
 
     results = []
 
+
+    # OCR configurations
     ocr_configs = [
-        '--psm 6',
-        '--psm 11'
+        "--psm 6",
+        "--psm 11"
     ]
 
+
     total_files = len(files)
 
+
+    # =====================================================
+    # LOOP THROUGH FILES
+    # =====================================================
+
     for index, file_obj in enumerate(files):
 
         progress(
@@ -222,25 +255,32 @@ def process_images(files, progress=gr.Progress()):
             desc=f"Processing {index + 1}/{total_files}"
         )
 
+
         # -------------------------------------------------
         # GET FILE PATH
         # -------------------------------------------------
 
         if isinstance(file_obj, str):
+
             filepath = file_obj
 
         elif hasattr(file_obj, "name"):
+
             filepath = file_obj.name
 
         else:
+
             filepath = str(file_obj)
 
 
         # -------------------------------------------------
-        # PREPROCESS
+        # PREPROCESS IMAGE
         # -------------------------------------------------
 
-        preprocessed_imgs = preprocess_image(filepath)
+        preprocessed_imgs = preprocess_image(
+            filepath
+        )
+
 
         final_awb = ""
         final_packet = ""
@@ -249,15 +289,16 @@ def process_images(files, progress=gr.Progress()):
         found_all = False
 
 
-        # -------------------------------------------------
+        # =================================================
         # OCR
-        # -------------------------------------------------
+        # =================================================
 
         for img in preprocessed_imgs:
 
             if found_all:
                 break
 
+
             for config in ocr_configs:
 
                 try:
@@ -268,34 +309,45 @@ def process_images(files, progress=gr.Progress()):
                     )
 
                 except Exception:
+
                     continue
 
 
-                awb, packet, ship = parse_ocr_text(text)
+                # Parse OCR
+                awb, packet, ship = parse_ocr_text(
+                    text
+                )
 
 
+                # Save AWB
                 if not final_awb and awb:
                     final_awb = awb
 
+
+                # Save Packet ID
                 if not final_packet and packet:
                     final_packet = packet
 
+
+                # Save Ship To
                 if not final_ship and ship:
                     final_ship = ship
 
 
+                # Stop if everything found
                 if (
                     final_awb
                     and final_packet
                     and final_ship
                 ):
+
                     found_all = True
                     break
 
 
-        # -------------------------------------------------
-        # SAVE RESULT
-        # -------------------------------------------------
+        # =================================================
+        # ADD RESULT
+        # =================================================
 
         results.append({
             "AWB": final_awb,
@@ -305,7 +357,7 @@ def process_images(files, progress=gr.Progress()):
 
 
     # =====================================================
-    # DATAFRAME
+    # CREATE DATAFRAME
     # =====================================================
 
     df = pd.DataFrame(
@@ -319,7 +371,7 @@ def process_images(files, progress=gr.Progress()):
 
 
     # =====================================================
-    # EXCEL OUTPUT
+    # CREATE EXCEL
     # =====================================================
 
     output_excel = os.path.abspath(
@@ -336,21 +388,28 @@ def process_images(files, progress=gr.Progress()):
 
 
 # =========================================================
-# GRADIO UI
+# GRADIO INTERFACE
 # =========================================================
 
 with gr.Blocks(
     title="RVP Label OCR Extractor"
 ) as interface:
 
+
+    # -----------------------------------------------------
+    # TITLE
+    # -----------------------------------------------------
+
     gr.Markdown(
         "# 📦 RVP Label OCR Extractor"
     )
 
+
     gr.Markdown(
         "Upload RVP label images to extract "
         "the last 3 digits of AWB, "
-        "last 3 digits of Packet ID, and Ship To."
+        "last 3 digits of Packet ID, "
+        "and Ship To."
     )
 
 
@@ -371,7 +430,7 @@ with gr.Blocks(
 
 
     # -----------------------------------------------------
-    # BUTTON
+    # EXTRACT BUTTON
     # -----------------------------------------------------
 
     extract_btn = gr.Button(
@@ -381,7 +440,7 @@ with gr.Blocks(
 
 
     # -----------------------------------------------------
-    # OUTPUT TABLE
+    # RESULTS
     # -----------------------------------------------------
 
     output_df = gr.Dataframe(