DatasetGenerator.py

import os
import numpy as np
from PIL import Image
import pydicom

import torch
from torch.utils.data import Dataset

#-------------------------------------------------------------------------------- 

# define a function to convert dicom images to png
def dcm_to_png(imgdcm):
    shape = imgdcm.pixel_array.shape
    # Convert to float to avoid overflow or underflow losses.
    image_2d = imgdcm.pixel_array.astype(float)
    # Rescaling grey scale between 0-255
    image_2d_scaled = (np.maximum(image_2d,0) / image_2d.max()) * 255.0
    # Convert to uint
    image_2d_scaled = np.uint8(image_2d_scaled)
    # load pil image from np array
    imgpng = Image.fromarray(image_2d_scaled)
    return imgpng

#-------------------------------------------------------------------------------- 

class DatasetGenerator (Dataset):
    
    #-------------------------------------------------------------------------------- 
    
    def __init__ (self, pathImageDirectory, pathDatasetFile, transform):
    
        self.listImagePaths = []
        self.listImageLabels = []
        self.transform = transform
    
        #---- Open file, get image paths and labels
    
        fileDescriptor = open(pathDatasetFile, "r")
        
        #---- get into the loop
        line = True
        
        while line:
                
            line = fileDescriptor.readline()
            
            #--- if not empty
            if line:
          
                lineItems = line.split()
                
                imagePath = os.path.join(pathImageDirectory, lineItems[0])
                imageLabel = lineItems[1:]
                imageLabel = [int(i) for i in imageLabel]
                
                self.listImagePaths.append(imagePath)
                self.listImageLabels.append(imageLabel)   
            
        fileDescriptor.close()
    
    #-------------------------------------------------------------------------------- 
    
    def __getitem__(self, index):
        
        imagePath = self.listImagePaths[index]
        if imagePath[-4:]=='.dcm':
            dcmimg = pydicom.dcmread(imagePath)
            imageData = dcm_to_png(dcmimg).convert('RGB')
        else:
            imageData = Image.open(imagePath).convert('RGB')

        imageLabel= torch.FloatTensor(self.listImageLabels[index])
        
        if self.transform != None: imageData = self.transform(imageData)
        
        return imageData, imageLabel
        
    #-------------------------------------------------------------------------------- 
    
    def __len__(self):
        
        return len(self.listImagePaths)
    
 #--------------------------------------------------------------------------------