Upload layer_diff_dataset/batch_mask.py with huggingface_hub

layer_diff_dataset/batch_mask.py

@@ -0,0 +1,52 @@
+import cv2
+import numpy as np
+import os
+from pathlib import Path
+
+# 定义文件夹路径
+image_folder = '/home/jsh/nas_jianchong/sihui.jsh/data/video_dataset/YoutubeVIS/train/JPEGImages'
+mask_folder = '/home/jsh/nas_jianchong/sihui.jsh/data/video_dataset/YoutubeVOS/train/Annotations'
+
+# 定义输出文件夹路径
+output_image_folder = '/home/jsh/nas_jianchong/sihui.jsh/data/video_dataset/YoutubeVOS/train/fg'
+output_mask_folder = '/home/jsh/nas_jianchong/sihui.jsh/data/video_dataset/YoutubeVOS/train/mask'
+
+# 创建输出文件夹如果它们不存在
+Path(output_image_folder).mkdir(parents=True, exist_ok=True)
+Path(output_mask_folder).mkdir(parents=True, exist_ok=True)
+
+# 遍历JPEGImages目录中的所有子目录
+for subdir in next(os.walk(image_folder))[1]:
+    print(subdir)
+    subdir_path = os.path.join(image_folder, subdir)
+    subdir_mask_path = os.path.join(mask_folder, subdir)
+    output_image_subdir = os.path.join(output_image_folder, subdir)
+    output_mask_subdir = os.path.join(output_mask_folder, subdir)
+    Path(output_image_subdir).mkdir(parents=True, exist_ok=True)
+    Path(output_mask_subdir).mkdir(parents=True, exist_ok=True)
+    # 遍历子目录里的所有文件
+    for filename in os.listdir(subdir_path):
+        if filename.lower().endswith(('.jpg', '.jpeg', '.png')):
+            # 读入图片和对应的掩码
+            img_path = os.path.join(subdir_path, filename)
+            mask_path = os.path.join(subdir_mask_path, filename.replace('.jpg', '.png'))
+            img = cv2.imread(img_path)
+            mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE).clip(0, 1)
+            
+            # 如果图片和掩码读取成功
+            if img is not None and mask is not None:
+                # 应用掩码
+                mask_image = mask * 255
+                result_img = img * mask[:, :, np.newaxis]
+                
+                # 构建输出文件名并保存图片和掩码
+                # output_img_name = f"{subdir}/{filename}"
+                # output_mask_name = f"{subdir}/{filename}"
+                output_img_path = os.path.join(output_image_subdir, filename.replace('.jpg', '.png'))
+                output_mask_path = os.path.join(output_mask_subdir, filename.replace('.jpg', '.png'))
+                # print(output_img_path)
+                # print(output_mask_path)
+                
+                cv2.imwrite(output_img_path, result_img)
+                cv2.imwrite(output_mask_path, mask_image)
+    # exit(0)