Spaces:
Sleeping
Sleeping
| import torch | |
| from PIL import Image | |
| from torchvision.transforms import v2 | |
| import torchvision.models as models | |
| import os | |
| import gradio as gr | |
| # Load the pre-trained model that was used in the training script. In this case it was ResNet18 model | |
| model = models.resnet18() | |
| # Modify the final fully connected layer to have 73 output classes, same as in the training script | |
| num_ftrs = model.fc.in_features | |
| model.fc = torch.nn.Linear(num_ftrs, 73) | |
| # Load the saved state dictionary (Device agnostic; inference for image classification is decent even on CPU) | |
| model.load_state_dict(torch.load('Dogrun2.pth', map_location=torch.device('cpu'))) | |
| # Set the model to evaluation mode | |
| model.eval() | |
| # Automatically detect the available device (CPU or GPU) | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| model.to(device) | |
| #Transforms the images to how they were tested for the model to read for inference. Keep Exactly the same as the transformation for the test and valid sets. No randomizing here! | |
| transforms_test = v2.Compose([ | |
| v2.Resize((224, 224), antialias=True), | |
| v2.CenterCrop((224, 224)), | |
| v2.ToImage(), | |
| v2.ToDtype(torch.float32, scale=True), | |
| v2.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) | |
| ]) | |
| labels = ['Afghan-Hound', 'Airedale-Terrier', 'Akita', 'Alaskan-Malamute', 'American-Foxhound', 'American-Hairless-Terrier', 'American-Water-Spaniel', 'Basenji', 'Basset-Hound', 'Beagle', 'Bearded-Collie', 'Belgian-Malinois', 'Belgian-Sheepdog', 'Bernese-Mountain-Dog', 'Bichon-Frise', 'Bloodhound', 'Bluetick-Coonhound', 'Border-Collie', 'Borzoi', 'Boston-Terrier', 'Boxer', 'Bull-Terrier', 'Bulldog', 'Bullmastiff', 'Cairn-Terrier', 'Cane-Corso', 'Cavalier-King-Charles-Spaniel', 'Chihuahua', 'Chinese-Crested', 'Chinese-Shar-Pei', 'Chow-Chow', 'Clumber-Spaniel', 'Cockapoo', 'Cocker-Spaniel', 'Collie', 'Dachshund', 'Dalmatian', 'Doberman-Pinscher', 'French-Bulldog', 'German-Shepherd', 'German-Shorthaired-Pointer', 'Golden-Retriever', 'Great-Dane', 'Great-Pyrenees', 'Greyhound', 'Irish-Water-Spaniel', 'Irish-Wolfhound', 'Japanese-Chin', 'Komondor', 'Labradoodle', 'Labrador-Retriever', 'Lhasa-Apso', 'Maltese', 'Miniature-Schnauzer', 'Newfoundland', 'Norwegian-Elkhound', 'Pekingese', 'Pembroke-Welsh-Corgi', 'Pomeranian', 'Poodle', 'Pug', 'Rhodesian-Ridgeback', 'Rottweiler', 'Saint-Bernard', 'Samoyed', 'Scottish-Terrier', 'Shiba-Inu', 'Shih-Tzu', 'Siberian-Husky', 'Staffordshire-Bull-Terrier', 'Vizsla', 'Xoloitzcuintli', 'Yorkshire-Terrier'] | |
| #I have the breed_nicknames dictionary set because some breeds arent recognized that much by the official breed, such as the ones below. | |
| breed_nicknames = { | |
| 'Xoloitzcuintli': ' (Mexican Hairless)', | |
| 'Staffordshire-Bull-Terrier': ' (Pitbull)', | |
| 'Pembroke-Welsh-Corgi': ' (Corgi)',} | |
| def predict(input_img): | |
| transformed_img = transforms_test(input_img) | |
| transformed_img = transformed_img.unsqueeze(0).to(device) | |
| logits = model(transformed_img) | |
| output_softmax = torch.nn.functional.softmax(logits, dim=1) | |
| topk_values, topk_indices = torch.topk(output_softmax, 3) | |
| topk_indices = topk_indices.tolist()[0] | |
| topk_labels = [labels[index] for index in topk_indices] | |
| return input_img, topk_labels | |
| gradio_app = gr.Interface( | |
| fn=predict, | |
| inputs=gr.Image(label="Select hot dog candidate", sources=['upload', 'webcam'], type="pil"), | |
| outputs=[gr.Image(label="Processed Image"), gr.Label(label="Result")], | |
| title="What is Your Dog Breed?", | |
| ) | |
| if __name__ == "__main__": | |
| gradio_app.launch() | |