### 1. Imports and class names setup ###
import gradio as gr
import os 
import torch

from model import create_transformer_model
from timeit import default_timer as timer
from typing import Tuple, Dict

# Setup class names
class_names = ['AFRICAN CROWNED CRANE',
 'AFRICAN FIREFINCH',
 'ALBATROSS',
 'ALEXANDRINE PARAKEET',
 'AMERICAN AVOCET',
 'AMERICAN BITTERN',
 'AMERICAN COOT',
 'AMERICAN GOLDFINCH',
 'AMERICAN KESTREL',
 'AMERICAN PIPIT',
 'AMERICAN REDSTART',
 'ANHINGA',
 'ANNAS HUMMINGBIRD',
 'ANTBIRD',
 'ARARIPE MANAKIN',
 'ASIAN CRESTED IBIS',
 'BALD EAGLE',
 'BALD IBIS',
 'BALI STARLING',
 'BALTIMORE ORIOLE',
 'BANANAQUIT',
 'BANDED BROADBILL',
 'BANDED PITA',
 'BAR-TAILED GODWIT',
 'BARN OWL',
 'BARN SWALLOW',
 'BARRED PUFFBIRD',
 'BAY-BREASTED WARBLER',
 'BEARDED BARBET',
 'BEARDED BELLBIRD',
 'BEARDED REEDLING',
 'BELTED KINGFISHER',
 'BIRD OF PARADISE',
 'BLACK & YELLOW bROADBILL',
 'BLACK BAZA',
 'BLACK FRANCOLIN',
 'BLACK SKIMMER',
 'BLACK SWAN',
 'BLACK TAIL CRAKE',
 'BLACK THROATED BUSHTIT',
 'BLACK THROATED WARBLER',
 'BLACK VULTURE',
 'BLACK-CAPPED CHICKADEE',
 'BLACK-NECKED GREBE',
 'BLACK-THROATED SPARROW',
 'BLACKBURNIAM WARBLER',
 'BLONDE CRESTED WOODPECKER',
 'BLUE COAU',
 'BLUE GROUSE',
 'BLUE HERON',
 'BLUE THROATED TOUCANET',
 'BOBOLINK',
 'BORNEAN BRISTLEHEAD',
 'BORNEAN LEAFBIRD',
 'BORNEAN PHEASANT',
 'BRANDT CORMARANT',
 'BROWN CREPPER',
 'BROWN NOODY',
 'BROWN THRASHER',
 'BULWERS PHEASANT',
 'CACTUS WREN',
 'CALIFORNIA CONDOR',
 'CALIFORNIA GULL',
 'CALIFORNIA QUAIL',
 'CANARY',
 'CAPE GLOSSY STARLING',
 'CAPE MAY WARBLER',
 'CAPPED HERON',
 'CAPUCHINBIRD',
 'CARMINE BEE-EATER',
 'CASPIAN TERN',
 'CASSOWARY',
 'CEDAR WAXWING',
 'CERULEAN WARBLER',
 'CHARA DE COLLAR',
 'CHESTNET BELLIED EUPHONIA',
 'CHIPPING SPARROW',
 'CHUKAR PARTRIDGE',
 'CINNAMON TEAL',
 'CLARKS NUTCRACKER',
 'COCK OF THE  ROCK',
 'COCKATOO',
 'COLLARED ARACARI',
 'COMMON FIRECREST',
 'COMMON GRACKLE',
 'COMMON HOUSE MARTIN',
 'COMMON LOON',
 'COMMON POORWILL',
 'COMMON STARLING',
 'COUCHS KINGBIRD',
 'CRESTED AUKLET',
 'CRESTED CARACARA',
 'CRESTED NUTHATCH',
 'CRIMSON SUNBIRD',
 'CROW',
 'CROWNED PIGEON',
 'CUBAN TODY',
 'CUBAN TROGON',
 'CURL CRESTED ARACURI',
 'D-ARNAUDS BARBET',
 'DARK EYED JUNCO',
 'DOUBLE BARRED FINCH',
 'DOUBLE BRESTED CORMARANT',
 'DOWNY WOODPECKER',
 'EASTERN BLUEBIRD',
 'EASTERN MEADOWLARK',
 'EASTERN ROSELLA',
 'EASTERN TOWEE',
 'ELEGANT TROGON',
 'ELLIOTS  PHEASANT',
 'EMPEROR PENGUIN',
 'EMU',
 'ENGGANO MYNA',
 'EURASIAN GOLDEN ORIOLE',
 'EURASIAN MAGPIE',
 'EVENING GROSBEAK',
 'FAIRY BLUEBIRD',
 'FIRE TAILLED MYZORNIS',
 'FLAME TANAGER',
 'FLAMINGO',
 'FRIGATE',
 'GAMBELS QUAIL',
 'GANG GANG COCKATOO',
 'GILA WOODPECKER',
 'GILDED FLICKER',
 'GLOSSY IBIS',
 'GO AWAY BIRD',
 'GOLD WING WARBLER',
 'GOLDEN CHEEKED WARBLER',
 'GOLDEN CHLOROPHONIA',
 'GOLDEN EAGLE',
 'GOLDEN PHEASANT',
 'GOLDEN PIPIT',
 'GOULDIAN FINCH',
 'GRAY CATBIRD',
 'GRAY KINGBIRD',
 'GRAY PARTRIDGE',
 'GREAT GRAY OWL',
 'GREAT KISKADEE',
 'GREAT POTOO',
 'GREATOR SAGE GROUSE',
 'GREEN BROADBILL',
 'GREEN JAY',
 'GREEN MAGPIE',
 'GREY PLOVER',
 'GROVED BILLED ANI',
 'GUINEA TURACO',
 'GUINEAFOWL',
 'GYRFALCON',
 'HARLEQUIN DUCK',
 'HARPY EAGLE',
 'HAWAIIAN GOOSE',
 'HELMET VANGA',
 'HIMALAYAN MONAL',
 'HOATZIN',
 'HOODED MERGANSER',
 'HOOPOES',
 'HORNBILL',
 'HORNED GUAN',
 'HORNED LARK',
 'HORNED SUNGEM',
 'HOUSE FINCH',
 'HOUSE SPARROW',
 'HYACINTH MACAW',
 'IMPERIAL SHAQ',
 'INCA TERN',
 'INDIAN BUSTARD',
 'INDIAN PITTA',
 'INDIAN ROLLER',
 'INDIGO BUNTING',
 'IWI',
 'JABIRU',
 'JAVA SPARROW',
 'KAGU',
 'KAKAPO',
 'KILLDEAR',
 'KING VULTURE',
 'KIWI',
 'KOOKABURRA',
 'LARK BUNTING',
 'LAZULI BUNTING',
 'LILAC ROLLER',
 'LONG-EARED OWL',
 'MAGPIE GOOSE',
 'MALABAR HORNBILL',
 'MALACHITE KINGFISHER',
 'MALAGASY WHITE EYE',
 'MALEO',
 'MALLARD DUCK',
 'MANDRIN DUCK',
 'MANGROVE CUCKOO',
 'MARABOU STORK',
 'MASKED BOOBY',
 'MASKED LAPWING',
 'MIKADO  PHEASANT',
 'MOURNING DOVE',
 'MYNA',
 'NICOBAR PIGEON',
 'NOISY FRIARBIRD',
 'NORTHERN CARDINAL',
 'NORTHERN FLICKER',
 'NORTHERN FULMAR',
 'NORTHERN GANNET',
 'NORTHERN GOSHAWK',
 'NORTHERN JACANA',
 'NORTHERN MOCKINGBIRD',
 'NORTHERN PARULA',
 'NORTHERN RED BISHOP',
 'NORTHERN SHOVELER',
 'OCELLATED TURKEY',
 'OKINAWA RAIL',
 'ORANGE BRESTED BUNTING',
 'ORIENTAL BAY OWL',
 'OSPREY',
 'OSTRICH',
 'OVENBIRD',
 'OYSTER CATCHER',
 'PAINTED BUNTIG',
 'PALILA',
 'PARADISE TANAGER',
 'PARAKETT  AKULET',
 'PARUS MAJOR',
 'PATAGONIAN SIERRA FINCH',
 'PEACOCK',
 'PELICAN',
 'PEREGRINE FALCON',
 'PHILIPPINE EAGLE',
 'PINK ROBIN',
 'POMARINE JAEGER',
 'PUFFIN',
 'PURPLE FINCH',
 'PURPLE GALLINULE',
 'PURPLE MARTIN',
 'PURPLE SWAMPHEN',
 'PYGMY KINGFISHER',
 'QUETZAL',
 'RAINBOW LORIKEET',
 'RAZORBILL',
 'RED BEARDED BEE EATER',
 'RED BELLIED PITTA',
 'RED BROWED FINCH',
 'RED FACED CORMORANT',
 'RED FACED WARBLER',
 'RED FODY',
 'RED HEADED DUCK',
 'RED HEADED WOODPECKER',
 'RED HONEY CREEPER',
 'RED NAPED TROGON',
 'RED TAILED HAWK',
 'RED TAILED THRUSH',
 'RED WINGED BLACKBIRD',
 'RED WISKERED BULBUL',
 'REGENT BOWERBIRD',
 'RING-NECKED PHEASANT',
 'ROADRUNNER',
 'ROBIN',
 'ROCK DOVE',
 'ROSY FACED LOVEBIRD',
 'ROUGH LEG BUZZARD',
 'ROYAL FLYCATCHER',
 'RUBY THROATED HUMMINGBIRD',
 'RUDY KINGFISHER',
 'RUFOUS KINGFISHER',
 'RUFUOS MOTMOT',
 'SAMATRAN THRUSH',
 'SAND MARTIN',
 'SANDHILL CRANE',
 'SATYR TRAGOPAN',
 'SCARLET CROWNED FRUIT DOVE',
 'SCARLET IBIS',
 'SCARLET MACAW',
 'SCARLET TANAGER',
 'SHOEBILL',
 'SHORT BILLED DOWITCHER',
 'SMITHS LONGSPUR',
 'SNOWY EGRET',
 'SNOWY OWL',
 'SORA',
 'SPANGLED COTINGA',
 'SPLENDID WREN',
 'SPOON BILED SANDPIPER',
 'SPOONBILL',
 'SPOTTED CATBIRD',
 'SRI LANKA BLUE MAGPIE',
 'STEAMER DUCK',
 'STORK BILLED KINGFISHER',
 'STRAWBERRY FINCH',
 'STRIPED OWL',
 'STRIPPED MANAKIN',
 'STRIPPED SWALLOW',
 'SUPERB STARLING',
 'SWINHOES PHEASANT',
 'TAIWAN MAGPIE',
 'TAKAHE',
 'TASMANIAN HEN',
 'TEAL DUCK',
 'TIT MOUSE',
 'TOUCHAN',
 'TOWNSENDS WARBLER',
 'TREE SWALLOW',
 'TROPICAL KINGBIRD',
 'TRUMPTER SWAN',
 'TURKEY VULTURE',
 'TURQUOISE MOTMOT',
 'UMBRELLA BIRD',
 'VARIED THRUSH',
 'VENEZUELIAN TROUPIAL',
 'VERMILION FLYCATHER',
 'VICTORIA CROWNED PIGEON',
 'VIOLET GREEN SWALLOW',
 'VULTURINE GUINEAFOWL',
 'WALL CREAPER',
 'WATTLED CURASSOW',
 'WHIMBREL',
 'WHITE BROWED CRAKE',
 'WHITE CHEEKED TURACO',
 'WHITE NECKED RAVEN',
 'WHITE TAILED TROPIC',
 'WHITE THROATED BEE EATER',
 'WILD TURKEY',
 'WILSONS BIRD OF PARADISE',
 'WOOD DUCK',
 'YELLOW BELLIED FLOWERPECKER',
 'YELLOW CACIQUE',
 'YELLOW HEADED BLACKBIRD']

