diff --git a/object-locator/train.py b/object-locator/train.py index 751c2bc..bed9de2 100644 --- a/object-locator/train.py +++ b/object-locator/train.py @@ -414,10 +414,11 @@ titles=['(Validation) Image w/ output heatmap'], window_ids=[5]) else: - # Send heatmap with a circle at the estimated centroids to Visdom + # Send heatmap with a cross at the estimated centroids to Visdom img_with_x = utils.paint_circles(img=orig_img_w_heatmap_origsize, - points=centroids_wrt_orig, - color='red') + points=centroids_wrt_orig, + color='red', + crosshair=True ) log.image(imgs=[img_with_x], titles=['(Validation) Image w/ output heatmap ' 'and point estimations'], diff --git a/object-locator/utils.py b/object-locator/utils.py index 1626ed9..3e7b1fe 100644 --- a/object-locator/utils.py +++ b/object-locator/utils.py @@ -283,7 +283,7 @@ def overlay_heatmap(img, map, colormap=matplotlib.cm.viridis): return img_w_heatmap -def paint_circles(img, points, color='red'): +def paint_circles(img, points, color='red', crosshair=False): """ Paint points as circles on top of an image. @@ -292,6 +292,9 @@ def paint_circles(img, points, color='red'): First dimension must be color. :param centroids: List of centroids in (y, x) format. :param color: String of the color used to paint centroids. + Default: 'red'. + :param crosshair: Paint crosshair instead of circle. + Default: False. :return: Image with painted circles centered on the points. First dimension is be color. """ @@ -306,8 +309,14 @@ def paint_circles(img, points, color='red'): points = points.round().astype(np.uint16) img = np.moveaxis(img, 0, 2).copy() - for y, x in points: - img = cv2.circle(img, (x, y), 3, color, -1) + if not crosshair: + for y, x in points: + img = cv2.circle(img, (x, y), 3, color, -1) + else: + for y, x in points: + img = cv2.drawMarker(img, + (x, y), + color, cv2.MARKER_TILTED_CROSS, 9, 3, cv2.LINE_AA) img = np.moveaxis(img, 2, 0) return img