app.py

import gradio as gr
from transformers import pipeline
from datasets import load_dataset
from PIL import Image
import requests
import io

# Load your custom dataset
def load_my_dataset():
    try:
        dataset = load_dataset("Avmromanov/tripoexamples")
        return dataset
    except Exception as e:
        print(f"Error loading dataset: {e}")
        return None

# Load car classification model
car_classifier = pipeline("image-classification", 
                         model="dima806/car_brand_classification")

def identify_car(image):
    if image is None:
        return "Please upload an image of a car"
    
    try:
        if image.mode != 'RGB':
            image = image.convert('RGB')
        
        predictions = car_classifier(image)
        
        result_text = "🚗 Car Identification Results:\n\n"
        top_5 = predictions[:5]
        
        for i, pred in enumerate(top_5, 1):
            label = pred['label'].replace('_', ' ').title()
            confidence = pred['score']
            result_text += f"{i}. {label}: {confidence:.2%}\n"
        
        result_text += f"\n🔍 Most likely: **{top_5[0]['label'].replace('_', ' ').title()}** " \
                      f"(confidence: {top_5[0]['score']:.2%})"
        
        return result_text
        
    except Exception as e:
        return f"Error processing image: {str(e)}"

def get_dataset_examples(dataset, num_examples=3):
    """Extract example images from the dataset"""
    examples = []
    
    if dataset is None:
        return examples
    
    try:
        # Adjust this based on your dataset structure
        train_data = dataset['train']
        
        for i in range(min(num_examples, len(train_data))):
            example = train_data[i]
            
            # The structure depends on your dataset - adjust accordingly
            if 'image' in example:
                # If images are stored in the dataset
                examples.append(example['image'])
            elif 'url' in example:
                # If URLs are provided
                examples.append(example['url'])
            elif 'path' in example:
                # If file paths are provided
                examples.append(example['path'])
                
    except Exception as e:
        print(f"Error extracting examples: {e}")
    
    return examples

# Load your dataset
my_dataset = load_my_dataset()
dataset_examples = get_dataset_examples(my_dataset, num_examples=4)

# Create the interface
with gr.Blocks() as demo:
    gr.Markdown("# 🚗 Car Identifier with My Dataset")
    gr.Markdown("Using examples from: **Avmromanov/tripoexamples**")
    
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(label="Upload Car Photo", type="pil")
            identify_btn = gr.Button("Identify Car", variant="primary")
        with gr.Column():
            output_text = gr.Textbox(label="Results", lines=10)
    

    gr.Markdown(f"### Dataset Examples (showing {len(dataset_examples)} samples)")
    gr.Examples(
        examples=dataset_examples,
        inputs=image_input,
        outputs=output_text,
        fn=identify_car,
        cache_examples=True
    )
    
    gr.Markdown("""
    **Dataset Information:**
    - Name: tripoexamples
    - Author: Avmromanov
    - Type: Car images for identification
    """)

demo.launch()