modelchanging

app.py

@@ -1,133 +1,80 @@
-# pylint: disable=import-error
 import gradio as gr
 import numpy as np
 import torch
-from torchvision import transforms
 from PIL import Image
 import matplotlib.pyplot as plt
-from transformers import ViTForImageClassification, ViTImageProcessor
+from transformers import AutoFeatureExtractor, AutoModelForImageClassification
 
-# Load pre-trained Vision Transformer model
-model_name = "google/vit-base-patch16-224"
-model = ViTForImageClassification.from_pretrained(model_name)
-processor = ViTImageProcessor.from_pretrained(model_name)
+# Use a smaller, more efficient model
+model_name = "microsoft/resnet-18"  # Smaller model that should work with Hugging Face constraints
 
-# Function to predict image class
+# Load model and feature extractor
+feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
+model = AutoModelForImageClassification.from_pretrained(model_name)
+
+# Function to classify image
 def classify_image(image):
     if image is None:
-        return None, None
-    
-    # Process image
-    inputs = processor(images=image, return_tensors="pt")
-    
-    # Make prediction
-    with torch.no_grad():
-        outputs = model(**inputs)
-        logits = outputs.logits
-    
-    # Get predicted class and probabilities
-    predicted_class_idx = logits.argmax(-1).item()
-    predicted_class = model.config.id2label[predicted_class_idx]
-    
-    # Get top 5 predictions
-    probs = torch.nn.functional.softmax(logits, dim=-1)[0]
-    top5_prob, top5_indices = torch.topk(probs, 5)
-    
-    # Create plot for visualization
-    fig, ax = plt.subplots(figsize=(10, 5))
-    
-    # Get class names and probabilities
-    classes = [model.config.id2label[idx.item()] for idx in top5_indices]
-    probabilities = [prob.item() * 100 for prob in top5_prob]
+        return "No image provided", None
     
-    # Create horizontal bar chart
-    bars = ax.barh(classes, probabilities, color='#4C72B0')
-    ax.set_xlabel('Probability (%)')
-    ax.set_title('Top 5 Predictions')
-    
-    # Add percentage labels
-    for i, bar in enumerate(bars):
-        width = bar.get_width()
-        ax.text(width + 1, bar.get_y() + bar.get_height()/2, 
-                f'{probabilities[i]:.1f}%', 
-                va='center', fontsize=10)
-    
-    # Improve layout
-    plt.tight_layout()
-    
-    return predicted_class, fig
-
-# Create Gradio interface
-with gr.Blocks(title="Image Classifier", theme=gr.themes.Soft()) as demo:
-    gr.Markdown(
-        """
-        # 🖼️ Image Classification Tool
+    try:
+        # Process image
+        inputs = feature_extractor(images=image, return_tensors="pt")
         
-        This application uses a Vision Transformer (ViT) model to classify images into 1,000 different categories.
+        # Make prediction
+        with torch.no_grad():
+            outputs = model(**inputs)
+            logits = outputs.logits
         
-        Upload an image or take a photo to see what the AI recognizes in it!
-        """
-    )
-    
-    with gr.Row():
-        with gr.Column():
-            image_input = gr.Image(
-                label="Upload or capture an image",
-                type="pil",
-                height=400
-            )
-            classify_btn = gr.Button("Classify Image", variant="primary")
+        # Get predicted class
+        predicted_class_idx = logits.argmax(-1).item()
+        predicted_class = model.config.id2label[predicted_class_idx]
         
-        with gr.Column():
-            prediction = gr.Textbox(label="Prediction")
-            confidence_plot = gr.Plot(label="Confidence Levels")
-    
-    # Add examples
-    example_images = [
-        "examples/dog.jpg",
-        "examples/cat.jpg",
-        "examples/coffee.jpg",
-        "examples/laptop.jpg",
-        "examples/beach.jpg"
-    ]
-    
-    gr.Examples(
-        examples=example_images,
-        inputs=image_input,
-        outputs=[prediction, confidence_plot],
-        fn=classify_image,
-        cache_examples=True
-    )
-    
-    # Set up the click event
-    classify_btn.click(
-        fn=classify_image,
-        inputs=image_input,
-        outputs=[prediction, confidence_plot]
-    )
-    
-    # Set up the input change event
-    image_input.change(
-        fn=classify_image,
-        inputs=image_input,
-        outputs=[prediction, confidence_plot]
-    )
-    
-    gr.Markdown("""
-    ### How it works
-    
-    This tool uses a Vision Transformer (ViT) model pre-trained on ImageNet, enabling it to recognize 1,000 
-    different object categories ranging from animals and plants to vehicles, household items, and more.
-    
-    ### Applications
-    
-    - **Content Categorization**: Automatically organize image libraries
-    - **Accessibility**: Help describe images for visually impaired users
-    - **Education**: Learn about objects in the world around you
-    - **Data Analysis**: Process and categorize large image datasets
+        # Get top 5 predictions
+        probs = torch.nn.functional.softmax(logits, dim=-1)[0]
+        top5_prob, top5_indices = torch.topk(probs, 5)
+        
+        # Create plot for visualization
+        fig, ax = plt.subplots(figsize=(10, 5))
+        
+        # Get class names and probabilities
+        classes = [model.config.id2label[idx.item()] for idx in top5_indices]
+        probabilities = [prob.item() * 100 for prob in top5_prob]
+        
+        # Create horizontal bar chart
+        bars = ax.barh(classes, probabilities, color='#4C72B0')
+        ax.set_xlabel('Probability (%)')
+        ax.set_title('Top 5 Predictions')
+        
+        # Add percentage labels
+        for i, bar in enumerate(bars):
+            width = bar.get_width()
+            ax.text(width + 1, bar.get_y() + bar.get_height()/2, 
+                    f'{probabilities[i]:.1f}%', 
+                    va='center', fontsize=10)
+        
+        # Improve layout
+        plt.tight_layout()
+        
+        return predicted_class, fig
     
-    Created by [Vinicius Guerra e Ribas](https://viniciusgribas.netlify.app/)
-    """)
+    except Exception as e:
+        return f"Error: {str(e)}", None
+
+# Create Gradio interface with simpler structure
+demo = gr.Interface(
+    fn=classify_image,
+    inputs=gr.Image(type="pil"),
+    outputs=[
+        gr.Textbox(label="Prediction"),
+        gr.Plot(label="Confidence Levels")
+    ],
+    title="🖼️ Image Classification Tool",
+    description="Upload an image to see what the AI recognizes in it!",
+    allow_flagging="never",
+    examples=[],  # No examples to avoid dependencies
+    theme=gr.themes.Soft()
+)
 
 # Launch the app
 if __name__ == "__main__":