1  import albumentations as A
 2  import numpy as np
 3  import pandas as pd
 4  
 5  def augment_patches(path_to_images, target_path,num_of_patches, onera=False):
 6      
 7      dataset = pd.DataFrame(columns=['pair1','pair2','overlap'])
 8      
 9      if onera == False:
10          images_df = pd.read_csv(path_to_images + "s2mtcp_set.csv")
11      else:
12          images_df = pd.read_csv(path_to_images + "onera_set.csv")
13          
14      transform_overlap = A.Compose(
15      [A.RandomCrop(width=96, height=96),
16       A.RandomRotate90(p=0.5),
17       A.HorizontalFlip(p=0.5),
18       A.VerticalFlip(p=0.5)
19      ],
20      additional_targets={'image0': 'image'}
21      )
22      
23      transform_non = A.Compose(
24      [A.RandomCrop(width=96, height=96),
25       A.RandomRotate90(p=0.5),
26       A.HorizontalFlip(p=0.5),
27       A.VerticalFlip(p=0.5)
28      ]
29      )
30      pos = 0
31      for index in range(len(images_df)):
32          img1 = np.load(path_to_images+images_df['pair1'][index])
33          img2 = np.load(path_to_images+images_df['pair2'][index])
34          #make num_of_patches overlaping patches for this image pair
35          for i in range(num_of_patches):
36              #left
37              transformed = transform_overlap(image=img1, image0=img2)
38              prefix = images_df['pair1'][index][:-4]
39              np.save( target_path + prefix + '_lo_'+ str(i) +'.npy',transformed['image'])
40              #save in a dataframe
41              dataset.loc[pos,'pair1'] = prefix + '_lo_'+ str(i) +'.npy'
42              dataset.loc[pos,'overlap'] = 0
43              #right
44              prefix = images_df['pair2'][index][:-4]
45              np.save( target_path + prefix + '_ro_'+ str(i) + '.npy',transformed['image0'])
46              #save in a dataframe
47              dataset.loc[pos,'pair2'] =  prefix + '_ro_'+ str(i) + '.npy'
48              pos += 1
49                  
50          #make num_of_patches non_overlaping patches for this image pair
51          for i in range(num_of_patches):
52              #left
53              transformed = transform_non(image=img1)
54              prefix = images_df['pair1'][index][:-4]
55              np.save( target_path + prefix + '_ln_'+ str(i) +'.npy',transformed['image'])
56              #save in a dataframe
57              dataset.loc[pos,'pair1'] = prefix + '_ln_'+ str(i) +'.npy'
58              dataset.loc[pos,'overlap'] = 1
59              #right
60              transformed = transform_non(image=img2)
61              prefix = images_df['pair2'][index][:-4]
62              np.save( target_path + prefix + '_rn_'+ str(i) + '.npy',transformed['image'])
63              #save in a dataframe
64              dataset.loc[pos,'pair2'] =  prefix + '_rn_'+ str(i) + '.npy'
65              pos+=1
66                 
67      dataset.to_csv(target_path + 'dataset.csv', index=False)