diff --git a/deepface/modules/preprocessing.py b/deepface/modules/preprocessing.py index 23edb5d..5a6ebf5 100644 --- a/deepface/modules/preprocessing.py +++ b/deepface/modules/preprocessing.py @@ -3,12 +3,13 @@ import os from typing import Union, Tuple import base64 from pathlib import Path -import imghdr +import io # 3rd party import numpy as np import cv2 import requests +from PIL import Image def load_image(img: Union[str, np.ndarray]) -> Tuple[np.ndarray, str]: @@ -86,7 +87,9 @@ def load_base64(uri: str) -> np.ndarray: encoded_data = encoded_data_parts[1] decoded_bytes = base64.b64decode(encoded_data) - file_type = imghdr.what(None, h=decoded_bytes) + + img = Image.open(io.BytesIO(decoded_bytes)) + file_type = img.format.lower() # similar to find functionality, we are just considering these extensions # content type is safer option than file extension diff --git a/deepface/modules/recognition.py b/deepface/modules/recognition.py index 18e1a93..7d11abf 100644 --- a/deepface/modules/recognition.py +++ b/deepface/modules/recognition.py @@ -3,12 +3,12 @@ import os import pickle from typing import List, Union, Optional, Dict, Any import time -import imghdr # 3rd party dependencies import numpy as np import pandas as pd from tqdm import tqdm +from PIL import Image # project dependencies from deepface.commons.logger import Logger @@ -298,7 +298,16 @@ def __list_images(path: str) -> List[str]: for r, _, f in os.walk(path): for file in f: exact_path = os.path.join(r, file) - file_type = imghdr.what(exact_path) + + _, ext = os.path.splitext(exact_path) + ext_lower = ext.lower() + + if ext_lower not in {".jpg", ".jpeg", ".png"}: + continue + + img = Image.open(exact_path) # lazy + + file_type = img.format.lower() if file_type in ["jpeg", "png"]: images.append(exact_path) return images