Upload layer_diff_dataset/test_inp_4 copy 4.py with huggingface_hub

layer_diff_dataset/test_inp_4 copy 4.py

@@ -0,0 +1,85 @@
+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'
+
+root_folder = '../data/video_dataset/YoutubeVOS/train'
+jpeg_folder = os.path.join(root_folder,'JPEGImages')
+mask_folder = os.path.join(root_folder,'mask_dilate')
+inp_folder = os.path.join(root_folder,'inp_image_256')
+os.makedirs(inp_folder,exist_ok=True)
+vid_list = os.listdir(jpeg_folder)
+# print('vid_list',vid_list)
+# exit(0)
+
+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:2',
+    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<=800:
+        continue
+    if i>1000:
+        break
+    if not vid_name == '45f8128b97':
+        continue
+    # 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)
+    # 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)):
+            print('here')
+            continue
+        last_idx = min(15,len(file_list)-1)
+        if not (i == 0 or i == last_idx):
+            print("there")
+            continue
+        # print('i',i)
+        # if i<1000:
+        #     continue
+        # if i>500:
+        #     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]
+        output = cv2.resize(output,(256,256))
+        cv2.imwrite(os.path.join(folder_path_,image_name), output)
+    # exit(0)