Skip to content

Commit

Permalink
Implement feature to draw KIE linking lines with arrowheads and Add c…
Browse files Browse the repository at this point in the history
…onfiguration to control whether the scores field is displayed (#426)
  • Loading branch information
CVHub520 committed Jul 27, 2024
1 parent 295d1c0 commit b8f98f1
Show file tree
Hide file tree
Showing 3 changed files with 109 additions and 13 deletions.
9 changes: 6 additions & 3 deletions anylabeling/configs/xanylabeling_config.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -9,11 +9,13 @@ keep_prev_contrast: false
auto_use_last_label: false
auto_highlight_shape: true
auto_switch_to_edit_mode: true
show_groups: true
show_groups: false
show_texts: true
show_labels: true
show_scores: true
show_degrees: false
show_shapes: true
show_linking: true
logger_level: info
save_mode: default

Expand All @@ -38,8 +40,8 @@ shape:
select_line_color: [255, 255, 255, 255]
select_fill_color: [0, 255, 0, 155]
hvertex_fill_color: [255, 255, 255, 255]
point_size: 8
line_width: 2
point_size: 10
line_width: 4

# main
flag_dock:
Expand Down Expand Up @@ -131,6 +133,7 @@ shortcuts:
show_overview: Ctrl+G
show_texts: Ctrl+T
show_labels: Ctrl+L
show_linking: Ctrl+K
toggle_auto_use_last_label: Ctrl+Y
toggle_visibility_shapes: Ctrl+H

Expand Down
35 changes: 35 additions & 0 deletions anylabeling/views/labeling/label_widget.py
Original file line number Diff line number Diff line change
Expand Up @@ -807,6 +807,15 @@ def __init__(
checked=self._config["show_labels"],
enabled=True,
)
show_scores = action(
self.tr("&Show Scores"),
self.enable_show_scores,
tip=self.tr("Show score inside shapes"),
icon=None,
checkable=True,
checked=self._config["show_scores"],
enabled=True,
)
show_degrees = action(
self.tr("&Show Degress"),
self.enable_show_degrees,
Expand All @@ -816,6 +825,16 @@ def __init__(
checked=self._config["show_degrees"],
enabled=True,
)
show_linking = action(
self.tr("&Show KIE Linking"),
self.enable_show_linking,
shortcut=shortcuts["show_linking"],
tip=self.tr("Show KIE linking between key and value"),
icon=None,
checkable=True,
checked=self._config["show_linking"],
enabled=True,
)

# Languages
select_lang_en = action(
Expand Down Expand Up @@ -1160,7 +1179,9 @@ def __init__(
show_groups=show_groups,
show_texts=show_texts,
show_labels=show_labels,
show_scores=show_scores,
show_degrees=show_degrees,
show_linking=show_linking,
zoom_actions=zoom_actions,
open_next_image=open_next_image,
open_prev_image=open_prev_image,
Expand Down Expand Up @@ -1366,7 +1387,9 @@ def __init__(
set_cross_line,
show_texts,
show_labels,
show_scores,
show_degrees,
show_linking,
show_groups,
hide_selected_polygons,
show_hidden_polygons,
Expand Down Expand Up @@ -3134,12 +3157,24 @@ def enable_show_labels(self, enabled):
self.canvas.set_show_labels(enabled)
save_config(self._config)

def enable_show_scores(self, enabled):
self._config["show_scores"] = enabled
self.actions.show_scores.setChecked(enabled)
self.canvas.set_show_scores(enabled)
save_config(self._config)

def enable_show_degrees(self, enabled):
self._config["show_degrees"] = enabled
self.actions.show_degrees.setChecked(enabled)
self.canvas.set_show_degrees(enabled)
# save_config(self._config)

def enable_show_linking(self, enabled):
self._config["show_linking"] = enabled
self.actions.show_linking.setChecked(enabled)
self.canvas.set_show_linking(enabled)
save_config(self._config)

def on_new_brightness_contrast(self, qimage):
self.canvas.load_pixmap(
QtGui.QPixmap.fromImage(qimage), clear_shapes=False
Expand Down
78 changes: 68 additions & 10 deletions anylabeling/views/labeling/widgets/canvas.py
Original file line number Diff line number Diff line change
Expand Up @@ -104,10 +104,12 @@ def __init__(self, *args, **kwargs):
# Set widget options.
self.setMouseTracking(True)
self.setFocusPolicy(QtCore.Qt.WheelFocus)
self.show_shape_groups = True
self.show_groups = False
self.show_texts = True
self.show_labels = True
self.show_shape_degrees = False
self.show_scores = True
self.show_degrees = False
self.show_linking = True

# Set cross line options.
self.cross_line_show = True
Expand Down Expand Up @@ -1037,7 +1039,7 @@ def paintEvent(self, event): # noqa: C901
return

# Draw groups
if self.show_shape_groups:
if self.show_groups:
pen = QtGui.QPen(QtGui.QColor("#AAAAAA"), 2, Qt.SolidLine)
p.setPen(pen)
grouped_shapes = {}
Expand Down Expand Up @@ -1099,6 +1101,50 @@ def paintEvent(self, event): # noqa: C901
)
p.drawRect(wrap_rect)

# Draw KIE linking
if self.show_linking:
pen = QtGui.QPen(QtGui.QColor("#AAAAAA"), 2, Qt.SolidLine)
p.setPen(pen)
gid2point = {}
linking_pairs = []
group_color = (255, 128, 0)
for shape in self.shapes:
try:
linking_pairs += shape.kie_linking
except:
pass

if (shape.group_id is None
or shape.shape_type not in
["rectangle", "polygon", "rotation"]):
continue
rect = shape.bounding_rect()
cx = rect.x() + (rect.width() / 2.)
cy = rect.y() + (rect.height() / 2.)
gid2point[shape.group_id] = (cx, cy)

for linking in linking_pairs:
pen.setStyle(Qt.SolidLine)
pen.setWidth(max(1, int(round(4.0 / Shape.scale))))
pen.setColor(QtGui.QColor(*group_color))
p.setPen(pen)
key, value = linking
# Adapt to the 'ungroup_selected_shapes' operation
if key not in gid2point or value not in gid2point:
continue
kp, vp = gid2point[key], gid2point[value]
# Draw a link from key point to value point
p.drawLine(QtCore.QPointF(*kp), QtCore.QPointF(*vp))
# Draw the triangle arrowhead
arrow_size = max(1, int(round(10.0 / Shape.scale))) # Size of the arrowhead
angle = math.atan2(vp[1] - kp[1], vp[0] - kp[0]) # Angle towards the value point
arrow_points = [
QtCore.QPointF(vp[0], vp[1]),
QtCore.QPointF(vp[0] - arrow_size * math.cos(angle - math.pi / 6), vp[1] - arrow_size * math.sin(angle - math.pi / 6)),
QtCore.QPointF(vp[0] - arrow_size * math.cos(angle + math.pi / 6), vp[1] - arrow_size * math.sin(angle + math.pi / 6)),
]
p.drawPolygon(arrow_points)

# Draw degrees
for shape in self.shapes:
if (
Expand All @@ -1119,7 +1165,7 @@ def paintEvent(self, event): # noqa: C901
(shape.points[0].x() + shape.points[2].x()) / 2,
(shape.points[0].y() + shape.points[2].y()) / 2,
)
if self.show_shape_degrees:
if self.show_degrees:
degrees = str(int(math.degrees(shape.direction))) + "°"
p.setFont(
QtGui.QFont(
Expand Down Expand Up @@ -1180,12 +1226,14 @@ def paintEvent(self, event): # noqa: C901

# Draw texts
if self.show_texts:
text_color = "#FFFFFF"
background_color = "#007BFF"
p.setFont(
QtGui.QFont(
"Arial", int(max(6.0, int(round(8.0 / Shape.scale))))
)
)
pen = QtGui.QPen(QtGui.QColor("#00FF00"), 8, Qt.SolidLine)
pen = QtGui.QPen(QtGui.QColor(background_color), 8, Qt.SolidLine)
p.setPen(pen)
for shape in self.shapes:
description = shape.description
Expand All @@ -1198,14 +1246,14 @@ def paintEvent(self, event): # noqa: C901
int(rect.y() + bbox.y()),
int(rect.width()),
int(rect.height()),
QtGui.QColor("#00FF00"),
QtGui.QColor(background_color),
)
p.drawText(
int(bbox.x()),
int(bbox.y()),
description,
)
pen = QtGui.QPen(QtGui.QColor("#000000"), 8, Qt.SolidLine)
pen = QtGui.QPen(QtGui.QColor(text_color), 8, Qt.SolidLine)
p.setPen(pen)
for shape in self.shapes:
description = shape.description
Expand Down Expand Up @@ -1289,7 +1337,7 @@ def paintEvent(self, event): # noqa: C901
) + (
f"{shape.label}"
) + (
f" {shape.score:.2f}" if shape.score is not None else ""
f" {shape.score:.2f}" if (shape.score is not None and self.show_scores) else ""
)
if not label:
continue
Expand Down Expand Up @@ -1701,7 +1749,7 @@ def set_cross_line(self, show, width, color, opacity):

def set_show_groups(self, enabled):
"""Set showing shape groups"""
self.show_shape_groups = enabled
self.show_groups = enabled
self.update()

def set_show_texts(self, enabled):
Expand All @@ -1714,9 +1762,19 @@ def set_show_labels(self, enabled):
self.show_labels = enabled
self.update()

def set_show_scores(self, enabled):
"""Set showing scores"""
self.show_scores = enabled
self.update()

def set_show_degrees(self, enabled):
"""Set showing degrees"""
self.show_shape_degrees = enabled
self.show_degrees = enabled
self.update()

def set_show_linking(self, enabled):
"""Set showing KIE linking"""
self.show_linking = enabled
self.update()

def gen_new_group_id(self):
Expand Down

0 comments on commit b8f98f1

Please sign in to comment.