From 4f0fa6ee22170ced65d78fd1e7ee35b756408694 Mon Sep 17 00:00:00 2001
From: "Samuel J. Woodward" <sam@woodward.fyi>
Date: Thu, 9 Jan 2025 10:35:47 -0500
Subject: [PATCH] load_image now accepts file objects that support being read

---
 deepface/commons/image_utils.py | 40 +++++++++++++++++++++++++++++----
 tests/test_represent.py         | 16 +++++++++++++
 2 files changed, 52 insertions(+), 4 deletions(-)

diff --git a/deepface/commons/image_utils.py b/deepface/commons/image_utils.py
index 868eaf25f..10d177f55 100644
--- a/deepface/commons/image_utils.py
+++ b/deepface/commons/image_utils.py
@@ -1,7 +1,7 @@
 # built-in dependencies
 import os
 import io
-from typing import Generator, List, Union, Tuple
+from typing import IO, Generator, List, Union, Tuple
 import hashlib
 import base64
 from pathlib import Path
@@ -77,11 +77,11 @@ def find_image_hash(file_path: str) -> str:
     return hasher.hexdigest()
 
 
-def load_image(img: Union[str, np.ndarray]) -> Tuple[np.ndarray, str]:
+def load_image(img: Union[str, np.ndarray, IO[bytes]]) -> Tuple[np.ndarray, str]:
     """
-    Load image from path, url, base64 or numpy array.
+    Load image from path, url, file object, base64 or numpy array.
     Args:
-        img: a path, url, base64 or numpy array.
+        img: a path, url, file object, base64 or numpy array.
     Returns:
         image (numpy array): the loaded image in BGR format
         image name (str): image name itself
@@ -91,6 +91,14 @@ def load_image(img: Union[str, np.ndarray]) -> Tuple[np.ndarray, str]:
     if isinstance(img, np.ndarray):
         return img, "numpy array"
 
+    # The image is an object that supports `.read`
+    if hasattr(img, 'read') and callable(img.read):
+        if isinstance(img, io.StringIO):
+            raise ValueError(
+                'img requires bytes and cannot be an io.StringIO object.'
+            )
+        return load_image_from_io_object(img), 'io object'
+
     if isinstance(img, Path):
         img = str(img)
 
@@ -120,6 +128,30 @@ def load_image(img: Union[str, np.ndarray]) -> Tuple[np.ndarray, str]:
     return img_obj_bgr, img
 
 
+def load_image_from_io_object(obj: IO[bytes]) -> np.ndarray:
+    """
+    Load image from an object that supports being read
+    Args:
+        obj: a file like object.
+    Returns:
+        img (np.ndarray): The decoded image as a numpy array (OpenCV format).
+    """
+    try:
+        _ = obj.seek(0)
+    except (AttributeError, TypeError, io.UnsupportedOperation):
+        seekable = False
+        obj = io.BytesIO(obj.read())
+    else:
+        seekable = True
+    try:
+        nparr = np.frombuffer(obj.read(), np.uint8)
+        img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+        return img
+    finally:
+        if not seekable:
+            obj.close()
+
+
 def load_image_from_base64(uri: str) -> np.ndarray:
     """
     Load image from base64 string.
diff --git a/tests/test_represent.py b/tests/test_represent.py
index 085dff213..91a24b68e 100644
--- a/tests/test_represent.py
+++ b/tests/test_represent.py
@@ -1,4 +1,5 @@
 # built-in dependencies
+import io
 import cv2
 
 # project dependencies
@@ -18,6 +19,21 @@ def test_standard_represent():
     logger.info("✅ test standard represent function done")
 
 
+def test_standard_represent_with_io_object():
+    img_path = "dataset/img1.jpg"
+    defualt_embedding_objs = DeepFace.represent(img_path)
+    io_embedding_objs = DeepFace.represent(open(img_path, 'rb'))
+    assert defualt_embedding_objs == io_embedding_objs
+
+    # Confirm non-seekable io objects are handled properly
+    io_obj = io.BytesIO(open(img_path, 'rb').read())
+    io_obj.seek = None
+    no_seek_io_embedding_objs = DeepFace.represent(io_obj)
+    assert defualt_embedding_objs == no_seek_io_embedding_objs
+
+    logger.info("✅ test standard represent with io object function done")
+
+
 def test_represent_for_skipped_detector_backend_with_image_path():
     face_img = "dataset/img5.jpg"
     img_objs = DeepFace.represent(img_path=face_img, detector_backend="skip")