File size: 3,121 Bytes
3a34b35 9d2f778 3a34b35 cc8f9e4 3a34b35 41d48c2 3a34b35 | 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 | import gradio as gr
import cv2
import numpy as np
import pytesseract
from ultralytics import YOLO
import easyocr
import os
os.environ['TESSDATA_PREFIX'] = '/usr/share/tesseract-ocr/5/tessdata'
# Load YOLO model
# yolo_model = YOLO("models/NER_66_37_50_V2.pt") # Your trained model Metz2_V2_Dataset_2_E_68_35best
yolo_model = YOLO("models/Metz2_V2_Dataset_2_E_68_35best.pt") # Your trained model
easyocr_readers = {} # Cache for readers
# Store global crops
detected_crops = []
# --- Detection Function ---
def detect_objects(image):
global detected_crops
results = yolo_model(image)[0]
annotated = image.copy()
detected_crops = []
for box in results.boxes.data.tolist():
x1, y1, x2, y2, conf, cls = box
x1, y1 = max(0, int(x1)), max(0, int(y1))
x2, y2 = int(x2), int(y2)
label = yolo_model.names[int(cls)]
cropped = image[y1:y2, x1:x2]
detected_crops.append((label, cropped))
cv2.rectangle(annotated, (x1, y1), (x2, y2), (0,255,0), 2)
cv2.putText(annotated, label, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255,0,0), 2)
crops = [crop for _, crop in detected_crops]
return annotated, crops
# --- OCR Function ---
def run_ocr(engine, lang = 'fr'):
results = []
for label, crop in detected_crops:
try:
if engine == "EasyOCR":
if lang not in easyocr_readers:
easyocr_readers[lang] = easyocr.Reader([lang], gpu=False)
reader = easyocr_readers[lang]
text = "\n".join(reader.readtext(crop, detail=0))
elif engine == "Tesseract":
config = f'--oem 3 --psm 6 -l {lang}'
text = pytesseract.image_to_string(crop, config=config)
else:
text = "[Unsupported Engine]"
except Exception as e:
text = f"[Error: {e}]"
results.append(f"{label}:\n{text}")
return "\n\n".join(results)
# --- Gradio UI ---
with gr.Blocks() as demo:
gr.Markdown("## 🧠 YOLO Detection + OCR Modular App - NER")
with gr.Row():
input_img = gr.Image(type="numpy", label="Upload Image")
detected_img = gr.Image(type="numpy", label="Annotated Image")
detect_button = gr.Button("🔍 Run Detection")
gallery = gr.Gallery(label="Detected Crops")
with gr.Row():
ocr_engine = gr.Radio(["EasyOCR", "Tesseract"], value="EasyOCR", label="OCR Engine")
lang_input = gr.Textbox("fr", label="Language Code")
ocr_button = gr.Button("📝 Run OCR on Detected Regions")
ocr_output = gr.Textbox(label="Extracted Text", lines=10)
detect_button.click(fn=detect_objects, inputs=input_img, outputs=[detected_img, gallery])
ocr_button.click(fn=run_ocr, inputs=[ocr_engine, lang_input], outputs=ocr_output)
# Examples
example_images = ["images/example1.jpg", "images/example2.jpg", "images/example3.jpg", "images/example4.jpg", "images/example5.jpg"]
gr.Examples(
examples=[[img] for img in example_images],
inputs=[input_img],
label="Example Images"
)
demo.launch() |