Spaces:
Running
Running
File size: 2,659 Bytes
8011742 b8048df 8011742 b8048df 8011742 b8048df 8011742 8601923 fc269ec 8601923 8011742 |
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 |
import gradio as gr
import pytesseract
from PIL import Image, ImageDraw
def calculate_iou(box1, box2):
xA = max(box1[0], box2[0])
yA = max(box1[1], box2[1])
xB = min(box1[0] + box1[2], box2[0] + box2[2])
yB = min(box1[1] + box1[3], box2[1] + box2[3])
intersection_area = max(0, xB - xA) * max(0, yB - yA)
box1_area = box1[2] * box1[3]
box2_area = box2[2] * box2[3]
iou = intersection_area / float(box1_area + box2_area - intersection_area)
return iou
def extract_text_and_boxes(image):
data = pytesseract.image_to_data(image, output_type=pytesseract.Output.DICT)
boxes_and_words = []
for i in range(len(data['text'])):
if data['text'][i].strip() != '':
x, y, w, h = data['left'][i], data['top'][i], data['width'][i], data['height'][i]
word = data['text'][i]
boxes_and_words.append({'box': (x, y, w, h), 'word': word})
# Remove overlapping boxes based on IoU
to_remove = set()
for i in range(len(boxes_and_words)):
for j in range(i + 1, len(boxes_and_words)):
box1 = boxes_and_words[i]['box']
box2 = boxes_and_words[j]['box']
iou = calculate_iou(box1, box2)
if iou > 0.5:
# Remove the box with the smaller area
area1 = box1[2] * box1[3]
area2 = box2[2] * box2[3]
if area1 > area2:
to_remove.add(j)
else:
to_remove.add(i)
filtered_boxes_and_words = [bw for i, bw in enumerate(boxes_and_words) if i not in to_remove]
# Draw the filtered boxes
draw = ImageDraw.Draw(image)
for bw in filtered_boxes_and_words:
x, y, w, h = bw['box']
draw.rectangle([x, y, x + w, y + h], outline='red', width=2)
return image, filtered_boxes_and_words
iface = gr.Interface(fn=extract_text_and_boxes,
inputs=gr.Image(type='pil'),
outputs=[gr.Image(type='pil', label="Image with Bounding Boxes"),
gr.JSON(label="Extracted Words and Boxes")],
title="Test Tesseract",
description="Test PyTesseract.",
examples=[
["examples/0OO7hzT.jpg"],
["examples/AMzONbY.jpg"],
["examples/eHKLg23.jpg"],
["examples/GXEadxE.jpg"],
["examples/HiuFcCq.jpg"],
["examples/aDsFny6.jpg"],
["examples/iF0qg3H.jpg"],
])
iface.launch()
|