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]