How to sample images belonging to particular classes

vision
9

Given a list classes of classes, I did the following:

def checkfun(args):
  return args.split("/")[-2] in classes and args.endswith(".jpg")


def ___find_classes(self, dir):
  return classes, {c: i for i, c in enumerate(classes)}

torchvision.datasets.ImageFolder._find_classes = ___find_classes
dataset = torchvision.datasets.ImageFolder(root="root_training_dir", is_valid_file=checkfun, )

Basically,

  • Iv’e created filter function to filter out images with the wrong class.
  • Torch _find_classes function just list directory, before the filtering, so Iv’e just replaced it with a new function. Probably the better thing to do here is to build some context manager to switch back to original function after I done.