Spaces:
Runtime error
Runtime error
import cv2, torch | |
import gradio as gr | |
import numpy as np | |
from PIL import Image | |
import torch.nn as nn | |
import torchvision.models as models | |
from torchvision import transforms as T | |
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
new_model = models.resnet18() | |
num_ftrs = new_model.fc.in_features | |
new_model.fc = nn.Linear(num_ftrs, 75) | |
checkpoint = torch.load('model_best_checkpoint.pth.tar', map_location=torch.device('cpu')) | |
new_model.load_state_dict(checkpoint['model']) | |
new_model.to(device) | |
we_are = ['INDRA SWALLOW', | |
'MALACHITE', | |
'COMMON BANDED AWL', | |
'DANAID EGGFLY', | |
'EASTERN PINE ELFIN', | |
'YELLOW SWALLOW TAIL', | |
'WOOD SATYR', | |
'ULYSES', | |
'MESTRA', | |
'MANGROVE SKIPPER', | |
'BECKERS WHITE', | |
'CRECENT', | |
'RED SPOTTED PURPLE', | |
'SOOTYWING', | |
'BLACK HAIRSTREAK', | |
'STRAITED QUEEN', | |
'ELBOWED PIERROT', | |
'ORANGE OAKLEAF', | |
'CHESTNUT', | |
'POPINJAY', | |
'COMMON WOOD-NYMPH', | |
'BROWN SIPROETA', | |
'QUESTION MARK', | |
'ADONIS', | |
'CLOUDED SULPHUR', | |
'TWO BARRED FLASHER', | |
'GOLD BANDED', | |
'BANDED ORANGE HELICONIAN', | |
'PURPLISH COPPER', | |
'VICEROY', | |
'RED CRACKER', | |
'SILVER SPOT SKIPPER', | |
'ZEBRA LONG WING', | |
'ORCHARD SWALLOW', | |
'RED POSTMAN', | |
'SOUTHERN DOGFACE', | |
'SCARCE SWALLOW', | |
'EASTERN COMA', | |
'CAIRNS BIRDWING', | |
'GREEN CELLED CATTLEHEART', | |
'METALMARK', | |
'LARGE MARBLE', | |
'AMERICAN SNOOT', | |
'COPPER TAIL', | |
'AN 88', | |
'AFRICAN GIANT SWALLOWTAIL', | |
'PAPER KITE', | |
'EASTERN DAPPLE WHITE', | |
'PEACOCK', | |
'ATALA', | |
'JULIA', | |
'RED ADMIRAL', | |
'GREAT JAY', | |
'GREAT EGGFLY', | |
'GREY HAIRSTREAK', | |
'PIPEVINE SWALLOW', | |
'PURPLE HAIRSTREAK', | |
'ORANGE TIP', | |
'BLUE SPOTTED CROW', | |
'TROPICAL LEAFWING', | |
'CLEOPATRA', | |
'APPOLLO', | |
'IPHICLUS SISTER', | |
'CABBAGE WHITE', | |
'BANDED PEACOCK', | |
'MONARCH', | |
'CRIMSON PATCH', | |
'BLUE MORPHO', | |
'MOURNING CLOAK', | |
'SLEEPY ORANGE', | |
'CLODIUS PARNASSIAN', | |
'MILBERTS TORTOISESHELL', | |
'PINE WHITE', | |
'CHECQUERED SKIPPER', | |
'PAINTED LADY'] | |
def classify(image_): | |
model = new_model.eval() | |
image = Image.open(image_) | |
image = image_transforms(image).float().to(device) | |
image = image.unsqueeze(0) | |
output = model(image) | |
_, predicted = torch.max(output, 1) | |
return we_are[predicted] | |
label = gr.outputs.Label(num_top_classes=75) | |
gr.Interface(fn=classify, inputs='image', outputs=label,interpretation='default', title = 'Butterfly Classification detection ', description = 'It will classify 75 different species ').launch() | |