-
Notifications
You must be signed in to change notification settings - Fork 0
/
utils.py
56 lines (44 loc) · 1.91 KB
/
utils.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
# import the necessary packages
import os
import numpy as np
import cv2
import torch
def preprocess_image(image):
# swap the color channels from BGR to RGB, resize it, and scale
# the pixel values to [0, 1] range
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image = cv2.resize(image, (224, 224))
image = image.astype("float32") / 255.0
# subtract ImageNet mean, divide by ImageNet standard deviation,
# set "channels first" ordering, and add a batch dimension
MEAN = [0.485, 0.456, 0.406]
STD = [0.229, 0.224, 0.225]
image -= MEAN
image /= STD
image = np.transpose(image, (2, 0, 1))
image = np.expand_dims(image, 0)
# return the preprocessed image
return image
"""
Reference: https://github.com/PyImageSearch/imutils/blob/master/imutils/paths.py
"""
image_types = (".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff")
def list_images(basePath, contains=None):
# return the set of files that are valid
return list_files(basePath, validExts=image_types, contains=contains)
def list_files(basePath, validExts=None, contains=None):
# loop over the directory structure
for (rootDir, dirNames, filenames) in os.walk(basePath):
# loop over the filenames in the current directory
for filename in filenames:
# if the contains string is not none and the filename does not contain
# the supplied string, then ignore the file
if contains is not None and filename.find(contains) == -1:
continue
# determine the file extension of the current file
ext = filename[filename.rfind("."):].lower()
# check to see if the file is an image and should be processed
if validExts is None or ext.endswith(validExts):
# construct the path to the image and yield it
imagePath = os.path.join(rootDir, filename)
yield imagePath