yonadab commited on
Commit
bdb83bf
·
verified ·
1 Parent(s): 339ccca

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +8 -5
app.py CHANGED
@@ -1,24 +1,27 @@
1
  import gradio as gr
2
  from transformers import AutoImageProcessor, AutoModel
3
  import torch
 
4
 
5
- # Cargar el modelo DINOv2 una sola vez
6
  processor = AutoImageProcessor.from_pretrained("facebook/dinov2-base")
7
  model = AutoModel.from_pretrained("facebook/dinov2-base")
 
8
 
9
- def get_embedding(image):
 
10
  inputs = processor(images=image, return_tensors="pt")
11
  with torch.no_grad():
12
  embeddings = model(**inputs).last_hidden_state[:, 0] # CLS token
13
  return embeddings.squeeze().tolist()
14
 
15
- # Gradio UI solo para aceptar imágenes y devolver JSON
16
  iface = gr.Interface(
17
  fn=get_embedding,
18
- inputs=gr.Image(type="pil"),
19
  outputs="json",
20
  description="Microservicio para extraer embeddings de imágenes usando DINOv2."
21
  )
22
 
23
  iface.launch()
24
- iface.queue() # 👈 Esta línea activa el sistema de event_id y polling
 
1
  import gradio as gr
2
  from transformers import AutoImageProcessor, AutoModel
3
  import torch
4
+ from PIL import Image
5
 
6
+ # Cargar el modelo solo una vez
7
  processor = AutoImageProcessor.from_pretrained("facebook/dinov2-base")
8
  model = AutoModel.from_pretrained("facebook/dinov2-base")
9
+ model.eval()
10
 
11
+ def get_embedding(image_file):
12
+ image = Image.open(image_file).convert("RGB")
13
  inputs = processor(images=image, return_tensors="pt")
14
  with torch.no_grad():
15
  embeddings = model(**inputs).last_hidden_state[:, 0] # CLS token
16
  return embeddings.squeeze().tolist()
17
 
18
+ # Interfaz Gradio para uso visual o programático (API)
19
  iface = gr.Interface(
20
  fn=get_embedding,
21
+ inputs=gr.Image(type="file"), # << cambia aquí
22
  outputs="json",
23
  description="Microservicio para extraer embeddings de imágenes usando DINOv2."
24
  )
25
 
26
  iface.launch()
27
+ iface.queue()