Upload layer_diff_dataset/test_inp_4_try.py with huggingface_hub

layer_diff_dataset/test_inp_4_try.py

@@ -0,0 +1,69 @@
+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
+
+# 进行图片 impainting
+
+# 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'
+
+root_folder = '../data/aim-500'
+jpeg_folder = os.path.join(root_folder,'original')
+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='cuda:0',
+    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<1000:
+    #     continue
+    # if i==10:
+    #     break
+    # folder_path_0 = 'YoutubeVOS/JPEGImages/0043f083b5'
+    folder_path_0 = os.path.join(jpeg_folder,vid_name)
+    folder_path = os.path.join(mask_folder,vid_name.replace('.jpg','.png'))
+    # folder_path = 'YoutubeVOS/mask_dilate/0043f083b5'
+    # folder_path_ = 'YoutubeVOS/inp/0043f083b5'
+    folder_path_ = os.path.join(inp_folder,vid_name)
+    # 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)
+
+    # for i, image_name in enumerate(file_list):
+    # if os.path.exists(os.path.join(folder_path_,image_name)):
+    #     continue
+    # last_idx = min(15,len(file_list)-1)
+    # if not (i == 0 or i == last_idx):
+    #     continue
+    
+    # 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': folder_path_0,
+        'mask': folder_path,
+        'prompt': prompt
+    }
+    output = image_inpainting(input)[OutputKeys.OUTPUT_IMG]
+    # output = cv2.resize(output,(256,256))
+    cv2.imwrite(folder_path_, output)
+    # exit(0)