def transform_points(self, points):
mode = self.REG_MODE
points = np.array(points)
center = points.mean(axis=0)
if mode == 'offsets-polar':
points = ut.organize_polygon_points(points)
return center, ut.cartesian2polar(points - center,
to_degree=True).reshape((-1))
if mode == 'fcos':
points = ut.bounding_rect(points).reshape((2, 2))
center = points.mean(axis=0)
return center, np.abs(points - center).reshape(-1)
if mode == 'centernet':
p1, p2 = points = ut.bounding_rect(points).reshape((2, 2))
w, h = abs(p1[0] - p2[0]), abs(p1[1] - p2[1])
center = points.mean(axis=0)
return center, np.array([w, h])
if mode == 'circle':
points = ut.organize_polygon_points(points)
center = points.mean(axis=0)
radius = 0.5 * ut.polygon_length(points) / len(points)
return center, radius
if mode == 'offsets':
points = ut.organize_polygon_points(points)
return center, (points - center).reshape((-1))
else:
raise
def data_augmentation(self, img, objs):
objs.sort(key=lambda obj: np.array(obj[0])[:, 0].mean())
shapes = [ut.organize_polygon_points(shape) for shape, cls in objs]
clses = [cls for shape, cls in objs]
if self.transform:
img, shapes = self.transform(img, shapes)
objs = list(zip(shapes, clses))
return img, objs
def transform_points(self, points):
points_transform = self.params['points_transform'].split(',')
points = ut.organize_polygon_points(points)
points = np.array(points)
center = points.mean(axis=0)
offsets = points - center
if 'none' in points_transform:
return points
if 'cartesian' in points_transform:
offsets = offsets.reshape((-1))
elif 'polar' in points_transform:
offsets = ut.cartesian2polar(offsets, to_degree=True).reshape((-1))
if 'abs' in points_transform:
offsets = np.abs(offsets)
return offsets
def process_batch_annots(self, batch_annots):
down = self.params['heatmap_downsample']
batch_size = len(batch_annots)
imw, imh = self.input_size
hm_h, hm_w = int(imh / down), int(imw / down)
batch_imgs = np.zeros((batch_size, 3, imh, imw), dtype=np.float32)
batch_center_heatmap = np.zeros(
(batch_size, self.num_classes, hm_h, hm_w), dtype=np.float32)
batch_corner_heatmap = np.zeros(
(batch_size, self.num_classes, hm_h, hm_w), dtype=np.float32)
batch_offsets = np.zeros((batch_size, self.num_points * 2, hm_h, hm_w),
dtype=np.float32)
batch_offsets_mask = np.zeros((batch_size, 1, hm_h, hm_w),
dtype=np.float32)
for i, (img_path, objs) in enumerate(batch_annots):
img = cv2.imread(img_path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img, objs = self.data_augmentation(img, objs)
img = np.array(img)
img = img / 255
img = ut.normalize(img)
batch_imgs[i, ...] = np.transpose(img, (2, 0, 1))
for points, cls in objs:
points = ut.organize_polygon_points(points)
points = np.array(points)
offsets = (points - points.mean(axis=0)).reshape((-1))
cx, cy = center = (points.mean(axis=0) / down).astype(np.int)
points = (points / down).astype(np.int)
ut.draw_points_heatmap(batch_center_heatmap[i, cls], [center])
batch_offsets[i, :, cy, cx] = offsets / down
batch_offsets_mask[i, 0, cy, cx] = 1
for pnt in points:
ut.draw_points_heatmap(batch_corner_heatmap[i, cls], [pnt])
batch_imgs = torch.from_numpy(batch_imgs).to(self.device)
batch_center_heatmap = torch.from_numpy(batch_center_heatmap).to(
self.device)
batch_corner_heatmap = torch.from_numpy(batch_corner_heatmap).to(
self.device)
batch_offsets = torch.from_numpy(batch_offsets).to(self.device)
batch_offsets_mask = torch.from_numpy(batch_offsets_mask).to(
self.device)
batch_labels = dict(
center_heatmap=batch_center_heatmap,
corner_heatmap=batch_corner_heatmap,
offsets=batch_offsets,
offsets_mask=batch_offsets_mask,
)
return batch_imgs, batch_labels
def test(cfg):
ut.set_default_font_path(cfg.FONT_PATH)
detector = cfg.get_detector()
num_params = 4
det_type = 'bbox'
dir = cfg.TEST_DIR
saver = wpcv.ImageSaver(dir + '_out',
remake_dir=True,
auto_make_subdir=True)
fs = glob.glob(dir + '/*.bmp')
fs.sort()
font = ut.get_default_font(32)
for i, f in enumerate(fs):
img = cv2.imread(f)
pim = wpcv.pilimg(img)
polygons = detector.predict(img)
if not len(polygons):
continue
polys = polygons[:, :num_params].astype(np.int)
# polys=list(filter(lambda poly:poly[num_params]>0.1,))
im = pim
for j, poly in enumerate(polys):
score = polygons[j][num_params]
label = '%.2f' % score
if det_type == 'bbox':
im = wpcv.draw_boxes_with_label(im, [(poly, label)],
line_width=2,
box_color='red',
font=font,
text_color='blue',
offset=(0, -18))
elif det_type == 'circle':
pass
elif False:
poly = np.array(poly).reshape((-1, 2))
box = wpcv.bounding_rect(poly)
im = wpcv.draw_boxes_with_label(im, [(box, label)],
line_width=2,
box_color='red',
font=font,
text_color='blue',
offset=(0, -18))
else:
assert det_type == 'polygon'
poly = np.array(poly).reshape((-1, 2))
poly = ut.organize_polygon_points(poly)
im = wpcv.draw_polygon(im,
poly,
color='red',
width=4,
label=label,
label_xy=(0, -18),
label_color='red',
font=font)
if hasattr(detector, 'get_middle_results'):
res = detector.get_middle_results()
saver.save(res['center_heatmap_img'])
saver.save(res['corner_heatmap_img'])
saver.save(im)
# im=visualize(pim,polygons)
# saver.save(im)
print(i, f, len(polygons))