luckyo87 commited on
Commit
f224c77
·
1 Parent(s): d5a5ca6

Add application file

Browse files
Files changed (1) hide show
  1. app.py +15 -10
app.py CHANGED
@@ -3,7 +3,7 @@ from transformers import ViTFeatureExtractor, ViTForImageClassification
3
  from PIL import Image
4
  import torch
5
 
6
- # Cargar el modelo y el extractor de características
7
  model = ViTForImageClassification.from_pretrained("akahana/vit-base-cats-vs-dogs")
8
  feature_extractor = ViTFeatureExtractor.from_pretrained("google/vit-base-patch16-224-in21k")
9
 
@@ -16,13 +16,18 @@ def classify_image(image):
16
  predicted_class = model.config.id2label[predicted_class_idx]
17
  return predicted_class
18
 
19
- # Crear la interfaz de Gradio con título y descripción
20
- interface = gr.Interface(
21
- fn=classify_image,
22
- inputs=gr.Image(type="pil"),
23
- outputs="text",
24
- title="Clasificador de Gatos vs Perros",
25
- description="Sube una imagen de un gato o un perro y este modelo basado en ViT te dirá qué animal aparece en la imagen."
26
- )
27
 
28
- interface.launch()
 
 
 
 
 
 
 
 
 
 
3
  from PIL import Image
4
  import torch
5
 
6
+ # Cargar modelo y extractor
7
  model = ViTForImageClassification.from_pretrained("akahana/vit-base-cats-vs-dogs")
8
  feature_extractor = ViTFeatureExtractor.from_pretrained("google/vit-base-patch16-224-in21k")
9
 
 
16
  predicted_class = model.config.id2label[predicted_class_idx]
17
  return predicted_class
18
 
19
+ # Interfaz mejorada
20
+ with gr.Blocks() as demo:
21
+ gr.Markdown("# 🐱🐶 Clasificador de Gatos vs Perros")
22
+ gr.Markdown("Sube una imagen de un gato o un perro. Este modelo basado en Vision Transformer (ViT) te dirá cuál es.")
 
 
 
 
23
 
24
+ with gr.Row():
25
+ with gr.Column():
26
+ image_input = gr.Image(label="📷 Sube tu imagen", type="pil")
27
+ submit_btn = gr.Button("🔍 Clasificar")
28
+ with gr.Column():
29
+ output_label = gr.Textbox(label="🔎 Resultado", interactive=False)
30
+
31
+ submit_btn.click(fn=classify_image, inputs=image_input, outputs=output_label)
32
+
33
+ demo.launch()