import gradio as gr
import torch
from PIL import Image
from torchvision import transforms
import huggingface_hub as hf

# HuggingFace model and Spaces
model = torch.hub.load('facebookresearch/deit:main', 'deit_tiny_patch16_224')
repo = "uploader"

# File upload component with local or remote option
file = gr.FileInput(type="file", label="Upload Image File", preview=True)

# Display image 
image = gr.Image(label="Uploaded Image")

# Upload file to HuggingFace Spaces
def upload_to_hf(filename):
  with open(filename, 'rb') as f:
    data = f.read()
  hf.upload_file(data, f"/{repo}/{filename}")
  return f"/{repo}/{filename}"

# Run model on image
def run(file):
  if file.startswith("http"): # remote file
    filename = file.split("/")[-1]
    filepath = upload_to_hf(filename)
  else: # local file
    filepath = file

  image = Image.open(filepath).convert('RGB')

  # preprocess, run model, return output
  transform = transforms.Compose([
    transforms.Resize(256), 
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize([0.5,0.5,0.5],[0.5,0.5,0.5])  
  ])

  tensor = transform(image)
  tensor = tensor.unsqueeze(0)

  with torch.no_grad():
    output = model(tensor)

  image.update(filepath)
  return image

# Launch app
app = gr.Interface(fn=run, inputs=[file], outputs=[image])
app.launch()