Upload layer_diff_dataset/test_inp_4_try_index.py with huggingface_hub

layer_diff_dataset/test_inp_4_try_index.py

@@ -0,0 +1,65 @@
+import cv2
+import torch
+import os
+from tqdm import tqdm
+from modelscope.outputs import OutputKeys
+from modelscope.pipelines import pipeline
+from modelscope.utils.constant import Tasks
+import argparse
+
+# 创建命令行参数解析器
+def parse_arguments():
+    parser = argparse.ArgumentParser(description="Image Inpainting with Pipeline")
+    parser.add_argument('--index', type=int, default=0,
+                        help='Index for selecting images to process (default: 0)')
+    return parser.parse_args()
+
+def main(args):
+    INDEX = args.index
+
+    # 进行图片 inpainting
+    prompt = 'hazy background with nothing on'
+
+    root_folder = '../data/P3M-10k/validation/P3M-500-NP'
+    jpeg_folder = os.path.join(root_folder, 'original_image')
+    mask_folder = os.path.join(root_folder, 'mask_dilate')
+    inp_folder = os.path.join(root_folder, 'inpainting')
+    os.makedirs(inp_folder, exist_ok=True)
+    vid_list = os.listdir(jpeg_folder)
+
+    image_inpainting = pipeline(
+        Tasks.image_inpainting,
+        model='/mnt/workspace/workgroup/sihui.jsh/alpha_work/diffusers/iic/cv_stable-diffusion-v2_image-inpainting_base',
+        device=f'cuda:{INDEX}',  # 使用 index 来选择 GPU
+        torch_dtype=torch.float32,
+        enable_attention_slicing=True
+    )
+    
+    pbar = tqdm(enumerate(vid_list), total=len(vid_list))
+    for i, vid_name in pbar:
+        if i < INDEX * 125:
+            continue
+        elif i >= INDEX * 125 + 125:
+            break
+        folder_path_0 = os.path.join(jpeg_folder, vid_name)
+        folder_path = os.path.join(mask_folder, vid_name.replace('.jpg', '.png'))
+        folder_path_ = os.path.join(inp_folder, vid_name)
+        
+        if os.path.exists(folder_path_):
+            continue
+        h, w = cv2.imread(folder_path_0).shape[:2]
+        if h*w > 2000000:
+            image_tmp = cv2.resize(cv2.imread(folder_path_0), (int(w*0.8), int(h*0.8)))
+            cv2.imwrite('tmp.jpg', image_tmp)
+            folder_path_0 = 'tmp.jpg'
+        input_data = {
+            'image': folder_path_0,
+            'mask': folder_path,
+            'prompt': prompt
+        }
+        output = image_inpainting(input_data)[OutputKeys.OUTPUT_IMG]
+        cv2.imwrite(folder_path_, output)
+
+if __name__ == "__main__":
+    args = parse_arguments()
+    main(args)