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
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.

View File

@ -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")