Upload layer_diff_dataset/test_inp_sd_3.py with huggingface_hub

layer_diff_dataset/test_inp_sd_3.py

@@ -0,0 +1,102 @@
+from diffusers import AutoPipelineForInpainting
+from diffusers.utils import load_image
+import torch
+import os 
+from tqdm import tqdm
+import cv2
+from PIL import Image
+import imageio
+import numpy as np
+
+pipe = AutoPipelineForInpainting.from_pretrained("../alpha_work/diffusers/stable-diffusion-xl-1.0-inpainting_", torch_dtype=torch.float16, variant="fp16").to("cuda")
+# print('pipe',pipe)
+# StableDiffusionXLInpaintPipeline
+
+preprocessed_root_path = '../data/video_dataset/YoutubeVOS/train/impainting_256'
+# folder_path_0 = '../codes/Inpaint-Anything/results/0b6f9105fc'
+mask_root_path = '../data/video_dataset/YoutubeVOS/train/mask'
+output_root_path = '../data/video_dataset/YoutubeVOS/train/inp_preprocess_sd_0.9_base_ppt_w'
+os.makedirs(output_root_path,exist_ok=True)
+# img_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png"
+# mask_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png"
+
+
+prompt = "hazy background with nothing on"
+generator = torch.Generator(device="cuda").manual_seed(22)
+# base_image = Image.open(base_image_path).resize((1024, 1024))
+
+vid_list = os.listdir(preprocessed_root_path)
+pbar = tqdm(enumerate(vid_list),total=len(vid_list))
+for i, vid_name in pbar:
+    # if not vid_name=='7d18074fef':
+    #     continue
+    # if i<=800:
+    #     continue
+    if i>550:
+        break
+    
+    output_folder = os.path.join(output_root_path, vid_name)
+    output_gif = os.path.join(output_folder, f'{vid_name}.gif')  # 输出GIF的文件名
+    if os.path.exists(output_gif):
+        continue
+    os.makedirs(output_folder,exist_ok=True)
+    
+    img_folder = os.path.join(preprocessed_root_path, vid_name)
+    img_list = os.listdir(img_folder)
+    img_list = [i for i in img_list if i.endswith('.jpg')]
+    img_list.sort()
+    # print(len(img_list))
+    mask_folder = os.path.join(mask_root_path, vid_name)
+    
+    
+    frames = []
+    for i,image_name in enumerate(img_list):
+        
+        # print('i',i)
+        # if not i%2 == 0:
+        #     continue
+        image_path = os.path.join(img_folder,image_name)
+        mask_path = os.path.join(mask_folder,image_name.split('.')[0]+'.png')
+        image = Image.open(image_path).resize((1024, 1024))
+        mask_image = Image.open(mask_path).resize((1024, 1024))
+        # image = cv2.resize(cv2.imread(image_path),(1024,1024))
+        # mask_image = cv2.resize(cv2.imread(mask_path,cv2.IMREAD_GRAYSCALE),(1024,1024))
+        # image = load_image(img_url).resize((1024, 1024))
+        # mask_image = load_image(mask_url).resize((1024, 1024))
+        if i==0:
+            base_image_1 = image
+            base_image = base_image_1
+            strength = 0.99
+        else:
+            strength = 0.5
+            if i==1:
+                base_image_2 = base_image_1
+                base_image_1 = image_out
+                image_array_1 = np.array(base_image_1.convert('RGB'))
+                image_array_2 = np.array(base_image_2.convert('RGB'))
+                base_image = (image_array_1*0.8+image_array_2*0.2).astype(np.uint8)
+                base_image = Image.fromarray(base_image)
+            else:
+                base_image_3 = base_image_2
+                base_image_2 = base_image_1
+                base_image_1 = image_out
+                image_array_1 = np.array(base_image_1.convert('RGB'))
+                image_array_2 = np.array(base_image_2.convert('RGB'))
+                image_array_3 = np.array(base_image_3.convert('RGB'))
+                base_image = (image_array_1*0.7+image_array_2*0.2+image_array_3*0.1).astype(np.uint8)
+                base_image = Image.fromarray(base_image)
+                
+        image_out = pipe(
+        prompt=prompt,
+        image=image,
+        base_image=base_image,
+        mask_image=mask_image,
+        guidance_scale=8.0,
+        num_inference_steps=20,  # steps between 15 and 30 work well for us
+        strength=strength,  # make sure to use `strength` below 1.0
+        generator=generator,
+        ).images[0]
+        image_out.save(os.path.join(output_folder,image_name))
+        frames.append(imageio.imread(os.path.join(output_folder,image_name)))
+    imageio.mimsave(output_gif, frames, fps=8, loop=0)
+    # exit(0)