using with clause to open img with pil to guarantee it is closed

This commit is contained in:
Sefik Ilkin Serengil 2024-04-07 10:18:47 +01:00
parent ba4d193535
commit 83c381e395
2 changed files with 7 additions and 10 deletions

View File

@ -88,11 +88,10 @@ def load_base64(uri: str) -> np.ndarray:
encoded_data = encoded_data_parts[1]
decoded_bytes = base64.b64decode(encoded_data)
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
with Image.open(io.BytesIO(decoded_bytes)) as img:
file_type = img.format.lower()
if file_type not in ["jpeg", "png"]:
raise ValueError(f"input image can be jpg or png, but it is {file_type}")

View File

@ -305,10 +305,8 @@ def __list_images(path: str) -> List[str]:
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"]:
with Image.open(exact_path) as img: # lazy
if img.format.lower() in ["jpeg", "png"]:
images.append(exact_path)
return images