Upload layer_diff_dataset/test_inp_1 copy 4.py with huggingface_hub

layer_diff_dataset/test_inp_1 copy 4.py

@@ -0,0 +1,58 @@
+import cv2
+import torch
+import os
+import json
+from tqdm import tqdm
+from modelscope.outputs import OutputKeys
+from modelscope.pipelines import pipeline
+from modelscope.utils.constant import Tasks
+
+# input_location = 'https://modelscope.oss-cn-beijing.aliyuncs.com/test/images/image_inpainting/image_inpainting_1.png'
+# input_mask_location = 'https://modelscope.oss-cn-beijing.aliyuncs.com/test/images/image_inpainting/image_inpainting_mask_1.png'
+prompt = 'hazy background with nothing on'
+
+folder_path_0 = '/mnt/workspace/workgroup/sihui.jsh/layer_diff_dataset/train/im_resized'
+folder_path = '/mnt/workspace/workgroup/sihui.jsh/layer_diff_dataset/train/gt_dilate'
+folder_path_ = '/mnt/workspace/workgroup/sihui.jsh/layer_diff_dataset/train/gt_inp_iic_gt_dilate'
+os.makedirs(folder_path_,exist_ok=True)
+file_list = os.listdir(folder_path)
+file_list = [i for i in file_list if i.endswith('.png')]
+file_list.sort()
+
+pbar = tqdm(enumerate(file_list),total=len(file_list))
+# with open('/mnt/workspace/workgroup/sihui.jsh/layer_diff_dataset/train/im_rgba.json', 'r') as file:
+#     data = json.load(file)
+image_inpainting = pipeline(
+    Tasks.image_inpainting,
+    model='/mnt/workspace/workgroup/sihui.jsh/alpha_work/diffusers/iic/cv_stable-diffusion-v2_image-inpainting_base',
+    device='cuda:1',
+    torch_dtype=torch.float16,
+    enable_attention_slicing=True)
+for i, image_name in pbar:
+    if os.path.exists(os.path.join(folder_path_,image_name)):
+        continue
+    # if i<1000:
+    #     continue
+    if i<=2000:
+        continue
+    elif i > 2500:
+        break
+
+    # if i==1:
+    #     break
+    # gt_name = data[i]["images"].split('/')[-1]
+    # print(image_name,gt_name)
+    # prompt = data[i]["prompt_fg"]
+    # print(prompt)
+    
+    # mask_image = load_image(os.path.join(folder_path,image_name)).resize((512, 512))
+    # image = load_image(os.path.join(folder_path_0,image_name.split('.png')[0]+'.jpg')).resize((512, 512))
+
+    input = {
+        'image': os.path.join(folder_path_0,image_name.split('.png')[0]+'.jpg'),
+        'mask': os.path.join(folder_path,image_name),
+        'prompt': prompt
+    }
+    
+    output = image_inpainting(input)[OutputKeys.OUTPUT_IMG]
+    cv2.imwrite(os.path.join(folder_path_,image_name), output)