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()