### 2. Model and transforms perparation ###
vit, vit_transforms = create_transformer_model(
    num_classes=325)

# Load save weights
vit.load_state_dict(
    torch.load(
        f="bird_classification_vit.pth",
        map_location=torch.device("cpu") # load the model to the CPU
    )
)

### 3. Predict function ### 

def predict(img) -> Tuple[Dict, float]:
  # Start a timer
  start_time = timer()

  # Transform the input image for use with EffNetB2
  img = vit_transforms(img).unsqueeze(0) # unsqueeze = add batch dimension on 0th index

  # Put model into eval mode, make prediction
  vit.eval()
  with torch.inference_mode():
    # Pass transformed image through the model and turn the prediction logits into probaiblities
    pred_probs = torch.softmax(effnetb2(img), dim=1)

  # Create a prediction label and prediction probability dictionary
  pred_labels_and_probs = {class_names[i]: float(pred_probs[0][i]) for i in range(len(class_names))}

  # Calculate pred time
  end_time = timer()
  pred_time = round(end_time - start_time, 4)

  # Return pred dict and pred time
  return pred_labels_and_probs, pred_time

### 4. Gradio app ### 

# Create title, description and article
title = "Bird Species Classification"
description = "A [Transformer] computer vision model to classify Bird species.."
article = "Created by Chidera Stanley [Transformer using PyTorch]"

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

# Create the Gradio demo
app = gr.Interface(fn=predict, # maps inputs to outputs
                    inputs=gr.Image(type="pil"),
                    outputs=[gr.Label(num_top_classes=3, label="Predictions"),
                             gr.Number(label="Prediction time (s)")],
                    examples=example_list,
                    title=title,
                    description=description,
                    article=article)

# Launch the demo!
app.launch()