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

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +7 -13
app.py CHANGED
@@ -1,30 +1,24 @@
1
  import gradio as gr
2
  from transformers import AutoImageProcessor, AutoModel
3
  import torch
4
- from PIL import Image
5
- import base64
6
- import io
7
-
8
 
 
9
  processor = AutoImageProcessor.from_pretrained("facebook/dinov2-base")
10
  model = AutoModel.from_pretrained("facebook/dinov2-base")
11
 
12
- def get_embedding(base64_str):
13
- header, encoded = base64_str.split(",", 1)
14
- image_data = base64.b64decode(encoded)
15
- image = Image.open(io.BytesIO(image_data)).convert("RGB")
16
-
17
  inputs = processor(images=image, return_tensors="pt")
18
  with torch.no_grad():
19
- embeddings = model(**inputs).last_hidden_state[:, 0]
20
  return embeddings.squeeze().tolist()
21
 
22
-
23
  iface = gr.Interface(
24
  fn=get_embedding,
25
- inputs="text", # ahora recibimos un string base64
26
  outputs="json",
27
- description="Microservicio para extraer embeddings desde base64."
28
  )
29
 
 
30
  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
 
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