hw3-test / app.py
Avmromanov's picture
fix4
59c7d1b
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()