Upload layer_diff_dataset/try_ani.py with huggingface_hub

layer_diff_dataset/try_ani.py

@@ -0,0 +1,89 @@
+import cv2
+import numpy as np
+import random
+import os
+
+image_folder = 'train/im'
+mask_folder = 'train/gt'
+image_list = os.listdir(image_folder)
+root_img_dir = 'train/jpeg'
+root_mask_dir = 'train/mask'
+os.makedirs(root_img_dir, exist_ok=True)
+os.makedirs(root_mask_dir, exist_ok=True)
+for idx, image in enumerate(image_list):
+    # if idx == 0:
+    #     continue
+    if idx == 10:
+        break
+    os.makedirs(os.path.join(root_img_dir, image.split('.jpg')[0]), exist_ok=True)
+    os.makedirs(os.path.join(root_mask_dir, image.split('.jpg')[0]), exist_ok=True)
+    image_name = os.path.join(image_folder, image)
+    mask_name = os.path.join(mask_folder, image.split('.jpg')[0]+'.png')
+    print(image_name, mask_name)
+    # image_name = 'train/im/1#Accessories#1#Bag#2339506821_83cf9f1d22_o.jpg'
+    # mask_name = 'train/gt/1#Accessories#1#Bag#2339506821_83cf9f1d22_o.png'
+    img = cv2.imread(image_name)
+    img = cv2.resize(img, (512, 512))
+    # cv2.imwrite('0.jpg', img)
+    H,W = img.shape[0], img.shape[1]
+    mask = cv2.imread(mask_name, cv2.IMREAD_GRAYSCALE)
+    mask = cv2.resize(mask, (512, 512))
+    # cv2.imwrite('1.jpg', mask)
+    # print(img.shape, mask.shape)
+    bboxes_y, bboxes_x = np.nonzero(mask)
+
+    sample_n_frames = 16
+    option = random.randint(0,2)
+    # option = 0  
+
+
+    # root_img = 'tmp_1/tmp_img'
+    # root_mask = 'tmp_1/tmp_mask'
+    # os.makedirs(root_img, exist_ok=True)
+    # os.makedirs(root_mask, exist_ok=True)
+    # root_img = 'tmp_1/tmp_img_1'
+    # root_mask = 'tmp_1/tmp_mask_1'
+    # os.makedirs(root_img, exist_ok=True)
+    # os.makedirs(root_mask, exist_ok=True)
+    # root_img = 'tmp_1/tmp_img_2'
+    # root_mask = 'tmp_1/tmp_mask_2'
+    # os.makedirs(root_img, exist_ok=True)
+    # os.makedirs(root_mask, exist_ok=True)
+    if option == 0:
+        small_ratio, large_ratio = 0.9, 1.2
+        for i in range(sample_n_frames):
+            
+            # x1,x2,y1,y2 = bboxes_x.min(), bboxes_x.max(), bboxes_y.min(), bboxes_y.max()
+            # print(x1,x2,y1,y2)
+            # x1,x2,y1,y2 = x1*(small_ratio-i*1.0/sample_n_frames*small_ratio), min(x2*(large_ratio+i/sample_n_frames*small_ratio), W), y1*(small_ratio-i/sample_n_frames*small_ratio), min(y2*(large_ratio+i/sample_n_frames*small_ratio), H)
+            x1,x2,y1,y2 = (1-i*1.0/sample_n_frames)*0.1*W, 0.9*W+i*1.0/sample_n_frames*0.1*W, (1-i*1.0/sample_n_frames)*0.1*H, 0.9*H+i*1.0/sample_n_frames*0.1*H
+            x1,x2,y1,y2 = int(x1),int(x2),int(y1),int(y2)
+            
+            # print(x1,x2,y1,y2)
+            img_ = cv2.resize(img[y1:y2, x1:x2, :], (H,W), cv2.INTER_CUBIC)
+            mask_ = cv2.resize(mask[y1:y2, x1:x2], (H,W), cv2.INTER_NEAREST)
+            cv2.imwrite(os.path.join(root_img_dir, image.split('.jpg')[0], '%05d'%i+'.jpg'), img_)
+            cv2.imwrite(os.path.join(root_mask_dir, image.split('.jpg')[0],  '%05d'%i+'.png'), mask_)
+    elif option == 1:
+        small_ratio, large_ratio = 0.2, 0.9
+        for i in range(sample_n_frames):
+            # x1,x2,y1,y2 = bboxes_x.min(), bboxes_x.max(), bboxes_y.min(), bboxes_y.max()
+            # x1,x2,y1,y2 = x1*(small_ratio+i*1.0/sample_n_frames*large_ratio), min(x2*(2.0-i/sample_n_frames*large_ratio), W), y1*(small_ratio+i/sample_n_frames*large_ratio), min(y2*(2.0-i/sample_n_frames*large_ratio), H)
+            x1,x2,y1,y2 = i*1.0/sample_n_frames*0.1*W, W-i*1.0/sample_n_frames*0.1*W, i*1.0/sample_n_frames*0.1*H, H-i*1.0/sample_n_frames*0.1*H
+            x1,x2,y1,y2 = int(x1),int(x2),int(y1),int(y2)
+            img_ = cv2.resize(img[y1:y2, x1:x2, :], (H,W), cv2.INTER_CUBIC)
+            mask_ = cv2.resize(mask[y1:y2, x1:x2], (H,W), cv2.INTER_NEAREST)
+            cv2.imwrite(os.path.join(root_img_dir, image.split('.jpg')[0], '%05d'%i+'.jpg'), img_)
+            cv2.imwrite(os.path.join(root_mask_dir, image.split('.jpg')[0], '%05d'%i+'.png'), mask_)
+    else:
+        pad = min(random.randint(200,400), bboxes_x.min())
+        for i in range(sample_n_frames):
+            x1,x2,y1,y2 = bboxes_x.min(), bboxes_x.max(), bboxes_y.min(), bboxes_y.max()
+            # print(x1,x2,y1,y2)
+            x1,x2,y1,y2 = x1-pad+pad//sample_n_frames*i,x2-pad+pad//sample_n_frames*i,0,img.shape[1]
+            x1,x2,y1,y2 = int(x1),int(x2),int(y1),int(y2)
+            # print(x1,x2)
+            img_ = cv2.resize(img[y1:y2, x1:x2, :], (H,W), cv2.INTER_CUBIC)
+            mask_ = cv2.resize(mask[y1:y2, x1:x2], (H,W), cv2.INTER_NEAREST)
+            cv2.imwrite(os.path.join(root_img_dir, image.split('.jpg')[0], '%05d'%i+'.jpg'), img_)
+            cv2.imwrite(os.path.join(root_mask_dir, image.split('.jpg')[0], '%05d'%i+'.png'), mask_)