dataset_bone.py

import os, torch
import numpy as np
from PIL import Image
from torch.utils.data import Dataset
from torchvision import transforms
import cv2
import random
import torchio as tio
import slicerio
import nrrd
import monai
import pickle
import nibabel as nib
from scipy.ndimage import zoom
from monai.transforms import OneOf
import einops
from funcs import *
from torchvision.transforms import InterpolationMode
#from .utils.transforms import ResizeLongestSide


class MRI_dataset(Dataset):
    def __init__(self,args, img_folder, mask_folder, img_list,phase='train',sample_num=50,channel_num=1,crop=False,crop_size=1024,targets=['femur','hip'],part_list=['all'],cls=1,if_prompt=True,prompt_type='point',region_type='largest_15',prompt_num=15,delete_empty_masks=False,if_attention_map=None):
        super(MRI_dataset, self).__init__()
        self.img_folder = img_folder
        self.mask_folder = mask_folder
        self.crop = crop
        self.crop_size = crop_size
        self.phase = phase
        self.channel_num=channel_num
        self.targets = targets
        self.segment_names_to_labels = []
        self.args = args
        self.cls = cls
        self.if_prompt = if_prompt
        self.region_type = region_type
        self.prompt_type = prompt_type
        self.prompt_num = prompt_num
        self.if_attention_map = if_attention_map
        
        for i,tag in enumerate(targets):
            self.segment_names_to_labels.append((tag,i))
            
        namefiles = open(img_list,'r')
        self.data_list = namefiles.read().split('\n')[:-1]

        if delete_empty_masks=='delete' or delete_empty_masks=='subsample':
            keep_idx = []
            for idx,data in enumerate(self.data_list):
                mask_path = data.split(' ')[1]
                if os.path.exists(os.path.join(self.mask_folder,mask_path)):
                    msk = Image.open(os.path.join(self.mask_folder,mask_path)).convert('L')
                else:
                    msk = Image.open(os.path.join(self.mask_folder.replace('2D-slices','2D-slices-generated'),mask_path)).convert('L')
                if 'all' in self.targets: # combine all targets as single target
                    mask_cls = np.array(np.array(msk,dtype=int)>0,dtype=int)
                else:
                    mask_cls = np.array(msk==self.cls,dtype=int)
                if part_list[0]=='all' and np.sum(mask_cls)>0:
                    keep_idx.append(idx) 
                elif np.sum(mask_cls)>0:
                    if_keep = False
                    for part in part_list:
                        if mask_path.find(part)>=0:
                            if_keep = True
                    if if_keep:
                        keep_idx.append(idx) 
            print('num with non-empty masks',len(keep_idx),'num with all masks',len(self.data_list))  
            if delete_empty_masks=='subsample':
                empty_idx = list(set(range(len(self.data_list)))-set(keep_idx))
                keep_empty_idx = random.sample(empty_idx, int(len(empty_idx)*0.1))
                keep_idx = empty_idx + keep_idx
            self.data_list = [self.data_list[i] for i in keep_idx] # keep the slices that contains target mask
  
        if phase == 'train':
            self.aug_img = [transforms.RandomEqualize(p=0.1),
                             transforms.ColorJitter(brightness=0.3, contrast=0.3,saturation=0.3,hue=0.3),
                             transforms.RandomAdjustSharpness(0.5, p=0.5),
                             ]
            self.transform_spatial = transforms.Compose([transforms.RandomResizedCrop(crop_size, scale=(0.8, 1.2)),
                     transforms.RandomRotation(45)])
            transform_img = [transforms.ToTensor()]
        else:
            transform_img = [
                         transforms.ToTensor(),
                             ]
        self.transform_img = transforms.Compose(transform_img)
            
    def __len__(self):
        return len(self.data_list)
        
    def __getitem__(self,index):
        # load image and the mask
        data = self.data_list[index]
        img_path = data.split(' ')[0]
        mask_path = data.split(' ')[1]
        slice_num = data.split(' ')[3] # total slice num for this object
        #print(img_path,mask_path)
        try:
            if os.path.exists(os.path.join(self.img_folder,img_path)):
                img = Image.open(os.path.join(self.img_folder,img_path)).convert('RGB')
            else:
                img = Image.open(os.path.join(self.img_folder.replace('2D-slices','2D-slices-generated'),img_path)).convert('RGB')
        except:
            # try to load image as numpy file
            img_arr = np.load(os.path.join(self.img_folder,img_path)) 
            img_arr = np.array((img_arr-img_arr.min())/(img_arr.max()-img_arr.min()+1e-8)*255,dtype=np.uint8)
            img_3c = np.tile(img_arr[:, :,None], [1, 1, 3])
            img = Image.fromarray(img_3c, 'RGB')
        if os.path.exists(os.path.join(self.mask_folder,mask_path)):
            msk = Image.open(os.path.join(self.mask_folder,mask_path)).convert('L')
        else:
            msk = Image.open(os.path.join(self.mask_folder.replace('2D-slices','2D-slices-generated'),mask_path)).convert('L')
                    
        if self.if_attention_map:
            slice_id = int(img_path.split('-')[-1].split('.')[0])
            slice_fraction = int(slice_id/int(slice_num)*4)
            img_id = '/'.join(img_path.split('-')[:-1]) +'_'+str(slice_fraction) + '.npy'
            attention_map = torch.tensor(np.load(os.path.join(self.if_attention_map,img_id)))
        else:
            attention_map = torch.zeros((64,64))
        
        img = transforms.Resize((self.args.image_size,self.args.image_size))(img)
        msk = transforms.Resize((self.args.image_size,self.args.image_size),InterpolationMode.NEAREST)(msk)
        
        state = torch.get_rng_state()
        if self.crop:
            im_w, im_h = img.size
            diff_w = max(0,self.crop_size-im_w)
            diff_h = max(0,self.crop_size-im_h)
            padding = (diff_w//2, diff_h//2, diff_w-diff_w//2, diff_h-diff_h//2)
            img = transforms.functional.pad(img, padding, 0, 'constant')
            torch.set_rng_state(state)
            t,l,h,w=transforms.RandomCrop.get_params(img,(self.crop_size,self.crop_size))
            img = transforms.functional.crop(img, t, l, h,w) 
            msk = transforms.functional.pad(msk, padding, 0, 'constant')
            msk = transforms.functional.crop(msk, t, l, h,w)
        if self.phase =='train':
            # add random optimazition
            aug_img_fuc = transforms.RandomChoice(self.aug_img)
            img = aug_img_fuc(img)

        img = self.transform_img(img)
        if self.phase == 'train':
            # It will randomly choose one
            random_transform = OneOf([monai.transforms.RandGaussianNoise(prob=0.5, mean=0.0, std=0.1),\
                                      monai.transforms.RandKSpaceSpikeNoise(prob=0.5, intensity_range=None, channel_wise=True),\
                                      monai.transforms.RandBiasField(degree=3),\
                                      monai.transforms.RandGibbsNoise(prob=0.5, alpha=(0.0, 1.0))
                                     ],weights=[0.3,0.3,0.2,0.2])
            img = random_transform(img).as_tensor()
        else:
            if img.mean()<0.05:
                img = min_max_normalize(img)
                img = monai.transforms.AdjustContrast(gamma=0.8)(img)

        
        if 'all' in self.targets: # combine all targets as single target
            msk = np.array(np.array(msk,dtype=int)>0,dtype=int)
        else:
            msk = np.array(msk,dtype=int)
            
        mask_cls = np.array(msk==self.cls,dtype=int)

        if self.phase=='train' and (not self.if_attention_map==None):
            mask_cls = np.repeat(mask_cls[np.newaxis,:, :], 3, axis=0)
            both_targets = torch.cat((img.unsqueeze(0), torch.tensor(mask_cls).unsqueeze(0)),0)
            transformed_targets = self.transform_spatial(both_targets)
            img = transformed_targets[0]
            mask_cls = np.array(transformed_targets[1][0].detach(),dtype=int)
        
        img = (img-img.min())/(img.max()-img.min()+1e-8)
        img = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])(img)
        
        # generate mask and prompt
        if self.if_prompt:
            if self.prompt_type =='point':
                prompt,mask_now = get_first_prompt(mask_cls,region_type=self.region_type,prompt_num=self.prompt_num)
                pc = torch.as_tensor(prompt[:,:2], dtype=torch.float)
                pl = torch.as_tensor(prompt[:, -1], dtype=torch.float)
                msk = torch.unsqueeze(torch.tensor(mask_now,dtype=torch.long),0)
                return {'image':img,
                    'mask':msk,
                    'point_coords': pc,
                    'point_labels':pl,
                    'img_name':img_path,
                    'atten_map':attention_map,
            }
            elif self.prompt_type =='box':
                prompt,mask_now = get_top_boxes(mask_cls,region_type=self.region_type,prompt_num=self.prompt_num)
                box = torch.as_tensor(prompt, dtype=torch.float)
                msk = torch.unsqueeze(torch.tensor(mask_now,dtype=torch.long),0)
                return {'image':img,
                    'mask':msk,
                    'boxes':box,
                    'img_name':img_path,
                    'atten_map':attention_map,
            }
        else:
            msk = torch.unsqueeze(torch.tensor(mask_cls,dtype=torch.long),0)
            return {'image':img,
                'mask':msk,
                'img_name':img_path,
                'atten_map':attention_map,
        }


