import gradio as gr import torch import torch.nn.functional as F import torchvision.transforms as T import numpy as np from PIL import Image import io import base64 from model import resnet18_cifar # CIFAR-100 class names CIFAR100_CLASSES = [ 'apple', 'aquarium_fish', 'baby', 'bear', 'beaver', 'bed', 'bee', 'beetle', 'bicycle', 'bottle', 'bowl', 'boy', 'bridge', 'bus', 'butterfly', 'camel', 'can', 'castle', 'caterpillar', 'cattle', 'chair', 'chimpanzee', 'clock', 'cloud', 'cockroach', 'couch', 'crab', 'crocodile', 'cup', 'dinosaur', 'dolphin', 'elephant', 'flatfish', 'forest', 'fox', 'girl', 'hamster', 'house', 'kangaroo', 'keyboard', 'lamp', 'lawn_mower', 'leopard', 'lion', 'lizard', 'lobster', 'man', 'maple_tree', 'motorcycle', 'mountain', 'mouse', 'mushroom', 'oak_tree', 'orange', 'orchid', 'otter', 'palm_tree', 'pear', 'pickup_truck', 'pine_tree', 'plain', 'plate', 'poppy', 'porcupine', 'possum', 'rabbit', 'raccoon', 'ray', 'road', 'rocket', 'rose', 'sea', 'seal', 'shark', 'shrew', 'skunk', 'skyscraper', 'snail', 'snake', 'spider', 'squirrel', 'streetcar', 'sunflower', 'sweet_pepper', 'table', 'tank', 'telephone', 'television', 'tiger', 'tractor', 'train', 'trout', 'tulip', 'turtle', 'wardrobe', 'whale', 'willow_tree', 'wolf', 'woman', 'worm' ] # CIFAR-100 normalization constants CIFAR100_MEAN = (0.5071, 0.4867, 0.4408) CIFAR100_STD = (0.2675, 0.2565, 0.2761) # Initialize model device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') def load_model(): """Load the model. Try to load from checkpoint if available, otherwise use initialized model.""" model = resnet18_cifar(num_classes=100, width=64) # Try to load a trained checkpoint checkpoint_paths = ['best.pth', 'checkpoints/demo_model.pth', 'last.pth'] for checkpoint_path in checkpoint_paths: try: if torch.cuda.is_available(): checkpoint = torch.load(checkpoint_path, map_location=device) else: checkpoint = torch.load(checkpoint_path, map_location='cpu') model.load_state_dict(checkpoint['model_state']) print(f"✅ Loaded model from {checkpoint_path}") break except FileNotFoundError: continue except Exception as e: print(f"⚠️ Error loading {checkpoint_path}: {e}") continue else: print("ℹ️ No trained checkpoint found. Using initialized model for demo purposes.") print("Note: For best results, train the model using train.py first.") model.eval() return model.to(device) # Load the model model = load_model() def preprocess_image(image): """Preprocess image for CIFAR-100 ResNet model.""" # Resize to 32x32 (CIFAR-100 size) if isinstance(image, str): # If it's a file path image = Image.open(image).convert('RGB') elif hasattr(image, 'convert'): # If it's already a PIL Image image = image.convert('RGB') # Resize to CIFAR-100 dimensions image = image.resize((32, 32), Image.Resampling.LANCZOS) # Apply transforms transform = T.Compose([ T.ToTensor(), T.Normalize(CIFAR100_MEAN, CIFAR100_STD) ]) return transform(image).unsqueeze(0) def predict(image): """Make prediction on uploaded image.""" try: # Preprocess the image input_tensor = preprocess_image(image).to(device) # Make prediction with torch.no_grad(): outputs = model(input_tensor) probabilities = F.softmax(outputs, dim=1) # Get top 5 predictions top5_prob, top5_idx = torch.topk(probabilities, 5) top5_prob = top5_prob.cpu().numpy()[0] top5_idx = top5_idx.cpu().numpy()[0] # Create results dictionary results = {} for i, (idx, prob) in enumerate(zip(top5_idx, top5_prob)): class_name = CIFAR100_CLASSES[idx] results[f"{class_name}"] = float(prob) return results except Exception as e: return {"Error": f"Prediction failed: {str(e)}"} def predict_and_explain(image): """Make prediction and provide explanation.""" prediction = predict(image) if "Error" in prediction: return prediction, "Error occurred during prediction." # Get the top prediction top_class = max(prediction.keys(), key=prediction.get) confidence = prediction[top_class] explanation = f""" **Model Architecture:** ResNet-18 adapted for CIFAR-100 - Input: 32×32 RGB images - Output: 100 classes (CIFAR-100 categories) - Architecture: Residual blocks with skip connections **Top Prediction:** {top_class} ({confidence:.2%} confidence) **About this model:** This ResNet-18 model is specifically designed for CIFAR-100 classification. The architecture uses: - 3×3 convolutions with stride 1 (no max pooling in the stem) - Residual blocks with skip connections for gradient flow - Batch normalization and ReLU activations - Adaptive average pooling before the final classifier **Note:** This is a demonstration model. For best results, the model should be trained on CIFAR-100 dataset using the provided training script. """ return prediction, explanation # Create Gradio interface def create_interface(): with gr.Blocks(title="CIFAR-100 ResNet Classifier", theme=gr.themes.Soft()) as demo: gr.Markdown(""" # 🖼️ CIFAR-100 ResNet-18 Image Classifier Upload an image to classify it into one of 100 CIFAR-100 categories using a ResNet-18 model. The model is optimized for small 32×32 images but can handle larger images (they will be resized). **Categories include:** animals, vehicles, household items, plants, and more! """) with gr.Row(): with gr.Column(): image_input = gr.Image( type="pil", label="Upload Image", height=300 ) predict_btn = gr.Button( "🔍 Classify Image", variant="primary", size="lg" ) gr.Markdown(""" ### 💡 Tips: - Images are resized to 32×32 pixels (CIFAR-100 format) - Works best with clear, centered objects - Try images of animals, vehicles, plants, or household items """) with gr.Column(): prediction_output = gr.Label( label="Top 5 Predictions", num_top_classes=5 ) explanation_output = gr.Markdown( label="Model Information", value="Upload an image to see predictions and model details." ) # Example images section gr.Markdown("### 📚 Try these example categories:") gr.Examples( examples=[ # We'll use placeholder text since we don't have actual example images ["Upload images of: animals (cats, dogs, bears)", ""], ["Vehicles (cars, bicycles, motorcycles)", ""], ["Plants (flowers, trees)", ""], ["Household items (chairs, tables, bottles)", ""], ], inputs=[gr.Textbox(visible=False), gr.Textbox(visible=False)], label="Common CIFAR-100 Categories" ) # Connect the interface predict_btn.click( fn=predict_and_explain, inputs=[image_input], outputs=[prediction_output, explanation_output] ) # Auto-predict on image upload image_input.change( fn=predict_and_explain, inputs=[image_input], outputs=[prediction_output, explanation_output] ) gr.Markdown(""" --- ### 🔧 Technical Details **Model:** ResNet-18 adapted for CIFAR-100 - **Parameters:** Configurable width (default: 64 channels) - **Training:** OneCycle learning rate scheduling with mixed precision - **Features:** Integrated Grad-CAM visualization support - **Optimization:** Label smoothing, gradient clipping, data augmentation **Architecture Features:** - Residual blocks with skip connections - Batch normalization for stable training - No max pooling in stem (optimized for 32×32 inputs) - Adaptive global average pooling **Dataset:** CIFAR-100 (100 classes, 32×32 color images) *Note: This demo uses an initialized model. For production use, train the model using the provided training script.* """) return demo # Create and launch the interface if __name__ == "__main__": demo = create_interface() demo.launch()