Spaces:
Running on Zero
Running on Zero
| import gradio as gr | |
| import torch | |
| import torchvision.models as models | |
| import torchvision.transforms as transforms | |
| from PIL import Image | |
| import math | |
| import spaces | |
| CLASS_NAMES = [ | |
| "Amphibia", # 0 | |
| "Animalia", # 1 | |
| "Arachnida", # 2 | |
| "Aves", # 3 | |
| "Fungi", # 4 | |
| "Insecta", # 5 | |
| "Mammalia", # 6 | |
| "Mollusca", # 7 | |
| "Plantae", # 8 | |
| "Reptilia" # 9 | |
| ] | |
| # 2. Setup the exact inference transform from your notebook | |
| val_transform = transforms.Compose([ | |
| transforms.Resize([224], interpolation=transforms.InterpolationMode.BICUBIC), | |
| transforms.CenterCrop([224]), | |
| transforms.ToTensor(), | |
| transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) | |
| ]) | |
| # 3. Create a helper function to load a specific model | |
| def load_single_model(model_name): | |
| import torch.nn as nn | |
| # Initialize with 10 classes | |
| if model_name in ["ResNet18 V1", "ResNet18 V2"]: | |
| model = models.resnet18(weights=None, num_classes=10) | |
| else: | |
| model = models.resnet34(weights=None, num_classes=10) | |
| model.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False) | |
| model.maxpool = nn.Identity() | |
| # Load the specific weights based on the dropdown | |
| if model_name == "ResNet18 V1": | |
| model.load_state_dict(torch.load("best_resnet18.pth", map_location="cpu")) | |
| elif model_name == "ResNet18 V2": | |
| model.load_state_dict(torch.load("best_resnet18_v2_1_clean.pth", map_location="cpu")) | |
| else: | |
| model.load_state_dict(torch.load("best_resnet34_clean.pth", map_location="cpu")) | |
| model.eval() | |
| return model | |
| # 4. The Prediction Function (Now with Ensembling!) | |
| def predict(image, model_choice): | |
| # Transform image and add batch dimension | |
| img_tensor = val_transform(image).unsqueeze(0) | |
| # If the user chose the Ensemble, run all 3 models! | |
| if model_choice == "Ensemble (All 3 Models - Geometric Mean)": | |
| model_v1 = load_single_model("ResNet18 V1") | |
| model_v2 = load_single_model("ResNet18 V2") | |
| model_v3 = load_single_model("ResNet 34") | |
| with torch.no_grad(): | |
| prob_v1 = torch.nn.functional.softmax(model_v1(img_tensor)[0], dim=0) | |
| prob_v2 = torch.nn.functional.softmax(model_v2(img_tensor)[0], dim=0) | |
| prob_v3 = torch.nn.functional.softmax(model_v3(img_tensor)[0], dim=0) | |
| # Geometric Mean formula from your notebook (with eps to prevent log(0) errors) | |
| eps = 1e-8 | |
| final_probs = ((prob_v1 + eps) * (prob_v2 + eps) * (prob_v3 + eps)) ** (1.0 / 3.0) | |
| # Normalize so probabilities sum to 1.0 for Gradio | |
| probabilities = final_probs / final_probs.sum() | |
| # Otherwise, just run the single chosen model | |
| else: | |
| model = load_single_model(model_choice) | |
| with torch.no_grad(): | |
| outputs = model(img_tensor) | |
| probabilities = torch.nn.functional.softmax(outputs[0], dim=0) | |
| # Return a dictionary of class confidences for Gradio | |
| confidences = {CLASS_NAMES[i]: float(probabilities[i]) for i in range(10)} | |
| return confidences | |
| # 5. Build the Gradio App | |
| # 5. Build the Gradio App | |
| interface = gr.Interface( | |
| fn=predict, | |
| inputs=[ | |
| gr.Image(type="pil", label="Upload an Image"), | |
| gr.Dropdown( | |
| choices=["ResNet18 V1", "ResNet18 V2", "ResNet 34", "Ensemble (All 3 Models - Geometric Mean)"], | |
| value="Ensemble (All 3 Models - Geometric Mean)", | |
| label="Select Model Architecture" | |
| ) | |
| ], | |
| outputs=gr.Label(num_top_classes=3, label="Predictions"), | |
| title="Week 9 Biological Image Classifier", | |
| description="Upload an image (e.g., a plant, insect, or animal) and let the trained ResNet models classify its biological kingdom/class! Try the Geometric Mean Ensemble for maximum accuracy.", | |
| examples=[ | |
| # Amphibia | |
| ["Amphibia_image_0001.jpg", "Ensemble (All 3 Models - Geometric Mean)"], | |
| ["Amphibia_image_0010.jpg", "Ensemble (All 3 Models - Geometric Mean)"], | |
| ["Amphibia_image_0017.jpg", "Ensemble (All 3 Models - Geometric Mean)"], | |
| # Animalia | |
| ["Animalia_image_0013.jpg", "Ensemble (All 3 Models - Geometric Mean)"], | |
| ["Animalia_image_0014.jpg", "Ensemble (All 3 Models - Geometric Mean)"], | |
| ["Animalia_image_0016.jpg", "Ensemble (All 3 Models - Geometric Mean)"], | |
| # Arachnida | |
| ["Arachnida_image_0003.jpg", "Ensemble (All 3 Models - Geometric Mean)"], | |
| ["Arachnida_image_0005.jpg", "Ensemble (All 3 Models - Geometric Mean)"], | |
| ["Arachnida_image_0006.jpg", "Ensemble (All 3 Models - Geometric Mean)"], | |
| # Aves | |
| ["Aves_image_0001.jpg", "Ensemble (All 3 Models - Geometric Mean)"], | |
| ["Aves_image_0008.jpg", "Ensemble (All 3 Models - Geometric Mean)"], | |
| ["Aves_image_0010.jpg", "Ensemble (All 3 Models - Geometric Mean)"], | |
| # Fungi | |
| ["Fungi_image_0001.jpg", "Ensemble (All 3 Models - Geometric Mean)"], | |
| ["Fungi_image_0012.jpg", "Ensemble (All 3 Models - Geometric Mean)"], | |
| ["Fungi_image_0024.jpg", "Ensemble (All 3 Models - Geometric Mean)"], | |
| # Insecta | |
| ["Insecta_image_0001.jpg", "Ensemble (All 3 Models - Geometric Mean)"], | |
| ["Insecta_image_0004.jpg", "Ensemble (All 3 Models - Geometric Mean)"], | |
| ["Insecta_image_0006.jpg", "Ensemble (All 3 Models - Geometric Mean)"], | |
| # Mammalia | |
| ["Mammalia_image_0001.jpg", "Ensemble (All 3 Models - Geometric Mean)"], | |
| ["Mammalia_image_0009.jpg", "Ensemble (All 3 Models - Geometric Mean)"], | |
| ["Mammalia_image_0018.jpg", "Ensemble (All 3 Models - Geometric Mean)"], | |
| # Mollusca | |
| ["Mollusca_image_0001.jpg", "Ensemble (All 3 Models - Geometric Mean)"], | |
| ["Mollusca_image_0010.jpg", "Ensemble (All 3 Models - Geometric Mean)"], | |
| ["Mollusca_image_0025.jpg", "Ensemble (All 3 Models - Geometric Mean)"], | |
| # Plantae | |
| ["Plantae_image_0001.jpg", "Ensemble (All 3 Models - Geometric Mean)"], | |
| ["Plantae_image_0007.jpg", "Ensemble (All 3 Models - Geometric Mean)"], | |
| ["Plantae_image_0020.jpg", "Ensemble (All 3 Models - Geometric Mean)"], | |
| # Reptilia | |
| ["Reptilia_image_0001.jpg", "Ensemble (All 3 Models - Geometric Mean)"], | |
| ["Reptilia_image_0011.jpg", "Ensemble (All 3 Models - Geometric Mean)"], | |
| ["Reptilia_image_0013.jpg", "Ensemble (All 3 Models - Geometric Mean)"], | |
| ["Reptilia_image_0014.jpg", "Ensemble (All 3 Models - Geometric Mean)"] | |
| ] | |
| ) | |
| interface.launch() |