import os
import gradio as gr
from PIL import Image
from app_predict import main

os.system('wget https://github.com/FanChiMao/Competition-2022-Pytorch-Orchid_Classification/releases/download/v0.0/beit_1.pth -P experiments/pretrained_models')
os.system('wget https://github.com/FanChiMao/Competition-2022-Pytorch-Orchid_Classification/releases/download/v0.0/convnext.pth -P experiments/pretrained_models')
os.system('wget https://github.com/FanChiMao/Competition-2022-Pytorch-Orchid_Classification/releases/download/v0.0/dmnfnet.pth -P experiments/pretrained_models')
os.system('wget https://github.com/FanChiMao/Competition-2022-Pytorch-Orchid_Classification/releases/download/v0.0/ecaresnet_50.pth -P experiments/pretrained_models')
os.system('wget https://github.com/FanChiMao/Competition-2022-Pytorch-Orchid_Classification/releases/download/v0.0/efficient.pth -P experiments/pretrained_models')
os.system('wget https://github.com/FanChiMao/Competition-2022-Pytorch-Orchid_Classification/releases/download/v0.0/regnet.pth -P experiments/pretrained_models')
os.system('wget https://github.com/FanChiMao/Competition-2022-Pytorch-Orchid_Classification/releases/download/v0.0/swin.pth -P experiments/pretrained_models')
os.system('wget https://github.com/FanChiMao/Competition-2022-Pytorch-Orchid_Classification/releases/download/v0.0/vit.pth -P experiments/pretrained_models')

def inference(img, model):
    os.system('mkdir test')
    img.save("test/1.png", "PNG")
    
    if model == 'Swin transformer':
        predict = main('swin')
    elif model == 'BEiT':
        predict = main('beit')
    elif model == 'NFNet':
        predict = main('dmnfnet')
    elif model == 'ECA-Resnet':
        predict = main('ecaresnet_50')
    elif model == 'EfficientNet':
        predict = main('efficient')
    elif model == 'Regnet':
        predict = main('regnet')
    elif model == 'ViT':
        predict = main('vit')
    elif model == 'ConvNext':
        predict = main('convnext') 
    print(predict )   
    return predict 


title = "[AICUP 2022] Orchid Image Classification (single image quick demo)"
description = ""
article = "<p style='text-align: center'><a href='https://github.com/FanChiMao/Competition-2022-Pytorch-Orchid_Classification' target='_blank'>Orchid image classification</a> | <a href='https://github.com/FanChiMao/Competition-2022-Pytorch-Orchid_Classification' target='_blank'>Github Repo</a></p> <center><img src='https://visitor-badge.glitch.me/badge?page_id=2022aicuphg' alt='visitor badge'></center>"

examples = [
['figures/1.jpg', 'ConvNext'], 
['figures/2.jpg', 'ConvNext'],
['figures/3.jpg', 'ConvNext'],
['figures/4.jpg', 'ConvNext'],
['figures/5.jpg', 'ConvNext'],
]
gr.Interface(
    inference,
    [gr.inputs.Image(type="pil", label="Input"), gr.inputs.Dropdown(choices=[
    'Swin transformer', 
    'BEiT',
    'NFNet',
    'ECA-Resnet',
    'EfficientNet',
    'Regnet',
    'ViT',
    'ConvNext'
    ], type="value", default='Swin transformer', label="model")],
    outputs="label",
    title=title,
    description=description,
    article=article,
    allow_flagging=False,
    allow_screenshot=False,
    examples=examples
).launch(debug=True)