Close tracked polygon by appending its initial point and generate poi…

…nt grid if 'point' or 'grid' mentioned in description
healthonrails · Mar 13, 2024 · dd5b272 · dd5b272
1 parent 320abe5
commit dd5b272
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/annolid/postprocessing/quality_control.py b/annolid/postprocessing/quality_control.py
@@ -63,14 +63,17 @@ def pred_dict_to_labelme(pred_row,
 
                 shape = Shape(label=label_name,
                               shape_type='polygon',
-                              flags={}
+                              flags={},
+                              visible=True,
                               )
                 all_points = np.array(
                     list(zip(polys[0::2], polys[1::2])))
                 for x, y in all_points:
                     # do not add 0,0 to the list
                     if x >= 1 and y >= 1:
                         shape.addPoint((x, y))
+                # To close the polygon
+                shape.points.append(shape.points[0])
                 label_list.append(shape)
             except IndexError:
                 pass
@@ -311,7 +314,7 @@ def instance_center_distances(self, old_instances, cur_instances):
                 if (ci['frame_number'] == oi['frame_number']
                         and int(ci['cx']) == int(oi['cx'])
                         and int(ci['cy']) == int(oi['cy'])
-                        ):
+                    ):
                     continue
                 dist = np.sqrt((ci['cx'] - oi['cx'])**2 +
                                (ci['cy']-oi['cy']) ** 2)

diff --git a/annolid/tracker/cotracker/track.py b/annolid/tracker/cotracker/track.py
@@ -104,7 +104,11 @@ def _process_shapes(self, shapes, frame_number):
                 processed_queries.append(self._process_point(
                     shape['points'][0], frame_number, label))
             elif (shape_type == 'polygon' and "description" in shape
-                  and shape["description"] is not None):
+                  and shape["description"] is not None
+                  and ('grid' in shape['description'] or
+                       'point' in shape['description']
+                       )
+                  ):
                 processed_queries.extend(self._process_polygon(
                     shape['points'], frame_number, label))