def __getitem__()

in part_selector.py [0:0]

• 26 lines of code
• 3 McCabe index (conditional complexity)

    def __getitem__(self, index):
        path = self.paths[index]
        part_id = self.parts_id[index]
        json_data = json.load(open(path))
        input_parts_json = json_data['input_parts']
        img_partial_test = []
        vector_input_part = []
        for i in range(self.n_part):
            key = self.id_to_part[i]
            vector_input_part += input_parts_json[key]
            img_partial_test.append(self.processed_part_to_raster(input_parts_json[key], side=self.image_size))
        img_partial_test.append(self.processed_part_to_raster(vector_input_part, side=self.image_size))
        # random affine
        theta = np.random.uniform(*self.rotate)
        trans_pixel = 512*self.trans
        translate_x = np.random.uniform(-trans_pixel, trans_pixel)
        translate_y = np.random.uniform(-trans_pixel, trans_pixel)
        scale = np.random.uniform(self.scale)
        # apply
        processed_img_partial = []
        affine_vector_input_part = []
        for i in range(self.n_part):
            key = self.id_to_part[i]
            affine_input_part_json = self.affine_trans(input_parts_json[key], theta, translate_x, translate_y, scale)
            affine_vector_input_part += affine_input_part_json
            processed_img_partial.append(self.processed_part_to_raster(affine_input_part_json, side=self.image_size))
        processed_img_partial.append(self.processed_part_to_raster(affine_vector_input_part, side=self.image_size))
        return part_id, torch.cat(processed_img_partial, 0), torch.cat(img_partial_test, 0)