1  import albumentations as A
 2  import numpy as np
 3  import pandas as pd
 4  
 5  def CDaugment_patches(path_to_images, path_to_train, path_to_test, target_path_train, target_path_test, num_of_patches):
 6      
 7      dataset_train = pd.DataFrame(columns=['pair1','pair2','change_mask'])
 8      dataset_test = pd.DataFrame(columns=['pair1','pair2','change_mask'])
 9      
10      transform = A.Compose(
11      [A.RandomCrop(width=96, height=96),
12       A.RandomRotate90(5),
13       A.HorizontalFlip(p=0.5),
14       A.VerticalFlip(p=1)
15      ],
16      additional_targets={'image0': 'image'}
17      )
18      
19      #train
20      train_setdf = pd.read_csv(path_to_train+'train_set.csv', dtype=str)
21      #test
22      test_setdf = pd.read_csv(path_to_test+'test_set.csv', dtype=str)
23      
24      ####################################train#############################################################
25      pos = 0
26      for index in range(len(train_setdf)):
27          img1 = np.load(path_to_images+train_setdf['pair1'][index])
28          img2 = np.load(path_to_images+train_setdf['pair2'][index])
29          mask = np.load(path_to_train+train_setdf['change_mask'][index])
30          
31          #make 100 patches per image pair
32          for i in range(num_of_patches):
33              
34              transformed = transform(image=img1, image0=img2, mask = mask)
35              #left
36              prefix_l = train_setdf['pair1'][index][:-4]
37              np.save(target_path_train + prefix_l + '_l_'+ str(i) +'.npy',transformed['image'])
38              #mask
39              prefix_m = train_setdf['change_mask'][index][:-4]
40              np.save(target_path_train + prefix_m +'_'+ str(i) +'.npy',transformed['mask'])
41              #right
42              prefix_r = train_setdf['pair2'][index][:-4]
43              np.save(target_path_train + prefix_r + '_r_'+ str(i) + '.npy',transformed['image0'])
44              #save in a dataframe
45              dataset_train.loc[pos,'pair1'] = prefix_l + '_l_'+ str(i) +'.npy'
46              dataset_train.loc[pos,'pair2'] =  prefix_r + '_r_'+ str(i) + '.npy'
47              dataset_train.loc[pos,'change_mask'] = prefix_m +'_'+ str(i) +'.npy'
48              pos += 1
49              
50      dataset_train.to_csv(target_path_train + 'dataset_train.csv', index=False)
51      #################################test#####################################################
52      pos = 0
53      for index in range(len(test_setdf)):
54          img1 = np.load(path_to_images+test_setdf['pair1'][index])
55          img2 = np.load(path_to_images+test_setdf['pair2'][index])
56          mask = np.load(path_to_test+test_setdf['change_mask'][index])
57          
58          #make 100 patches per image pair
59          for i in range(num_of_patches):
60              
61              transformed = transform(image=img1, image0=img2, mask = mask)
62              #left
63              prefix_l = test_setdf['pair1'][index][:-4]
64              np.save(target_path_test + prefix_l + '_l_'+ str(i) +'.npy',transformed['image'])
65              #mask
66              prefix_m = test_setdf['change_mask'][index][:-4]
67              np.save(target_path_test + prefix_m +'_'+ str(i) +'.npy',transformed['mask'])
68              #right
69              prefix_r = test_setdf['pair2'][index][:-4]
70              np.save(target_path_test + prefix_r + '_r_'+ str(i) + '.npy',transformed['image0'])
71              #save in a dataframe
72              dataset_test.loc[pos,'pair1'] = prefix_l + '_l_'+ str(i) +'.npy'
73              dataset_test.loc[pos,'pair2'] =  prefix_r + '_r_'+ str(i) + '.npy'
74              dataset_test.loc[pos,'change_mask'] = prefix_m +'_'+ str(i) +'.npy'
75              pos += 1
76              
77      dataset_test.to_csv(target_path_test + 'dataset_test.csv', index=False)
78  
79  
80  # Paths and parameters
81  path_to_images = '/home/dvalsamis/Documents/data/sysu/SYSU_NPY/total_NPY/'
82  path_to_train = '/home/dvalsamis/Documents/data/sysu/SYSU_NPY/train_labels_NPY/'
83  path_to_test = '/home/dvalsamis/Documents/data/sysu/SYSU_NPY/test_labels_NPY/'
84  target_path_train = '/home/dvalsamis/Documents/data/sysu/SYSU_NPY/aug_train_data/'
85  target_path_test = '/home/dvalsamis/Documents/data/sysu/SYSU_NPY/aug_test_data/'
86  num_of_patches = 5
87  
88  # Call the function
89  CDaugment_patches(path_to_images, path_to_train, path_to_test, target_path_train, target_path_test, num_of_patches)
90  
91  print("Done")