Spaces:
Sleeping
Sleeping
File size: 5,395 Bytes
8808a46 847202a 8808a46 95e0f84 b12bb6f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 | import gradio as gr
from ultralytics import YOLO
from PIL import Image
import tempfile
import os
import shutil
import numpy as np
import easyocr
from collections import defaultdict
import cv2
# Load YOLO model
print("[INFO] Loading Custom YOLO Segmentation model...")
model = YOLO('trained_seg_model/seg_v1.pt')
print("[INFO] Custom YOLO Segmentation model loaded successfully.")
# Load EasyOCR model
print("[INFO] Loading EasyOCR model...")
reader = easyocr.Reader(['fr'])
print("[INFO] EasyOCR model initialized.")
# OCR helper function
def run_easyocr_on_yolo_segments(results, original_img: Image.Image):
print("[INFO] Running EasyOCR on segmented regions...")
original_img_np = np.array(original_img.convert("RGB"))
ocr_dict = defaultdict(list)
class_counts = defaultdict(int)
cropped_images = []
# for result in results:
# if result.masks is None or result.boxes is None:
# continue
for result_index, result in enumerate(results):
print(f"[DEBUG] Processing result index: {result_index}")
if result.masks is None or result.boxes is None:
print("[WARN] Result has no masks or boxes.")
continue
# for seg, class_id in zip(result.masks.xy, result.boxes.cls):
for seg_index, (seg, class_id) in enumerate(zip(result.masks.xy, result.boxes.cls)):
class_name = model.names[int(class_id)]
print(f"[DEBUG] Segment {seg_index}: Class = {class_name}")
poly = np.array(seg, dtype=np.int32)
x, y, w, h = cv2.boundingRect(poly)
crop = original_img_np[y:y+h, x:x+w]
if crop.size == 0:
print(f"[WARN] Empty crop for class: {class_name}")
text = ""
else:
ocr_result = reader.readtext(crop, detail=0)
text = " ".join(ocr_result).strip()
print(f"[DEBUG] OCR result for {class_name}: {text[:60]}...")
# Always use a suffix, starting from 1
class_counts[class_name] += 1
class_key = f"{class_name}_{class_counts[class_name]}"
ocr_dict[class_key] = text
# Convert cropped numpy image to PIL for gallery
cropped_pil = Image.fromarray(crop)
cropped_images.append((class_key, cropped_pil))
print(f"[INFO] OCR completed. Total regions processed: {len(cropped_images)}")
return cropped_images, dict(ocr_dict)
# Main function for Gradio
def segment_and_ocr(img: Image.Image):
print("[INFO] Starting segmentation + OCR pipeline...")
with tempfile.TemporaryDirectory() as tmpdir:
input_path = os.path.join(tmpdir, "input.jpg")
img.save(input_path)
print(f"[INFO] Saved input image to: {input_path}")
print("[INFO] Running YOLO prediction...")
results = model.predict(
source=input_path,
save=True,
save_txt=True,
project=tmpdir,
name="predict",
exist_ok=True
)
print("[INFO] YOLO prediction completed.")
output_img_path = os.path.join(tmpdir, "predict", "input.jpg")
label_txt_path = os.path.join(tmpdir, "predict", "labels", "input.txt")
# segmented_img = Image.open(output_img_path) if os.path.exists(output_img_path) else None
# label_data = open(label_txt_path).read() if os.path.exists(label_txt_path) else "No labels generated."
segmented_img = None
if os.path.exists(output_img_path):
segmented_img = Image.open(output_img_path)
print(f"[INFO] Segmented image found: {output_img_path}")
else:
print("[ERROR] Segmented image not found!")
if os.path.exists(label_txt_path):
with open(label_txt_path) as f:
label_data = f.read()
print(f"[INFO] Label file found with length: {len(label_data)} chars")
else:
label_data = "No labels generated."
print("[WARN] Label file not found.")
print("[INFO] Extracting OCR results from segments...")
cropped_images_with_labels, ocr_dict = run_easyocr_on_yolo_segments(results, img)
# Format gallery for Gradio [(image, label)]
gallery_output = [(image, label) for label, image in cropped_images_with_labels]
print(f"[INFO] Returning {len(gallery_output)} cropped images and {len(ocr_dict)} OCR entries.")
return segmented_img, label_data, gallery_output, ocr_dict
# Sample images for "Examples" section
examples = [
["sample_images/seg_v1.jpg"],
["sample_images/seg_v2.jpg"],
["sample_images/seg_v3.jpg"],
["sample_images/seg_v4.jpg"]
]
# Gradio UI
print("[INFO] Launching Gradio interface...")
gr.Interface(
fn=segment_and_ocr,
inputs=gr.Image(type="pil", label="Upload Image"),
outputs=[
gr.Image(type="pil", label="Segmented Image"),
gr.Textbox(label="YOLO Labels (.txt Output)"),
gr.Gallery(label="Segmented Crops (OCR Regions)"),
gr.JSON(label="OCR Output by Class (with suffix)")
],
examples=examples,
cache_examples=False, # <- this disables automatic execution
title="AI-Powered YOLO Segmentation + OCR with EasyOCR",
description="Upload or choose an image to segment using YOLO and extract text using EasyOCR from each region."
).launch() |