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

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +9 -11
app.py CHANGED
@@ -5,25 +5,23 @@ 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
29
- iface.launch()
 
5
  import base64
6
  import io
7
 
8
+ # Modelo
9
  processor = AutoImageProcessor.from_pretrained("facebook/dinov2-base")
10
  model = AutoModel.from_pretrained("facebook/dinov2-base")
11
 
12
+ # Recibe string base64
13
  def get_embedding(base64_str):
14
+ # Separar encabezado
15
  header, encoded = base64_str.split(",", 1)
16
+ # Decodificar
17
+ image_bytes = base64.b64decode(encoded)
18
+ image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
19
 
20
+ # Procesar
21
  inputs = processor(images=image, return_tensors="pt")
22
  with torch.no_grad():
23
  embeddings = model(**inputs).last_hidden_state[:, 0]
24
  return embeddings.squeeze().tolist()
25
 
26
+ # Gradio Interface
27
+ gr.Interface(fn=get_embedding, inputs="text", outputs="json").launch()