File size: 2,981 Bytes
1be9dc9 abb83b0 1be9dc9 |
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 |
# Import libraries
import cv2 # for reading images, draw bounding boxes
from ultralytics import YOLO
import gradio as gr
# Define constants
ENTITIES_COLORS = {
"Caption": (191, 100, 21),
"Footnote": (2, 62, 115),
"Formula": (140, 80, 58),
"List-item": (168, 181, 69),
"Page-footer": (2, 69, 84),
"Page-header": (83, 115, 106),
"Picture": (255, 72, 88),
"Section-header": (0, 204, 192),
"Table": (116, 127, 127),
"Text": (0, 153, 221),
"Title": (196, 51, 2)
}
BOX_PADDING = 2
# Load models
DETECTION_MODEL = YOLO("models/dla-model.pt")
def detect(image_path):
"""
Output inference image with bounding box
Args:
- image: to check for checkboxes
Return: image with bounding boxes drawn
"""
image = cv2.imread(image_path)
if image is None:
return image
# Predict on image
results = DETECTION_MODEL.predict(source=image, conf=0.2, iou=0.8) # Predict on image
boxes = results[0].boxes # Get bounding boxes
if len(boxes) == 0:
return image
# Get bounding boxes
for box in boxes:
detection_class_conf = round(box.conf.item(), 2)
cls = list(ENTITIES_COLORS)[int(box.cls)]
# Get start and end points of the current box
start_box = (int(box.xyxy[0][0]), int(box.xyxy[0][1]))
end_box = (int(box.xyxy[0][2]), int(box.xyxy[0][3]))
# 01. DRAW BOUNDING BOX OF OBJECT
line_thickness = round(0.002 * (image.shape[0] + image.shape[1]) / 2) + 1
image = cv2.rectangle(img=image,
pt1=start_box,
pt2=end_box,
color=ENTITIES_COLORS[cls],
thickness = line_thickness) # Draw the box with predefined colors
# 02. DRAW LABEL
text = cls + " " + str(detection_class_conf)
# Get text dimensions to draw wrapping box
font_thickness = max(line_thickness - 1, 1)
(text_w, text_h), _ = cv2.getTextSize(text=text, fontFace=2, fontScale=line_thickness/3, thickness=font_thickness)
# Draw wrapping box for text
image = cv2.rectangle(img=image,
pt1=(start_box[0], start_box[1] - text_h - BOX_PADDING*2),
pt2=(start_box[0] + text_w + BOX_PADDING * 2, start_box[1]),
color=ENTITIES_COLORS[cls],
thickness=-1)
# Put class name on image
start_text = (start_box[0] + BOX_PADDING, start_box[1] - BOX_PADDING)
image = cv2.putText(img=image, text=text, org=start_text, fontFace=0, color=(255,255,255), fontScale=line_thickness/3, thickness=font_thickness)
return image
iface = gr.Interface(fn=detect,
inputs=gr.Image(label="Upload scanned document", type="filepath"),
outputs="image")
iface.launch()
|