diff --git a/rt_utils/image_helper.py b/rt_utils/image_helper.py index 5770b42..c5991b6 100644 --- a/rt_utils/image_helper.py +++ b/rt_utils/image_helper.py @@ -271,15 +271,34 @@ def get_slice_mask_from_slice_contour_data( ): # Go through all contours in a slice, create polygons in correct space and with a correct format # and append to polygons array (appropriate for fillPoly) + slice_mask = create_empty_slice_mask(series_slice).astype(np.uint8) polygons = [] + for contour_coords in slice_contour_data: reshaped_contour_data = np.reshape(contour_coords, [len(contour_coords) // 3, 3]) translated_contour_data = apply_transformation_to_3d_points(reshaped_contour_data, transformation_matrix) polygon = [np.around([translated_contour_data[:, :2]]).astype(np.int32)] polygon = np.array(polygon).squeeze() - polygons.append(polygon) - slice_mask = create_empty_slice_mask(series_slice).astype(np.uint8) - cv.fillPoly(img=slice_mask, pts = polygons, color = 1) + + num_points = len(reshaped_contour_data) + + if num_points == 1: + # Single point contour - draw a point marker + point = tuple(polygon.flatten()[:2]) + cv.circle(slice_mask, point, radius=1, color=1, thickness=-1) + elif num_points == 2: + # Two point contour - draw a line + pt1 = tuple(polygon[0]) + pt2 = tuple(polygon[1]) + cv.line(slice_mask, pt1, pt2, color=1, thickness=1) + else: + # Three or more points - use fillPoly + polygons.append(polygon) + + # Fill all valid polygons (3+ points) at once + if polygons: + cv.fillPoly(img=slice_mask, pts=polygons, color=1) + return slice_mask def create_empty_series_mask(series_data):