detectMonument / app.py
Ujjwal123's picture
First working version of the app
0bf004a
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()