user-agent commited on
Commit
3fd704a
·
verified ·
1 Parent(s): 28b424d

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +14 -7
app.py CHANGED
@@ -6,12 +6,18 @@ import gradio as gr
6
  from io import BytesIO
7
 
8
  @spaces.GPU
9
- def get_face_embedding(base64_image):
10
- # Decode the base64 image
11
- img_data = base64.b64decode(base64_image)
12
- np_arr = np.frombuffer(img_data, np.uint8)
13
- image = face_recognition.load_image_file(BytesIO(img_data))
14
-
 
 
 
 
 
 
15
  # Get the face encodings for all faces in the image
16
  face_encodings = face_recognition.face_encodings(image)
17
 
@@ -22,13 +28,14 @@ def get_face_embedding(base64_image):
22
  # Return the first face encoding as a list
23
  return face_encodings[0].tolist()
24
 
 
25
  # Define the Gradio interface
26
  interface = gr.Interface(
27
  fn=get_face_embedding,
28
  inputs="text",
29
  outputs="json",
30
  title="Face Embedding Extractor",
31
- description="Input a base64 encoded image to get a 128-dimensional face embedding vector. If no face is detected, an empty list is returned."
32
  )
33
 
34
  if __name__ == "__main__":
 
6
  from io import BytesIO
7
 
8
  @spaces.GPU
9
+ def get_face_embedding(image_input):
10
+ # Check if the input is a URL
11
+ if isinstance(image_input, str) and (image_input.startswith("http://") or image_input.startswith("https://")):
12
+ response = requests.get(image_input)
13
+ image = face_recognition.load_image_file(BytesIO(response.content))
14
+ else:
15
+ # Assume input is a base64 encoded string
16
+ if ',' in image_input:
17
+ image_input = image_input.split(',')[1] # Remove the prefix
18
+ img_data = base64.b64decode(image_input)
19
+ image = face_recognition.load_image_file(BytesIO(img_data))
20
+
21
  # Get the face encodings for all faces in the image
22
  face_encodings = face_recognition.face_encodings(image)
23
 
 
28
  # Return the first face encoding as a list
29
  return face_encodings[0].tolist()
30
 
31
+
32
  # Define the Gradio interface
33
  interface = gr.Interface(
34
  fn=get_face_embedding,
35
  inputs="text",
36
  outputs="json",
37
  title="Face Embedding Extractor",
38
+ description="Input a base64 encoded image or an image link to get a 128-dimensional face embedding vector. If no face is detected, an empty list is returned."
39
  )
40
 
41
  if __name__ == "__main__":