import gradio as gr
from transformers import pipeline
from datasets import load_dataset


def get_dataset_examples():
    dataset = load_dataset("Avmromanov/tripoexamples")
    train_data = dataset['train']
    example_ids = [0, 3, 6]

    examples = []
    for i in example_ids:
        example = train_data[i]
        examples.append(example['image'])
    
    return examples

def identify_car(image):
    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"\nMost likely: **{top_5[0]['label'].replace('_', ' ').title()}** " \
                    f"(confidence: {top_5[0]['score']:.2%})"
    
    return result_text


car_classifier = pipeline("image-classification", model="dima806/car_models_image_detection")
dataset_examples = get_dataset_examples()

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.Examples(
        examples=dataset_examples,
        inputs=image_input,
        outputs=output_text,
        fn=identify_car,
        cache_examples=True
    )

demo.launch()