Upload layer_diff_dataset/test_inp_2.py with huggingface_hub

layer_diff_dataset/test_inp_2.py

@@ -0,0 +1,72 @@
+import cv2
+import torch
+from modelscope.outputs import OutputKeys
+from modelscope.pipelines import pipeline
+from modelscope.utils.constant import Tasks
+import PIL
+import numpy as np
+import os
+from tqdm import tqdm
+import random
+
+# 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'#'background'
+data_dir = '/mnt/workspace/workgroup/sihui.jsh/layer_diff_dataset'
+split = 'try'
+input_dir = f'{data_dir}/{split}/im'
+input_mask_dir = f'{data_dir}/{split}/gt'
+output_image_path = f'{data_dir}/{split}/bg_inpaint/'
+if not os.path.exists(output_image_path):
+    os.mkdir(output_image_path)
+image_inpainting = pipeline(
+        Tasks.image_inpainting,
+        model='/mnt/workspace/workgroup/sihui.jsh/alpha_work/diffusers/iic/cv_stable-diffusion-v2_image-inpainting_base',
+        device='gpu',
+        torch_dtype=torch.float16,
+        enable_attention_slicing=True,
+        )
+
+all_images = os.listdir(input_dir)
+all_images.sort()
+# random.shuffle(all_images)
+for img in tqdm(all_images):
+    img_name = img[:-4]
+    input_location = os.path.join(input_dir,img_name+'.jpg')
+    input_mask_location = os.path.join(input_mask_dir,img_name+'.png')
+    input_image = PIL.Image.open(input_location).convert('RGB')
+    mask_image = PIL.Image.open(input_mask_location).convert('L')
+    h,w = input_image.size
+    # print(h,w)
+    target_size = 1024
+    if h>target_size:
+        input_image = input_image.resize((target_size,int(input_image.size[-1]*(target_size/input_image.size[0]))))
+        mask_image = mask_image.resize((target_size,int(input_image.size[-1]*(target_size/input_image.size[0]))))
+    input_image.save('input_image.png')
+    mask_image = np.array(mask_image)
+
+    kernel = np.ones((40, 40), np.uint8)
+    mask_image = cv2.dilate(mask_image,kernel,20)
+    # mask_indice = np.argwhere(mask_image>0)
+    # x0,x1 = mask_indice[:,0].min(),mask_indice[:,0].max(),
+    # y0,y1 = mask_indice[:,1].min(),mask_indice[:,1].max(),
+
+    # mask_new = np.zeros_like(mask_image)
+    # mask_new[x0:x1,y0:y1] = 1
+    
+    mask_image = mask_image / 255.0
+    mask_image = (mask_image>0)
+    PIL.Image.fromarray((mask_image*255.0).astype(np.uint8)).save('mask_dilate.png')
+
+    input = {
+        'image': 'input_image.png',
+        'mask': 'mask_dilate.png',
+        'prompt': prompt
+    }
+    
+    output = image_inpainting(input)[OutputKeys.OUTPUT_IMG]
+    # print(output.shape)
+    input_image = input_image.resize((output.shape[1],output.shape[0]))
+    res = np.concatenate([np.array(input_image)[:,:,::-1],output],axis=1)
+    cv2.imwrite(os.path.join(output_image_path,img_name+'.png'),res)
+# print('pipeline: the output image path is {}'.format(output_image_path))