def on_batch_end(self):
super().on_batch_end()
# log cls & box loss. I'm using a very dirty way to do it by accessing
# protected attributes of loss class
self.cls_loss_meter.update(
utils.to_numpy(self.state.criterion._cls_loss))
self.box_loss_meter.update(
utils.to_numpy(self.state.criterion._box_loss))
# skip creating coco res for train
if self.state.is_train:
return
cls_out, box_out = self.state.output
res = decode(cls_out, box_out, self.anchors, **self.decode_params)
# rescale to image size. don't really need to clip after that
res[..., :4] *= self.batch_ratios.view(-1, 1, 1).to(res)
# xyxy -> xywh
res[..., 2:4] = res[..., 2:4] - res[..., :2]
for batch in range(res.size(0)):
for one_res in res[batch]:
if one_res[4].tolist(
) < 0.001: # stop when below this threshold, scores in descending order
break
coco_result = dict(
image_id=self.batch_img_ids[batch, 0].tolist(),
bbox=one_res[:4].tolist(),
score=one_res[4].tolist(),
category_id=int(one_res[5].tolist()),
)
self.all_results.append(coco_result)
def predict(self, x):
"""Run forward on given images and decode raw prediction into bboxes
Returns: bboxes, scores, classes
"""
class_outputs, box_outputs = self.forward(x)
anchors = box_utils.generate_anchors_boxes(x.shape[-2:])[0]
return box_utils.decode(class_outputs, box_outputs, anchors)
def predict(self, x):
"""
Run forward on given images and decode raw prediction into bboxes
Returns:
torch.Tensor with bboxes, scores and classes. bboxes in `lrtb` format
shape [BS, MAX_DETECTION_PER_IMAGE, 6]
"""
class_outputs, box_outputs = self.forward(x)
anchors = box_utils.generate_anchors_boxes(x.shape[-2:])[0]
return box_utils.decode(class_outputs, box_outputs, anchors)
def on_batch_end(self):
if self.state.is_train:
return
cls_out, box_out = self.state.output
res = decode(cls_out, box_out,
self.anchors) # TODO: return **self.decode_params
# rescale to image size. don't really need to clip after that
res[..., :4] *= self.batch_ratios.view(-1, 1, 1).to(res)
# xyxy -> xywh
res[..., 2:4] = res[..., 2:4] - res[..., :2]
for batch in range(res.size(0)):
for one_res in res[batch]:
self.all_img_ids.append(self.batch_img_ids[batch, 0].tolist())
coco_result = dict(
image_id=self.batch_img_ids[batch, 0].tolist(),
bbox=one_res[:4].tolist(),
score=one_res[4].tolist(),
category_id=int(one_res[5].tolist()),
)
self.all_results.append(coco_result)