load_image now accepts file objects that support being read

This commit is contained in:
Samuel J. Woodward 2025-01-09 10:35:47 -05:00
parent 5bab888411
commit 4f0fa6ee22
2 changed files with 52 additions and 4 deletions

View File

@ -1,7 +1,7 @@
# built-in dependencies # built-in dependencies
import os import os
import io import io
from typing import Generator, List, Union, Tuple from typing import IO, Generator, List, Union, Tuple
import hashlib import hashlib
import base64 import base64
from pathlib import Path from pathlib import Path
@ -77,11 +77,11 @@ def find_image_hash(file_path: str) -> str:
return hasher.hexdigest() 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: Args:
img: a path, url, base64 or numpy array. img: a path, url, file object, base64 or numpy array.
Returns: Returns:
image (numpy array): the loaded image in BGR format image (numpy array): the loaded image in BGR format
image name (str): image name itself 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): if isinstance(img, np.ndarray):
return img, "numpy array" 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): if isinstance(img, Path):
img = str(img) img = str(img)
@ -120,6 +128,30 @@ def load_image(img: Union[str, np.ndarray]) -> Tuple[np.ndarray, str]:
return img_obj_bgr, img 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: def load_image_from_base64(uri: str) -> np.ndarray:
""" """
Load image from base64 string. Load image from base64 string.

View File

@ -1,4 +1,5 @@
# built-in dependencies # built-in dependencies
import io
import cv2 import cv2
# project dependencies # project dependencies
@ -18,6 +19,21 @@ def test_standard_represent():
logger.info("✅ test standard represent function done") 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(): def test_represent_for_skipped_detector_backend_with_image_path():
face_img = "dataset/img5.jpg" face_img = "dataset/img5.jpg"
img_objs = DeepFace.represent(img_path=face_img, detector_backend="skip") img_objs = DeepFace.represent(img_path=face_img, detector_backend="skip")