def __getitem__(self, index):
img_id = self.images[index]
img_path = os.path.join(
self.img_dir,
self.coco.loadImgs(ids=[img_id])[0]['file_name'])
ann_ids = self.coco.getAnnIds(imgIds=[img_id])
annotations = self.coco.loadAnns(ids=ann_ids)
img = self.coco.loadImgs(ids=[img_id])[0]
w_img = int(img['width'])
h_img = int(img['height'])
labels = []
bboxes = []
shapes = []
for anno in annotations:
if anno['iscrowd'] == 1: # Excludes crowd objects
continue
polygons = get_connected_polygon_using_mask(
anno['segmentation'], (h_img, w_img),
n_vertices=self.n_vertices,
closing_max_kernel=50)
gt_x1, gt_y1, gt_w, gt_h = anno['bbox']
contour = np.array(polygons).reshape((-1, 2))
# Downsample the contour to fix number of vertices
if len(contour) > self.n_vertices:
fixed_contour = resample(contour, num=self.n_vertices)
else:
fixed_contour = turning_angle_resample(contour,
self.n_vertices)
fixed_contour[:, 0] = np.clip(fixed_contour[:, 0], gt_x1,
gt_x1 + gt_w)
fixed_contour[:, 1] = np.clip(fixed_contour[:, 1], gt_y1,
gt_y1 + gt_h)
contour_std = np.sqrt(np.sum(np.std(fixed_contour, axis=0)**2))
if contour_std < 1e-6 or contour_std == np.inf or contour_std == np.nan: # invalid shapes
continue
updated_bbox = [
np.min(fixed_contour[:, 0]),
np.min(fixed_contour[:, 1]),
np.max(fixed_contour[:, 0]),
np.max(fixed_contour[:, 1])
]
shapes.append(np.ndarray.flatten(fixed_contour).tolist())
labels.append(self.cat_ids[anno['category_id']])
bboxes.append(updated_bbox)
labels = np.array(labels)
bboxes = np.array(bboxes, dtype=np.float32)
shapes = np.array(shapes, dtype=np.float32)
if len(bboxes) == 0:
bboxes = np.array([[0., 0., 0., 0.]], dtype=np.float32)
labels = np.array([[0]])
shapes = np.zeros((1, self.n_vertices * 2), dtype=np.float32)
# bboxes[:, 2:] += bboxes[:, :2] # xywh to xyxy
img = cv2.imread(img_path)
height, width = img.shape[0], img.shape[1]
center = np.array([width / 2., height / 2.],
dtype=np.float32) # center of image
scale = max(height, width) * 1.0
flipped = False
if self.split == 'train':
scale = scale * np.random.choice(self.rand_scales)
w_border = get_border(160, width)
h_border = get_border(160, height)
center[0] = np.random.randint(low=w_border, high=width - w_border)
center[1] = np.random.randint(low=h_border, high=height - h_border)
if np.random.random() < 0.5:
flipped = True
img = img[:, ::-1, :]
center[0] = width - center[0] - 1
trans_img = get_affine_transform(
center, scale, 0, [self.img_size['w'], self.img_size['h']])
img = cv2.warpAffine(img, trans_img,
(self.img_size['w'], self.img_size['h']))
# -----------------------------------debug---------------------------------
# image_show = img.copy()
# for bbox, label in zip(bboxes, labels):
# if flipped:
# bbox[[0, 2]] = width - bbox[[2, 0]] - 1
# bbox[:2] = affine_transform(bbox[:2], trans_img)
# bbox[2:] = affine_transform(bbox[2:], trans_img)
# bbox[[0, 2]] = np.clip(bbox[[0, 2]], 0, self.img_size['w'] - 1)
# bbox[[1, 3]] = np.clip(bbox[[1, 3]], 0, self.img_size['h'] - 1)
# cv2.rectangle(image_show, (int(bbox[0]), int(bbox[1])), (int(bbox[2]), int(bbox[3])), (255, 0, 0), 2)
# cv2.putText(image_show, self.class_name[label + 1], (int(bbox[0]), int(bbox[1])),
# cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
# cv2.imshow('img', image_show)
# cv2.waitKey()
# -----------------------------------debug---------------------------------
img = img.astype(np.float32) / 255.
if self.split == 'train':
color_aug(self.data_rng, img, self.eig_val, self.eig_vec)
img -= self.mean
img /= self.std
img = img.transpose(2, 0, 1) # from [H, W, C] to [C, H, W]
trans_fmap = get_affine_transform(
center, scale, 0, [self.fmap_size['w'], self.fmap_size['h']])
hmap = np.zeros(
(self.num_classes, self.fmap_size['h'], self.fmap_size['w']),
dtype=np.float32) # heatmap
w_h_ = np.zeros((self.max_objs, 2),
dtype=np.float32) # width and height of bboxes
shapes_ = np.zeros((self.max_objs, self.n_vertices * 2),
dtype=np.float32) # gt amodal segmentation polygons
center_offsets = np.zeros(
(self.max_objs, 2),
dtype=np.float32) # gt mass centers to bbox center
codes_ = np.zeros((self.max_objs, self.n_codes), dtype=np.float32)
contour_std_ = np.zeros(
(self.max_objs, 1),
dtype=np.float32) # keep track of codes that is activated
regs = np.zeros(
(self.max_objs, 2),
dtype=np.float32) # regression for offsets of shape center
inds = np.zeros((self.max_objs, ), dtype=np.int64)
ind_masks = np.zeros((self.max_objs, ), dtype=np.uint8)
for k, (bbox, label, shape) in enumerate(zip(bboxes, labels, shapes)):
if flipped:
bbox[[0, 2]] = width - bbox[[2, 0]] - 1
# Flip the contour x-axis
for m in range(self.n_vertices):
shape[2 * m] = width - shape[2 * m] - 1
bbox[:2] = affine_transform(bbox[:2], trans_fmap)
bbox[2:] = affine_transform(bbox[2:], trans_fmap)
bbox[[0, 2]] = np.clip(bbox[[0, 2]], 0, self.fmap_size['w'] - 1)
bbox[[1, 3]] = np.clip(bbox[[1, 3]], 0, self.fmap_size['h'] - 1)
h, w = bbox[3] - bbox[1], bbox[2] - bbox[0]
# generate gt shape mean and std from contours
for m in range(self.n_vertices
): # apply scale and crop transform to shapes
shape[2 * m:2 * m + 2] = affine_transform(
shape[2 * m:2 * m + 2], trans_fmap)
shape_clipped = np.reshape(shape, (self.n_vertices, 2))
shape_clipped[:, 0] = np.clip(shape_clipped[:, 0], 0,
self.fmap_size['w'] - 1)
shape_clipped[:, 1] = np.clip(shape_clipped[:, 1], 0,
self.fmap_size['h'] - 1)
clockwise_flag = check_clockwise_polygon(shape_clipped)
if not clockwise_flag:
fixed_contour = np.flip(shape_clipped, axis=0)
else:
fixed_contour = shape_clipped.copy()
# Indexing from the left-most vertex, argmin x-axis
idx = np.argmin(fixed_contour[:, 0])
indexed_shape = np.concatenate(
(fixed_contour[idx:, :], fixed_contour[:idx, :]), axis=0)
mass_center = np.mean(indexed_shape, axis=0)
contour_std = np.std(indexed_shape, axis=0) + 1e-4
if h < 1e-6 or w < 1e-6: # remove small bboxes
continue
# centered_shape = indexed_shape - mass_center
norm_shape = (indexed_shape - mass_center) / np.sqrt(
np.sum(contour_std**2))
if h > 0 and w > 0:
obj_c = np.array([(bbox[0] + bbox[2]) / 2,
(bbox[1] + bbox[3]) / 2],
dtype=np.float32)
# obj_c = mass_center
obj_c_int = obj_c.astype(np.int32)
radius = max(
0,
int(
gaussian_radius((math.ceil(h), math.ceil(w)),
self.gaussian_iou)))
draw_umich_gaussian(hmap[label], obj_c_int, radius)
shapes_[k] = norm_shape.reshape((1, -1))
center_offsets[k] = mass_center - obj_c
codes_[k], _ = fast_ista(norm_shape.reshape((1, -1)),
self.dictionary,
lmbda=self.sparse_alpha,
max_iter=60)
contour_std_[k] = np.sqrt(np.sum(contour_std**2))
w_h_[k] = 1. * w, 1. * h
# w_h_[k] = mass_center[1] - bbox[1], bbox[3] - mass_center[1], \
# mass_center[0] - bbox[0], bbox[2] - mass_center[0] # [top, bottom, left, right] distance
regs[k] = obj_c - obj_c_int # discretization error
inds[k] = obj_c_int[1] * self.fmap_size['w'] + obj_c_int[0]
ind_masks[k] = 1
return {
'image': img,
'shapes': shapes_,
'codes': codes_,
'offsets': center_offsets,
'std': contour_std_,
'hmap': hmap,
'w_h_': w_h_,
'regs': regs,
'inds': inds,
'ind_masks': ind_masks,
'c': center,
's': scale,
'img_id': img_id
}
def main():
cfg.device = torch.device('cuda')
torch.backends.cudnn.benchmark = False
max_per_image = 100
num_classes = 80 if cfg.dataset == 'coco' else 4
dictionary = np.load(cfg.dictionary_file)
colors = COCO_COLORS if cfg.dataset == 'coco' else DETRAC_COLORS
names = COCO_NAMES if cfg.dataset == 'coco' else DETRAC_NAMES
for j in range(len(names)):
col_ = [c * 255 for c in colors[j]]
colors[j] = tuple(col_)
print('Creating model and recover from checkpoint ...')
if 'hourglass' in cfg.arch:
model = exkp(n=5,
nstack=2,
dims=[256, 256, 384, 384, 384, 512],
modules=[2, 2, 2, 2, 2, 4],
num_classes=num_classes)
else:
model = get_pose_net(num_layers=int(cfg.arch.split('_')[-1]),
num_classes=80)
# raise NotImplementedError
model = load_demo_model(model, cfg.ckpt_dir)
model = model.to(cfg.device)
model.eval()
# Loading COCO validation images
annotation_file = '{}/annotations/instances_{}.json'.format(
cfg.data_dir, cfg.data_type)
coco = COCO(annotation_file)
# Load all annotations
cats = coco.loadCats(coco.getCatIds())
nms = [cat['name'] for cat in cats]
catIds = coco.getCatIds(catNms=nms)
# imgIds = coco.getImgIds(catIds=catIds)
imgIds = coco.getImgIds()
# annIds = coco.getAnnIds(catIds=catIds)
# all_anns = coco.loadAnns(ids=annIds)
# print(len(imgIds), imgIds)
for id in imgIds:
annt_ids = coco.getAnnIds(imgIds=[id])
annotations_per_img = coco.loadAnns(ids=annt_ids)
# print('All annots: ', len(annotations_per_img), annotations_per_img)
img = coco.loadImgs(id)[0]
image_path = '%s/images/%s/%s' % (cfg.data_dir, cfg.data_type,
img['file_name'])
w_img = int(img['width'])
h_img = int(img['height'])
if w_img < 1 or h_img < 1:
continue
img_original = cv2.imread(image_path)
img_connect = cv2.imread(image_path)
img_recon = cv2.imread(image_path)
print('Image id: ', id)
for annt in annotations_per_img:
if annt['iscrowd'] == 1 or type(annt['segmentation']) != list:
continue
polygons = get_connected_polygon_using_mask(
annt['segmentation'], (h_img, w_img),
n_vertices=cfg.num_vertices,
closing_max_kernel=60)
gt_bbox = annt['bbox']
gt_x1, gt_y1, gt_w, gt_h = gt_bbox
contour = np.array(polygons).reshape((-1, 2))
# Downsample the contour to fix number of vertices
if len(contour) > cfg.num_vertices:
resampled_contour = resample(contour, num=cfg.num_vertices)
else:
resampled_contour = turning_angle_resample(
contour, cfg.num_vertices)
resampled_contour[:, 0] = np.clip(resampled_contour[:, 0], gt_x1,
gt_x1 + gt_w)
resampled_contour[:, 1] = np.clip(resampled_contour[:, 1], gt_y1,
gt_y1 + gt_h)
clockwise_flag = check_clockwise_polygon(resampled_contour)
if not clockwise_flag:
fixed_contour = np.flip(resampled_contour, axis=0)
else:
fixed_contour = resampled_contour.copy()
# Indexing from the left-most vertex, argmin x-axis
idx = np.argmin(fixed_contour[:, 0])
indexed_shape = np.concatenate(
(fixed_contour[idx:, :], fixed_contour[:idx, :]), axis=0)
x1, y1, x2, y2 = gt_x1, gt_y1, gt_x1 + gt_w, gt_y1 + gt_h
# bbox_width, bbox_height = x2 - x1, y2 - y1
# bbox = [x1, y1, bbox_width, bbox_height]
# bbox_center = np.array([(x1 + x2) / 2., (y1 + y2) / 2.])
bbox_center = np.mean(indexed_shape, axis=0)
centered_shape = indexed_shape - bbox_center
# visualize resampled points with multiple parts in image side by side
for cnt in range(len(annt['segmentation'])):
polys = np.array(annt['segmentation'][cnt]).reshape((-1, 2))
cv2.polylines(img_original, [polys.astype(np.int32)],
True, (10, 10, 255),
thickness=2)
# cv2.drawContours(img_original, [polys.astype(np.int32)], contourIdx=-1, color=(10, 10, 255), thickness=-1)
cv2.polylines(img_connect, [indexed_shape.astype(np.int32)],
True, (10, 10, 255),
thickness=2)
# cv2.drawContours(img_connect, [indexed_shape.astype(np.int32)], contourIdx=-1, color=(10, 10, 255), thickness=-1)
learned_val_codes, _ = fast_ista(centered_shape.reshape((1, -1)),
dictionary,
lmbda=0.1,
max_iter=60)
recon_contour = np.matmul(learned_val_codes, dictionary).reshape(
(-1, 2))
recon_contour = recon_contour + bbox_center
cv2.polylines(img_recon, [recon_contour.astype(np.int32)],
True, (10, 10, 255),
thickness=2)
# cv2.drawContours(img_recon, [recon_contour.astype(np.int32)], contourIdx=-1, color=(10, 10, 255), thickness=-1)
# plot gt mean and std
# image = cv2.imread(image_path)
# # cv2.ellipse(image, center=(int(contour_mean[0]), int(contour_mean[1])),
# # axes=(int(contour_std[0]), int(contour_std[1])),
# # angle=0, startAngle=0, endAngle=360, color=(0, 255, 0),
# # thickness=2)
# cv2.rectangle(image, pt1=(int(contour_mean[0] - contour_std[0] / 2.), int(contour_mean[1] - contour_std[1] / 2.)),
# pt2=(int(contour_mean[0] + contour_std[0] / 2.), int(contour_mean[1] + contour_std[1] / 2.)),
# color=(0, 255, 0), thickness=2)
# cv2.polylines(image, [fixed_contour.astype(np.int32)], True, (0, 0, 255))
# cv2.rectangle(image, pt1=(int(min(fixed_contour[:, 0])), int(min(fixed_contour[:, 1]))),
# pt2=(int(max(fixed_contour[:, 0])), int(max(fixed_contour[:, 1]))),
# color=(255, 0, 0), thickness=2)
# cv2.imshow('GT segments', image)
# if cv2.waitKey() & 0xFF == ord('q'):
# break
image = cv2.imread(image_path)
original_image = image.copy()
height, width = image.shape[0:2]
padding = 127 if 'hourglass' in cfg.arch else 31
imgs = {}
for scale in cfg.test_scales:
new_height = int(height * scale)
new_width = int(width * scale)
if cfg.img_size > 0:
img_height, img_width = cfg.img_size, cfg.img_size
center = np.array([new_width / 2., new_height / 2.],
dtype=np.float32)
scaled_size = max(height, width) * 1.0
scaled_size = np.array([scaled_size, scaled_size],
dtype=np.float32)
else:
img_height = (new_height | padding) + 1
img_width = (new_width | padding) + 1
center = np.array([new_width // 2, new_height // 2],
dtype=np.float32)
scaled_size = np.array([img_width, img_height],
dtype=np.float32)
img = cv2.resize(image, (new_width, new_height))
trans_img = get_affine_transform(center, scaled_size, 0,
[img_width, img_height])
img = cv2.warpAffine(img, trans_img, (img_width, img_height))
img = img.astype(np.float32) / 255.
img -= np.array(
COCO_MEAN if cfg.dataset == 'coco' else DETRAC_MEAN,
dtype=np.float32)[None, None, :]
img /= np.array(COCO_STD if cfg.dataset == 'coco' else DETRAC_STD,
dtype=np.float32)[None, None, :]
img = img.transpose(
2, 0, 1)[None, :, :, :] # from [H, W, C] to [1, C, H, W]
# if cfg.test_flip:
# img = np.concatenate((img, img[:, :, :, ::-1].copy()), axis=0)
imgs[scale] = {
'image': torch.from_numpy(img).float(),
'center': np.array(center),
'scale': np.array(scaled_size),
'fmap_h': np.array(img_height // 4),
'fmap_w': np.array(img_width // 4)
}
with torch.no_grad():
segmentations = []
predicted_codes = []
start_time = time.time()
for scale in imgs:
imgs[scale]['image'] = imgs[scale]['image'].to(cfg.device)
output = model(imgs[scale]['image'])[-1]
# segms, codes_ = ctsegm_scaled_decode_debug(*output, torch.from_numpy(dictionary.astype(np.float32)).to(cfg.device),
# K=cfg.test_topk)
segms = ctsegm_code_n_offset_decode(
*output,
torch.from_numpy(dictionary.astype(np.float32)).to(
cfg.device),
K=cfg.test_topk)
segms = segms.detach().cpu().numpy().reshape(
1, -1, segms.shape[2])[0]
# codes_ = codes_.detach().cpu().numpy().reshape(1, -1, codes_.shape[2])[0]
top_preds = {}
code_preds = {}
for j in range(cfg.num_vertices):
segms[:, 2 * j:2 * j + 2] = transform_preds(
segms[:, 2 * j:2 * j + 2], imgs[scale]['center'],
imgs[scale]['scale'],
(imgs[scale]['fmap_w'], imgs[scale]['fmap_h']))
segms[:, cfg.num_vertices * 2:cfg.num_vertices * 2 +
2] = transform_preds(
segms[:,
cfg.num_vertices * 2:cfg.num_vertices * 2 + 2],
imgs[scale]['center'], imgs[scale]['scale'],
(imgs[scale]['fmap_w'], imgs[scale]['fmap_h']))
segms[:, cfg.num_vertices * 2 + 2:cfg.num_vertices * 2 +
4] = transform_preds(
segms[:, cfg.num_vertices * 2 +
2:cfg.num_vertices * 2 + 4],
imgs[scale]['center'], imgs[scale]['scale'],
(imgs[scale]['fmap_w'], imgs[scale]['fmap_h']))
clses = segms[:, -1]
for j in range(num_classes):
inds = (clses == j)
top_preds[j + 1] = segms[inds, :cfg.num_vertices * 2 +
5].astype(np.float32)
top_preds[j + 1][:, :cfg.num_vertices * 2 + 4] /= scale
# code_preds[j + 1] = codes_[inds, :]
segmentations.append(top_preds)
predicted_codes.append(code_preds)
segms_and_scores = {
j: np.concatenate([d[j] for d in segmentations], axis=0)
for j in range(1, num_classes + 1)
} # a Dict label: segments
# codes_and_scores = {j: np.concatenate([d[j] for d in predicted_codes], axis=0)
# for j in range(1, num_classes + 1)} # a Dict label: segments
scores = np.hstack([
segms_and_scores[j][:, cfg.num_vertices * 2 + 4]
for j in range(1, num_classes + 1)
])
if len(scores) > max_per_image:
kth = len(scores) - max_per_image
thresh = np.partition(scores, kth)[kth]
for j in range(1, num_classes + 1):
keep_inds = (segms_and_scores[j][:, cfg.num_vertices * 2 +
4] >= thresh)
segms_and_scores[j] = segms_and_scores[j][keep_inds]
# codes_and_scores[j] = codes_and_scores[j][keep_inds]
# Use opencv functions to output a video
output_image = original_image
blend_mask = np.zeros(shape=output_image.shape, dtype=np.uint8)
# print(blend_mask.shape)
for lab in segms_and_scores:
for idx in range(len(segms_and_scores[lab])):
res = segms_and_scores[lab][idx]
# c_ = codes_and_scores[lab][idx]
# for res in segms_and_scores[lab]:
contour, bbox, score = res[:-5], res[-5:-1], res[-1]
bbox[0] = np.clip(bbox[0], 0, w_img)
bbox[1] = np.clip(bbox[1], 0, h_img)
bbox[2] = np.clip(bbox[2], 0, w_img)
bbox[3] = np.clip(bbox[3], 0, h_img)
if score > cfg.detect_thres:
text = names[lab] # + ' %.2f' % score
# label_size = cv2.getTextSize(text, cv2.FONT_HERSHEY_COMPLEX, thickness=2, fontScale=0.5)
polygon = contour.reshape((-1, 2))
# print('Shape: Poly -- ', polygon.shape)
# print(polygon)
polygon[:, 0] = np.clip(polygon[:, 0], 0, w_img - 1)
polygon[:, 1] = np.clip(polygon[:, 1], 0, h_img - 1)
# use bb tools to draw predictions
color = random.choice(COLOR_WORLD)
bb.add(output_image, bbox[0], bbox[1], bbox[2],
bbox[3], text, color)
cv2.polylines(output_image, [polygon.astype(np.int32)],
True,
RGB_DICT[color],
thickness=1)
cv2.drawContours(blend_mask,
[polygon.astype(np.int32)],
contourIdx=-1,
color=RGB_DICT[color],
thickness=-1)
# color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
# contour_mean = np.mean(polygon, axis=0)
# contour_std = np.std(polygon, axis=0)
# center_x, center_y = np.mean(polygon, axis=0).astype(np.int32)
# text_location = [bbox[0] + 1, bbox[1] + 1,
# bbox[1] + label_size[0][0] + 1,
# bbox[0] + label_size[0][1] + 1]
# cv2.rectangle(output_image, pt1=(int(bbox[0]), int(bbox[1])),
# pt2=(int(bbox[2]), int(bbox[3])),
# color=color, thickness=1)
# cv2.rectangle(output_image, pt1=(int(np.min(polygon[:, 0])), int(np.min(polygon[:, 1]))),
# pt2=(int(np.max(polygon[:, 0])), int(np.max(polygon[:, 1]))),
# color=(0, 255, 0), thickness=1)
# cv2.polylines(output_image, [polygon.astype(np.int32)], True, color, thickness=2)
# cv2.putText(output_image, text, org=(int(text_location[0]), int(text_location[3])),
# fontFace=cv2.FONT_HERSHEY_COMPLEX, thickness=2, fontScale=0.5,
# color=(255, 0, 0))
# cv2.putText(output_image, text, org=(int(bbox[0]), int(bbox[1])),
# fontFace=cv2.FONT_HERSHEY_COMPLEX, thickness=1, fontScale=0.5,
# color=color)
# show the histgram for predicted codes
# fig = plt.figure()
# plt.plot(np.arange(cfg.n_codes), c_.reshape((-1,)), color='green',
# marker='o', linestyle='dashed', linewidth=2, markersize=6)
# plt.ylabel('Value of each coefficient')
# plt.xlabel('All predicted {} coefficients'.format(cfg.n_codes))
# plt.title('Distribution of the predicted coefficients for {}'.format(text))
# plt.show()
value = [255, 255, 255]
dst_img = cv2.addWeighted(output_image, 0.5, blend_mask, 0.5, 0)
dst_img[blend_mask == 0] = output_image[blend_mask == 0]
img_original = cv2.copyMakeBorder(img_original, 0, 0, 0, 10,
cv2.BORDER_CONSTANT, None, value)
img_connect = cv2.copyMakeBorder(img_connect, 0, 0, 10, 10,
cv2.BORDER_CONSTANT, None, value)
img_recon = cv2.copyMakeBorder(img_recon, 0, 0, 10, 10,
cv2.BORDER_CONSTANT, None, value)
dst_img = cv2.copyMakeBorder(dst_img, 0, 0, 10, 0,
cv2.BORDER_CONSTANT, None, value)
im_cat = np.concatenate(
(img_original, img_connect, img_recon, dst_img), axis=1)
# im_cat = np.concatenate((img_original, img_connect, img_recon), axis=1)
cv2.imshow('GT:Resample:Recons:Predict', im_cat)
if cv2.waitKey() & 0xFF == ord('q'):
break
def __getitem__(self, index):
img_id = self.images[index]
img_path = os.path.join(
self.img_dir,
self.coco.loadImgs(ids=[img_id])[0]['file_name'])
ann_ids = self.coco.getAnnIds(imgIds=[img_id])
annotations = self.coco.loadAnns(ids=ann_ids)
img = self.coco.loadImgs(ids=[img_id])[0]
w_img = int(img['width'])
h_img = int(img['height'])
labels = []
bboxes = []
a_bboxes = []
shapes = []
a_shapes = []
for anno in annotations:
if anno['category_id'] not in KINS_IDS:
continue # excludes 3: person-sitting class for evaluation
a_polygons = anno['segmentation'][
0] # only one mask for each instance
polygons = anno['i_segm'][0]
# gt_x1, gt_y1, gt_w, gt_h = anno['a_bbox'] # this is used to clip resampled polygons
a_contour = np.array(a_polygons).reshape((-1, 2))
contour = np.array(polygons).reshape((-1, 2))
# Downsample the contour to fix number of vertices
if cv2.contourArea(contour.astype(
np.int32)) < 5: # remove tiny objects
continue
fixed_contour = uniformsample(a_contour, self.n_vertices)
i_contour = uniformsample(contour, self.n_vertices)
# fixed_contour[:, 0] = np.clip(fixed_contour[:, 0], gt_x1, gt_x1 + gt_w)
# fixed_contour[:, 1] = np.clip(fixed_contour[:, 1], gt_y1, gt_y1 + gt_h)
# contour_std = np.sqrt(np.sum(np.std(fixed_contour, axis=0) ** 2))
# if contour_std < 1e-6 or contour_std == np.inf or contour_std == np.nan: # invalid shapes
# continue
shapes.append(np.ndarray.flatten(i_contour).tolist())
a_shapes.append(np.ndarray.flatten(fixed_contour).tolist())
labels.append(self.cat_ids[anno['category_id']])
bboxes.append(anno['bbox'])
a_bboxes.append(anno['a_bbox'])
labels = np.array(labels)
bboxes = np.array(bboxes, dtype=np.float32)
a_bboxes = np.array(a_bboxes, dtype=np.float32)
shapes = np.array(shapes, dtype=np.float32)
a_shapes = np.array(a_shapes, dtype=np.float32)
if len(bboxes) == 0:
bboxes = np.array([[0., 0., 0., 0.]], dtype=np.float32)
a_bboxes = np.array([[0., 0., 0., 0.]], dtype=np.float32)
labels = np.array([[0]])
shapes = np.zeros((1, self.n_vertices * 2), dtype=np.float32)
a_shapes = np.zeros((1, self.n_vertices * 2), dtype=np.float32)
bboxes[:, 2:] += bboxes[:, :2] # xywh to xyxy
a_bboxes[:, 2:] += a_bboxes[:, :2]
img = cv2.imread(img_path)
height, width = img.shape[0], img.shape[1]
center = np.array([width / 2., height / 2.],
dtype=np.float32) # center of image
scale = max(height, width) * 1.0
flipped = False
if self.split == 'train':
scale = scale * np.random.choice(self.rand_scales)
w_border = get_border(360, width)
h_border = get_border(160, height)
center[0] = np.random.randint(low=w_border, high=width - w_border)
center[1] = np.random.randint(low=h_border, high=height - h_border)
if np.random.random() < 0.5:
flipped = True
img = img[:, ::-1, :]
center[0] = width - center[0] - 1
trans_img = get_affine_transform(
center, scale, 0, [self.img_size['w'], self.img_size['h']])
# -----------------------------------debug---------------------------------
# image_show = img.copy()
img = cv2.warpAffine(img, trans_img,
(self.img_size['w'], self.img_size['h']))
img = img.astype(np.float32) / 255.
if self.split == 'train':
color_aug(self.data_rng, img, self.eig_val, self.eig_vec)
img -= self.mean
img /= self.std
img = img.transpose(2, 0, 1) # from [H, W, C] to [C, H, W]
trans_fmap = get_affine_transform(
center, scale, 0, [self.fmap_size['w'], self.fmap_size['h']])
# -----------------------------------debug---------------------------------
# image_show = cv2.warpAffine(image_show, trans_fmap, (self.fmap_size['w'], self.fmap_size['h']))
hmap = np.zeros(
(self.num_classes, self.fmap_size['h'], self.fmap_size['w']),
dtype=np.float32) # heatmap of centers
occ_map = np.zeros(
(1, self.fmap_size['h'], self.fmap_size['w']),
dtype=np.float32) # grayscale map for occlusion levels
w_h_ = np.zeros((self.max_objs, 2),
dtype=np.float32) # width and height of inmodal bboxes
shapes_ = np.zeros((self.max_objs, self.n_vertices * 2),
dtype=np.float32) # gt amodal segmentation polygons
center_offsets = np.zeros(
(self.max_objs, 2),
dtype=np.float32) # gt amodal mass centers to inmodal bbox center
codes_ = np.zeros((self.max_objs, self.n_codes),
dtype=np.float32) # gt amodal coefficients
regs = np.zeros((self.max_objs, 2),
dtype=np.float32) # regression for quantization error
inds = np.zeros((self.max_objs, ), dtype=np.int64)
ind_masks = np.zeros((self.max_objs, ), dtype=np.uint8)
votes_ = np.zeros((self.max_objs, self.vote_length),
dtype=np.float32) # voting for heatmaps
for k, (bbox, a_bbox, label, shape, a_shape) in enumerate(
zip(bboxes, a_bboxes, labels, shapes, a_shapes)):
if flipped:
bbox[[0, 2]] = width - bbox[[2, 0]] - 1
a_bbox[[0, 2]] = width - a_bbox[[2, 0]] - 1
# Flip the contour x-axis
for m in range(self.n_vertices):
a_shape[2 * m] = width - a_shape[2 * m] - 1
shape[2 * m] = width - shape[2 * m] - 1
bbox[:2] = affine_transform(bbox[:2], trans_fmap)
bbox[2:] = affine_transform(bbox[2:], trans_fmap)
bbox[[0, 2]] = np.clip(bbox[[0, 2]], 0, self.fmap_size['w'] - 1)
bbox[[1, 3]] = np.clip(bbox[[1, 3]], 0, self.fmap_size['h'] - 1)
h, w = bbox[3] - bbox[1], bbox[2] - bbox[
0] # This box is the inmodal boxes
a_bbox[:2] = affine_transform(a_bbox[:2], trans_fmap)
a_bbox[2:] = affine_transform(a_bbox[2:], trans_fmap)
a_bbox[[0, 2]] = np.clip(a_bbox[[0, 2]], 0,
self.fmap_size['w'] - 1)
a_bbox[[1, 3]] = np.clip(a_bbox[[1, 3]], 0,
self.fmap_size['h'] - 1)
# generate gt shape mean and std from contours
for m in range(self.n_vertices
): # apply scale and crop transform to shapes
a_shape[2 * m:2 * m + 2] = affine_transform(
a_shape[2 * m:2 * m + 2], trans_fmap)
shape[2 * m:2 * m + 2] = affine_transform(
shape[2 * m:2 * m + 2], trans_fmap)
shape_clipped = np.reshape(a_shape, (self.n_vertices, 2))
shape_clipped[:, 0] = np.clip(shape_clipped[:, 0], 0,
self.fmap_size['w'] - 1)
shape_clipped[:, 1] = np.clip(shape_clipped[:, 1], 0,
self.fmap_size['h'] - 1)
i_shape_clipped = np.reshape(shape, (self.n_vertices, 2))
i_shape_clipped[:, 0] = np.clip(i_shape_clipped[:, 0], 0,
self.fmap_size['w'] - 1)
i_shape_clipped[:, 1] = np.clip(i_shape_clipped[:, 1], 0,
self.fmap_size['h'] - 1)
clockwise_flag = check_clockwise_polygon(shape_clipped)
if not clockwise_flag:
fixed_contour = np.flip(shape_clipped, axis=0)
else:
fixed_contour = shape_clipped.copy()
# Indexing from the left-most vertex, argmin x-axis
idx = np.argmin(fixed_contour[:, 0])
indexed_shape = np.concatenate(
(fixed_contour[idx:, :], fixed_contour[:idx, :]), axis=0)
mass_center = np.mean(indexed_shape, axis=0)
if h < 1e-6 or w < 1e-6: # remove small bboxes
continue
centered_shape = indexed_shape - mass_center # these are amodal mask shapes
if h > 0 and w > 0:
obj_c = np.array([(bbox[0] + bbox[2]) / 2,
(bbox[1] + bbox[3]) / 2],
dtype=np.float32)
obj_c_int = obj_c.astype(np.int32)
radius = max(
0,
int(
gaussian_radius((math.ceil(h), math.ceil(w)),
self.gaussian_iou)))
draw_umich_gaussian(hmap[label], obj_c_int, radius)
shapes_[k] = centered_shape.reshape((1, -1))
center_offsets[k] = mass_center - obj_c
codes_[k], _ = fast_ista(centered_shape.reshape((1, -1)),
self.dictionary,
lmbda=self.sparse_alpha,
max_iter=60)
a_shifted_poly = indexed_shape - np.array([
a_bbox[0], a_bbox[1]
]) # crop amodal shapes to the amodal bboxes
amodal_obj_mask = self.polys_to_mask(
[np.ndarray.flatten(a_shifted_poly, order='C').tolist()],
a_bbox[3], a_bbox[2])
i_shifted_poly = i_shape_clipped - np.array([
a_bbox[0], a_bbox[1]
]) # crop inmodal shapes to the same amodal bboxes
inmodal_obj_mask = self.polys_to_mask(
[np.ndarray.flatten(i_shifted_poly, order='C').tolist()],
a_bbox[3], a_bbox[2])
obj_mask = (
amodal_obj_mask + inmodal_obj_mask
) * 255. / 2 # convert to float type in image scale
obj_mask = cv2.resize(
obj_mask.astype(np.uint8),
dsize=(self.vote_vec_dim, self.vote_vec_dim),
interpolation=cv2.INTER_LINEAR) * 1.
votes_[k] = obj_mask.reshape((1, -1)) / 255.
w_h_[k] = 1. * w, 1. * h
regs[k] = obj_c - obj_c_int # discretization error
inds[k] = obj_c_int[1] * self.fmap_size['w'] + obj_c_int[0]
ind_masks[k] = 1
# occlusion level map gt
occ_map[0] += self.polys_to_mask(
[np.ndarray.flatten(indexed_shape).tolist()],
self.fmap_size['h'], self.fmap_size['w']) * 1.
occ_map = np.clip(occ_map, 0, self.max_occ) / self.max_occ
# -----------------------------------debug---------------------------------
# for bbox, label, shape in zip(bboxes, labels, shapes_):
# # cv2.rectangle(image_show, (int(bbox[0]), int(bbox[1])), (int(bbox[2]), int(bbox[3])), (0, 255, 0), 1)
# cv2.putText(image_show, str(self.reverse_labels[label]), (int(bbox[0]), int(bbox[1])), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
# # print(shape, shape.shape)
# cv2.polylines(image_show, [shape.reshape(self.n_vertices, 2).astype(np.int32)], True, (0, 0, 255),
# thickness=1)
# # cv2.imshow('img', image_show)
# # cv2.imshow('occ', occ_map.astype(np.uint8).reshape(occ_map.shape[1], occ_map.shape[2]) * 255)
# m_img = cv2.cvtColor((occ_map * 255).astype(np.uint8).reshape(occ_map.shape[1], occ_map.shape[2]),
# code=cv2.COLOR_GRAY2BGR)
# cat_img = np.concatenate([m_img, image_show], axis=0)
# cv2.imshow('segm', cat_img)
# cv2.waitKey()
# -----------------------------------debug---------------------------------
return {
'image': img,
'shapes': shapes_,
'codes': codes_,
'offsets': center_offsets,
'occ_map': occ_map,
'hmap': hmap,
'w_h_': w_h_,
'regs': regs,
'inds': inds,
'ind_masks': ind_masks,
'votes': votes_,
'c': center,
's': scale,
'img_id': img_id
}
def main():
cfg.device = torch.device('cuda')
torch.backends.cudnn.benchmark = False
max_per_image = 100
num_classes = 80 if cfg.dataset == 'coco' else 4
dictionary = np.load(cfg.dictionary_file)
colors = COCO_COLORS if cfg.dataset == 'coco' else DETRAC_COLORS
names = COCO_NAMES if cfg.dataset == 'coco' else DETRAC_NAMES
for j in range(len(names)):
col_ = [c * 255 for c in colors[j]]
colors[j] = tuple(col_)
print('Creating model and recover from checkpoint ...')
if 'hourglass' in cfg.arch:
model = exkp(n=5, nstack=2, dims=[256, 256, 384, 384, 384, 512],
modules=[2, 2, 2, 2, 2, 4], num_classes=num_classes)
elif 'resdcn' in cfg.arch:
model = get_pose_resdcn(num_layers=int(cfg.arch.split('_')[-1]), head_conv=64,
num_classes=num_classes, num_codes=cfg.n_codes)
else:
raise NotImplementedError
model = load_demo_model(model, cfg.ckpt_dir)
model = model.to(cfg.device)
model.eval()
# Loading COCO validation images
if 'train' in cfg.data_type:
annotation_file = '{}/annotations/instances_train2017.json'.format(cfg.data_dir)
cfg.data_type = 'train2017'
else:
annotation_file = '{}/annotations/instances_val2017.json'.format(cfg.data_dir)
cfg.data_type = 'val2017'
coco = COCO(annotation_file)
# Load all annotations
cats = coco.loadCats(coco.getCatIds())
# nms = [cat['name'] for cat in cats]
nms = ['giraffe']
catIds = coco.getCatIds(catNms=nms)
imgIds = coco.getImgIds(catIds=catIds)
annIds = coco.getAnnIds(catIds=catIds)
all_anns = coco.loadAnns(ids=annIds)
for annotation in all_anns:
if annotation['iscrowd'] == 1 or type(annotation['segmentation']) != list or len(
annotation['segmentation']) > 1:
continue
img = coco.loadImgs(annotation['image_id'])[0]
image_path = '%s/images/%s/%s' % (cfg.data_dir, cfg.data_type, img['file_name'])
w_img = int(img['width'])
h_img = int(img['height'])
if w_img < 350 or h_img < 350:
continue
polygons = annotation['segmentation'][0]
gt_bbox = annotation['bbox']
gt_x1, gt_y1, gt_w, gt_h = gt_bbox
contour = np.array(polygons).reshape((-1, 2))
if cv2.contourArea(contour.astype(np.int32)) < 200:
continue
# Downsample the contour to fix number of vertices
fixed_contour = resample(contour, num=cfg.num_vertices)
clockwise_flag = check_clockwise_polygon(fixed_contour)
if not clockwise_flag:
fixed_contour = np.flip(fixed_contour, axis=0)
# else:
# fixed_contour = indexed_shape.copy()
# Indexing from the left-most vertex, argmin x-axis
idx = np.argmin(fixed_contour[:, 0])
indexed_shape = np.concatenate((fixed_contour[idx:, :], fixed_contour[:idx, :]), axis=0)
indexed_shape[:, 0] = np.clip(indexed_shape[:, 0], gt_x1, gt_x1 + gt_w)
indexed_shape[:, 1] = np.clip(indexed_shape[:, 1], gt_y1, gt_y1 + gt_h)
updated_bbox = [np.min(indexed_shape[:, 0]), np.min(indexed_shape[:, 1]),
np.max(indexed_shape[:, 0]), np.max(indexed_shape[:, 1])]
w, h = updated_bbox[2] - updated_bbox[0], updated_bbox[3] - updated_bbox[1]
contour_mean = np.mean(indexed_shape, axis=0)
# contour_std = np.std(indexed_shape, axis=0)
# if contour_std < 1e-6 or contour_std == np.inf or contour_std == np.nan: # invalid shapes
# continue
norm_shape = (indexed_shape - contour_mean) / np.array([w / 2., h / 2.])
gt_codes, _ = fast_ista(norm_shape.reshape((1, -1)), dictionary, lmbda=0.005, max_iter=80)
recon_contour = np.matmul(gt_codes, dictionary).reshape((-1, 2)) * np.array([w / 2., h / 2.])
recon_contour = recon_contour + contour_mean
image = cv2.imread(image_path, cv2.IMREAD_COLOR)
if image is None:
continue
original_image = image.copy()
height, width = image.shape[0:2]
padding = 127 if 'hourglass' in cfg.arch else 31
imgs = {}
for scale in cfg.test_scales:
new_height = int(height * scale)
new_width = int(width * scale)
if cfg.img_size > 0:
img_height, img_width = cfg.img_size, cfg.img_size
center = np.array([new_width / 2., new_height / 2.], dtype=np.float32)
scaled_size = max(height, width) * 1.0
scaled_size = np.array([scaled_size, scaled_size], dtype=np.float32)
else:
img_height = (new_height | padding) + 1
img_width = (new_width | padding) + 1
center = np.array([new_width // 2, new_height // 2], dtype=np.float32)
scaled_size = np.array([img_width, img_height], dtype=np.float32)
img = cv2.resize(image, (new_width, new_height))
trans_img = get_affine_transform(center, scaled_size, 0, [img_width, img_height])
img = cv2.warpAffine(img, trans_img, (img_width, img_height))
img = img.astype(np.float32) / 255.
img -= np.array(COCO_MEAN if cfg.dataset == 'coco' else DETRAC_MEAN, dtype=np.float32)[None, None, :]
img /= np.array(COCO_STD if cfg.dataset == 'coco' else DETRAC_STD, dtype=np.float32)[None, None, :]
img = img.transpose(2, 0, 1)[None, :, :, :] # from [H, W, C] to [1, C, H, W]
# if cfg.test_flip:
# img = np.concatenate((img, img[:, :, :, ::-1].copy()), axis=0)
imgs[scale] = {'image': torch.from_numpy(img).float(),
'center': np.array(center),
'scale': np.array(scaled_size),
'fmap_h': np.array(img_height // 4),
'fmap_w': np.array(img_width // 4)}
with torch.no_grad():
segmentations = []
predicted_codes = []
mass_centers = []
start_time = time.time()
print('Start running model ......')
for scale in imgs:
imgs[scale]['image'] = imgs[scale]['image'].to(cfg.device)
hmap, regs, w_h_, _, _, codes, offsets = model(imgs[scale]['image'])[-1]
output = [hmap, regs, w_h_, codes, offsets]
# segms = ctsegm_scale_decode(*output,
# torch.from_numpy(dictionary.astype(np.float32)).to(cfg.device),
# K=cfg.test_topk)
# print(len(output))
segms, pred_codes, pred_center = ctsegm_scale_decode_debug(*output,
torch.from_numpy(dictionary.astype(np.float32)).to(cfg.device),
K=cfg.test_topk)
segms = segms.detach().cpu().numpy().reshape(1, -1, segms.shape[2])[0]
pred_codes = pred_codes.detach().cpu().numpy().reshape(-1, pred_codes.shape[-1])
pred_center = pred_center.detach().cpu().numpy().reshape(-1, 2)
top_preds = {}
code_preds = {}
center_preds = {}
for j in range(cfg.num_vertices):
segms[:, 2 * j:2 * j + 2] = transform_preds(segms[:, 2 * j:2 * j + 2],
imgs[scale]['center'],
imgs[scale]['scale'],
(imgs[scale]['fmap_w'], imgs[scale]['fmap_h']))
segms[:, cfg.num_vertices * 2:cfg.num_vertices * 2 + 2] = transform_preds(
segms[:, cfg.num_vertices * 2:cfg.num_vertices * 2 + 2],
imgs[scale]['center'],
imgs[scale]['scale'],
(imgs[scale]['fmap_w'], imgs[scale]['fmap_h']))
segms[:, cfg.num_vertices * 2 + 2:cfg.num_vertices * 2 + 4] = transform_preds(
segms[:, cfg.num_vertices * 2 + 2:cfg.num_vertices * 2 + 4],
imgs[scale]['center'],
imgs[scale]['scale'],
(imgs[scale]['fmap_w'], imgs[scale]['fmap_h']))
# For mass center
pred_center = transform_preds(pred_center,
imgs[scale]['center'],
imgs[scale]['scale'],
(imgs[scale]['fmap_w'], imgs[scale]['fmap_h']))
clses = segms[:, -1]
for j in range(num_classes):
inds = (clses == j)
top_preds[j + 1] = segms[inds, :cfg.num_vertices * 2 + 5].astype(np.float32)
top_preds[j + 1][:, :cfg.num_vertices * 2 + 4] /= scale
center_preds[j + 1] = pred_center[inds, :] / scale
code_preds[j + 1] = pred_codes[inds, :]
segmentations.append(top_preds)
predicted_codes.append(code_preds)
mass_centers.append(center_preds)
segms_and_scores = {j: np.concatenate([d[j] for d in segmentations], axis=0)
for j in range(1, num_classes + 1)} # a Dict label: segments
segms_and_codes = {j: np.concatenate([d[j] for d in predicted_codes], axis=0)
for j in range(1, num_classes + 1)}
segms_and_centers = {j: np.concatenate([d[j] for d in mass_centers], axis=0)
for j in range(1, num_classes + 1)}
scores = np.hstack(
[segms_and_scores[j][:, cfg.num_vertices * 2 + 4] for j in range(1, num_classes + 1)])
print('Image processing time {:.4f} sec, preparing output image ......'.format(time.time() - start_time))
if len(scores) > max_per_image:
kth = len(scores) - max_per_image
thresh = np.partition(scores, kth)[kth]
for j in range(1, num_classes + 1):
keep_inds = (segms_and_scores[j][:, cfg.num_vertices * 2 + 4] >= thresh)
segms_and_scores[j] = segms_and_scores[j][keep_inds]
segms_and_codes[j] = segms_and_codes[j][keep_inds]
segms_and_centers[j] = segms_and_centers[j][keep_inds]
# Use opencv functions to output
# output_image = original_image
# blend_mask = np.zeros(shape=output_image.shape, dtype=np.uint8)
counter = 1
for lab in segms_and_scores:
output_image = original_image.copy()
# if cfg.dataset == 'coco':
# if names[lab] not in display_cat and cfg.dataset != 'kins':
# continue
for idx in range(len(segms_and_scores[lab])):
res = segms_and_scores[lab][idx]
p_code = segms_and_codes[lab][idx]
p_center = segms_and_centers[lab][idx]
contour, bbox, score = res[:-5], res[-5:-1], res[-1]
bbox[0] = np.clip(bbox[0], 0, width - 1)
bbox[1] = np.clip(bbox[1], 0, height - 1)
bbox[2] = np.clip(bbox[2], 0, width - 1)
bbox[3] = np.clip(bbox[3], 0, height - 1)
polygon = contour.reshape((-1, 2))
polygon[:, 0] = np.clip(polygon[:, 0], 0, width - 1)
polygon[:, 1] = np.clip(polygon[:, 1], 0, height - 1)
if score > cfg.detect_thres:
# text = names[lab] + ' %.2f' % score
# label_size = cv2.getTextSize(text, cv2.FONT_HERSHEY_COMPLEX, 0.3, 1)
# text_location = [int(bbox[0]) + 2, int(bbox[1]) + 2,
# int(bbox[0]) + 2 + label_size[0][0],
# int(bbox[1]) + 2 + label_size[0][1]]
# cv2.rectangle(output_image, pt1=(int(bbox[0]), int(bbox[1])),
# pt2=(int(bbox[2]), int(bbox[3])),
# color=colors[lab], thickness=2)
# cv2.rectangle(output_image, pt1=(int(bbox[0]), int(bbox[1])),
# pt2=(int(bbox[2]), int(bbox[3])),
# color=nice_colors[names[lab]], thickness=2)
# cv2.putText(output_image, text, org=(int(text_location[0]), int(text_location[3])),
# fontFace=cv2.FONT_HERSHEY_COMPLEX, thickness=1, fontScale=0.3,
# color=nice_colors[names[lab]])
# use_color_key = COLOR_WORLD[random.randint(1, len(COLOR_WORLD)) - 1]
# cv2.polylines(output_image, [polygon.astype(np.int32)], True,
# color=switch_tuple(RGB_DICT[use_color_key]),
# thickness=2)
# cv2.drawContours(blend_mask, [polygon.astype(np.int32)], contourIdx=-1,
# color=switch_tuple(RGB_DICT[use_color_key]),
# thickness=-1)
# plot the polygons/contours
cv2.polylines(output_image, [recon_contour.astype(np.int32)], True,
color=switch_tuple(RGB_DICT['green']), thickness=2)
cv2.polylines(output_image, [polygon.astype(np.int32)], True,
color=switch_tuple(RGB_DICT['red']), thickness=2)
# plot the mass center location
cv2.circle(output_image, tuple(contour_mean.astype(np.int32).tolist()),
radius=9, color=switch_tuple(RGB_DICT['green']), thickness=-1)
cv2.circle(output_image, tuple(p_center.astype(np.int32).tolist()),
radius=9, color=switch_tuple(RGB_DICT['red']), thickness=-1)
# dst_img = cv2.addWeighted(output_image, 0.4, blend_mask, 0.6, 0)
# dst_img[blend_mask == 0] = output_image[blend_mask == 0]
# output_image = dst_img
cv2.imshow('Frames', output_image)
if cv2.waitKey() & 0xFF == ord('q'):
break
counter += 1
# show histogram
fig, (ax1, ax2) = plt.subplots(1, 2)
# plot 1
bins = np.linspace(-2, 2, 30)
ax1.hist(gt_codes.reshape((-1,)).tolist(), bins=bins, color='g', density=False, alpha=0.5)
ax1.hist(p_code.reshape((-1,)).tolist(), bins=bins, color='r', density=False, alpha=0.5)
ax1.legend(['GT Coeffs', 'Pred Coeffs'])
ax1.set_xlabel('Sparse Coefficients')
ax1.set_ylabel('Counts')
ax1.set_title('Histogram of Coefficients')
# plot 2
ax2.plot(gt_codes.reshape((-1,)), 'g*-', linewidth=2, markersize=6)
ax2.plot(p_code.reshape((-1,)), 'ro--', linewidth=1, markersize=5)
ax2.legend(['GT Coeffs', 'Pred Coeffs'])
ax2.set_xlabel('Coefficients Index')
ax2.set_ylabel('Value')
ax2.set_title('Coefficients')
plt.show()
plt.close()
def __getitem__(self, index):
img_id = self.images[index]
img_path = os.path.join(self.img_dir, self.coco.loadImgs(ids=[img_id])[0]['file_name'])
ann_ids = self.coco.getAnnIds(imgIds=[img_id])
annotations = self.coco.loadAnns(ids=ann_ids)
# img = self.coco.loadImgs(ids=[img_id])[0]
# w_img = int(img['width'])
# h_img = int(img['height'])
# if w_img < 2 or h_img < 2:
# continue
labels = []
bboxes = []
shapes = []
for anno in annotations:
if anno['iscrowd'] == 1 or type(anno['segmentation']) != list: # Excludes crowd objects
continue
if len(anno['segmentation']) > 1:
obj_contours = [np.array(s).reshape((-1, 2)).astype(np.int32) for s in anno['segmentation']]
obj_contours = sorted(obj_contours, key=cv2.contourArea)
polygons = obj_contours[-1]
else:
polygons = anno['segmentation'][0]
gt_x1, gt_y1, gt_w, gt_h = anno['bbox']
if gt_w < 5 or gt_h < 5:
continue
contour = np.array(polygons).reshape((-1, 2))
# Downsample the contour to fix number of vertices
if cv2.contourArea(contour.astype(np.int32)) < 35:
continue
fixed_contour = uniformsample(contour, self.n_vertices)
# fixed_contour[:, 0] = np.clip(fixed_contour[:, 0], gt_x1, gt_x1 + gt_w)
# fixed_contour[:, 1] = np.clip(fixed_contour[:, 1], gt_y1, gt_y1 + gt_h)
contour_std = np.sqrt(np.sum(np.std(fixed_contour, axis=0) ** 2))
if contour_std < 1e-6 or contour_std == np.inf or contour_std == np.nan: # invalid shapes
continue
updated_bbox = [np.min(fixed_contour[:, 0]), np.min(fixed_contour[:, 1]),
np.max(fixed_contour[:, 0]), np.max(fixed_contour[:, 1])]
shapes.append(np.ndarray.flatten(fixed_contour).tolist())
labels.append(self.cat_ids[anno['category_id']])
bboxes.append(updated_bbox)
labels = np.array(labels)
bboxes = np.array(bboxes, dtype=np.float32)
shapes = np.array(shapes, dtype=np.float32)
if len(bboxes) == 0:
bboxes = np.array([[0., 0., 0., 0.]], dtype=np.float32)
labels = np.array([[0]])
shapes = np.zeros((1, self.n_vertices * 2), dtype=np.float32)
# bboxes[:, 2:] += bboxes[:, :2] # xywh to xyxy
img = cv2.imread(img_path)
height, width = img.shape[0], img.shape[1]
center = np.array([width / 2., height / 2.], dtype=np.float32) # center of image
scale = max(height, width) * 1.0
flipped = False
if self.split == 'train':
scale = scale * np.random.choice(self.rand_scales)
w_border = get_border(150, width)
h_border = get_border(150, height)
center[0] = np.random.randint(low=w_border, high=width - w_border)
center[1] = np.random.randint(low=h_border, high=height - h_border)
if np.random.random() < 0.5:
flipped = True
img = img[:, ::-1, :]
center[0] = width - center[0] - 1
trans_img = get_affine_transform(center, scale, 0, [self.img_size['w'], self.img_size['h']])
img = cv2.warpAffine(img, trans_img, (self.img_size['w'], self.img_size['h']))
img = img.astype(np.float32) / 255.
if self.split == 'train':
color_aug(self.data_rng, img, self.eig_val, self.eig_vec)
img -= self.mean
img /= self.std
img = img.transpose(2, 0, 1) # from [H, W, C] to [C, H, W]
trans_fmap = get_affine_transform(center, scale, 0, [self.fmap_size['w'], self.fmap_size['h']])
hmap = np.zeros((self.num_classes, self.fmap_size['h'], self.fmap_size['w']), dtype=np.float32) # heatmap
votes_ = np.zeros((self.max_objs, self.vote_length), dtype=np.float32) # votes for hmap and code
w_h_ = np.zeros((self.max_objs, 2), dtype=np.float32) # width and height of bboxes
shapes_ = np.zeros((self.max_objs, self.n_vertices * 2), dtype=np.float32) # gt amodal segmentation polygons
center_offsets = np.zeros((self.max_objs, 2), dtype=np.float32) # gt mass centers to bbox center
codes_ = np.zeros((self.max_objs, self.n_codes), dtype=np.float32)
regs = np.zeros((self.max_objs, 2), dtype=np.float32) # regression for offsets of shape center
inds = np.zeros((self.max_objs,), dtype=np.int64)
ind_masks = np.zeros((self.max_objs,), dtype=np.uint8)
for k, (bbox, label, shape) in enumerate(zip(bboxes, labels, shapes)):
if flipped:
bbox[[0, 2]] = width - bbox[[2, 0]] - 1
# Flip the contour x-axis
for m in range(self.n_vertices):
shape[2 * m] = width - shape[2 * m] - 1
bbox[:2] = affine_transform(bbox[:2], trans_fmap)
bbox[2:] = affine_transform(bbox[2:], trans_fmap)
bbox[[0, 2]] = np.clip(bbox[[0, 2]], 0, self.fmap_size['w'] - 1)
bbox[[1, 3]] = np.clip(bbox[[1, 3]], 0, self.fmap_size['h'] - 1)
h, w = bbox[3] - bbox[1], bbox[2] - bbox[0]
# generate gt shape mean and std from contours
for m in range(self.n_vertices): # apply scale and crop transform to shapes
shape[2 * m:2 * m + 2] = affine_transform(shape[2 * m:2 * m + 2], trans_fmap)
shape_clipped = np.reshape(shape, (self.n_vertices, 2))
shape_clipped[:, 0] = np.clip(shape_clipped[:, 0], 0, self.fmap_size['w'] - 1)
shape_clipped[:, 1] = np.clip(shape_clipped[:, 1], 0, self.fmap_size['h'] - 1)
clockwise_flag = check_clockwise_polygon(shape_clipped)
if not clockwise_flag:
fixed_contour = np.flip(shape_clipped, axis=0)
else:
fixed_contour = shape_clipped.copy()
# Indexing from the left-most vertex, argmin x-axis
idx = np.argmin(fixed_contour[:, 0])
indexed_shape = np.concatenate((fixed_contour[idx:, :], fixed_contour[:idx, :]), axis=0)
mass_center = np.mean(indexed_shape, axis=0)
# contour_std = np.std(indexed_shape, axis=0) + 1e-4
if h < 1e-6 or w < 1e-6: # remove small bboxes
continue
# centered_shape = indexed_shape - mass_center
norm_shape = (indexed_shape - mass_center) / np.array([w / 2., h / 2.])
if h > 0 and w > 0:
obj_c = np.array([(bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2], dtype=np.float32)
obj_c_int = obj_c.astype(np.int32)
radius = max(0, int(gaussian_radius((math.ceil(h), math.ceil(w)), self.gaussian_iou)))
draw_umich_gaussian(hmap[label], obj_c_int, radius)
shapes_[k] = norm_shape.reshape((1, -1))
center_offsets[k] = mass_center - obj_c
codes_[k], _ = fast_ista(norm_shape.reshape((1, -1)), self.dictionary,
lmbda=self.sparse_alpha, max_iter=80)
w_h_[k] = 1. * w, 1. * h
regs[k] = obj_c - obj_c_int # discretization error
inds[k] = obj_c_int[1] * self.fmap_size['w'] + obj_c_int[0]
ind_masks[k] = 1
# getting the gt votes
shifted_poly = indexed_shape - np.array([bbox[0], bbox[1]]) + 1 # crop to the bbox, add padding 1
# obj_mask = polys_to_mask([np.ndarray.flatten(shifted_poly, order='C').tolist()], h + 2, w + 2) * 255
obj_mask = np.zeros((int(h) + 3, int(w) + 3), dtype=np.uint8)
cv2.drawContours(obj_mask, shifted_poly[None, :, :].astype(np.int32), color=255, contourIdx=-1,
thickness=-1)
# instance = obj_mask.copy()
# obj_mask = cv2.resize(obj_mask.astype(np.uint8), dsize=(self.vote_vec_dim, self.vote_vec_dim),
# interpolation=cv2.INTER_LINEAR) * 1.
# votes_[k] = obj_mask.reshape((1, -1)) / 255.
# votes_[k] = (obj_mask.reshape((1, -1)) > 255 * 0.4) * 1.0
# show debug masks
obj_mask = cv2.resize(obj_mask.astype(np.uint8), dsize=(self.vote_vec_dim, self.vote_vec_dim),
interpolation=cv2.INTER_LINEAR) # INTER_AREA
# obj_mask = cv2.resize(obj_mask.astype(np.uint8), dsize=(self.vote_vec_dim, self.vote_vec_dim),
# interpolation=cv2.INTER_AREA)
votes_[k] = (obj_mask.reshape((1, -1)) > 0.2 * 255) * 1.0
# cv2.imshow('obj_mask', instance.astype(np.uint8))
# cv2.waitKey()
# cv2.imshow('votes', obj_mask.astype(np.uint8))
# cv2.waitKey()
return {'image': img, 'shapes': shapes_, 'codes': codes_, 'offsets': center_offsets, 'votes': votes_,
'hmap': hmap, 'w_h_': w_h_, 'regs': regs, 'inds': inds, 'ind_masks': ind_masks,
'c': center, 's': scale, 'img_id': img_id}
def main():
cfg.device = torch.device('cuda')
torch.backends.cudnn.benchmark = False
max_per_image = 100
num_classes = 80 if cfg.dataset == 'coco' else 4
dictionary = np.load(cfg.dictionary_file)
colors = COCO_COLORS if cfg.dataset == 'coco' else DETRAC_COLORS
names = COCO_NAMES if cfg.dataset == 'coco' else DETRAC_NAMES
for j in range(len(names)):
col_ = [c * 255 for c in colors[j]]
colors[j] = tuple(col_)
print('Creating model and recover from checkpoint ...')
if 'hourglass' in cfg.arch:
model = exkp(n=5, nstack=2, dims=[256, 256, 384, 384, 384, 512],
modules=[2, 2, 2, 2, 2, 4], num_classes=num_classes)
else:
raise NotImplementedError
model = load_demo_model(model, cfg.ckpt_dir)
model = model.to(cfg.device)
model.eval()
# Loading COCO validation images
annotation_file = '{}/annotations/instances_{}.json'.format(cfg.data_dir, cfg.data_type)
coco = COCO(annotation_file)
# Load all annotations
cats = coco.loadCats(coco.getCatIds())
nms = [cat['name'] for cat in cats]
catIds = coco.getCatIds(catNms=nms)
imgIds = coco.getImgIds(catIds=catIds)
annIds = coco.getAnnIds(catIds=catIds)
all_anns = coco.loadAnns(ids=annIds)
for annotation in all_anns:
if annotation['iscrowd'] == 1 or type(annotation['segmentation']) != list:
continue
img = coco.loadImgs(annotation['image_id'])[0]
image_path = '%s/images/%s/%s' % (cfg.data_dir, cfg.data_type, img['file_name'])
w_img = int(img['width'])
h_img = int(img['height'])
if w_img < 1 or h_img < 1:
continue
polygons = annotation['segmentation'][0]
gt_bbox = annotation['bbox']
gt_x1, gt_y1, gt_w, gt_h = gt_bbox
contour = np.array(polygons).reshape((-1, 2))
# Downsample the contour to fix number of vertices
fixed_contour = resample(contour, num=cfg.num_vertices)
# Indexing from the left-most vertex, argmin x-axis
idx = np.argmin(fixed_contour[:, 0])
indexed_shape = np.concatenate((fixed_contour[idx:, :], fixed_contour[:idx, :]), axis=0)
clockwise_flag = check_clockwise_polygon(indexed_shape)
if not clockwise_flag:
fixed_contour = np.flip(indexed_shape, axis=0)
else:
fixed_contour = indexed_shape.copy()
fixed_contour[:, 0] = np.clip(fixed_contour[:, 0], gt_x1, gt_x1 + gt_w)
fixed_contour[:, 1] = np.clip(fixed_contour[:, 1], gt_y1, gt_y1 + gt_h)
contour_mean = np.mean(fixed_contour, axis=0)
contour_std = np.std(fixed_contour, axis=0)
# norm_shape = (fixed_contour - contour_mean) / np.sqrt(np.sum(contour_std ** 2.))
# plot gt mean and std
# image = cv2.imread(image_path)
# # cv2.ellipse(image, center=(int(contour_mean[0]), int(contour_mean[1])),
# # axes=(int(contour_std[0]), int(contour_std[1])),
# # angle=0, startAngle=0, endAngle=360, color=(0, 255, 0),
# # thickness=2)
# cv2.rectangle(image, pt1=(int(contour_mean[0] - contour_std[0] / 2.), int(contour_mean[1] - contour_std[1] / 2.)),
# pt2=(int(contour_mean[0] + contour_std[0] / 2.), int(contour_mean[1] + contour_std[1] / 2.)),
# color=(0, 255, 0), thickness=2)
# cv2.polylines(image, [fixed_contour.astype(np.int32)], True, (0, 0, 255))
# cv2.rectangle(image, pt1=(int(min(fixed_contour[:, 0])), int(min(fixed_contour[:, 1]))),
# pt2=(int(max(fixed_contour[:, 0])), int(max(fixed_contour[:, 1]))),
# color=(255, 0, 0), thickness=2)
# cv2.imshow('GT segments', image)
# if cv2.waitKey() & 0xFF == ord('q'):
# break
image = cv2.imread(image_path)
original_image = image.copy()
height, width = image.shape[0:2]
padding = 127 if 'hourglass' in cfg.arch else 31
imgs = {}
for scale in cfg.test_scales:
new_height = int(height * scale)
new_width = int(width * scale)
if cfg.img_size > 0:
img_height, img_width = cfg.img_size, cfg.img_size
center = np.array([new_width / 2., new_height / 2.], dtype=np.float32)
scaled_size = max(height, width) * 1.0
scaled_size = np.array([scaled_size, scaled_size], dtype=np.float32)
else:
img_height = (new_height | padding) + 1
img_width = (new_width | padding) + 1
center = np.array([new_width // 2, new_height // 2], dtype=np.float32)
scaled_size = np.array([img_width, img_height], dtype=np.float32)
img = cv2.resize(image, (new_width, new_height))
trans_img = get_affine_transform(center, scaled_size, 0, [img_width, img_height])
img = cv2.warpAffine(img, trans_img, (img_width, img_height))
img = img.astype(np.float32) / 255.
img -= np.array(COCO_MEAN if cfg.dataset == 'coco' else DETRAC_MEAN, dtype=np.float32)[None, None, :]
img /= np.array(COCO_STD if cfg.dataset == 'coco' else DETRAC_STD, dtype=np.float32)[None, None, :]
img = img.transpose(2, 0, 1)[None, :, :, :] # from [H, W, C] to [1, C, H, W]
# if cfg.test_flip:
# img = np.concatenate((img, img[:, :, :, ::-1].copy()), axis=0)
imgs[scale] = {'image': torch.from_numpy(img).float(),
'center': np.array(center),
'scale': np.array(scaled_size),
'fmap_h': np.array(img_height // 4),
'fmap_w': np.array(img_width // 4)}
with torch.no_grad():
segmentations = []
start_time = time.time()
for scale in imgs:
imgs[scale]['image'] = imgs[scale]['image'].to(cfg.device)
output = model(imgs[scale]['image'])[-1]
segms = ctsegm_decode(*output, torch.from_numpy(dictionary.astype(np.float32)).to(cfg.device),
K=cfg.test_topk)
segms = segms.detach().cpu().numpy().reshape(1, -1, segms.shape[2])[0]
top_preds = {}
for j in range(cfg.num_vertices):
segms[:, 2 * j:2 * j + 2] = transform_preds(segms[:, 2 * j:2 * j + 2],
imgs[scale]['center'],
imgs[scale]['scale'],
(imgs[scale]['fmap_w'], imgs[scale]['fmap_h']))
clses = segms[:, -1]
for j in range(num_classes):
inds = (clses == j)
top_preds[j + 1] = segms[inds, :cfg.num_vertices * 2 + 1].astype(np.float32)
top_preds[j + 1][:, :cfg.num_vertices * 2] /= scale
segmentations.append(top_preds)
segms_and_scores = {j: np.concatenate([d[j] for d in segmentations], axis=0)
for j in range(1, num_classes + 1)} # a Dict label: segments
scores = np.hstack(
[segms_and_scores[j][:, cfg.num_vertices * 2] for j in range(1, num_classes + 1)])
if len(scores) > max_per_image:
kth = len(scores) - max_per_image
thresh = np.partition(scores, kth)[kth]
for j in range(1, num_classes + 1):
keep_inds = (segms_and_scores[j][:, cfg.num_vertices * 2] >= thresh)
segms_and_scores[j] = segms_and_scores[j][keep_inds]
# Use opencv functions to output a video
output_image = original_image
for lab in segms_and_scores:
for res in segms_and_scores[lab]:
contour, score = res[:-1], res[-1]
if score > cfg.detect_thres:
text = names[lab] + ' %.2f' % score
label_size = cv2.getTextSize(text, cv2.FONT_HERSHEY_COMPLEX, 0.5, 1)
polygon = contour.reshape((-1, 2))
contour_mean = np.mean(polygon, axis=0)
contour_std = np.std(polygon, axis=0)
center_x, center_y = np.mean(polygon, axis=0).astype(np.int32)
text_location = [center_x, center_y,
center_x + label_size[0][0],
center_y + label_size[0][1]]
# cv2.rectangle(output_image, pt1=(int(contour_mean[0] - contour_std[0] / 2.), int(contour_mean[1] - contour_std[1] / 2.)),
# pt2=(int(contour_mean[0] + contour_std[0] / 2.), int(contour_mean[1] + contour_std[1] / 2.)),
# color=(0, 255, 0), thickness=1)
cv2.polylines(output_image, [polygon.astype(np.int32)], True, (255, 0, 0), thickness=2)
# cv2.putText(output_image, text, org=(int(text_location[0]), int(text_location[3])),
# fontFace=cv2.FONT_HERSHEY_COMPLEX, thickness=1, fontScale=0.5,
# color=(0, 0, 255))
cv2.imshow('Results', output_image)
if cv2.waitKey() & 0xFF == ord('q'):
break
def __getitem__(self, index):
img_id = self.images[index]
img_path = os.path.join(
self.img_dir,
self.coco.loadImgs(ids=[img_id])[0]['file_name'])
ann_ids = self.coco.getAnnIds(imgIds=[img_id])
annotations = self.coco.loadAnns(ids=ann_ids)
img = self.coco.loadImgs(ids=[img_id])[0]
w_img = int(img['width'])
h_img = int(img['height'])
labels = []
bboxes = []
shapes = []
for anno in annotations:
if anno['iscrowd'] == 1: # Excludes crowd objects
continue
# polygons = anno['segmentation'][0]
polygons = anno['segmentation']
if len(polygons) > 1:
bg = np.zeros((h_img, w_img, 1), dtype=np.uint8)
for poly in polygons:
len_poly = len(poly)
vertices = np.zeros((1, len_poly // 2, 2), dtype=np.int32)
for i in range(len_poly // 2):
vertices[0, i, 0] = int(poly[2 * i])
vertices[0, i, 1] = int(poly[2 * i + 1])
# cv2.fillPoly(bg, vertices, color=(255))
cv2.drawContours(bg,
vertices,
color=(255),
contourIdx=-1,
thickness=-1)
pads = 5
while True:
kernel = np.ones((pads, pads), np.uint8)
bg_closed = cv2.morphologyEx(bg, cv2.MORPH_CLOSE, kernel)
obj_contours, _ = cv2.findContours(bg_closed,
cv2.RETR_TREE,
cv2.CHAIN_APPROX_SIMPLE)
if len(obj_contours) > 1:
pads += 5
else:
polygons = obj_contours[0]
break
else:
# continue
polygons = anno['segmentation'][0]
gt_x1, gt_y1, gt_w, gt_h = anno['bbox']
contour = np.array(polygons).reshape((-1, 2))
# Downsample the contour to fix number of vertices
fixed_contour = resample(contour, num=self.n_vertices)
fixed_contour[:, 0] = np.clip(fixed_contour[:, 0], gt_x1,
gt_x1 + gt_w)
fixed_contour[:, 1] = np.clip(fixed_contour[:, 1], gt_y1,
gt_y1 + gt_h)
# contour_mean = np.mean(fixed_contour, axis=0)
contour_std = np.sqrt(np.sum(np.std(fixed_contour, axis=0)**2))
if contour_std < 1e-6 or contour_std == np.inf or contour_std == np.nan: # invalid shapes
continue
shapes.append(np.ndarray.flatten(fixed_contour).tolist())
labels.append(self.cat_ids[anno['category_id']])
bboxes.append(anno['bbox'])
labels = np.array(labels)
bboxes = np.array(bboxes, dtype=np.float32)
shapes = np.array(shapes, dtype=np.float32)
if len(bboxes) == 0:
bboxes = np.array([[0., 0., 0., 0.]], dtype=np.float32)
labels = np.array([[0]])
shapes = np.zeros((1, self.n_vertices * 2), dtype=np.float32)
bboxes[:, 2:] += bboxes[:, :2] # xywh to xyxy
# if img_id in self.all_annotations.keys():
# annotations = self.all_annotations[img_id]
# shape_annots = self.all_shapes[img_id]
# labels = annotations['cat_id']
# bboxes = annotations['bbox'] # xyxy format
# shapes = shape_annots['shape'] # polygonal vertices format xyxyxyxyxy...
# codes = annotations['codes']
# labels = np.array(labels)
# bboxes = np.array(bboxes, dtype=np.float32)
# codes = np.array(codes, dtype=np.float32)
# shapes = np.array(shapes, dtype=np.float32)
# else:
# bboxes = np.array([[0., 0., 0., 0.]], dtype=np.float32)
# labels = np.array([[0]])
# codes = np.zeros(shape=(1, self.n_codes), dtype=np.float32)
# shapes = np.zeros(shape=(1, self.n_vertices * 2), dtype=np.float32)
img = cv2.imread(img_path)
height, width = img.shape[0], img.shape[1]
center = np.array([width / 2., height / 2.],
dtype=np.float32) # center of image
scale = max(height, width) * 1.0
flipped = False
if self.split == 'train':
scale = scale * np.random.choice(self.rand_scales)
w_border = get_border(128, width)
h_border = get_border(128, height)
center[0] = np.random.randint(low=w_border, high=width - w_border)
center[1] = np.random.randint(low=h_border, high=height - h_border)
if np.random.random() < 0.5:
flipped = True
img = img[:, ::-1, :]
center[0] = width - center[0] - 1
trans_img = get_affine_transform(
center, scale, 0, [self.img_size['w'], self.img_size['h']])
img = cv2.warpAffine(img, trans_img,
(self.img_size['w'], self.img_size['h']))
# -----------------------------------debug---------------------------------
# image_show = img.copy()
# for bbox, label, shape in zip(bboxes, labels, shapes):
# if flipped:
# bbox[[0, 2]] = width - bbox[[2, 0]] - 1
# # Flip the contour
# for m in range(self.n_vertices):
# shape[2 * m] = width - shape[2 * m] - 1
# bbox[:2] = affine_transform(bbox[:2], trans_img)
# bbox[2:] = affine_transform(bbox[2:], trans_img)
# bbox[[0, 2]] = np.clip(bbox[[0, 2]], 0, self.img_size['w'] - 1)
# bbox[[1, 3]] = np.clip(bbox[[1, 3]], 0, self.img_size['h'] - 1)
#
# # generate gt shape mean and std from contours
# for m in range(self.n_vertices): # apply scale and crop transform to shapes
# shape[2 * m:2 * m + 2] = affine_transform(shape[2 * m:2 * m + 2], trans_img)
#
# contour = np.reshape(shape, (self.n_vertices, 2))
# # Indexing from the left-most vertex, argmin x-axis
# idx = np.argmin(contour[:, 0])
# indexed_shape = np.concatenate((contour[idx:, :], contour[:idx, :]), axis=0)
#
# clockwise_flag = check_clockwise_polygon(indexed_shape)
# if not clockwise_flag:
# fixed_contour = np.flip(indexed_shape, axis=0)
# else:
# fixed_contour = indexed_shape
#
# contour[:, 0] = np.clip(fixed_contour[:, 0], 0, self.img_size['w'] - 1)
# contour[:, 1] = np.clip(fixed_contour[:, 1], 0, self.img_size['h'] - 1)
#
# # cv2.rectangle(image_show, (int(bbox[0]), int(bbox[1])), (int(bbox[2]), int(bbox[3])), (255, 0, 0), 2)
# # cv2.polylines(image_show, [contour.astype(np.int32)], True, (0, 0, 255), thickness=2)
# cv2.drawContours(image_show, [contour.astype(np.int32)],
# color=(random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)),
# contourIdx=-1, thickness=-1)
#
# cv2.imshow('img', image_show)
# cv2.waitKey()
# -----------------------------------debug---------------------------------
img = img.astype(np.float32) / 255.
if self.split == 'train':
color_aug(self.data_rng, img, self.eig_val, self.eig_vec)
img -= self.mean
img /= self.std
img = img.transpose(2, 0, 1) # from [H, W, C] to [C, H, W]
trans_fmap = get_affine_transform(
center, scale, 0, [self.fmap_size['w'], self.fmap_size['h']])
hmap = np.zeros(
(self.num_classes, self.fmap_size['h'], self.fmap_size['w']),
dtype=np.float32) # heatmap
# w_h_ = np.zeros((self.max_objs, 2), dtype=np.float32) # width and height of the shape
w_h_std = np.zeros((self.max_objs, 2),
dtype=np.float32) # width and height of the shape
codes_ = np.zeros((self.max_objs, self.n_codes),
dtype=np.float32) # gt coefficients/codes for shapes
regs = np.zeros(
(self.max_objs, 2),
dtype=np.float32) # regression for offsets of shape center
inds = np.zeros((self.max_objs, ), dtype=np.int64)
ind_masks = np.zeros((self.max_objs, ), dtype=np.uint8)
# detections = []
for k, (bbox, label, shape) in enumerate(zip(bboxes, labels, shapes)):
if flipped:
bbox[[0, 2]] = width - bbox[[2, 0]] - 1
# Flip the contour
for m in range(self.n_vertices):
shape[2 * m] = width - shape[2 * m] - 1
bbox[:2] = affine_transform(bbox[:2], trans_fmap)
bbox[2:] = affine_transform(bbox[2:], trans_fmap)
bbox[[0, 2]] = np.clip(bbox[[0, 2]], 0, self.fmap_size['w'] - 1)
bbox[[1, 3]] = np.clip(bbox[[1, 3]], 0, self.fmap_size['h'] - 1)
h, w = bbox[3] - bbox[1], bbox[2] - bbox[0]
# generate gt shape mean and std from contours
for m in range(self.n_vertices
): # apply scale and crop transform to shapes
shape[2 * m:2 * m + 2] = affine_transform(
shape[2 * m:2 * m + 2], trans_fmap)
contour = np.reshape(shape, (self.n_vertices, 2))
# Indexing from the left-most vertex, argmin x-axis
idx = np.argmin(contour[:, 0])
indexed_shape = np.concatenate(
(contour[idx:, :], contour[:idx, :]), axis=0)
clockwise_flag = check_clockwise_polygon(indexed_shape)
if not clockwise_flag:
fixed_contour = np.flip(indexed_shape, axis=0)
else:
fixed_contour = indexed_shape.copy()
contour[:, 0] = np.clip(fixed_contour[:, 0], 0,
self.fmap_size['w'] - 1)
contour[:, 1] = np.clip(fixed_contour[:, 1], 0,
self.fmap_size['h'] - 1)
contour_mean = np.mean(contour, axis=0)
contour_std = np.std(contour, axis=0)
if np.sqrt(np.sum(contour_std**2)) <= 1e-6:
continue
else:
norm_shape = (contour - contour_mean) / np.sqrt(
np.sum(contour_std**2))
if h > 0 and w > 0 and np.sqrt(np.sum(contour_std**2)) > 1e-6:
obj_c = contour_mean
obj_c_int = obj_c.astype(np.int32)
radius = max(
0,
int(
gaussian_radius((math.ceil(h), math.ceil(w)),
self.gaussian_iou)))
draw_umich_gaussian(hmap[label], obj_c_int, radius)
w_h_std[k] = contour_std
temp_codes, _ = fast_ista(norm_shape.reshape((1, -1)),
self.dictionary,
lmbda=self.sparse_alpha,
max_iter=80)
codes_[k] = np.exp(temp_codes)
regs[k] = obj_c - obj_c_int # discretization error
inds[k] = obj_c_int[1] * self.fmap_size['w'] + obj_c_int[0]
ind_masks[k] = 1
# groundtruth bounding box coordinate with class
# detections.append([obj_c[0] - w / 2, obj_c[1] - h / 2,
# obj_c[0] + w / 2, obj_c[1] + h / 2, 1, label])
# detections = np.array(detections, dtype=np.float32) \
# if len(detections) > 0 else np.zeros((1, 6), dtype=np.float32)
# -----------------------------------debug---------------------------------
# canvas = np.zeros((self.fmap_size['h'] * 2, self.fmap_size['w'] * 2, 3), dtype=np.float32)
# canvas[0:self.fmap_size['h'], 0:self.fmap_size['w'], :] = np.tile(np.expand_dims(hmap[0], 2), (1, 1, 3))
# canvas[0:self.fmap_size['h'], self.fmap_size['w']:, :] = np.tile(np.expand_dims(hmap[1], 2), (1, 1, 3))
# canvas[self.fmap_size['h']:, 0:self.fmap_size['w'], :] = np.tile(np.expand_dims(hmap[2], 2), (1, 1, 3))
# canvas[self.fmap_size['h']:, self.fmap_size['w']:, :] = np.tile(np.expand_dims(hmap[3], 2), (1, 1, 3))
# print(w_h_[0], regs[0])
# cv2.imshow('hmap', canvas)
# cv2.waitKey()
# -----------------------------------debug---------------------------------
# -----------------------------------debug---------------------------------
# image_show = img.copy()
# for bbox, label, shape in zip(bboxes, labels, shapes):
# cv2.rectangle(image_show, (int(bbox[0]), int(bbox[1])), (int(bbox[2]), int(bbox[3])), (255, 0, 0), 2)
# cv2.polylines(image_show, [contour.astype(np.int32)], True, (0, 0, 255), thickness=2)
# cv2.imshow('img', image_show)
# cv2.waitKey()
# -----------------------------------debug---------------------------------
return {
'image': img,
'codes': codes_,
'hmap': hmap,
'w_h_std': w_h_std,
'regs': regs,
'inds': inds,
'ind_masks': ind_masks,
'c': center,
's': scale,
'img_id': img_id
}
