/ 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)