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])