Skip to content

Commit ba4d193

Browse files
authored
Merge pull request #1173 from serengil/feat-task-0704-find-image-type-with-pil
checking file type with PIL
2 parents ec0b833 + 6c51489 commit ba4d193

File tree

2 files changed

+16
-4
lines changed

2 files changed

+16
-4
lines changed

deepface/modules/preprocessing.py

+5-2
Original file line numberDiff line numberDiff line change
@@ -3,12 +3,13 @@
33
from typing import Union, Tuple
44
import base64
55
from pathlib import Path
6-
import imghdr
6+
import io
77

88
# 3rd party
99
import numpy as np
1010
import cv2
1111
import requests
12+
from PIL import Image
1213

1314

1415
def load_image(img: Union[str, np.ndarray]) -> Tuple[np.ndarray, str]:
@@ -86,7 +87,9 @@ def load_base64(uri: str) -> np.ndarray:
8687

8788
encoded_data = encoded_data_parts[1]
8889
decoded_bytes = base64.b64decode(encoded_data)
89-
file_type = imghdr.what(None, h=decoded_bytes)
90+
91+
img = Image.open(io.BytesIO(decoded_bytes))
92+
file_type = img.format.lower()
9093

9194
# similar to find functionality, we are just considering these extensions
9295
# content type is safer option than file extension

deepface/modules/recognition.py

+11-2
Original file line numberDiff line numberDiff line change
@@ -3,12 +3,12 @@
33
import pickle
44
from typing import List, Union, Optional, Dict, Any
55
import time
6-
import imghdr
76

87
# 3rd party dependencies
98
import numpy as np
109
import pandas as pd
1110
from tqdm import tqdm
11+
from PIL import Image
1212

1313
# project dependencies
1414
from deepface.commons.logger import Logger
@@ -298,7 +298,16 @@ def __list_images(path: str) -> List[str]:
298298
for r, _, f in os.walk(path):
299299
for file in f:
300300
exact_path = os.path.join(r, file)
301-
file_type = imghdr.what(exact_path)
301+
302+
_, ext = os.path.splitext(exact_path)
303+
ext_lower = ext.lower()
304+
305+
if ext_lower not in {".jpg", ".jpeg", ".png"}:
306+
continue
307+
308+
img = Image.open(exact_path) # lazy
309+
310+
file_type = img.format.lower()
302311
if file_type in ["jpeg", "png"]:
303312
images.append(exact_path)
304313
return images

0 commit comments

Comments
 (0)