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()