def __getitem__(self, index): index = index % len(self.imnames) try: sample = self.raw_dataset['data'][self.imnames[index]] except Exception as ex: print('Exeption:', ex) print('At index:', index) print('Sample id:', self.imnames[index]) print( 'WARNING: h5py do not support index from multiple thread, please use num_worker=0') exit(0) image = sample[()] charBB = sample.attrs['charBB'] txt = [each.decode('utf-8') for each in sample.attrs['txt']] # print(txt) # Handle line-break all_words = [] for line in txt: if '\n' in line: all_words.extend(line.split('\n')) else: all_words.append(line) # Remove blank word for index, line in enumerate(all_words): all_words[index] = [word for word in line.strip().split(' ') if word not in ['', ' ']] # Split word to char for index, line in enumerate(all_words): new_line = [] for word in line: if len(word) >= 2: new_line.extend([char for char in word]) else: new_line.append(word) all_words[index] = new_line # print('--------') # print(all_words) # print('--------') # if len(image.shape) == 2: # image = np.repeat(image[:, :, None], repeats=3, axis=2) # elif image.shape[2] == 1: # image = np.repeat(image, repeats=3, axis=2) # else: # image = image[:, :, 0: 3] # Resize the image to (768, 768) image, character = resize(image, charBB.copy()) image = normalize_mean_variance(image).transpose(2, 0, 1) weight_character = generate_target( image.shape, character.copy()) # Generate character heatmap # print("WEIGHT CHARACTER: {} - {}".format(type(weight_character), weight_character)) weight_affinity = generate_affinity(image.shape, character.copy( ), all_words.copy()) # Generate affinity heatmap # print("WEIGHT affinity: {} - {}".format(type(weight_affinity), weight_affinity)) return \ image.astype(np.float32), \ weight_character.astype(np.float32), \ weight_affinity[0].astype(np.float32)
def __getitem__(self, item): item = item % len(self.imnames) image = plt.imread(self.base_path+'/' + self.imnames[item][0]) # Read the image if len(image.shape) == 2: image = np.repeat(image[:, :, None], repeats=3, axis=2) elif image.shape[2] == 1: image = np.repeat(image, repeats=3, axis=2) else: image = image[:, :, 0: 3] # Resize the image to (768, 768) image, character = resize(image, self.charBB[item].copy()) normal_image = image.astype(np.uint8).copy() image = normalize_mean_variance(image).transpose(2, 0, 1) # Generate character heatmap weight_character = generate_target( image.shape, character.copy()) # Generate affinity heatmap weight_affinity, affinity_bbox = generate_affinity( image.shape, character.copy(), self.txt[item].copy()) cv2.drawContours( normal_image, np.array(affinity_bbox).reshape([len(affinity_bbox), 4, 1, 2]).astype(np.int64), -1, (0, 255, 0), 2) enlarged_affinity_bbox = [] for i in affinity_bbox: center = np.mean(i, axis=0) i = i - center[None, :] i = i*60/25 i = i + center[None, :] enlarged_affinity_bbox.append(i) cv2.drawContours( normal_image, np.array(enlarged_affinity_bbox).reshape( [len(affinity_bbox), 4, 1, 2]).astype(np.int64), -1, (0, 0, 255), 2 ) return \ image.astype(np.float32), \ weight_character.astype(np.float32), \ weight_affinity.astype(np.float32), \ normal_image
def __getitem__(self, item): item = item % len(self.imnames) image = plt.imread(self.base_path + '/' + self.imnames[item][0]) # Read the image image, character = resize( image, self.charBB[item].copy()) # Resize the image to (768, 768) image = normalize_mean_variance(image).transpose(2, 0, 1) weight_character = generate_target( image.shape, character.copy()) # Generate character heatmap weight_affinity = generate_affinity( image.shape, character.copy(), self.txt[item].copy()) # Generate affinity heatmap return \ image.astype(np.float32), \ weight_character.astype(np.float32), \ weight_affinity.astype(np.float32)
def __getitem__(self, item_i): # noinspection PyArgumentList np.random.seed() check = np.random.uniform() if check < config.prob_synth and self.type_ == 'train': # probability of picking a Synth-Text image vs Image from dataset random_item = np.random.randint(len(self.imnames)) character = self.charBB[random_item].copy() image = plt.imread(self.base_path_synth + '/' + self.imnames[random_item][0]) # Read the image if len(image.shape) == 2: image = np.repeat(image[:, :, None], repeats=3, axis=2) elif image.shape[2] == 1: image = np.repeat(image, repeats=3, axis=2) else: image = image[:, :, 0:3] height, width, channel = image.shape image, character = resize( image, character) # Resize the image to (768, 768) image = normalize_mean_variance(image).transpose(2, 0, 1) # Generate character heatmap with weights weight_character, weak_supervision_char = generate_target( image.shape, character.copy(), weight=1) # Generate affinity heatmap with weights weight_affinity, weak_supervision_affinity = generate_affinity( image.shape, character.copy(), self.txt[random_item].copy(), weight=1) dataset_name = 'SYNTH' text_target = '' else: random_item = np.random.randint(len(self.gt)) image = plt.imread(self.base_path_other_images + '/' + self.gt[random_item][0]) # Read the image if len(image.shape) == 2: image = np.repeat(image[:, :, None], repeats=3, axis=2) elif image.shape[2] == 1: image = np.repeat(image, repeats=3, axis=2) else: image = image[:, :, 0:3] height, width, channel = image.shape character = [ np.array(word_i).reshape([len(word_i), 4, 1, 2]) for word_i in self.gt[random_item][1]['characters'].copy() ] affinity = [ np.array(word_i).reshape([len(word_i), 4, 1, 2]) for word_i in self.gt[random_item][1]['affinity'].copy() ] assert len(character) == len( affinity), 'word length different in character and affinity' # Resize the image to (768, 768) image, character, affinity = resize_generated( image, character.copy(), affinity.copy()) image = normalize_mean_variance(image).transpose(2, 0, 1) weights = [i for i in self.gt[random_item][1]['weights'].copy()] text_target = '#@#@#@'.join(self.gt[random_item][1]['text']) assert len(self.gt[random_item][1]['text']) == len(self.gt[random_item][1]['word_bbox']), \ 'Length of word_bbox != Length of text' # assert len(text_target.split('#@#@#@')) == len(self.gt[random_item][1]['word_bbox']), \ # 'Some error in splitting' # Generate character heatmap with weights weight_character, weak_supervision_char = generate_target_others( image.shape, character.copy(), weights.copy()) # Generate affinity heatmap with weights weight_affinity, weak_supervision_affinity = generate_target_others( image.shape, affinity.copy(), weights.copy()) # Get original word_bbox annotations dataset_name = 'ICDAR' return \ image.astype(np.float32), \ weight_character.astype(np.float32), \ weight_affinity.astype(np.float32), \ weak_supervision_char.astype(np.float32), \ weak_supervision_affinity.astype(np.float32), \ dataset_name, \ text_target, \ random_item, \ np.array([height, width])
def __getitem__(self, item_i): height, width, channel = 0, 0, 0 # noinspection PyArgumentList np.random.seed() check = np.random.uniform() if check < config.prob_synth and self.type_ == 'train': # probability of picking a Synth-Text image vs Image from dataset random_item = np.random.choice(self.imnames) sample = self.raw_dataset['data'][random_item] image = sample[()] charBB = sample.attrs['charBB'] txt = [each.decode('utf-8') for each in sample.attrs['txt']] # print(txt) # Handle line-break all_words = [] for line in txt: if '\n' in line: all_words.extend(line.split('\n')) else: all_words.append(line) # Remove blank word for index, line in enumerate(all_words): all_words[index] = [word for word in line.strip().split(' ') if word not in ['', ' ']] # Split word to char for index, line in enumerate(all_words): new_line = [] for word in line: if len(word) >= 2: new_line.extend([char for char in word]) else: new_line.append(word) all_words[index] = new_line # print('--------') # print(all_words) # print('--------') # Resize the image to (768, 768) image, character = resize(image, charBB.copy()) image = normalize_mean_variance(image).transpose(2, 0, 1) # Generate character heatmap with weights weight_character, weak_supervision_char = generate_target( image.shape, character.copy(), weight=1) # Generate affinity heatmap with weights weight_affinity, weak_supervision_affinity = generate_affinity( image.shape, character.copy(), all_words.copy(), weight=1) dataset_name = 'SYNTH' text_target = '' else: random_item = np.random.randint(len(self.gt)) image = plt.imread(os.path.join( self.base_path_other_images, self.gt[random_item][0])) # Read the image if len(image.shape) == 2: image = np.repeat(image[:, :, None], repeats=3, axis=2) elif image.shape[2] == 1: image = np.repeat(image, repeats=3, axis=2) else: image = image[:, :, 0: 3] height, width, channel = image.shape character = [ np.array(word_i).reshape([len(word_i), 4, 1, 2]) for word_i in self.gt[random_item][1]['characters'].copy()] affinity = [ np.array(word_i).reshape([len(word_i), 4, 1, 2]) for word_i in self.gt[random_item][1]['affinity'].copy()] assert len(character) == len( affinity), 'word length different in character and affinity' # Resize the image to (768, 768) image, character, affinity = resize_generated( image, character.copy(), affinity.copy()) image = normalize_mean_variance(image).transpose(2, 0, 1) weights = [i for i in self.gt[random_item][1]['weights'].copy()] text_target = '#@#@#@'.join(self.gt[random_item][1]['text']) assert len(self.gt[random_item][1]['text']) == len(self.gt[random_item][1]['word_bbox']), \ 'Length of word_bbox != Length of text' # assert len(text_target.split('#@#@#@')) == len(self.gt[random_item][1]['word_bbox']), \ # 'Some error in splitting' # Generate character heatmap with weights weight_character, weak_supervision_char = generate_target_others( image.shape, character.copy(), np.array(weights)[:, 0]) # Generate affinity heatmap with weights weight_affinity, weak_supervision_affinity = generate_target_others( image.shape, affinity.copy(), np.array(weights)[:, 0]) # Get original word_bbox annotations dataset_name = 'datapile' return \ image.astype(np.float32), \ weight_character.astype(np.float32), \ weight_affinity.astype(np.float32), \ weak_supervision_char.astype(np.float32), \ weak_supervision_affinity.astype(np.float32), \ dataset_name, \ text_target, \ str(random_item), \ np.array([height, width])