Update run/transform.py

run/transform.py

@@ -1,46 +1,48 @@
-import os
-import cv2
-import numpy as np
-
-
-def rgb_to_binary_mask(img_path, output_path):
-    '''
-    Convert an RGB image to a binary mask where non-black pixels are set to 1, black pixels are 0.
-    '''
-
-    img_bgr = cv2.imread(img_path)
-
-    img_binary_mask = np.zeros(img_bgr.shape[:2], dtype=np.uint8)
-
-    non_black_pixels = np.any(img_bgr != [0, 0, 0], axis=-1)
-    img_binary_mask[non_black_pixels] = 1
-
-    cv2.imwrite(output_path, img_binary_mask)
-
-
-def convert_rgb_to_binary(Dataset_Path):
-    '''
-    Convert all RGB images in the dataset to binary mask images.
-    '''
-
-    img_dir = os.path.join(Dataset_Path, 'annotation')
-
-    ann_target_dir = os.path.join(Dataset_Path, 'ann_dir')
-
-
-    for img_name in os.listdir(img_dir):
-        try:
-            img_path = os.path.join(img_dir, img_name)
-            mask_path = os.path.join(ann_target_dir, f'{os.path.splitext(img_name)[0]}.png')
-
-            rgb_to_binary_mask(img_path, mask_path)
-
-        except Exception as e:
-            print(f"Failed to convert {img_name}: {e}")
-
-    print("Conversion completed.")
-
-
-if __name__ == '__main__':
-    Dataset_Path = 'J:/dataset_panicle'
-    convert_rgb_to_binary(Dataset_Path)
+import os
+import cv2
+import numpy as np
+
+
+def rgb_to_binary_mask(img_path, output_path):
+    '''
+    Convert an RGB image to a binary mask where non-black pixels are set to 1, black pixels are 0.
+    '''
+
+    img_bgr = cv2.imread(img_path)
+
+    img_binary_mask = np.zeros(img_bgr.shape[:2], dtype=np.uint8)
+
+    non_black_pixels = np.any(img_bgr != [0, 0, 0], axis=-1)
+    img_binary_mask[non_black_pixels] = 1
+
+    cv2.imwrite(output_path, img_binary_mask)
+
+
+def convert_rgb_to_binary(Dataset_Path):
+    '''
+    Convert all RGB images in the dataset to binary mask images.
+    '''
+
+    img_dir = os.path.join(Dataset_Path, 'annotation')
+
+    ann_target_dir = os.path.join(Dataset_Path, 'ann_dir')
+        if not os.path.exists(ann_target_dir):
+            os.mkdir(ann_target_dir)
+
+
+    for img_name in os.listdir(img_dir):
+        try:
+            img_path = os.path.join(img_dir, img_name)
+            mask_path = os.path.join(ann_target_dir, f'{os.path.splitext(img_name)[0]}.png')
+
+            rgb_to_binary_mask(img_path, mask_path)
+
+        except Exception as e:
+            print(f"Failed to convert {img_name}: {e}")
+
+    print("Conversion completed.")
+
+
+if __name__ == '__main__':
+    Dataset_Path = 'CVRPDataset'
+    convert_rgb_to_binary(Dataset_Path)