|
import gradio as gr
|
|
from PIL import Image
|
|
import pytesseract
|
|
|
|
import yolov5
|
|
|
|
|
|
choices = pytesseract.get_languages(config='')
|
|
|
|
def text_inference(filepath, languages):
|
|
|
|
return pytesseract.image_to_string(Image.open(filepath), lang=', '.join(languages))
|
|
|
|
|
|
model = yolov5.load('nickgambirasi/yolov5s-recycling')
|
|
|
|
model.conf = 0.2
|
|
model.iou = 0.45
|
|
model.agnostic = False
|
|
model.multi_label = False
|
|
model.max_det = 1000
|
|
|
|
def yolo_inference(image):
|
|
|
|
results = model(Image.open(image), size=640)
|
|
predictions = results.pred[0]
|
|
|
|
scores = predictions[:, 4]
|
|
|
|
return scores
|
|
|
|
def full_inference(image, languages):
|
|
|
|
text_infer = text_inference(image, languages)
|
|
symbol_infer = yolo_inference(image)
|
|
|
|
output = {}
|
|
|
|
if text_infer:
|
|
|
|
output.update({'language_id': 'all_pass'})
|
|
output.update({'text': text_infer})
|
|
|
|
|
|
else:
|
|
|
|
output.update({'language_id': 'fail'})
|
|
|
|
if any(symbol_infer > model.conf):
|
|
|
|
output.update({'symbols': 'all_pass'})
|
|
|
|
else:
|
|
|
|
output.update({'symbols': 'fail'})
|
|
|
|
return output
|
|
|
|
title = "Hyperintelligent Art Parser"
|
|
description = "Gradio deployment of models for text and symbol detection on product labels, powered by Tesseract and YOLOv5"
|
|
gr.Interface(
|
|
full_inference,
|
|
[gr.inputs.Image(type="filepath", label="Upload Your Label"), gr.inputs.CheckboxGroup(choices, type="value", default=['eng'], label="Language of Text")],
|
|
'text',
|
|
title=title,
|
|
description=description,
|
|
article=None,
|
|
examples=None
|
|
).launch(enable_queue=True) |