Skip to content

Commit ada444a

Browse files
stephenyan1231facebook-github-bot
authored andcommitted
minor changes to Visualizer: enforce int type label
Summary: Minor changes - enforce int type label by doing **int(label)** - support draw_text and edge_color arguments in method **draw_sem_seg()** Differential Revision: D77811205 fbshipit-source-id: c1036e46eb94f2788a2d311173987b982a801c70
1 parent b15f64e commit ada444a

File tree

1 file changed

+13
-5
lines changed

1 file changed

+13
-5
lines changed

detectron2/utils/visualizer.py

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -445,7 +445,14 @@ def draw_instance_predictions(self, predictions, jittering: bool = True):
445445
)
446446
return self.output
447447

448-
def draw_sem_seg(self, sem_seg, area_threshold=None, alpha=0.8):
448+
def draw_sem_seg(
449+
self,
450+
sem_seg,
451+
area_threshold=None,
452+
alpha=0.8,
453+
draw_text: bool = True,
454+
edge_color: tuple[float, float, float] = _OFF_WHITE,
455+
):
449456
"""
450457
Draw semantic segmentation predictions/labels.
451458
@@ -454,7 +461,8 @@ def draw_sem_seg(self, sem_seg, area_threshold=None, alpha=0.8):
454461
Each value is the integer label of the pixel.
455462
area_threshold (int): segments with less than `area_threshold` are not drawn.
456463
alpha (float): the larger it is, the more opaque the segmentations are.
457-
464+
draw_text (bool): if True, draw class name text over the mask
465+
edge_color: color of polygon edge
458466
Returns:
459467
output (VisImage): image object with visualizations.
460468
"""
@@ -465,16 +473,16 @@ def draw_sem_seg(self, sem_seg, area_threshold=None, alpha=0.8):
465473
labels = labels[sorted_idxs]
466474
for label in filter(lambda l: l < len(self.metadata.stuff_classes), labels):
467475
try:
468-
mask_color = [x / 255 for x in self.metadata.stuff_colors[label]]
476+
mask_color = [x / 255 for x in self.metadata.stuff_colors[int(label)]]
469477
except (AttributeError, IndexError):
470478
mask_color = None
471479

472480
binary_mask = (sem_seg == label).astype(np.uint8)
473-
text = self.metadata.stuff_classes[label]
481+
text = self.metadata.stuff_classes[int(label)] if draw_text else None
474482
self.draw_binary_mask(
475483
binary_mask,
476484
color=mask_color,
477-
edge_color=_OFF_WHITE,
485+
edge_color=edge_color,
478486
text=text,
479487
alpha=alpha,
480488
area_threshold=area_threshold,

0 commit comments

Comments
 (0)