src/data_loader/mammo_data.py

import pandas as pd
import numpy as np
import torch
from torch.utils.data import Dataset
import cv2
import os
import pydicom
import einops

class Mammography(Dataset):

    def __init__(self, excel_file, data_folder, category, task, height, width, background_crop = True, use_clahe= False, transform=None):
        """
        Args:
            excel_file (string): Path to the excel file with annotations.
            category (string) : 'Classification' for Benign and Malignant. 'Subtypes' for Subtype Classification
            transform (callable, optional): Optional transform to be applied
        """
        self.mammography = pd.read_csv(excel_file, dtype = str)
        self.data_folder = data_folder
        self.category = category
        self.task = task
        self.height = height
        self.width = width
        self.background_crop  = background_crop
        self.use_clahe = use_clahe
        self.transform = transform

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

    def class_name_to_labels(self, idx):

        if self.category == 'No Defects':
            class_abnormality = self.mammography.iloc[idx, 5]
            if class_abnormality in ['1', 1] : # No Defect   = True
                labels = 0.0
            elif class_abnormality in ['0', 0] : # No Defect   = False
                labels = 1.0
                
        if self.category == 'Skinfolds': ## Either 2i or 3i
            class_abnormality = self.mammography.iloc[idx, 6]
            if class_abnormality in ['0', 0 ] :
                labels = 0.0
            elif class_abnormality in ['1', 1]:
                labels = 1.0
                
        if self.category == '2i':
            class_abnormality = self.mammography.iloc[idx, 7]
            if class_abnormality in ['0', 0 ] :
                labels = 0.0
            elif class_abnormality in ['1', 1] :
                labels = 1.0
                
        if self.category == '3i':
            class_abnormality = self.mammography.iloc[idx, 8]
            if class_abnormality in ['0', 0 ] :
                labels = 0.0
            elif class_abnormality in ['1', 1] :
                labels = 1.0

        if self.category == 'Skinfold_MultiLabel':
            class_abnormality_skinfold = self.mammography.iloc[idx, 6]
            class_abnormality_defect = self.mammography.iloc[idx, 5]

            if class_abnormality_skinfold in ['0', 0 ] and  class_abnormality_defect in ['1', 1] :
                labels = [0., 0.]
            elif class_abnormality_skinfold in ['0', 0] and  class_abnormality_defect in ['0', 0 ]:
                labels = [0., 1.]
            elif class_abnormality_skinfold in ['1', 1] and  class_abnormality_defect in ['1', 1 ]:
                labels = [1., 0.]
            elif class_abnormality_skinfold in ['1', 1] and  class_abnormality_defect in ['0', 0 ]:
                labels = [1., 1.]

        if self.category == 'Skinfold_Defect_MultiLabel':
            class_abnormality_2i = self.mammography.iloc[idx, 7]
            class_abnormality_3i = self.mammography.iloc[idx, 8]
            class_abnormality_defect = self.mammography.iloc[idx, 5]

            if class_abnormality_2i in ['0', 0 ] and  class_abnormality_3i in ['0', 0 ]  and  class_abnormality_defect in ['1', 1 ]:
                labels = [0., 0., 0.]
            elif class_abnormality_2i in ['0', 0] and  class_abnormality_3i in ['1', 1 ] and  class_abnormality_defect in ['0', 0 ]:
                labels = [0., 1., 1.]
            elif class_abnormality_2i in ['1', 1] and  class_abnormality_3i in ['0', 0 ] and  class_abnormality_defect in ['0', 0 ]:
                labels = [1., 0., 1.]
            elif class_abnormality_2i in ['1', 1] and  class_abnormality_3i in ['1', 1 ] and  class_abnormality_defect in ['0', 0 ]:
                labels = [1., 1., 1.]

        if self.category == 'Calc_Mass_Malignant_MultiLabel':
            abnormality = self.mammography.iloc[idx, 3]
            classification = self.mammography.iloc[idx ,4]

            if abnormality == 'calcification' and classification == 'Benign':
                labels = [1. , 0., 0.]
            elif abnormality == 'both' and classification == 'Benign':
                labels = [1. , 1., 0.]
            elif abnormality == 'mass' and classification == 'Benign':
                labels = [0. , 1., 0.]
            elif abnormality == 'calcification' and classification == 'Malignant':
                labels = [1. , 0., 1.]
            elif abnormality == 'both' and classification == 'Malignant':
                labels = [1. , 1., 1.]
            elif abnormality == 'mass' and classification == 'Malignant':
                labels = [0. , 1., 1.]
            return labels

        return labels

    def detect_nonzero_regions(self, numbers):
        regions = []
        start = None
        
        for i, num in enumerate(numbers):
            if num != 0:
                if start is None:
                    start = i
            elif start is not None:
                regions.append((start, i-1))
                start = None
        
        if start is not None:
            regions.append((start, len(numbers)-1))
        return regions    
    
    def crop_images(self, data):
        columns = data.shape[1]
        count_all=[]
        for i in range(columns):
            count_non_zeros = np.count_nonzero(data[:, i][20:-20])
            count_all.append(count_non_zeros)
        column_indices = self.detect_nonzero_regions(count_all)
        differences = []
        for index, i in enumerate(column_indices):
            difference  = i[1]-i[0]
            differences.append(difference)

        index_max = np.argmax(differences)
        val = column_indices[index_max]
        if index_max == 0 and len(differences)!=1:
            column_indices = list(range(val[1], data.shape[1]))
        elif index_max == 0 and len(differences)==1:
            zero_columns = np.all(data == 0, axis=0)
            column_indices = np.where(zero_columns)[0]
        else:
            column_indices = list(range(0, val[0]))
        data = np.delete(data, column_indices, axis=1)
        zero_rows = np.all(data == 0, axis=1)
        row_indices = np.where(zero_rows)[0]
        data = np.delete(data, row_indices, axis=0)
        return data

    def image_load(self, idx, column):
        img_name = self.mammography.iloc[idx, column]
        
        if self.task == 'classification_dcm':
            image_name = os.path.join(self.data_folder, img_name +'.dcm')
            image = pydicom.dcmread(image_name)
            image_array = image.pixel_array
        elif self.task == 'classification_png':
            image_name = os.path.join(self.data_folder, img_name + '.png')
            image_array = cv2.imread(image_name)
        elif self.task == 'classification_jpeg':
            #filename = os.path.basename(img_name) ## Check why doesnt it work
            filename = img_name.split("\\")[-1] 
            image_name = os.path.join(self.data_folder, filename)
            image_array = cv2.imread(image_name)
            
        if self.task == 'classification_dcm' and self.background_crop:
            image_array = self.crop_images(image_array)
        image = cv2.resize(image_array, (self.width, self.height))

        
        if self.use_clahe:
            clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8)) ## Using standard values; Extent of impact unknown.
            image = clahe.apply(image)
        image = image * 1.0 / image.max()
        image = torch.from_numpy(image)
        if self.task == 'classification_dcm':
            image = image[None, :, :]
            image = einops.repeat(image, 'b h w -> (repeat b) h w', repeat=3)
        elif self.task in ['classification_png', 'classification_jpeg']:
            image = image.permute(2, 0, 1)
        return image, image_name

    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()

        image, image_name = self.image_load(idx, 9) 

        if self.transform:
            image = self.transform(image)
        labels = self.class_name_to_labels(idx)
        labels = torch.from_numpy(np.array(labels))
        
        if self.category in ['Calc_Mass_Malignant_MultiLabel']:
            ## For displaying the labels on the Output Image / GradCAM plots
            abnormality = self.mammography.iloc[idx, 3]
            classification = self.mammography.iloc[idx ,4]
            display_label = str(abnormality) + ' , ' + str(classification)

        else: 
            ## For displaying the labels on the Output Image / GradCAM plots
            two_i = self.mammography.iloc[idx, 7]
            three_i = self.mammography.iloc[idx, 8]
            defect = self.mammography.iloc[idx, 5]
            display_label = '2i:' + str(two_i)+ ' 3i:' + str(three_i) + ' Defect: ' + str(defect) 

        sample = {"image_name":image_name, "image": image, "label" : labels, "display_label":display_label}
        return sample