novos files

app.py

@@ -0,0 +1,134 @@
+# 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
+
+# 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)
+
+# Function to predict image class
+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]
+    
+    # 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
+        
+        This application uses a Vision Transformer (ViT) model to classify images into 1,000 different categories.
+        
+        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")
+        
+        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
+    
+    Created by [Vinicius Guerra e Ribas](https://viniciusgribas.netlify.app/)
+    """)
+
+# Launch the app
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1,7 @@
+gradio>=3.32.0
+torch>=1.7.0
+torchvision>=0.8.0
+transformers>=4.26.0
+matplotlib>=3.5.0
+numpy>=1.20.0
+Pillow>=8.0.0