yonadab commited on
Commit
bce21ce
·
verified ·
1 Parent(s): cef5090

nuevo update

Browse files
Files changed (1) hide show
  1. app.py +11 -7
app.py CHANGED
@@ -1,24 +1,28 @@
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
- print("Tipo de imagen recibido:", type(image))
 
 
 
11
  inputs = processor(images=image, return_tensors="pt")
12
  with torch.no_grad():
13
- embeddings = model(**inputs).last_hidden_state[:, 0] # CLS token
14
  return embeddings.squeeze().tolist()
15
 
16
- # Gradio UI solo para aceptar imágenes y devolver JSON
17
  iface = gr.Interface(
18
  fn=get_embedding,
19
- inputs=gr.Image(type="pil"),
20
  outputs="json",
21
- description="Microservicio para extraer embeddings de imágenes usando DINOv2."
22
  )
23
 
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
+ import base64
6
+ import io
7
 
 
8
  processor = AutoImageProcessor.from_pretrained("facebook/dinov2-base")
9
  model = AutoModel.from_pretrained("facebook/dinov2-base")
10
 
11
+ def get_embedding(base64_str):
12
+ header, encoded = base64_str.split(",", 1)
13
+ image_data = base64.b64decode(encoded)
14
+ image = Image.open(io.BytesIO(image_data)).convert("RGB")
15
+
16
  inputs = processor(images=image, return_tensors="pt")
17
  with torch.no_grad():
18
+ embeddings = model(**inputs).last_hidden_state[:, 0]
19
  return embeddings.squeeze().tolist()
20
 
 
21
  iface = gr.Interface(
22
  fn=get_embedding,
23
+ inputs="text", # ahora recibimos un string base64
24
  outputs="json",
25
+ description="Microservicio para extraer embeddings desde base64."
26
  )
27
 
28
  iface.queue() # 👈 Esta línea activa el sistema de event_id y polling