obj / app.py
aarav-kushwaha's picture
Update app.py
6ed7da3 verified
Raw
History Blame Contribute Delete
3.2 kB
import gradio as gr
import torch
from ultralytics import YOLO
from PIL import Image
import numpy as np
import easyocr
import cv2
# Load YOLOv8 Model (Nano for speed)
model = YOLO('yolov8n.pt')
# Load EasyOCR (English)
# This will download the OCR models on first run
reader = easyocr.Reader(['en'], gpu=torch.cuda.is_available())
def detect_and_read(image, threshold):
# 1. YOLO Detection
results = model.predict(source=image, conf=threshold)
# 2. EasyOCR Full Page Read
# Convert PIL to numpy for EasyOCR
img_np = np.array(image)
ocr_results = reader.readtext(img_np)
detections = []
# Process YOLO boxes (UI Components)
for r in results:
for box in r.boxes:
b = box.xyxy[0].cpu().numpy()
c = box.cls.cpu().item()
label = model.names[int(c)]
# Box coords [xmin, ymin, xmax, ymax]
xmin, ymin, xmax, ymax = b
# Link OCR text to this specific YOLO box
box_text = ""
for (ocr_box, text, prob) in ocr_results:
# Calculate OCR center
ox = (ocr_box[0][0] + ocr_box[2][0]) / 2
oy = (ocr_box[0][1] + ocr_box[2][1]) / 2
# If text is inside the UI box
if xmin <= ox <= xmax and ymin <= oy <= ymax:
box_text += text + " "
# Normalize coordinates for RapnssZ
norm_box = [
float(ymin / image.height), # ymin
float(xmin / image.width), # xmin
float(ymax / image.height), # ymax
float(xmax / image.width) # xmax
]
detections.append({
"box": norm_box,
"label": label,
"text": box_text.strip(),
"score": float(box.conf.cpu().item())
})
# 3. Handle Floating Text (OCR text not inside a YOLO box)
for (ocr_box, text, prob) in ocr_results:
already_matched = False
for d in detections:
if text in d["text"]:
already_matched = True
break
if not already_matched:
# Add as a 'text' element
# EasyOCR box format: [[x,y], [x,y], [x,y], [x,y]]
norm_box = [
float(ocr_box[0][1] / image.height),
float(ocr_box[0][0] / image.width),
float(ocr_box[2][1] / image.height),
float(ocr_box[2][0] / image.width)
]
detections.append({
"box": norm_box,
"label": "text",
"text": text,
"score": float(prob)
})
return detections
# Gradio Interface
iface = gr.Interface(
fn=detect_and_read,
inputs=[
gr.Image(type="pil", label="Screen Capture"),
gr.Slider(0, 1, value=0.2, label="Radar Sensitivity (Threshold)")
],
outputs=gr.JSON(label="OCR-Radar Analysis"),
title="RapnssZ OCR-Radar // Mega v4.5",
description="Zero-latency UI mapping using YOLOv8 and EasyOCR."
)
iface.launch()