Commit 20913655 authored by Jon Crall's avatar Jon Crall Committed by Soumith Chintala
Browse files

Improved image extension robustness (#322)

is_image_file is now case insensitive
Added docstring to is_image_file
is_image_file now returns true for pgm files
parent 15f6a225
...@@ -4,14 +4,20 @@ from PIL import Image ...@@ -4,14 +4,20 @@ from PIL import Image
import os import os
import os.path import os.path
IMG_EXTENSIONS = [ IMG_EXTENSIONS = ['.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm']
'.jpg', '.JPG', '.jpeg', '.JPEG',
'.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP',
]
def is_image_file(filename): def is_image_file(filename):
return any(filename.endswith(extension) for extension in IMG_EXTENSIONS) """Checks if a file is an image.
Args:
filename (string): path to a file
Returns:
bool: True if the filename ends with a known image extension
"""
filename_lower = filename.lower()
return any(filename_lower.endswith(ext) for ext in IMG_EXTENSIONS)
def find_classes(dir): def find_classes(dir):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment