How to sample images belonging to particular classes

shivangi: <p>I have an ImageFolder. I see that the  <code>__getitem__(index)</code>  method of ImageFolder can get both the tensor and the class of an index. But what if I want to use ImageFolder to sample the images from only a particular class of my choice.<br>
How would this be done?</p>

ptrblck: <p>If could provide a <code>WeightedRandomSampler</code> with all weights set to zero for elements of the unwanted classes.<br>
Alternatively, you could write your custom <code>Dataset</code> and just read all images from a particular folder, which will also result in a single class being sampled.</p>

shivangi: <aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="43776">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>WeightedRandomSampler</p>
</blockquote>
</aside>
<p>So currently I am using custom dataset and the class looks like this:</p>
<pre><code class="lang-auto">class CustomData(Dataset):
    """
    CustomData dataset
    """

    def __init__(self, name, dirpath, transform=None, should_invert=False):
        super(Dataset, self).__init__()
        self.dirpath = dirpath
        self.imageFolderDataset = dset.ImageFolder(root=self.dirpath)
        self.transform = transform
        self.should_invert = should_invert
        self.to_tensor = transforms.ToTensor()

    def __getitem__(self, index):
        # Training images
        images = self.imageFolderDataset.imgs
        img = cv2.imread(images[index][0])

        if self.should_invert:
            img = PIL.ImageOps.invert(img)

        if self.transform is not None:
            img = self.transform(img)

        img = np.array(img, dtype='uint8')
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img_as_tensor = self.to_tensor(img)
        return img_as_tensor, images[index][1]

    def __len__(self):
        return len(self.imageFolderDataset.imgs)
</code></pre>
<p>How can I incorporate reading from particular folder because ImageFolder gives error if there is no directory inside?</p>

ptrblck: <p>An easy solution would be to just copy one folder into the <code>root</code> directory, but that would probably include some data moving, which you might want to avoid.</p>
<p>I would suggest to pass the image paths of the single class you want and to just sample from them:</p>
<pre><code class="lang-python">def __init__(...):
    self.image_paths = image_paths  # Should contain a list of image paths of your desired class: e.g. ['./data/class0/img0.png', './data/class0/img1.png', ...]
    ...

def __getitem__(self, index):
    img = cv2.imread(self.image_paths[index])
    ...
    return img_as_tensor, torch.tensor([0])
</code></pre>

shivangi: <p>But would it not be possible without using passing image_paths. I was kind of avoiding this approach.</p>

ptrblck: <p>In that case, I would just use a <code>SubsetRandomSampler</code> based on the class indices.<br>
Here is a small example getting the class indices for class0 from an <code>ImageFolder</code> dataset and creating the <code>SubsetRandomSampler</code>:</p>
<pre><code class="lang-python">targets = torch.tensor(dataset.targets)
target_idx = (targets==0).nonzero()

sampler = torch.utils.data.sampler.SubsetRandomSampler(target_idx)

loader = DataLoader(
    dataset,
    sampler=sampler)

for data, target in loader:
    print(target)  # should only print zeros
</code></pre>
<p>You could create a member method inside <code>CustomData</code> to return the class indices and pass it to the sampler.</p>

shivangi: <p><a class="mention" href="/u/ptrblck">@ptrblck</a><br>
Again I have to specify <code>dataset.targets</code> somewhere. I just want to do something like passing the directory path and then loading images from there. Is it possible to do this without using a list if paths like <code>ImageFolder</code> does it.</p>

ptrblck: <p>Internally <code>ImageFolder</code> creates these paths, so the one approach would be to only have one subfolder inside <code>root</code> containing your desired class images.</p>
<p>Every other approach would make the usage of <code>ImageFolder</code> not really useful, as you would have to filter out the other unwanted classes.<br>
A kind of dirty hack would be to use <code>ImageFolder</code>, and increase the <code>index</code> inside <code>__getitem__</code> until you sample an image from the class you want.</p>

Ronen_Abravanel: <p>Given a list <code>classes</code> of classes, I did the following:</p>
<pre><code class="lang-auto">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, )
</code></pre>
<p>Basically,</p>
<ul>
<li>Iv’e created filter function to filter out images with the wrong class.</li>
<li>Torch <code>_find_classes</code> 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.</li>
</ul>

NicoHambauer: <p>I just posted on another thread, but i think this one answers this question as well and may answer it best:</p>
<p>I would also consider going one level above the ImageFolder Class which inherits from DatasetFolder.<br>
DatasetFolder uses a method to index the folder subdirectories for each class:</p>
<pre><code class="lang-auto">def find_classes(directory: str) -&gt; Tuple[List[str], Dict[str, int]]:
    """Finds the class folders in a dataset.

    See :class:`DatasetFolder` for details.
    """
    classes = sorted(entry.name for entry in os.scandir(directory) if entry.is_dir())
    if not classes:
        raise FileNotFoundError(f"Couldn't find any class folder in {directory}.")

    class_to_idx = {cls_name: i for i, cls_name in enumerate(classes)}
    return classes, class_to_idx
</code></pre>
<p>See here: <a href="https://pytorch.org/vision/stable/_modules/torchvision/datasets/folder.html#DatasetFolder" class="inline-onebox" rel="noopener nofollow ugc">torchvision.datasets.folder — Torchvision 0.10.0 documentation</a></p>
<p>You may simply create your own DatasetFolder (which inherits from VisionDataset, don’t forget to inherit from that) and then let your own Image Folder class inherit from DatasetFolder (if even needed).<br>
By creating your own DatasetFolder, create a new  find_classes method, which only scans for subdirectories in your dir, with your desired class name</p>
<pre><code class="lang-auto">def find_classes(directory: str, desired_class_names: List) -&gt; Tuple[List[str], Dict[str, int]]:
    """Finds the class folders in a dataset."""
    classes = sorted(entry.name for entry in os.scandir(directory) if entry.is_dir())
    classes = classes [desired_class_names] # TODO: do something like this line! Not tested it yet!
    if not classes:
        raise FileNotFoundError(f"Couldn't find any class folder in {directory}.")

    class_to_idx = {cls_name: i for i, cls_name in enumerate(classes)}
    return classes, class_to_idx
</code></pre>
<p>Hope that makes sense and helps! Also just came across this question and tomorrow I am going to solve it!</p>
<p>Maybe I should request a pull or contribute to the official pytorch repo, by adding a feature that filters for classes as a parameter in ImageFolder? Does this make any sense <a class="mention" href="/u/ptrblck">@ptrblck</a> ?</p>

ptrblck: <p>I think your approach is great and the right way to implement this filtering.<br>
You could surely create a feature request on GitHub and discuss the suggestion with the code owners.<br>
I don’t know how often this feature of skipping specific folders is used, but maybe having the ability to pass a custom <code>find_classes</code> or <code>make_dataset</code> function might make sense (same as passing a custom <code>loader</code>).</p>

NicoHambauer: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> thanks for your Evaluation, since I have less experience with submitting such suggestions!</p>

ptrblck: <p>Your suggestion is great and thanks for the feedback! Could you CC me on your feature request in GitHub, too?</p>

NicoHambauer: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> Sure I will do that! I am going to create a request soon and tag you on GitHub! Edit: Done! <a href="https://github.com/pytorch/vision/issues/4633" class="inline-onebox" rel="noopener nofollow ugc">Feature Request for torchvision ImageFolder using/inheriting DatasetFolder · Issue #4633 · pytorch/vision · GitHub</a></p>