import gradio as gr
import time
from PIL import Image
import torch
from torch import nn
import timm
import os


title = "Age-PT"
description = "ViT(medium clip) based model. transfer trained with custom dataset"
article = "Through bunch of fine tuning and experiments. REMEMBER! This model can be wrong."

MODEL_NAME = "vit_medium_patch16_clip_224.tinyclip_yfcc15m"
FILE_NAME = "pretrained_weight/vit_medium_patch16_clip_224.tinyclip_yfcc15m(trainable 0.00) (eval Score 0.9067, Loss 29.465482).pth"
DEVICE = "cpu"




torch.set_default_device(DEVICE)

model = timm.create_model(MODEL_NAME, pretrained=True, num_classes=0, drop_rate=0.7)
model_classifier = nn.Sequential(nn.Linear(512, 512),
                                nn.BatchNorm1d(512),
                                nn.GELU(),
                                nn.Linear(512, 1))
model = nn.Sequential(model, model_classifier)
model.load_state_dict(state_dict=torch.load(FILE_NAME, weights_only=True, map_location=torch.device(DEVICE)))

test_transform = timm.data.create_transform(input_size=224, is_training=False, interpolation="bicubic")




def predict(img):
    start_time = time.time()
    model.eval()
    with torch.inference_mode():
        img = test_transform(img).unsqueeze(dim=0).to(DEVICE)
        pred_age = model(img).item()
    
    end_time = time.time()
    
    elapsed_time = end_time - start_time
    fps = 1 / elapsed_time
    return pred_age, fps

example_list = [["examples/" + example] for example in os.listdir("examples")]


demo = gr.Interface(fn=predict, 
                    inputs=gr.Image(type="pil"),
                    outputs=[gr.Number(label="Age Prediction"),
                             gr.Number(label="Prediction speed (fps)")], 
                    examples=example_list,
                    title=title,
                    description=description,
                    article=article)
demo.launch()