import torch
import cv2
from inference_utils import get_predictions
import os
import gradio as gr


weights_path = os.path.join('best.pt')
yolo_path = os.path.join('yolov5')
model = torch.hub.load(yolo_path, 'custom', path = weights_path, source = 'local',device='cpu',force_reload=True)

def show_preds_image(image_path):
    """
    This function takes an image path as input and returns an image with bounding boxes around the detected monuments.
    Args:
        image_path (str): The path to the input image.
    Returns:
        numpy.ndarray: The image with bounding boxes around the detected monuments.
        prediction (dict) : The yolo predictions with bounding box coordinates
    """
        
    try:
        img = cv2.imread(image_path)
        if img is None:
            raise ValueError("Invalid image file")

        img_color = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        bbox_img, predictions = get_predictions(model, img_color, return_raw=True)
        return bbox_img, predictions
    except Exception as e:
        return None, f"Error: {e}"
    
    
    

inputs_image = [
    gr.components.Image(type="filepath", label="Input Image"),
]
outputs_image = [
    gr.components.Image(type="numpy", label="Output Image"),
    gr.components.Textbox()
]

interface_image = gr.Interface(
    fn=show_preds_image,
    inputs=inputs_image,
    outputs=outputs_image,
    title="Monument detector",
    cache_examples=False,
)

if __name__ == "__main__":
    gr.TabbedInterface(
        [interface_image],
        tab_names=['Image inference']
    ).queue().launch()