From 723a77285af8770c02a918a4a6f3c4fbf624cd7e Mon Sep 17 00:00:00 2001 From: cvhub Date: Fri, 27 Sep 2024 01:32:01 +0800 Subject: [PATCH] revert: Reset SAM2 Video tracking to align with original version (#651) --- .../auto_labeling/segment_anything_2_video.py | 59 +++++++++---------- 1 file changed, 29 insertions(+), 30 deletions(-) diff --git a/anylabeling/services/auto_labeling/segment_anything_2_video.py b/anylabeling/services/auto_labeling/segment_anything_2_video.py index c2ef7f6..95c5ca9 100644 --- a/anylabeling/services/auto_labeling/segment_anything_2_video.py +++ b/anylabeling/services/auto_labeling/segment_anything_2_video.py @@ -262,18 +262,17 @@ def post_process(self, masks, index=None): point[0] = int(point[0]) point[1] = int(point[1]) shape.add_point(QtCore.QPointF(point[0], point[1])) - break - # Create Polygon shape - shape.shape_type = "polygon" - shape.group_id = ( - self.group_ids[index] if index is not None else None - ) - shape.closed = True - shape.label = ( - "AUTOLABEL_OBJECT" if index is None else self.labels[index] - ) - shape.selected = False - shapes.append(shape) + # Create Polygon shape + shape.shape_type = "polygon" + shape.group_id = ( + self.group_ids[index] if index is not None else None + ) + shape.closed = True + shape.label = ( + "AUTOLABEL_OBJECT" if index is None else self.labels[index] + ) + shape.selected = False + shapes.append(shape) elif self.output_mode in ["rectangle", "rotation"]: x_min = 100000000 y_min = 100000000 @@ -293,24 +292,24 @@ def post_process(self, masks, index=None): y_min = min(y_min, point[1]) x_max = max(x_max, point[0]) y_max = max(y_max, point[1]) - # Create Rectangle shape - shape = Shape(flags={}) - shape.add_point(QtCore.QPointF(x_min, y_min)) - shape.add_point(QtCore.QPointF(x_max, y_min)) - shape.add_point(QtCore.QPointF(x_max, y_max)) - shape.add_point(QtCore.QPointF(x_min, y_max)) - shape.shape_type = ( - "rectangle" if self.output_mode == "rectangle" else "rotation" - ) - shape.closed = True - shape.group_id = ( - self.group_ids[index] if index is not None else None - ) - shape.label = ( - "AUTOLABEL_OBJECT" if index is None else self.labels[index] - ) - shape.selected = False - shapes.append(shape) + # Create Rectangle shape + shape = Shape(flags={}) + shape.add_point(QtCore.QPointF(x_min, y_min)) + shape.add_point(QtCore.QPointF(x_max, y_min)) + shape.add_point(QtCore.QPointF(x_max, y_max)) + shape.add_point(QtCore.QPointF(x_min, y_max)) + shape.shape_type = ( + "rectangle" if self.output_mode == "rectangle" else "rotation" + ) + shape.closed = True + shape.group_id = ( + self.group_ids[index] if index is not None else None + ) + shape.label = ( + "AUTOLABEL_OBJECT" if index is None else self.labels[index] + ) + shape.selected = False + shapes.append(shape) return shapes