Fixed the bug that it did not support capital image extensions like .JPG, .JPEG

This commit is contained in:
KVZN 2023-03-26 17:58:41 +08:00
parent 070c7eb728
commit 3f86e8f7df
4 changed files with 17 additions and 17 deletions

View File

@ -358,11 +358,10 @@ def train_model(
num_images = len(
[
f
for f in os.listdir(os.path.join(train_data_dir, folder))
if f.endswith('.jpg')
or f.endswith('.jpeg')
or f.endswith('.png')
or f.endswith('.webp')
for f, lower_f in (
(file, file.lower()) for file in os.listdir(os.path.join(train_data_dir, folder))
)
if lower_f.endswith(('.jpg', '.jpeg', '.png', '.webp'))
]
)

View File

@ -368,8 +368,10 @@ def train_model(
image_num = len(
[
f
for f in os.listdir(image_folder)
if f.endswith('.jpg') or f.endswith('.png') or f.endswith('.webp')
for f, lower_f in (
(file, file.lower()) for file in os.listdir(image_folder)
)
if lower_f.endswith(('.jpg', '.jpeg', '.png', '.webp'))
]
)
print(f'image_num = {image_num}')

View File

@ -419,14 +419,14 @@ def train_model(
num_images = len(
[
f
for f in os.listdir(os.path.join(train_data_dir, folder))
if f.endswith('.jpg')
or f.endswith('.jpeg')
or f.endswith('.png')
or f.endswith('.webp')
for f, lower_f in (
(file, file.lower()) for file in os.listdir(os.path.join(train_data_dir, folder))
)
if lower_f.endswith(('.jpg', '.jpeg', '.png', '.webp'))
]
)
print(f'Folder {folder}: {num_images} images found')
# Calculate the total number of steps for this folder

View File

@ -375,11 +375,10 @@ def train_model(
num_images = len(
[
f
for f in os.listdir(os.path.join(train_data_dir, folder))
if f.endswith('.jpg')
or f.endswith('.jpeg')
or f.endswith('.png')
or f.endswith('.webp')
for f, lower_f in (
(file, file.lower()) for file in os.listdir(os.path.join(train_data_dir, folder))
)
if lower_f.endswith(('.jpg', '.jpeg', '.png', '.webp'))
]
)