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)