panda1835's picture
Update app.py
346c324
from PIL import Image
import numpy as np
import gradio as gr
import pandas as pd
import torch
import json
from ultralytics import YOLO
from datetime import datetime
print(f"Is CUDA available: {torch.cuda.is_available()}")
# print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
# Load the model
# detect_model = torch.hub.load('ultralytics/yolov5', 'custom', path='best.pt', force_reload=True, _verbose=False)
detect_model = YOLO('yolov8m_2023-10-23_best.pt')
def detect(image):
results = detect_model(image, conf=0.1)
# Get the current time
current_time = datetime.now()
# Format the current time as a string
formatted_time = current_time.strftime("%Y-%m-%d %H:%M:%S")
print(formatted_time)
try:
results = results[0].boxes.xyxy[0].cpu().numpy()
top = int(results[1])
left = int(results[0])
width = int(results[2] - results[0])
height = int(results[3] - results[1])
return {
"top": top,
"left": left,
"width": width,
"height": height
}
except:
return {
"top": 0,
"left": 0,
"width": 0,
"height": 0
}
title = "🐢"
gr.Interface(
fn=detect,
inputs=gr.Image(type="pil", label="Input Image"),
outputs=[gr.JSON()],
# live=True,
title=title,
).launch()