How to improve custom Dataset class for reading DICOM images?

username: sguerin
cooked: <p>Hi everyone!<br>
I’m very new to PyTorch or python although I know basics of programming. I’m trying to process some MR images in DICOM format to classify them into two classes. I’ve created a custom dataset class (code bellow) and I would like to know if I’m thinking it right.</p>
<p>So, my questions are:</p>
<p>How can I improve my code? Is it a good idea to create a list (image pixel_data, label) or should I use another approach? In __getitem__ is it a good practice to use a dictionary?</p>
<pre><code class="lang-python">import os
import torch
import pydicom
import numpy as np
import torchvision.transforms as transforms
from torch.utils.data import Dataset
"""
# Train images: ImageData/train/0; ImageData/train/1
# Test images: ImageData/test/0; ImageData/test/1
# Classes: "not_operable": 0; "operable":1
#
# set_file_matrix():
# Count total items in sub-folders of root/image_dir:
# Create a list with all items from root/image_dir "(pixel_array, label)"
# Takes label from sub-folder name "0" or "1"
"""
class DicomDataset(Dataset):
    def __init__(self, root, image_dir):
        self.image_dir = os.path.join(root, image_dir)  # ImageData/train or ImageData/test
        self.data = self.set_file_matrix()
        self.transform = transforms.Compose([
            transforms.Resize(128),
            transforms.CenterCrop(128)
        ])

    def __len__(self):
        return len(self.data)

    def set_file_matrix(self):
        # count elements
        total = 0
        root = self.image_dir
        folders = ([name for name in os.listdir(root)
                    if os.path.isdir(os.path.join(root, name))])
        for folder in folders:
            new_path = os.path.join(root, folder)
            contents = len([name for name in os.listdir(new_path) if os.path.isfile(os.path.join(new_path, name))])
            total += contents

        # create list(img_name, label)
        files = []
        labels = ([name for name in os.listdir(root)
                   if os.path.isdir(os.path.join(root, name))])
        for label in labels:
            new_path = os.path.join(root, label)
            file_list = os.listdir(new_path)
            for file in file_list:
                files.append([file, label])

        return files

    def __getitem__(self, index):
        image_file = pydicom.dcmread(os.path.join(self.image_dir, self.data[index][1], self.data[index][0]))
        image = np.array(image_file.pixel_array, dtype=np.float32)[np.newaxis]  # Add channel dimension
        image = torch.from_numpy(image)
        label = float(self.data[index][1])
        if self.transform:
            image = self.transform(image)

        return {'image': image, 'label': label}
</code></pre>
<p><em><strong>Thanks in advance!</strong></em></p>

username: ptrblck
cooked: <p>Your code looks generally fine.<br>
This line of code looks a bit strange:</p>
<pre><code class="lang-python">label = float(self.data[index][1])
</code></pre>
<p>as I would assume <code>self.data</code> contains (part of) a path.</p>

username: sguerin
cooked: <p>Thanks <a class="mention" href="/u/ptrblck">@ptrblck</a> !</p>
<p>About that strange line maybe there is something that I’m not doing right.</p>
<p>This is the way I load the dataset:</p>
<pre><code class="lang-python">train_set = DicomDataset(ROOT_PATH, 'train')
test_set = DicomDataset(ROOT_PATH, 'test')
train_set_loader = torch.utils.data.DataLoader(train_set, batch_size=5, shuffle=True)
test_set_loader = torch.utils.data.DataLoader(test_set, batch_size=5, shuffle=True)
</code></pre>
<p>And this is the way I iterate over it in my model:</p>
<pre><code class="lang-python"># more code before
 n_total_steps = len(self.train_set_loader)
        for epoch in range(self.num_epochs):
            i = 0
            for item in enumerate(self.train_set_loader):
                images = item[1]["image"].to(device)
                labels = item[1]["label"].to(device)

# some more code after
</code></pre>
<p>If I don’t use <code>float()</code> I get this error : <code>AttributeError: 'list' object has no attribute 'to'</code> or <code>AttributeError: 'str' object has no attribute 'to'</code></p>
<p>Is it a good practice use <code>float()</code> or should I use another approach to read/save/load/iterate “labels” ?</p>

username: ptrblck
cooked: <p>You shouldn’t be able to directly convert a string to a floating point number:</p>
<pre><code class="lang-python">float('a')
&gt; ValueError: could not convert string to float: 'a'
</code></pre>
<p>so I’m not sure what exactly is stored in <code>self.data[index][1]</code>.<br>
Based on the error messages it seems it might be a <code>list</code> or a <code>str</code>.</p>
<p>Generally, you would have to make sure the targets are created as valid tensors (e.g. as a class index).<br>
Depending how you are loading the dataset, you might need to create the class indices from the number of folders or any other mapping.</p>

username: JoeS
cooked: <p>I haven’t done much with DICOMs but <a href="https://github.com/JDSobek/PublicMLCode/blob/main/3dResnetUnet/NiftiCTDataset.py#L73" rel="noopener nofollow ugc">here</a> is a Dataset I made to read NIfTI files, although it was intended for segmentation not classification.  I haven’t given it a proper test as when I tried it I had a mistake with my segmentation masks, but I’m reasonably sure that was why my results were bad.</p>
<p>Depending on what I was trying to do I’ve used lists and dictionaries when assembling my dataset, depending on how my data was saved out.  Neither method has given me issues, but generally the training process works much, much faster and the datasets are a bit simpler if you do most of the preprocessing for your data and save it out as torch tensors beforehand <a href="https://github.com/JDSobek/PublicMLCode/blob/main/3dResnetUnet/dataset.py#L71" rel="noopener nofollow ugc">(contrast above to this dataset)</a>.  Doubly so if your files need normalization applied.</p>

username: sguerin
cooked: <p><code>self.data[index][1]</code> stores labels as <code>'0'</code>or <code>'1'</code>, I think that’s the reason why <code>float( )</code> gives no error.  About the recommendation, I think I’m not understanding what “class index” is. Is it a specific class like <code>class class_name() ... </code> or something else?</p>

username: sguerin
cooked: <p>Thanks! I’ll take a look to your code.</p>

username: ptrblck
cooked: <p>Ah, OK, that makes sense then and should be the right approach.</p>