Modify Python Class to filter files

New to python. How can I modify the class to filter files in the folder with a string. Right now it returns all files in folder_containing_the_content_folder which could be millions of items. I would like to isolate files that contain a specific string, for example, isolate all files that contain ‘v_1234_frame’:

# Image loader
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Lambda(lambda x: x.mul(255))
])
image_dataset = utils.ImageFolderWithPaths(folder_containing_the_content_folder, transform=transform)
image_loader = torch.utils.data.DataLoader(image_dataset, batch_size=batch_size)

The class that works requires a modification to filter file names that contain ‘v_1234_frame’:

class ImageFolderWithPaths(datasets.ImageFolder):
"""Custom dataset that includes image file paths.
Extends torchvision.datasets.ImageFolder()
Reference: https://discuss.pytorch.org/t/dataloader-filenames-in-each-batch/4212/2
"""

# override the __getitem__ method. this is the method dataloader calls
def __getitem__(self, index):
    # this is what ImageFolder normally returns
    original_tuple = super(ImageFolderWithPaths, self).__getitem__(index)

    # the image file path
    path = self.imgs[index][0]

    # make a new tuple that includes original and the path
    tuple_with_path = (*original_tuple, path)
    return tuple_with_path

I am learning python and just can’t seem to come up with the solution. Hope you can help/suggest a change to the class or calling method.

You could e.g. create a custom find_classes method and pass it to DatasetFolder as described here.

Very new to python and pytorch. In the FindClasses example how do you pass that result of that to my ImageFolderWithPaths class? Basically how can I pass anything to the ImageFolderWithPaths class in addition to the folder?

One approach would be to write a custom Dataset by deriving from DatasetFolder and override the find_classes method. You could also alternatively create an ImageFolder dataset and try to reassign the new find_classes method to the internal method (haven’t checked this approach).

Built my own data loader to isolate files via a wildcard pattern in glob and then loop through those to create a tensor for each image, passing that to my model which required it to be converted to a float. Extract the base name (image name) from the path (ex. img_frame1.jpg). Save result to my style folder. This method gives me total control over the files via the wildcard. I have included the other functions used in the solution. Note: I am not using a gpu to process these so I can run it on a standard python web server. Hopefully this helps someone in the future. Simple is sometimes better :slight_smile:

# Load image file
# def load_image(path):
#     # Images loaded as BGR
#     img = cv2.imread(path)
#     return img

# def itot(img, max_size=None):
#     # Rescale the image
#     if (max_size == None):
#         itot_t = transforms.Compose([
#             # transforms.ToPILImage(),
#             transforms.ToTensor(),
#             transforms.Lambda(lambda x: x.mul(255))
#         ])
#     else:
#         H, W, C = img.shape
#         image_size = tuple([int((float(max_size) / max([H, W])) * x) for x in [H, W]])
#         itot_t = transforms.Compose([
#             transforms.ToPILImage(),
#             transforms.Resize(image_size),
#             transforms.ToTensor(),
#             transforms.Lambda(lambda x: x.mul(255))
#         ])
# 
#     # Convert image to tensor
#     tensor = itot_t(img)
# 
#     # Add the batch_size dimension
#     tensor = tensor.unsqueeze(dim=0)
#     return tensor

folder_data = glob.glob(folder_containing_the_content_folder + "content_folder/" +  video_token + "_frame*.jpg")

# image_dataset = utils.ImageFolderWithPaths(folder_containing_the_content_folder, transform)
# image_loader = torch.utils.data.DataLoader(image_dataset, batch_size=batch_size)

# Load Transformer Network
net = transformer.TransformerNetwork()
net.load_state_dict(torch.load(style_path))
net = net.to(device)

torch.cuda.empty_cache()

with torch.no_grad():
    for image_name in folder_data:
        img = utils.load_image(image_name)
        img = img / 255.0  # standardize the data/transform
        img_tensor = utils.itot(img)

        # style image tensor
        generated_tensor = net(img_tensor.float())
        
        # convert image the model modified tensor back to an image
        generated_image = utils.ttoi(generated_tensor)
        image_name = os.path.basename(image_name)
        
        # save generated image to folder
        utils.saveimg(generated_image, save_folder + image_name)
1 Like