import torch
import gradio as gr
from PIL import Image
from torchvision import transforms
import timm

classes = ['la_badèche', 'la_blennie_coiffée', 'la_blennie_de_caneva', 'la_blennie_de_roux', 'la_blennie_de_zvonimir', 'la_blennie_diabolo', 'la_blennie_gattorugine', 'la_blennie_palmicorne', 'la_blennie_paon', 'la_blennie_pilicorne', 'la_blennie_sphynx', 'la_bogue', 'la_bécune_bouche_jaune', 'la_canthare', 'la_carangue_dentue', 'la_castagnole', 'la_coquette', 'la_dorade_royale', 'la_girelle', 'la_girelle-paon', 'la_gonnelle', 'la_grande_seriole', 'la_grande_vive', 'la_morue', 'la_mostelle_de_roche', 'la_murène_commune', 'la_palomine', 'la_petite_rascasse_rouge', 'la_rascasse_brune', 'la_rascasse_de_madère', 'la_saupe', 'la_sole_commune', 'la_vielle', 'le_baliste_commun', 'le_bar_commun', 'le_barbier_commun', 'le_chabot_buffle', 'le_chabot_commun', 'le_chapon', 'le_congre', 'le_corb', 'le_crénilabre_cendré', 'le_crénilabre_de_melops', 'le_crénilabre_méditerranéen', 'le_crénilabre_ocellé', 'le_crénilabre_tanche', 'le_crénilabre_à_5_taches', 'le_cténolabre', 'le_denté_commun', 'le_dragonnet_lyre', 'le_flet', 'le_gobie_de_sarato', 'le_gobie_marbré', 'le_gobie_nageur', 'le_gobie_noir', 'le_gobie_paganel', 'le_gobie_svelte', 'le_gobie_tacheté', 'le_gobie_varié', 'le_gobie_à_bouche_rouge', 'le_gobie_à_grosse_tête', 'le_gobie_à_tête_jaune', 'le_grondin-perlon', 'le_grondin_camard', 'le_joel', 'le_marbré', 'le_merle', 'le_mordocet', 'le_mulet_labeon', 'le_mulet_lippu', 'le_mérou_brun', 'le_nérophis_lombricoïde', 'le_poisson-lézard_rayé', 'le_poisson-perroquet', 'le_poisson_lapin_à_queue_tronquée', 'le_poisson_lapin_à_ventre_strié', 'le_rason', 'le_rombou', 'le_rouget-barbet_de_roche', 'le_rouget-barbet_de_vase', 'le_sar_commun', 'le_sar_à_grosses_lèvres', 'le_sar_à_museau_pointu', 'le_sar_à_tête_noire', 'le_serran_-chevrette', 'le_serran_ecriture', 'le_siphonostome_atlantique', 'le_sparaillon', 'le_sublet', 'le_syngnathe_aiguille', 'le_tripterygion_rouge', 'le_triptérygion_jaune', 'le_tryptérigion_nain', 'l’anguille', 'l’apogon_commun', 'l’hippocampe_moucheté', 'l’hippocampe_à_nez_court', 'l’hippocampe_à_ramules', 'l’oblade', 'l’orphie', 'l’épinoche']

# Load the model
model_path = "model.pth"
model = timm.create_model('resnet34', pretrained=False, num_classes=101)
model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
model = model.to('cpu')
model.eval()


# Preprocessing function
def preprocess(image):
    _preprocess = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
    return _preprocess(image)


# Inference function
def fish_classifier(img_input):
    img = Image.fromarray(img_input.astype('uint8'), 'RGB')
    input_tensor = preprocess(img).unsqueeze(0) 

    with torch.no_grad():
        pred = model(input_tensor)[0]
        pred = torch.nn.functional.softmax(pred, dim=0)
        return {classes[i]: float(pred[i]) for i in range(len(classes))}


# Gradio interface
gr.Interface(
    fn=fish_classifier,
    inputs=gr.Image(shape=(224, 224)),
    outputs=gr.Label(num_top_classes=5),
    examples=["dorade.jpg"]
).launch()