/ inputs_processing / augment_task_one_colorshifts.py
augment_task_one_colorshifts.py
1 #!/usr/bin/env python3 2 # -*- coding: utf-8 -*- 3 """ 4 Created on Wed Apr 13 20:49:24 2022 5 6 @author: aleoikon 7 """ 8 import albumentations as A 9 import numpy as np 10 import pandas as pd 11 12 13 def scaleMinMax(x): 14 return ((x - np.nanmin(x))/(np.nanmax(x) - np.nanmin(x))) 15 16 def create_rgb(x): 17 r = x[:,:,1] 18 g = x[:,:,2] 19 b = x[:,:,3] 20 r = scaleMinMax(r) 21 g = scaleMinMax(g) 22 b = scaleMinMax(b) 23 #vnir = x[:,:,8] 24 #rgbvnir = np.stack((r,g,b,vnir),axis=2).astype('float') 25 rgb = np.dstack((r,g,b)) 26 return(rgb) 27 28 29 def augment_patches(path_to_images, target_path,num_of_patches, onera=False): 30 31 dataset = pd.DataFrame(columns=['pair1','pair2','overlap']) 32 33 if onera == False: 34 images_df = pd.read_csv(path_to_images + "s2mtcp_set.csv") 35 else: 36 images_df = pd.read_csv(path_to_images + "onera_set.csv") 37 38 transform_overlap = A.Compose( 39 [A.RandomCrop(width=96, height=96), 40 A.RandomRotate90(p=0.5), 41 A.HorizontalFlip(p=0.5), 42 A.VerticalFlip(p=0.5), 43 A.ColorJitter(brightness=0.9, contrast=0.9, saturation=0.9, hue=0.9, always_apply=False, p=0.5), 44 ], 45 additional_targets={'image0': 'image'} 46 ) 47 48 transform_non = A.Compose( 49 [A.RandomCrop(width=96, height=96), 50 A.RandomRotate90(p=0.5), 51 A.HorizontalFlip(p=0.5), 52 A.VerticalFlip(p=0.5), 53 A.ColorJitter(brightness=0.9, contrast=0.9, saturation=0.9, hue=0.9, always_apply=False, p=0.5), 54 ] 55 ) 56 57 pos = 0 58 59 for index in range(len(images_df)): 60 img1 = np.load(path_to_images+images_df['pair1'][index]) 61 img2 = np.load(path_to_images+images_df['pair2'][index]) 62 if onera==False: 63 img1 = create_rgb(img1) 64 img2 = create_rgb(img2) 65 img1 = np.float32(img1) 66 img2 = np.float32(img2) 67 #make 5 overlaping patches for this image pair 68 for i in range(num_of_patches): 69 #left 70 transformed = transform_overlap(image=img1, image0=img2) 71 prefix = images_df['pair1'][index][:-4] 72 np.save( target_path + prefix + '_lo_'+ str(i) +'.npy',transformed['image']) 73 #save in a dataframe 74 dataset.loc[pos,'pair1'] = prefix + '_lo_'+ str(i) +'.npy' 75 dataset.loc[pos,'overlap'] = 0 76 #right 77 prefix = images_df['pair2'][index][:-4] 78 np.save( target_path + prefix + '_ro_'+ str(i) + '.npy',transformed['image0']) 79 #save in a dataframe 80 dataset.loc[pos,'pair2'] = prefix + '_ro_'+ str(i) + '.npy' 81 pos += 1 82 83 84 #make 5 non_overlaping patches for this image pair 85 for i in range(num_of_patches): 86 #left 87 transformed = transform_non(image=img1) 88 prefix = images_df['pair1'][index][:-4] 89 np.save( target_path + prefix + '_ln_'+ str(i) +'.npy',transformed['image']) 90 #save in a dataframe 91 dataset.loc[pos,'pair1'] = prefix + '_ln_'+ str(i) +'.npy' 92 dataset.loc[pos,'overlap'] = 1 93 #right 94 transformed = transform_non(image=img2) 95 prefix = images_df['pair2'][index][:-4] 96 np.save( target_path + prefix + '_rn_'+ str(i) + '.npy',transformed['image']) 97 #save in a dataframe 98 dataset.loc[pos,'pair2'] = prefix + '_rn_'+ str(i) + '.npy' 99 pos+=1 100 101 102 dataset.to_csv(target_path + 'dataset.csv', index=False)