PYTHON file type dir detector

#detect all with the file extension of jpg and png(file type dir detector)
import os

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
image_dir = os.path.join(BASE_DIR, "images")

#D:\main folder\software\PROGRAMMING\Python\Python Sample program\Images

for root, dirs, files in os.walk(image_dir):
    for file in files:
        if file.endswith("jpg") or file.endswith("png"):
            path = os.path.join(root, file)
            label = os.path.basename(os.path.dirname(path)).replace(" ", "-").lower()

            print(label, path)