Spaces:
Sleeping
Sleeping
| import os | |
| import re | |
| import cv2 | |
| import pandas as pd | |
| import gradio as gr | |
| import pytesseract | |
| import spaces | |
| from PIL import Image | |
| # ========================================================= | |
| # ZERO GPU STARTUP CHECK | |
| # ========================================================= | |
| 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" | |
| } | |
| for k, v in replacements.items(): | |
| text = text.replace(k, v) | |
| return text | |
| # ========================================================= | |
| # EXTRACT LAST 3 DIGITS | |
| # ========================================================= | |
| def extract_last_3_digits(val): | |
| if not val: | |
| return "" | |
| val = correct_ocr_mistakes(val) | |
| digits = re.findall(r"\d", val) | |
| if len(digits) >= 3: | |
| return "".join(digits[-3:]) | |
| return "" | |
| # ========================================================= | |
| # PARSE OCR TEXT | |
| # ========================================================= | |
| def parse_ocr_text(text): | |
| awb_val = "" | |
| packet_val = "" | |
| ship_to_val = "" | |
| if not text: | |
| return "", "", "" | |
| lines = [ | |
| line.strip() | |
| for line in text.split("\n") | |
| if line.strip() | |
| ] | |
| for i, line in enumerate(lines): | |
| # ------------------------------------------------- | |
| # AWB | |
| # ------------------------------------------------- | |
| if ( | |
| 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_]+)", | |
| line, | |
| re.IGNORECASE | |
| ) | |
| if m and len(m.group(1)) > 4: | |
| awb_val = m.group(1) | |
| elif i + 1 < len(lines): | |
| awb_val = lines[i + 1] | |
| # ------------------------------------------------- | |
| # PACKET ID | |
| # ------------------------------------------------- | |
| if re.search( | |
| r"Packet\s*[I1][D0]", | |
| line, | |
| re.IGNORECASE | |
| ): | |
| m = re.search( | |
| r"Packet\s*[I1][D0]\s*[:\-]?\s*([A-Z0-9_]+)", | |
| line, | |
| re.IGNORECASE | |
| ) | |
| if m and len(m.group(1)) > 4: | |
| packet_val = m.group(1) | |
| elif i + 1 < len(lines): | |
| packet_val = lines[i + 1] | |
| # ------------------------------------------------- | |
| # SHIP TO | |
| # ------------------------------------------------- | |
| if re.search( | |
| r"Ship\s*T[o0]", | |
| line, | |
| re.IGNORECASE | |
| ): | |
| m = re.search( | |
| r"Ship\s*T[o0]\s*[:\-]?\s*(.+)", | |
| line, | |
| re.IGNORECASE | |
| ) | |
| if m and m.group(1).strip(): | |
| ship_to_val = m.group(1).strip() | |
| elif i + 1 < len(lines): | |
| ship_to_val = lines[i + 1].strip() | |
| return ( | |
| extract_last_3_digits(awb_val), | |
| extract_last_3_digits(packet_val), | |
| ship_to_val | |
| ) | |
| # ========================================================= | |
| # IMAGE PREPROCESSING | |
| # ========================================================= | |
| def preprocess_image(image_path): | |
| img = cv2.imread(image_path) | |
| if img is None: | |
| return [] | |
| # Grayscale | |
| gray = cv2.cvtColor( | |
| img, | |
| cv2.COLOR_BGR2GRAY | |
| ) | |
| # Resize | |
| gray = cv2.resize( | |
| gray, | |
| None, | |
| fx=2, | |
| fy=2, | |
| interpolation=cv2.INTER_CUBIC | |
| ) | |
| # Denoise | |
| denoised = cv2.fastNlMeansDenoising( | |
| gray, | |
| None, | |
| h=10, | |
| templateWindowSize=7, | |
| searchWindowSize=21 | |
| ) | |
| # Adaptive threshold | |
| adaptive_thresh = cv2.adaptiveThreshold( | |
| denoised, | |
| 255, | |
| cv2.ADAPTIVE_THRESH_GAUSSIAN_C, | |
| cv2.THRESH_BINARY, | |
| 31, | |
| 2 | |
| ) | |
| return [ | |
| Image.fromarray(adaptive_thresh), | |
| Image.fromarray(gray) | |
| ] | |
| # ========================================================= | |
| # PROCESS IMAGES | |
| # ========================================================= | |
| def process_images(files, progress=gr.Progress()): | |
| if not files: | |
| return ( | |
| pd.DataFrame( | |
| columns=[ | |
| "AWB", | |
| "Packet ID", | |
| "Ship To" | |
| ] | |
| ), | |
| None | |
| ) | |
| results = [] | |
| # OCR configurations | |
| ocr_configs = [ | |
| "--psm 6", | |
| "--psm 11" | |
| ] | |
| total_files = len(files) | |
| # ===================================================== | |
| # LOOP THROUGH FILES | |
| # ===================================================== | |
| for index, file_obj in enumerate(files): | |
| progress( | |
| (index + 1) / total_files, | |
| 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 IMAGE | |
| # ------------------------------------------------- | |
| preprocessed_imgs = preprocess_image( | |
| filepath | |
| ) | |
| final_awb = "" | |
| final_packet = "" | |
| final_ship = "" | |
| found_all = False | |
| # ================================================= | |
| # OCR | |
| # ================================================= | |
| for img in preprocessed_imgs: | |
| if found_all: | |
| break | |
| for config in ocr_configs: | |
| try: | |
| text = pytesseract.image_to_string( | |
| img, | |
| config=config | |
| ) | |
| except Exception: | |
| continue | |
| # 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 | |
| # ================================================= | |
| # ADD RESULT | |
| # ================================================= | |
| results.append({ | |
| "AWB": final_awb, | |
| "Packet ID": final_packet, | |
| "Ship To": final_ship | |
| }) | |
| # ===================================================== | |
| # CREATE DATAFRAME | |
| # ===================================================== | |
| df = pd.DataFrame( | |
| results, | |
| columns=[ | |
| "AWB", | |
| "Packet ID", | |
| "Ship To" | |
| ] | |
| ) | |
| # ===================================================== | |
| # CREATE EXCEL | |
| # ===================================================== | |
| output_excel = os.path.abspath( | |
| "RVP_Data.xlsx" | |
| ) | |
| df.to_excel( | |
| output_excel, | |
| index=False | |
| ) | |
| return df, output_excel | |
| # ========================================================= | |
| # 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." | |
| ) | |
| # ----------------------------------------------------- | |
| # FILE UPLOAD | |
| # ----------------------------------------------------- | |
| file_input = gr.File( | |
| file_count="multiple", | |
| file_types=[ | |
| ".jpg", | |
| ".jpeg", | |
| ".png" | |
| ], | |
| label="Upload RVP Photos", | |
| type="filepath" | |
| ) | |
| # ----------------------------------------------------- | |
| # EXTRACT BUTTON | |
| # ----------------------------------------------------- | |
| extract_btn = gr.Button( | |
| "Extract Data", | |
| variant="primary" | |
| ) | |
| # ----------------------------------------------------- | |
| # RESULTS | |
| # ----------------------------------------------------- | |
| output_df = gr.Dataframe( | |
| headers=[ | |
| "AWB", | |
| "Packet ID", | |
| "Ship To" | |
| ], | |
| datatype=[ | |
| "str", | |
| "str", | |
| "str" | |
| ], | |
| label="Results", | |
| interactive=False | |
| ) | |
| # ----------------------------------------------------- | |
| # EXCEL DOWNLOAD | |
| # ----------------------------------------------------- | |
| output_file = gr.File( | |
| label="Download Excel File" | |
| ) | |
| # ----------------------------------------------------- | |
| # BUTTON EVENT | |
| # ----------------------------------------------------- | |
| extract_btn.click( | |
| fn=process_images, | |
| inputs=file_input, | |
| outputs=[ | |
| output_df, | |
| output_file | |
| ] | |
| ) | |
| # ========================================================= | |
| # LAUNCH | |
| # ========================================================= | |
| if __name__ == "__main__": | |
| interface.launch( | |
| server_name="0.0.0.0", | |
| server_port=7860, | |
| ssr_mode=False | |
| ) |