import gradio as gr
import torch
from huggingface_hub import hf_hub_download
import json
from omegaconf import OmegaConf
import sys
import os
from PIL import Image
import torchvision.transforms as transforms

# Pobierz model i config
repo_id = "Kiwinicki/sat2map-generator"
generator_path = hf_hub_download(repo_id=repo_id, filename="generator.pth")
config_path = hf_hub_download(repo_id=repo_id, filename="config.json")
model_path = hf_hub_download(repo_id=repo_id, filename="model.py")

# Dodaj ścieżkę do modelu
sys.path.append(os.path.dirname(model_path))
from model import Generator

# Załaduj konfigurację
with open(config_path, "r") as f:
    config_dict = json.load(f)
cfg = OmegaConf.create(config_dict)

# Inicjalizacja modelu
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
generator = Generator(cfg).to(device)
generator.load_state_dict(torch.load(generator_path, map_location=device))
generator.eval()

# Transformacje
transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])

def process_image(image):
    # Konwersja do tensora
    image_tensor = transform(image).unsqueeze(0).to(device)
    
    # Inferencja
    with torch.no_grad():
        output_tensor = generator(image_tensor)
    
    # Przygotowanie wyjścia
    output_image = output_tensor.squeeze(0).cpu()
    output_image = output_image * 0.5 + 0.5  # Denormalizacja
    output_image = transforms.ToPILImage()(output_image)
    
    return output_image

iface = gr.Interface(
    fn=process_image,
    inputs=gr.Image(type="pil"),
    outputs="image",
    title="Satellite to Map Generator"
)

iface.launch()