import gradio as gr
from PIL import Image
import pytesseract

import yolov5

# Pytesseract init
choices = pytesseract.get_languages(config='')

def text_inference(filepath, languages):

    return pytesseract.image_to_string(Image.open(filepath), lang=', '.join(languages))

# YOLOv5 Init
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)