class MRI_dataset_multicls(Dataset):
    def __init__(self, args, img_folder, mask_folder, img_list, phase='train', sample_num=50, channel_num=1,
                 crop=False, crop_size=1024, targets=['combine_all'], part_list=['all'], if_prompt=True, 
                 prompt_type='point', if_spatial = True, region_type='largest_20', prompt_num=20, delete_empty_masks=False, 
                 label_mapping=None, reference_slice_num=0, if_attention_map=None,label_frequency_path=None):
        super(MRI_dataset_multicls, self).__init__()
        self.initialize_parameters(args, img_folder, mask_folder, img_list, phase, sample_num, channel_num,
                                   crop, crop_size, targets, part_list, if_prompt, prompt_type, if_spatial, region_type,
                                   prompt_num, delete_empty_masks, label_mapping, reference_slice_num, if_attention_map,label_frequency_path)
        self.load_label_mapping()
        self.prepare_data_list()
        self.filter_data_list()
        if phase == 'train':
            self.setup_transformations_train(crop_size)
        else:
            self.setup_transformations_other()

    def initialize_parameters(self, args, img_folder, mask_folder, img_list, phase, sample_num, channel_num,
                              crop, crop_size, targets, part_list, if_prompt, prompt_type, if_spatial, region_type,
                              prompt_num, delete_empty_masks, label_mapping, reference_slice_num, if_attention_map,label_frequency_path):
        self.args = args
        self.img_folder = img_folder
        self.mask_folder = mask_folder
        self.img_list = img_list
        self.phase = phase
        self.sample_num = sample_num
        self.channel_num = channel_num
        self.crop = crop
        self.crop_size = crop_size
        self.targets = targets
        self.part_list = part_list
        self.if_prompt = if_prompt
        self.prompt_type = prompt_type
        self.if_spatial = if_spatial
        self.region_type = region_type
        self.prompt_num = prompt_num
        self.delete_empty_masks = delete_empty_masks
        self.label_mapping = label_mapping
        self.reference_slice_num = reference_slice_num
        self.if_attention_map = if_attention_map
        self.label_dic = {}
        self.label_frequency_path = label_frequency_path

    def load_label_mapping(self):
        # Load the basic label mappings from a pickle file
        if self.label_mapping:
            with open(self.label_mapping, 'rb') as handle:
                self.segment_names_to_labels = pickle.load(handle)
            self.label_dic = {seg[1]: seg[0] for seg in self.segment_names_to_labels}
            self.label_name_list = [seg[0] for seg in self.segment_names_to_labels]
            print(self.label_dic)
        else:
            self.label_dic = {value: 'all' for value in range(1, 256)}
        
        # Load frequency data and remap classes if required
        if 'remap_frequency' in self.targets:
            self.load_and_remap_classes_based_on_frequency()

    def load_and_remap_classes_based_on_frequency(self):
        if self.label_frequency_path:
            with open(self.label_frequency_path, 'r') as file:
                all_label_frequencies = json.load(file)
            all_label_frequencies = all_label_frequencies['train']
            
            
            # Example to select the target region dynamically based on some condition or configuration
            target_region = self.part_list[0] 
            if target_region in all_label_frequencies:
                label_frequencies = all_label_frequencies[target_region]
                self.label_frequencies = label_frequencies
                #print(label_frequencies)
                self.remap_classes_based_on_frequency(label_frequencies)
            else:
                print(f"Warning: No frequency data found for the target region '{target_region}'. No remapping applied.")
    
    def remap_classes_based_on_frequency(self, label_frequencies):
        # Determine the frequency threshold for high vs. low frequency classes
        total = max(label_frequencies.values())
        high_freq_threshold = total * 0.5  # Adjust this threshold as needed
        
        # Initialize dictionaries to hold new class mappings
        high_freq_classes = {}
        low_freq_classes = {}
        
        # Assign classes to high or low frequency based on the threshold
        for label, freq in label_frequencies.items():
            if freq >= high_freq_threshold:
                high_freq_classes[label] = freq
            else:
                low_freq_classes[label] = freq
    
        # Update label dictionary based on the frequency classification
        #self.label_dic: {old_cls: old_name}
        new_label_dic = {}
        for cls, name in self.label_dic.items():
            if name in high_freq_classes:
                new_label_dic[cls] = name  # Retain original name for high frequency classes
            elif name in low_freq_classes:
                new_label_dic[cls] = 'combined_low_freq'  # Combine low frequency classes into one
    
        self.updated_label_dic = new_label_dic
        #new_label_dic: {old_cls: new_name}
        #print("Updated label dictionary with frequency remapping:", new_label_dic)
        
        #print('new_label_dic:',new_label_dic)
    
        # Sort high frequency keys by their frequency in descending order
        sorted_high_freq_labels = sorted(high_freq_classes.items(), key=lambda item: item[1], reverse=True)
        
        # Create a mapping for high frequency classes based on the sorted order
        original_to_new = {label: idx + 1 for idx, (label, _) in enumerate(sorted_high_freq_labels)}

        
        combined_low_freq_class_id = len(original_to_new) + 1
        # Ensure combined low frequency class is mapped correctly
        if 'combined_low_freq' in new_label_dic.values():
            for cls in low_freq_classes.keys():
                original_to_new[cls] = combined_low_freq_class_id
                
        # orignal_to_new {old_name:new_cls} 
        #print('original_to_new:',original_to_new)

        
        # Create additional dictionaries
        self.old_name_to_new_name = {self.label_dic[cls]: new_label for cls, new_label in new_label_dic.items()}
        self.old_cls_to_new_cls = {cls: original_to_new[self.label_dic[cls]] for cls in self.label_dic.keys() if self.label_dic[cls] in original_to_new}

        print('remapped label dic:',self.old_name_to_new_name)
        print('remapped cls dic:',self.old_cls_to_new_cls)
            
    def prepare_data_list(self):
        with open(self.img_list, 'r') as namefiles:
            self.data_list = namefiles.read().split('\n')[:-1]
        self.sp_symbol = ',' if ',' in self.data_list[0] else ' '

    def filter_data_list(self):
        keep_idx = []
        for idx, data in enumerate(self.data_list):
            img_path, mask_path = self.extract_paths(data)
            msk = Image.open(os.path.join(self.mask_folder, mask_path)).convert('L')
            mask_cls = self.determine_mask_class(msk)

            if self.should_keep(mask_cls, mask_path):
                keep_idx.append(idx)
                if self.reference_slice_num > 1:
                    self.add_reference_slice(img_path, mask_path, data)

        self.data_list = [self.data_list[i] for i in keep_idx]
        print('num with non-empty masks', len(keep_idx), 'num with all masks', len(self.data_list))

    def extract_paths(self, data):
        img_path = data.split(self.sp_symbol)[0]
        mask_path = data.split(self.sp_symbol)[1]
        return img_path.lstrip('/'), mask_path.lstrip('/')

    def determine_mask_class(self, msk):
        if 'combine_all' in self.targets:
            return np.array(msk, dtype=int) > 0
        elif self.targets[0] in self.label_name_list:
            return np.array(msk, dtype=int) == self.cls
        return np.array(msk, dtype=int)

    def should_keep(self, mask_cls, mask_path):
        if self.delete_empty_masks:
            has_mask = np.any(mask_cls > 0)
            if has_mask:
                if self.part_list[0] == 'all':
                    return True
                return any(mask_path.find(part) >= 0 for part in self.part_list)
            return False
        return True


    def add_reference_slice(self, img_path, mask_path, data):
        volume_name = ''.join(img_path.split('-')[:-1])  # get volume name
        slice_num = data.split(self.sp_symbol)[2]
        if volume_name not in self.reference_slices:
            self.reference_slices[volume_name] = []
        self.reference_slices[volume_name].append((img_path, mask_path, slice_num))

    def setup_transformations_train(self, crop_size):
        self.transform_img = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])
        self.aug_img = transforms.RandomChoice([
            transforms.RandomEqualize(p=0.1),
            transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3, hue=0.3),
            transforms.RandomAdjustSharpness(0.5, p=0.5),
        ])
        if self.if_spatial:
                self.transform_spatial = transforms.Compose([transforms.RandomResizedCrop(self.crop_size, scale=(0.5, 1.5), interpolation=InterpolationMode.NEAREST),
                         transforms.RandomRotation(45, interpolation=InterpolationMode.NEAREST)])

    def setup_transformations_other(self):
        self.transform_img = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])

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

    def __getitem__(self, index):
        # Load image and mask, handle missing files
        data = self.data_list[index]
        img, msk, img_path, mask_path, slice_num = self.load_image_and_mask(data)
        
        # Optional: Load attention map
        attention_map = self.load_attention_map(img_path, slice_num) if self.if_attention_map else torch.zeros((64, 64))
    
        # Handle reference slices if necessary
        if self.reference_slice_num > 1:
            img, msk = self.handle_reference_slices(img_path, mask_path, slice_num)
        
        # Apply transformations
        img, msk = self.apply_transformations(img, msk)
    
        # Generate and process masks and prompts
        output_dict = self.prepare_output(img, msk, img_path, mask_path,attention_map)
        
    
        return output_dict
        
    def load_image_and_mask(self, data):
        img_path, mask_path = self.extract_paths(data)
        slice_num = data.split(self.sp_symbol)[3]  # Extract total slice number for this object
        
        img_folder = self.img_folder
        msk_folder = self.mask_folder
        
        img = Image.open(os.path.join(img_folder, img_path)).convert('RGB')
        msk = Image.open(os.path.join(msk_folder, mask_path)).convert('L')
    
        # Resize images for processing
        img = transforms.Resize((self.args.image_size, self.args.image_size))(img)
        msk = transforms.Resize((self.args.image_size, self.args.image_size), InterpolationMode.NEAREST)(msk)
    
        return img, msk, img_path, mask_path, int(slice_num)

    def load_attention_map(self, img_path, slice_num):
        slice_id = int(img_path.split('-')[-1].split('.')[0])
        slice_fraction = int(slice_id / slice_num * 4)
        img_id = '/'.join(img_path.split('-')[:-1]) + '_' + str(slice_fraction) + '.npy'
        attention_map = torch.tensor(np.load(os.path.join(self.if_attention_map, img_id)))
        return attention_map


    def apply_crop(self, img, msk):
        im_w, im_h = img.size
        diff_w = max(0, self.crop_size - im_w)
        diff_h = max(0, self.crop_size - im_h)
        padding = (diff_w // 2, diff_h // 2, diff_w - diff_w // 2, diff_h - diff_h // 2)
        img = transforms.functional.pad(img, padding, 0, 'constant')
        msk = transforms.functional.pad(msk, padding, 0, 'constant')
        t, l, h, w = transforms.RandomCrop.get_params(img, (self.crop_size, self.crop_size))
        img = transforms.functional.crop(img, t, l, h, w)
        msk = transforms.functional.crop(msk, t, l, h, w)
        return img, msk
        
    def apply_transformations(self, img, msk):
        if self.crop:
            img, msk = self.apply_crop(img, msk)
        if self.phase == 'train':
            img = self.aug_img(img)
        img = self.transform_img(img)
        if self.phase =='train' and self.if_spatial:
            mask_cls = np.array(msk,dtype=int)
            mask_cls = np.repeat(mask_cls[np.newaxis,:, :], 3, axis=0)
            both_targets = torch.cat((img.unsqueeze(0), torch.tensor(mask_cls).unsqueeze(0)),0)
            transformed_targets = self.transform_spatial(both_targets)
            img = transformed_targets[0]
            mask_cls = np.array(transformed_targets[1][0].detach(),dtype=int)
            msk = torch.tensor(mask_cls)
        return img, msk

    def handle_reference_slices(self, img_path, mask_path, slice_num):
        volume_name = ''.join(img_path.split('-')[:-1])
        ref_slices, ref_msks = [], []
        reference_slices = self.reference_slices.get(volume_name, [])
        for ref_slice in reference_slices:
            ref_img_path, ref_msk_path, _ = ref_slice
            ref_img = Image.open(os.path.join(self.img_folder, ref_img_path)).convert('RGB')
            ref_img = transforms.Resize((self.args.image_size, self.args.image_size))(ref_img)
            ref_img = self.transform_img(ref_img)
            ref_img = torch.unsqueeze(ref_img, 0)
            
            ref_msk = Image.open(os.path.join(self.mask_folder, ref_msk_path)).convert('L')
            ref_msk = transforms.Resize((self.args.image_size, self.args.image_size), InterpolationMode.NEAREST)(ref_msk)
            ref_msk = torch.tensor(ref_msk, dtype=torch.long)
            ref_msks.append(torch.unsqueeze(ref_msk, 0))
    
        img = torch.cat(ref_slices, dim=0)
        msk = torch.cat(ref_msks, dim=0)
        return img, msk
    
    def remap_classes_sequentially(self, mask, label_frequencies):
        # Apply the mapping to the mask
        remapped_mask = mask.copy()
        for old_cls, new_cls in  self.old_cls_to_new_cls.items():
            remapped_mask[mask == old_cls] = new_cls
        return remapped_mask


    def prepare_output(self, img, msk, img_path, mask_path, attention_map):
        # Normalize the image
        img = (img - img.min()) / (img.max() - img.min() + 1e-8)
        img = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])(img)
    
        msk = np.array(msk, dtype=int)
        #print('ori_msk:',np.unique(msk))
        if self.label_frequency_path:
            msk = self.remap_classes_sequentially(msk,self.label_frequencies)  # Assuming msk is already using updated IDs
            #print('new_msk------------------------:',self.old_cls_to_new_cls)
            # Prepare one-hot encoding for the remapped classes
        
        unique_classes = np.unique(msk).tolist()
        if 0 in unique_classes:
            unique_classes.remove(0)
    
        if len(unique_classes) > 0:
            selected_dic = {k: self.label_dic[k] for k in unique_classes if k in self.label_dic}
        else:
            selected_dic = {}
    
        if self.targets[0] == 'random':
            mask_cls, selected_label, cls_one_hot = self.handle_random_target(msk, unique_classes, selected_dic)
        elif self.targets[0] in self.label_name_list:
            selected_label = self.targets[0]
            mask_cls = np.array(msk == self.cls, dtype=int)
            cls_one_hot = torch.zeros(len(self.label_dic), dtype=torch.long)
            cls_one_hot[self.cls - 1] = 1
        else:
            selected_label = self.targets[0]
            mask_cls = msk
            cls_one_hot = torch.zeros(len(self.label_dic), dtype=torch.long)
    
        # Handle prompts
        if self.if_prompt:
            prompt, mask_now, mask_cls = self.generate_prompt(mask_cls)
            ref_msk,_ = torch.max(mask_now>0,dim=0)
            return_dict = {'image': img, 'mask': mask_now, 'selected_label_name': selected_label,
                           'cls_one_hot': cls_one_hot, 'prompt': prompt, 'img_name': img_path,
                           'mask_ori': msk, 'mask_cls': mask_cls, 'all_label_dic': selected_dic,'ref_mask':ref_msk}
        else:
            if len(mask_cls.shape)==2:
                msk = torch.unsqueeze(torch.tensor(mask_cls,dtype=torch.long),0)
            elif len(mask_cls.shape)==4:
                msk = torch.squeeze(torch.tensor(mask_cls,dtype=torch.long))
            else:
                msk = torch.tensor(mask_cls,dtype=torch.long)
            ref_msk,_ = torch.max(msk>0,dim=0)
            #print('unique mask values:',msk.unique())
            return_dict = {'image': img, 'mask': msk, 'selected_label_name': selected_label,
                           'cls_one_hot': cls_one_hot, 'img_name': img_path, 'mask_ori': msk,'ref_mask':ref_msk}
    
        return return_dict
        
    def generate_prompt(self, mask_cls):
        if self.prompt_type == 'point':
            prompt, mask_now = get_first_prompt(mask_cls, region_type=self.region_type, prompt_num=self.prompt_num)
        elif self.prompt_type == 'box':
            prompt, mask_now = get_top_boxes(mask_cls, region_type=self.region_type, prompt_num=self.prompt_num)
        else:
            prompt = mask_now = None
        
        # Handling the shape of mask_now for return
        if mask_now is not None:
            if len(mask_now.shape) == 2:
                mask_now = torch.unsqueeze(torch.tensor(mask_now, dtype=torch.long), 0)
                mask_cls = torch.unsqueeze(torch.tensor(mask_cls, dtype=torch.long), 0)
            elif len(mask_now.shape) == 4:
                mask_now = torch.squeeze(torch.tensor(mask_now, dtype=torch.long))
            else:
                mask_now = torch.tensor(mask_now, dtype=torch.long)
                mask_cls = torch.tensor(mask_cls, dtype=torch.long)
    
        return prompt, mask_now, mask_cls


    def handle_random_target(self, msk, unique_classes, selected_dic):
        if len(unique_classes) > 0:
            random_selected_cls = random.choice(unique_classes)
            selected_label = selected_dic[random_selected_cls]
            mask_cls = np.array(msk == random_selected_cls, dtype=int)
            
            cls_one_hot = torch.zeros(len(self.label_dic), dtype=torch.long)
            cls_one_hot[random_selected_cls - 1] = 1
        else:
            selected_label = None
            mask_cls = torch.zeros_like(msk)  # assuming msk is already a numpy array
            cls_one_hot = torch.zeros(len(self.label_dic), dtype=torch.long)
    
        return mask_cls, selected_label, cls_one_hot