from pathlib import Path
import gradio as gr
import spaces
import torch
import os
from PIL import Image
import subprocess
import tempfile

zero = torch.Tensor([0]).cuda()
print(zero.device) # <-- 'cpu' 🤔

CHECKPOINT="./D-FINE/weight/dfine-n.pth"

@spaces.GPU
def inference(image):
    temp_input = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False)
    image.save(temp_input.name)
    temp_input.close()

    subprocess.run([
        "python", "D-FINE/tools/inference/torch_inf.py",
        "-c", "D-FINE/configs/dfine/custom/dfine_hgnetv2_n_custom.yml",  
        "-r", CHECKPOINT,
        "--input", temp_input.name,
        "--device", "cuda:0"
    ], check=True)

    output_path = temp_input.name  
    output_image = Image.open(output_path)

    return output_image

def get_default_image_paths(folder_path):
    image_extensions = ('.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff')
    image_paths = [[os.path.join(folder_path, file)] for file in os.listdir(folder_path) 
                   if file.lower().endswith(image_extensions)]
    return image_paths

default_images = get_default_image_paths(Path("examples/"))

# Create Gradio Interface with title and description
iface = gr.Interface(
    fn=inference,
    inputs=[
        gr.Image(label="Upload Image", type="pil"),  # File input as PIL Image
    ],
    outputs=gr.Image(type="pil", label="Output (Image)"),  # Show output as an image
    examples=default_images,
    cache_examples=False,
    title="Strawberry Disease Detection DFINE S",
    description="This application detects diseases in strawberries using a trained D-FINE N model. " \
    "Upload an image use your webcam for analysis."
)

iface.launch()


# demo = gr.Interface(fn=greet, inputs=gr.Number(), outputs=gr.Text())
# demo.launch()