RadDataset.py

# Defines a custom PyTorch Dataset for radiologic questions

from torch.utils.data import Dataset
from utilFunctions import readCSV
import CONFIG
import os.path

RAD_DATASET_SUBTOPICS = [
    "Angiography",
    "ComputedTomography",
    "MagneticResonanceImaging",
    "Radiography",
]

SUBTOPIC_PHRASE_DICT = {
    "Angiography":"angiographic image",
    "ComputedTomography": "CT image",
    "MagneticResonanceImaging": "MR image",
    "Radiography": "radiograph",
}

class RadDataset(Dataset):
    def __init__(self,questionFolder = CONFIG.RAD_QUESTION_FOLDER,imgFolder=CONFIG.RAD_IMG_FOLDER,withContext = False):
        self.withContext = withContext
        self.sampleList = list()
        for subtopic in RAD_DATASET_SUBTOPICS:
            inpFile = questionFolder + subtopic + ".csv"
            head,*data = readCSV(inpFile)
            toImgName = lambda num: [imgFolder + num[0] + subtopic + ".jpg"]
            lst = [toImgName(x[:1])+ [subtopic] + x for x in data]
            self.sampleList.extend(lst)
        for el in self.sampleList: assert os.path.exists(el[0])
    
    def __len__(self):
        return len(self.sampleList)
    
    def __getitem__(self,idx):
        sample = self.sampleList[idx]
        context = ""
        if self.withContext:
            age = sample[3]
            sex = {"f":"female","m":"male"}[sample[4]]
            modalityPhrase = SUBTOPIC_PHRASE_DICT[sample[1]]
            context = f"Attached is a representative {modalityPhrase} of "

            complaint = sample[6]
            if complaint == "asymptomatic":
                context += f"an asymptomatic {age}-year-old {sex}."
            else:
                context += f"a {age}-year-old {sex} who complains of {complaint}."
            context += " "

        retDict = {
            "question": context + "What is the most likely diagnosis?",
            "answer": sample[5],
            "subtopic": sample[1],
            "imgPath": sample[0],
            }
        return retDict

class VQAMed2019Dataset(Dataset):
    # Loads and provides questions, answers, and metadata
    def __init__(self,datasetFile,imgFolder=CONFIG.VQA2019_IMG_FOLDER):
        # Initializes the dataset from the specified CSV file
        self.data = readCSV(datasetFile,"|")
        self.imgFolder = imgFolder
    
    def __len__(self):
        # Returns the total number of data samples
        return len(self.data)
    
    def __getitem__(self,idx):
        # Retrieves a single sample by index
        line = self.data[idx] 
        imgPath = self.imgFolder + line[0] + ".jpg"
        retDict = {
            "question": line[2],
            "answer": line[3],
            "subtopic": line[1],
            "imgPath": imgPath,
            }
        return retDict