StampDetection / app.py
HimankJ's picture
Update app.py
a4de4e1 verified
from object_detector import DocumentObjects
import cv2
import gradio as gr
import uuid,os
from PIL import Image, ImageDraw, ImageFont
detector = DocumentObjects()
print('Detector ready')
def showStamp(path,bounding_boxes,output_path=None):
image = cv2.imread(path)
for box in bounding_boxes:
x1, y1, x2, y2 = map(int, box)
cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
cv2.putText(image, 'stamp', (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
if output_path:
input_img = Image.fromarray(image)
input_img.save(output_path)
else:
return image
print('Going for inference')
def inference(input_img):
input_img = Image.fromarray(input_img)
file_id = str(uuid.uuid4())
file_name = f'Img_{file_id}.png'
input_img.save(file_name)
resp = detector.detect_objects(file_name)
found = resp['stamp']['found']
coordinates = resp['stamp']['loc']
annotated_img = showStamp(file_name,coordinates)
os.remove(file_name)
if found:
return annotated_img, 'Stamp found'
else:
return annotated_img, 'Stamp not found'
title = "Stamp Detection usong YOLOV9"
description = "A simple Gradio interface to infer on Yolo model trained on custom data for stamp detection"
examples = [
'example1.png',
'example2.jpeg'
]
demo = gr.Interface(
inference,
inputs = [
gr.Image(label="Input Image"),
],
outputs = [
gr.Image(label="Output Image"),
gr.Label()
],
title = title,
description = description,
examples = examples,
)
demo.launch(debug=True)