def __getitem__(self, item):
im_name = os.path.basename(self.image_lists[item])
#print(self.image_lists[item])
img = Image.open(self.image_lists[item]).convert("RGB")
width, height = img.size
if self.gts_dir is not None:
gt_path = os.path.join(self.gts_dir, im_name + '.txt')
#ipdb.set_trace()
words, boxes, charsbbs, segmentations = self.load_gt_from_txt(
gt_path, height, width)
if words[0] == '':
use_char_ann = False
else:
use_char_ann = True
if not self.use_charann:
use_char_ann = False
target = BoxList(boxes[:, :4],
img.size,
mode="xyxy",
use_char_ann=use_char_ann)
classes = torch.ones(len(boxes))
target.add_field("labels", classes)
masks = SegmentationMask(segmentations, img.size)
target.add_field("masks", masks)
char_masks = SegmentationCharMask(charsbbs,
words=words,
use_char_ann=use_char_ann,
size=img.size)
target.add_field("char_masks", char_masks)
else:
target = None
if self.transforms is not None:
img, target = self.transforms(img, target)
if self.vis:
new_im = img.numpy().copy().transpose(
[1, 2, 0]) + [102.9801, 115.9465, 122.7717]
new_im = Image.fromarray(new_im.astype(np.uint8)).convert('RGB')
mask = target.extra_fields['masks'].polygons[0].convert('mask')
mask = Image.fromarray(
(mask.numpy() * 255).astype(np.uint8)).convert('RGB')
if self.use_charann:
m, _ = target.extra_fields['char_masks'].chars_boxes[
0].convert('char_mask')
color = self.creat_color_map(37, 255)
color_map = color[m.numpy().astype(np.uint8)]
char = Image.fromarray(color_map.astype(
np.uint8)).convert('RGB')
char = Image.blend(char, new_im, 0.5)
else:
char = new_im
new = Image.blend(char, mask, 0.5)
img_draw = ImageDraw.Draw(new)
#ipdb.set_trace()
for box in target.bbox.numpy():
box = list(box)
box = box[:2] + [box[2], box[1]] + box[2:] + [box[0], box[3]
] + box[:2]
img_draw.line(box, fill=(255, 0, 0), width=2)
new.save('./vis/char_' + im_name)
return img, target, self.image_lists[item]
def __getitem__(self, item):
while True:
img_path = self.image_lists[item]
try:
img = Image.open(img_path).convert("RGB")
break
except BaseException:
item += 1
im_name = os.path.basename(img_path)
width, height = img.size
gt_path = self.gt_lists[item]
words, boxes, charsbbs, segmentations = self.load_gt_from_txt(
gt_path, height, width)
target = BoxList(boxes[:, :4],
img.size,
mode="xyxy",
use_char_ann=self.use_charann)
classes = torch.ones(len(boxes))
target.add_field("labels", classes)
masks = SegmentationMask(segmentations, img.size)
target.add_field("masks", masks)
if words[0] == "":
use_char_ann = False
else:
use_char_ann = True
if not self.use_charann:
use_char_ann = False
char_masks = SegmentationCharMask(charsbbs,
words=words,
use_char_ann=use_char_ann,
size=img.size,
char_num_classes=len(
self.char_classes))
target.add_field("char_masks", char_masks)
if self.transforms is not None:
img, target = self.transforms(img, target)
if self.vis:
new_im = img.numpy().copy().transpose([1, 2, 0]) + [
102.9801,
115.9465,
122.7717,
]
new_im = Image.fromarray(new_im.astype(np.uint8)).convert("RGB")
mask = target.extra_fields["masks"].polygons[0].convert("mask")
mask = Image.fromarray(
(mask.numpy() * 255).astype(np.uint8)).convert("RGB")
if self.use_charann:
m, _ = (target.extra_fields["char_masks"].chars_boxes[0].
convert("char_mask"))
color = self.creat_color_map(37, 255)
color_map = color[m.numpy().astype(np.uint8)]
char = Image.fromarray(color_map.astype(
np.uint8)).convert("RGB")
char = Image.blend(char, new_im, 0.5)
else:
char = new_im
new = Image.blend(char, mask, 0.5)
img_draw = ImageDraw.Draw(new)
for box in target.bbox.numpy():
box = list(box)
box = box[:2] + [box[2], box[1]] + box[2:] + [box[0], box[3]
] + box[:2]
img_draw.line(box, fill=(255, 0, 0), width=2)
new.save("./vis/char_" + im_name)
return img, target, self.image_lists[item]