def get_ratio_by_est(self, boxlist, is_train):
return_boxlist = []
device = boxlist[0].bbox.device
for target in boxlist:
target_bbox = target.bbox
n = target_bbox.shape[0]
img_size = target.size
regvals = target.get_field("reg_vals")
reg_vals = est_decode(regvals)
new_bbox = []
for k in range(n):
p_bbox = target_bbox[k, :]
h = p_bbox[3] - p_bbox[1]
new_h = h * (1. / (1. - reg_vals[k]))
p_bbox[3] = p_bbox[1] + new_h
new_bbox.append(p_bbox.tolist())
new_bboxlist = BoxList(new_bbox, img_size, mode="xyxy")
new_bboxlist._copy_extra_fields(target)
new_bboxlist.add_field("pad_ratio", reg_vals)
return_boxlist.append(new_bboxlist)
return_boxlist = [
return_box.to(device) for return_box in return_boxlist
]
return return_boxlist
def __call__(self, image, target=None):
if target is None:
return image
if not self.do:
return image, target
keypoints = target.get_field("keypoints")
polygons = list(
map(lambda x: x.polygons[0].numpy(), keypoints.instances.polygons))
# polygons = np.stack( polygons, axis=0 ).reshape( (-1, 4, 2) )
# rrects = list( map( self._to_rrect, polygons ) )
rrects_np = np.array(rrects, dtype=np.float32).reshape((-1, 8))
xmins = np.min(rrects_np[:, ::2], axis=1)
ymins = np.min(rrects_np[:, 1::2], axis=1)
xmaxs = np.max(rrects_np[:, ::2], axis=1)
ymaxs = np.max(rrects_np[:, 1::2], axis=1)
xyxy = np.vstack([xmins, ymins, xmaxs, ymaxs]).transpose()
boxes = torch.from_numpy(xyxy).reshape(-1, 4) # guard against no boxes
new_target = BoxList(boxes, image.size, mode="xyxy")
new_target._copy_extra_fields(target)
new_keypoints = SegmentationMask(rrects_np.reshape(
(-1, 1, 8)).tolist(),
image.size,
mode='poly')
new_target.add_field("keypoints", new_keypoints)
return image, new_target
def exchange_box(self, boxlists1, boxlists2):
result = []
for boxlist1, boxlist2 in zip(boxlists1, boxlists2):
result_box = BoxList(boxlist1.bbox, boxlist1.size, mode="xyxy")
result_box._copy_extra_fields(boxlist2)
result.append(result_box)
return result
def __call__(self, image, target=None):
if target == None:
return image
if not self.do:
return image, target
w, h = image.size
if w != h:
return image, target
assert w == h
cx = w / 2
cy = cx
degree = random.uniform(0, 360)
radian = degree * math.pi / 180
new_image = image.rotate(-degree)
sin = math.sin(radian)
cos = math.cos(radian)
keypoints = target.get_field("keypoints")
polygons = list(
map(lambda x: x.polygons[0], keypoints.instances.polygons))
polygons = torch.stack(polygons, 0).reshape((-1, 2)).t()
M = torch.Tensor([[cos, -sin], [sin, cos]])
b = torch.Tensor([[(1 - cos) * cx + cy * sin],
[(1 - cos) * cy - cx * sin]])
new_points = M.mm(polygons) + b
new_points = new_points.t().reshape((-1, 8))
xmins, _ = torch.min(new_points[:, ::2], 1)
minx_idx = xmins < 1
xmins[minx_idx] = 1
ymins, _ = torch.min(new_points[:, 1::2], 1)
miny_idx = ymins < 1
ymins[miny_idx] = 1
xmaxs, _ = torch.max(new_points[:, ::2], 1)
xmax_idx = xmaxs > 1024
xmaxs[xmax_idx] = 1024
ymaxs, _ = torch.max(new_points[:, 1::2], 1)
ymax_idx = ymaxs > 1024
ymaxs[ymax_idx] = 1024
boxes = torch.stack([xmins, ymins, xmaxs, ymaxs], 1).reshape((-1, 4))
new_target = BoxList(boxes, image.size, mode="xyxy")
new_target._copy_extra_fields(target)
new_keypoints = SegmentationMask(new_points.reshape(
(-1, 1, 8)).tolist(),
image.size,
mode='poly')
new_target.add_field("keypoints", new_keypoints)
return new_image, new_target
def generate_right_proposals(self, boxes):
boxes_right = []
for boxes_per_image in boxes:
if boxes_per_image.mode != "xyxy": boxes_per_image.convert("xyxy")
bbox_right = boxes_per_image.bbox.clone()
disp_unity = boxes_per_image.get_field("depths")
disparity = disp_unity * boxes_per_image.size[0] * 0.209313 # TODO: embed them in ground truth
# print(bbox_right)
bbox_right[:,0::4] -= disparity
bbox_right[:,2::4] -= disparity
bbox_right = BoxList(bbox_right, boxes_per_image.size, mode=boxes_per_image.mode)
bbox_right._copy_extra_fields(boxes_per_image)
boxes_right.append(bbox_right)
return boxes_right
def get_ratio(self, boxlist, is_train):
"""
boxlist: [Bbox, Bbox, ...]
"""
"""
for those without keypoints:
global, not partition
"""
return_boxlist = []
device = boxlist[0].bbox.device
for target in boxlist:
target_bbox = target.bbox
keypoint = target.get_field("keypoints")
kp = keypoint.keypoints
n, _, _ = kp.shape
bbox = target.bbox
img_size = target.size
new_bbox = []
new_pad = []
for k in range(n):
p_kp = kp[k]
pad = 1.0 if is_train else 0.
if p_kp.sum().item() > 0:
pad = 0.
for iteration, i in enumerate(self._idx[::-1][:-1]):
# assume thorax exists
vis = False
store_y = None
for j in i:
if p_kp[j][2] > self.INVIS_THRSH:
vis = True
if vis:
store_y = max(p_kp[i[0]][1], p_kp[i[1]][1])
break
if not vis:
# hips, knees, ankles not visible
pad += sum(self._pratio[2:])
res = F.relu(target_bbox[k, 3]-p_kp[self.thrx_idx, 1])
known = F.relu(p_kp[self.thrx_idx, 1]-target_bbox[k, 1])
tmp = F.relu((self.r_thrx2hip/self.r_head2thrx)*known-res) # pixel
pad += (self.r_thrx2hip*tmp/(tmp+res)).item()
if p_kp[self.thrx_idx, 1].item() == 0:
pad = 1.0
elif iteration == 0:
pad = 0.
else:
pad += sum(self._pratio[::-1][:iteration])
res = F.relu(target_bbox[k, 3]-store_y)
known = F.relu(p_kp[self.thrx_idx, 1]-target_bbox[k, 1])
tmp = F.relu((self._pratio[::-1][iteration]/self.r_head2thrx)*known-res)
pad += (self._pratio[::-1][iteration]*tmp/(tmp+res)).item()
if p_kp[self.thrx_idx, 1].item() == 0:
pad = 1.0
p_bbox = 1.*bbox[k, :]
h = p_bbox[3] - p_bbox[1]
if pad == 1.0:
new_h = h
if not is_train:
pad = 0.
p_bbox[3] = p_bbox[1] + new_h
new_bbox.append(p_bbox.tolist())
new_pad.append(pad)
else:
if not is_train:
curr_aug_per = 1
else:
curr_aug_per = self.aug_per + 1
p_bbox_repeat = p_bbox.repeat(curr_aug_per, 1)
for ap in range(curr_aug_per):
if ap == 0:
random_cut = 0.
else:
random_cut = self.rand_cut*random.random()
# 0-0.3
update_pad = pad + (1.-pad)*random_cut
p_bbox_repeat_ = p_bbox_repeat[ap]
new_h = h*(1./(1.-update_pad))
p_bbox_repeat_[3] = p_bbox_repeat_[1] + new_h
new_bbox.append(p_bbox_repeat_.tolist())
new_pad.append(update_pad)
new_bboxlist = BoxList(new_bbox, img_size, mode="xyxy")
new_bboxlist._copy_extra_fields(target)
new_bboxlist.add_field("pad_ratio", torch.tensor(new_pad))
return_boxlist.append(new_bboxlist)
return_boxlist = [return_box.to(device) for return_box in return_boxlist]
return return_boxlist
def forward(self, features_, proposals, targets=None, query=False):
"""
Arguments:
features (list[Tensor]): feature-maps from possibly several levels
proposals (list[BoxList]): proposal boxes
targets (list[BoxList], optional): the ground-truth targets.
Returns:
x (Tensor): the result of the feature extractor
proposals (list[BoxList]): during training, the subsampled proposals
are returned. During testing, the predicted boxlists are returned
losses (dict[Tensor]): During training, returns the losses for the
head. During testing, returns an empty dict.
"""
tmpp = [feature.shape[2:] for feature in features_]
features = [
F.pad(feature, (0, 0, 0, size[0]), "constant", value=0.0)
for size, feature in zip(tmpp, features_)
]
if self.training:
targets = self.ratio_estimator.get_ratio(targets, self.training)
proposals = targets
elif query:
for target, proposal in zip(targets, proposals):
target.add_field("embeds", proposal.get_field("embeds"))
if self.padreg:
for target, proposal in zip(targets, proposals):
target.add_field("reg_vals",
proposal.get_field("reg_vals"))
proposals = targets
query_get_ratio = self.ratio_estimator.get_ratio \
if self.query_by_gt else self.ratio_estimator.get_ratio_by_est
proposals = query_get_ratio(proposals, self.training)
else:
proposals = targets
elif self.padreg:
old_proposals = proposals
device_ = proposals[0].bbox.device
return_bboxlist = []
for proposal in proposals:
p_bbox = 1.0 * proposal.bbox
new_bbox = []
n_proposal = p_bbox.shape[0]
regvals = proposal.get_field("reg_vals")
reg_vals = est_decode(regvals)
# reg_valss = (reg_vals+1)/2
for j in range(n_proposal):
bbox = p_bbox[j, :]
h = bbox[3] - bbox[1]
new_h = h * (1. / (1. - reg_vals[j]))
bbox[3] = bbox[1] + new_h
new_bbox.append(bbox.tolist())
if n_proposal == 0:
new_bbox = torch.tensor([]).view(0, 4)
new_bboxlist = BoxList(new_bbox, proposal.size, mode="xyxy")
new_bboxlist._copy_extra_fields(proposal)
return_bboxlist.append(new_bboxlist)
return_bboxlist = [
return_box.to(device_) for return_box in return_bboxlist
]
proposals = return_bboxlist
else:
pass
# keep the proposals
x = self.feature_extractor(features, proposals)
# final classifier that converts the features into predictions
part_feat = self.predictor(x)
if not self.training:
# when no training
# for query, proposals are ground truth
# for gallery, proposals are results, just add part_embeds on it
if not query and self.padreg:
proposals = self.exchange_box(old_proposals, proposals)
result = self.post_processor(part_feat, proposals)
return x, result, {}
loss_part_oim = self.loss_evaluator(part_feat, targets)
loss_dict = dict(
zip([
"loss_reid_p" + str(i) for i in range(1,
len(loss_part_oim) + 1)
], loss_part_oim))
return (
x,
proposals,
loss_dict,
)
