@@ -77,85 +77,66 @@ def choose_biggest_detection(result: ultralytics.engine.results.Results, trackin
7777 mask = yolo_mask
7878 return box , mask
7979
80- def convert_segment_masks_to_yolo_segmentation_labels (masks_dir , output_dir , pixel_to_class_mapping ):
80+ def _get_unique_pixel_values (mask : Mask ) -> list [int ]:
81+ # get unique values except background (0)
82+ unique_values = np .unique (mask ).tolist ()
83+ if 0 in unique_values : unique_values .remove (0 ) # remove background class
84+ return unique_values
85+
86+ def convert_segment_masks_to_yolo_labels (masks_dir , output_dir_segmentation_labels , output_dir_detection_labels , pixel_to_class_mapping ):
8187 """
8288 pixel_to_class_mapping is a dict providing a mapping from pixel value to class id.
8389 e.g. if you only have a single class with id 0 and binary masks use pixel value 255 then this would be:
8490 pixel_to_class_mapping = {255: 0}
8591
86- source : ultralytics.data.converter.convert_segment_masks_to_yolo_seg
92+ Based of : ultralytics.data.converter.convert_segment_masks_to_yolo_seg
8793 """
94+ def get_yolo_box (contour ) -> tuple [float ]:
95+ x , y , w , h = cv2 .boundingRect (contour )
96+ h , w = mask .shape [:2 ]
97+ center_x = x + w / 2
98+ center_y = y + h / 2
99+ yolo_box = center_x / w , center_y / h , w / w , h / h
100+ return yolo_box
101+
88102 for mask_path in Path (masks_dir ).iterdir ():
89103 if mask_path .suffix in {".png" , ".jpg" }:
90104 mask = cv2 .imread (str (mask_path ), cv2 .IMREAD_GRAYSCALE )
91105 img_height , img_width = mask .shape
92106
93- unique_values = np .unique (mask ) # Get unique pixel values representing different classes
94- yolo_format_data = []
107+ unique_values = _get_unique_pixel_values (mask )
108+ yolo_segmentation_format_data = []
109+ yolo_detection_format_data = []
95110
96111 for value in unique_values :
97- if value == 0 :
98- continue # Skip background
99112 class_index = pixel_to_class_mapping .get (value , - 1 )
100113 if class_index == - 1 :
101114 print (f"Unknown class for pixel value { value } { mask_path } )
102115 continue
103116
104117 # Create a binary mask for the current class and find contours
105- contours , _ = cv2 .findContours (
106- (mask == value ).astype (np .uint8 ), cv2 .RETR_EXTERNAL , cv2 .CHAIN_APPROX_SIMPLE
107- ) # Find contours
118+ binary_mask_for_current_class = (mask == value ).astype (np .uint8 )
119+ contours , _ = cv2 .findContours (binary_mask_for_current_class , cv2 .RETR_EXTERNAL , cv2 .CHAIN_APPROX_SIMPLE )
108120
109121 for contour in contours :
110122 if len (contour ) >= 3 : # YOLO requires at least 3 points for a valid segmentation
111123 contour = contour .squeeze () # Remove single-dimensional entries
112- yolo_format = [class_index ]
124+ yolo_segmentation_format = [class_index ]
113125 for point in contour :
114126 # Normalize the coordinates
115- yolo_format .append (round (point [0 ] / img_width , 6 )) # Rounding to 6 decimal places
116- yolo_format .append (round (point [1 ] / img_height , 6 ))
117- yolo_format_data .append (yolo_format )
127+ yolo_segmentation_format .append (round (point [0 ] / img_width , 6 )) # Rounding to 6 decimal places
128+ yolo_segmentation_format .append (round (point [1 ] / img_height , 6 ))
129+ yolo_segmentation_format_data .append (yolo_segmentation_format )
130+ yolo_detection_format_data .append (get_yolo_box (contour ))
131+
118132 # Save Ultralytics YOLO format data to file
119- output_path = Path (output_dir ) / f"{ mask_path .stem }
133+ output_path = Path (output_dir_segmentation_labels ) / f"{ mask_path .stem }
120134 with open (output_path , "w" , encoding = "utf-8" ) as file :
121- for item in yolo_format_data :
135+ for item in yolo_segmentation_format_data :
136+ line = " " .join (map (str , item ))
137+ file .write (line + "\n " )
138+ output_path = Path (output_dir_detection_labels ) / f"{ mask_path .stem }
139+ with open (output_path , "w" , encoding = "utf-8" ) as file :
140+ for item in yolo_detection_format_data :
122141 line = " " .join (map (str , item ))
123142 file .write (line + "\n " )
124-
125-
126- def convert_binary_mask_to_yolo_detection_labels (masks_dir , output_dir , pixel_to_class_mapping ):
127- """
128- pixel_to_class_mapping is a dict providing a mapping from pixel value to class id.
129- e.g. if you only have a single class with id 0 and binary masks use pixel value 255 then this would be:
130- pixel_to_class_mapping = {255: 0}
131-
132- """
133-
134- def _convert_binary_mask_to_yolo_detection_labels (mask : Mask ) -> tuple [float ]:
135- t , l , b , r = mask_utils .get_box (mask )
136- h , w = mask .shape [:2 ]
137- box_width = r - l
138- box_height = b - t
139- box_center_x = l + box_width / 2
140- box_center_y = t + box_height / 2
141- yolo_box = box_center_x / w , box_center_y / h , box_width / w , box_height / h
142- return yolo_box
143-
144- def _get_class_id (mask : Mask ) -> int :
145- unique_values = np .unique (mask ).tolist ()
146- if 0 in unique_values : unique_values .remove (0 ) # remove background class
147- assert len (unique_values ) == 1 , f"only single class / binary segmentation mask supported but found these values: { unique_values }
148- mask_val = unique_values [0 ]
149-
150- class_id = pixel_to_class_mapping .get (mask_val , - 1 )
151- assert class_id != - 1 , f"Unknown class for pixel value { mask_val } { mask_path }
152- return class_id
153-
154- for mask_path in Path (masks_dir ).iterdir ():
155- if mask_path .suffix in {".png" , ".jpg" }:
156- mask = cv2 .imread (str (mask_path ), cv2 .IMREAD_GRAYSCALE )
157- class_id = _get_class_id (mask )
158- yolo_box = _convert_binary_mask_to_yolo_detection_labels (mask )
159- label_file_path = Path (output_dir ).joinpath (Path (mask_path ).with_suffix ('.txt' ).name )
160- with open (label_file_path , 'a' ) as file :
161- file .write (f"{ class_id } { yolo_box [0 ]} { yolo_box [1 ]} { yolo_box [2 ]} { yolo_box [3 ]} )
0 commit comments