def __getitem__(self, index):
# if _DEBUG:
# index = 0
# img_id = self.ids[index]
im_path = self.annobase[index]['image']# os.path.join(self.root, img_id + '.jpg')
img = Image.open(im_path).convert("RGB")
# im = cv2.imread(im_path)
anno = self.annobase[index]
target = RBoxList(torch.from_numpy(anno["boxes"]), (anno['width'], anno['height']), mode="xywha")
target.add_field("labels", torch.from_numpy(anno["gt_classes"]))
target.add_field("difficult", torch.Tensor([0 for i in range(len(anno["gt_classes"]))]))
target = target.clip_to_image(remove_empty=True)
# print('target:', target, im_path)
if self.transforms is not None:
# off = int(self.num_samples * np.random.rand())
# mix_index = (off + index) % self.num_samples
# img_mix = Image.open(self.annobase[mix_index]['image']).convert("RGB")
# img, target = self.mixup(img, img_mix, target)
img, target = self.transforms(img, target)
if _DEBUG:
if not target is None:
self.show_boxes(img, target)
return img, target, index
def __getitem__(self, index):
if _DEBUG:
index = 0
anno = self.annobase[index % self.database_num][
int(index / self.database_num) %
len(self.annobase[index % self.database_num])]
im_path = anno['image']
img = Image.open(im_path).convert("RGB")
# print('im_path:', im_path)
text, text_len = self.wk_converter.encode(anno['gt_words'])
text_label_split = []
off_cnt = 0
mx_len = np.max(text_len)
word_num = len(text_len)
for i in range(len(text_len)):
text_label_split.append(text[off_cnt:off_cnt + text_len[i]])
off_cnt += text_len[i]
padding_words = np.zeros((word_num, mx_len))
for i in range(word_num):
padding_words[i][:text_len[i]] = text_label_split[i]
if anno["boxes"].shape[0] > 0:
target = RBoxList(torch.from_numpy(anno["boxes"]),
(anno['width'], anno['height']),
mode="xywha")
target.add_field("labels", torch.from_numpy(anno["gt_classes"]))
target.add_field(
"difficult",
torch.tensor([0 for i in range(len(anno["gt_classes"]))]))
target.add_field("words", torch.from_numpy(padding_words))
target.add_field("word_length", torch.tensor(text_len))
target = target.clip_to_image(remove_empty=True)
else:
target = torch.from_numpy(padding_words)
if self.transforms is not None:
img, target = self.transforms(img, target)
if _DEBUG:
self.show_boxes(img, target)
return img, target, index
def forward_for_single_feature_map(self, anchors, objectness_,
box_regression_, scale):
"""
Arguments:
anchors: list[BoxList]
objectness: tensor of size N, A, H, W
box_regression: tensor of size N, A * 5, H, W
"""
device = objectness_.device
N, A, H, W = objectness_.shape
width, height = anchors[0].size
# scale = width / W
# put in the same format as anchors
objectness = objectness_.permute(0, 2, 3, 1)
objectness = objectness.reshape(N, -1)
# get the first 5 channels
box_regression = box_regression_[:, :5].view(N, -1, 5, H,
W).permute(0, 3, 4, 1, 2)
box_regression = box_regression.reshape(N, -1, 5)
all_proposals = eastbox2rbox(box_regression, self.base_size, (H, W),
scale)
num_anchors = A * H * W
pre_nms_top_n = min(self.pre_nms_top_n, num_anchors)
objectness, topk_idx = objectness.topk(pre_nms_top_n,
dim=1,
sorted=True)
batch_idx = torch.arange(N, device=device)[:, None]
proposals = all_proposals.view(N, -1, 5)[batch_idx, topk_idx]
image_shapes = [box.size for box in anchors]
result = []
for proposal, score, im_shape in zip(proposals, objectness,
image_shapes):
if not self.training:
# print("score:", score.shape)
# print("proposal:", proposal.shape)
proposal = proposal[score > self.score_thresh]
score = score[score > self.score_thresh]
# print("score:", score.shape, score)
# print("proposal:", proposal.shape)
# print("score:", score)
boxlist = RBoxList(proposal, im_shape, mode="xywha")
boxlist.add_field("objectness", score)
boxlist = boxlist.clip_to_image(remove_empty=False)
boxlist = remove_small_boxes(boxlist, self.min_size)
boxlist = self.nms_fn(
boxlist,
self.nms_thresh,
max_proposals=self.post_nms_top_n,
score_field="objectness",
)
result.append(boxlist)
return result