from dataset.baseset import BaseSet
import random, cv2
import numpy as np


class iNaturalist(BaseSet):
    def __init__(self, mode='train', cfg=None, transform=None):
        super(iNaturalist, self).__init__(mode, cfg, transform)
        random.seed(0)
        self.class_dict = self._get_class_dict()

    def __getitem__(self, index):
        if self.cfg.TRAIN.SAMPLER.TYPE == "weighted sampler" and self.mode == 'train':
            assert self.cfg.TRAIN.SAMPLER.WEIGHTED_SAMPLER.TYPE in ["balance", 'square', 'progressive']
            if  self.cfg.TRAIN.SAMPLER.WEIGHTED_SAMPLER.TYPE == "balance":
                sample_class = random.randint(0, self.num_classes - 1)
            elif self.cfg.TRAIN.SAMPLER.WEIGHTED_SAMPLER.TYPE == "square":
                sample_class = np.random.choice(np.arange(self.num_classes), p=self.square_p)
            else:
                sample_class = np.random.choice(np.arange(self.num_classes), p=self.progress_p)
            sample_indexes = self.class_dict[sample_class]
            index = random.choice(sample_indexes)
        now_info = self.data[index]
        img = self._get_image(now_info)
        image = self.transform(img)
        meta = dict()
        image_label = now_info['category_id']  # 0-index
        return image, image_label, meta













[((9, 8, 9, 22), 'random.seed', 'random.seed', ({(9, 20, 9, 21): '(0)'}, {}), '(0)', False, 'import random, cv2\n'), ((22, 20, 22, 49), 'random.choice', 'random.choice', ({(22, 34, 22, 48): 'sample_indexes'}, {}), '(sample_indexes)', False, 'import random, cv2\n'), ((16, 31, 16, 70), 'random.randint', 'random.randint', ({(16, 46, 16, 47): '0', (16, 49, 16, 69): 'self.num_classes - 1'}, {}), '(0, self.num_classes - 1)', False, 'import random, cv2\n'), ((18, 48, 18, 75), 'numpy.arange', 'np.arange', ({(18, 58, 18, 74): 'self.num_classes'}, {}), '(self.num_classes)', True, 'import numpy as np\n'), ((20, 48, 20, 75), 'numpy.arange', 'np.arange', ({(20, 58, 20, 74): 'self.num_classes'}, {}), '(self.num_classes)', True, 'import numpy as np\n')]