preprocessing.py

# -*- coding: utf-8 -*-
"""preprocessing.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/19F6KeMIolMuRiPMEcHQWYidBeZDqHQnA
"""

import pandas as pd
import numpy as np
import json
import copy
import cv2
import math
import torch
import matplotlib.pyplot as plt
from argparse import ArgumentParser
from PIL import Image
import os
from sklearn.model_selection import KFold
import torch.nn.functional as F
from skimage import io
from sklearn.preprocessing import LabelEncoder
import configparser

def kfold_generation(folds,train_csv,train_csv_with_flip,destination_folder):
  '''
  This function is used to generate 4 folds for training data. flipped images of a fold are added to the respective fold.
  folds :- number of folds required
  train_csv :- csv file with train data
  train_csv_with_parts :- csv file with train data and flipped images
  destination_folder:- destination file path to store the folds csv.
  '''
  print('Genertaing folds')
  # reading training and flipped images csv file
  df_flipped=pd.read_csv(train_csv_with_flip)

  # reading train data csv
  df=pd.read_csv(train_csv,header=None,names=["ids","image"])

  # encoding labels to 0-107
  lb = LabelEncoder()
  df['encoded_labels'] = lb.fit_transform(df['ids'])
  df.head()

  #initialize with 4 folds
  kf =KFold(n_splits=folds, shuffle=True, random_state=42)
  index=0
  for train_index,test_index in kf.split(df):
    df_train=df.iloc[train_index]
    df_test=df.iloc[test_index]
    
    train_names=df_train.image.values
    train_names='flipped'+train_names
    test_names=df_test.image.values
    test_names='flipped'+test_names

    # adding flipped images to the fold
    temp=df_flipped.loc[df_flipped['image'].isin(train_names)]
    temp=temp.reset_index(drop=True)
    df_train=df_train.append(temp)
    temp=df_flipped.loc[df_flipped['image'].isin(test_names)]
    temp=temp.reset_index(drop=True)
    df_test=df_test.append(temp)

    #saving the folds to the destination folder
    df_train.to_csv(destination_folder+'fold_train'+str(index)+'.csv',index=False)
    df_test.to_csv(destination_folder+'fold_test'+str(index)+'.csv',index=False)
    index=index+1
  print('Folds generated')

def remove0(x):
    while len(x) != 0:
        if 0 in x:
            x.remove(0)
        else:
            break
    return x
def masked(img_path, x1, x2, x3, x4, y1, y2, y3, y4):
    """
    reference:https://github.com/LcenArthas/CVWC2019-Amur-Tiger-Re-ID
    """
    img = cv2.imread(img_path)
    (h, w, c) = img.shape
    mask = np.zeros((h, w, 3), np.uint8)                                
    triangle = np.array([[x1, y1], [x2, y2], [x3, y3], [x4, y4]])       
    cv2.fillConvexPoly(mask, triangle, (255, 255, 255))                  
    img_mask = cv2.bitwise_and(img, mask)                               
    # plt.imshow(img_mask)
    # plt.show()
    xmin = min(x1, x2, x3, x4)
    xmax = max(x1, x2, x3, x4)
    ymin = min(y1, y2, y3, y4)
    ymax = max(y1, y2, y3, y4)

    if xmin <= 0:
        xmin = 0
    if ymin <= 0:
        ymin = 0
    if xmax > w:
        xmax = w
    if ymax > h:
        ymax = h
    if xmax <= 0 or ymax <=0 or xmin >= xmax or ymin >= ymax:
        print('exception')
        return 1
    else:
        img_crop = img_mask[ymin: ymax, xmin: xmax]
        return img_crop
def find_bbox(x1,y1,x2,y2,root_dir,image_name,ind):
  '''
   """
    reference:https://github.com/LcenArthas/CVWC2019-Amur-Tiger-Re-ID
    """
  This function gets bounding box around the limb parts.
  x1,y1,x2,y2:-coordinates of the rectangle
  root_dir:- directory of images
  image_name:- image name
  ind :- specifies the location of limb, for eg: front limb, hind limb etc.
  '''
  img = Image.open(os.path.join(root_dir, image_name))
  img=np.asarray(img)
  h,w,c=img.shape
  L = math.sqrt(pow(x2 - x1, 2) + pow(y2 - y1, 2))
  # L=distance.euclidean(image[y1,x1],image[y2,x2])
  tha = math.atan((x2 -x1)/((y2 -y1)+0.000000000001))
  if ind == 1 or ind == 2:
      deltax = math.cos(tha) * 0.3 * L
      deltay = math.sin(tha) * 0.3 * L
  if ind == 3 or ind == 5:
      deltax = math.cos(tha) * 0.45 * L
      deltay = math.sin(tha) * 0.45 * L
  else:
      deltax = math.cos(tha) * 0.3 * L
      deltay = math.sin(tha) * 0.3 * L
  xx1=int(x1-deltax)
  xx2=int(x1+deltax)
  yy1=int(y1+deltay)
  yy2=int(y1-deltay)
  xx3=int(x2-deltax)
  xx4=int(x2+deltax)
  yy3=int(y2-deltay)
  yy4=int(y2+deltay)
  img1=masked(os.path.join(root_dir,image_name), xx1, xx2, xx4, xx3, yy1, yy2, yy3, yy4)
  return img1
  
def extract_paw(pos,image_name,root_dir,dest_dir):
  '''
   """
    reference:https://github.com/LcenArthas/CVWC2019-Amur-Tiger-Re-ID
    """
  This function extracts limbs from the tiger image
  pos:- keypoints of the tiger
  image_name:- image_name to be extracted
  root_dir:- directory of the images
  dest_dir:- destination directory to store the limb parts
  '''
  pairs = [[5, 6], [3, 4], [7, 8], [8, 9], [10, 11], [11, 12]]
  for index,pair in enumerate(pairs):
    index+=1
    i,j=pair[0],pair[1]
    if 0 not in [pos[i, 2], pos[j, 2]]: 
      
      x1=pos[i,0]
      y1=pos[i,1]
      x2=pos[j,0]
      y2=pos[j,1]
      img1=find_bbox(x1,y1,x2,y2,root_dir,image_name,index)
      if type(img1)==int:
        break
      else:
          img1 = Image.fromarray(cv2.cvtColor(img1, cv2.COLOR_BGR2RGB))   
          # plt.imshow(img1)
          # plt.show()
                   
          img1.save(dest_dir+ image_name.split('.')[0] +'_'+ str(index) +'part.jpg')
          
          
def extract_body(pos,image_name,root_dir,dest_dir):
  '''
   """
    reference:https://github.com/LcenArthas/CVWC2019-Amur-Tiger-Re-ID
    """
  This function extracts trunk part of the tiger
  pos:- keypoints of the tiger
  image_name:- image_name to be extracted
  root_dir:- directory of the images
  dest_dir:- destination directory to store the trunk parts
  '''
  img = Image.open(os.path.join(root_dir, image_name))
  
  # extracting x,y coordinates of root of the tail,left ear, right ear, left hip, right hip, nose to get the rectangle of trunk
  x1=[pos[13,0],pos[1,0],pos[0,0],pos[10,0],pos[7,0],pos[2,0]]
  x2=[pos[13,0],pos[1,0],pos[0,0],pos[10,0],pos[7,0],pos[2,0]]
  y1=[pos[13,1],pos[1,1],pos[0,1]]
  y2=[pos[7,1],pos[10,1],pos[3,1],pos[5,1],pos[2,1]]
  
  xmin=min(x1)
  xmax=max(x1)
  ymin=min(y1)
  ymax=max(y2)
  img_body = img.crop((xmin, ymin, xmax, ymax))
  x1 = [pos[2, 0], pos[13, 0], pos[0, 0], pos[1, 0], pos[10, 0], pos[7, 0]]     
  x2 = [pos[2, 0], pos[13, 0], pos[0, 0], pos[1, 0], pos[10, 0], pos[7, 0]]     
  y1 = [pos[1, 1], pos[0, 1], pos[13, 1]]                                       
  y2 = [pos[10, 1], pos[7, 1], pos[3, 1], pos[5, 1]]                            
  x1 = remove0(x1)
  x2 = remove0(x2)
  y1 = remove0(y1)
  y2 = remove0(y2)
  if x1 and x2 and y1 and y2:       
      
      if len(x1) >1:
          xmin = min(x1)
          xmax = max(x2)
          ymin = min(y1)
          ymax = max(y2)
          
          img_body = img.crop((xmin, ymin, xmax, ymax))
          # plt.imshow(img_body)
          # plt.show()
          
          img_body.save(dest_dir+ image_name.split('.')[0] +'_body.jpg')
      else:
          pass
  else:
      pass

def extract_parts(train_csv_with_flipped,keypoints_json_file,images_folder,destination_folder):
  '''
   """
    reference:https://github.com/LcenArthas/CVWC2019-Amur-Tiger-Re-ID
    """
  This function extract trunk part and limbs from the tiger image.
  train_csv_with_flipped:- train csv which includes flipped images
  keypoints_json_file:- json file which includes keypoints for train and flipped images
  images_folder:- images folder
  destination folder:- destination folder to store all the part images
  '''
  train_df=pd.read_csv(train_csv_with_flipped)
  df=json.load(open(keypoints_json_file))
  if not os.path.exists(destination_folder):
    os.mkdir(destination_folder)
  print(train_df.head())
  print('Extracting parts.....')
  for i in range(len(train_df)):
    pos = np.array(df[train_df.iloc[i,1]]).reshape((15, 3))
    extract_body(pos,train_df.iloc[i,1],images_folder,destination_folder)
    extract_paw(pos,train_df.iloc[i,1],images_folder,destination_folder)
  print('Parts extracted.')

def flip_images(train_csv_file,keypoints_json_file,root_dir,final_keypoints_json_file,final_keypoints_train_csv):
  '''
   """
    reference:https://github.com/LcenArthas/CVWC2019-Amur-Tiger-Re-ID
    """
  This function flips images and generate keypoints for flipped images
  train_csv_file:- csv file with train data
  keypoints_json_file:- json file with keypoints
  root_dir:- root directory of images
  final_keypoints_json_file:- destination file to store keypoints of all images including flipped images
  final_keypoints_train_csv:- csv file to store all images including flipped images.
  '''
  print('Flipping images....')
  train_df=pd.read_csv(train_csv_file,header=None,names=["ids","image"])
  df=json.load(open(keypoints_json_file))
  image_path=root_dir
  for i in train_df["image"]:
    if not os.path.exists(image_path+"flipped"+i):
      print(i)
      im = plt.imread(image_path+i)
      im = np.flip(im,1)
      im = Image.fromarray(im)
      im.save(image_path+"flipped"+i)
  for iid in train_df["image"]:
    df_flipped={}
    im = plt.imread(image_path+iid)
    h,w,c=im.shape
    im = np.flip(im,1)
    points=copy.deepcopy(df[iid])
    for i in range(int(len(points)/3)):
      if points[i*3]!=0 or points[i*3+1]!=0:
        points[i*3]=w-points[i*3]
    df["flipped"+iid]=points
  with open(final_keypoints_json_file, 'w') as f:
    json.dump(df, f) 
  train_df["image"]=train_df["image"].apply(lambda x: "flipped"+x)
  train_df1=pd.read_csv(train_csv_file,header=None,names=["ids","image"])
  train_df1=train_df1.append(train_df)
  lb = LabelEncoder()
  train_df1['encoded_labels'] = lb.fit_transform(train_df1['ids'])
  train_df1.head()
  train_df1.to_csv(final_keypoints_train_csv,index=False)
  print('Images flipped.')

def preprocess(train_csv,keypoint_json,images_folder,final_json_file,final_train_csv,parts_dir,fold_directory):
  
  flip_images(train_csv,keypoint_json,images_folder,final_json_file,final_train_csv)
  extract_parts(final_train_csv,final_json_file,images_folder,parts_dir)
  kfold_generation(int(config["test_parameters"]["no_of_folds"]),train_csv,final_train_csv,fold_directory)

if  __name__=="__main__":                
  config = configparser.ConfigParser()
  config.sections()
  config.read('config.ini')  

  
  train_csv=config["train_images"]["train_csv_path"]
  keypoint_json=config["train_images"]["keypoint_json_file"]
  images_folder=config["train_images"]["train_image_path"]
  final_train_csv=config["train_images"]["final_train_csv"]
  final_json_file=config["train_images"]["final_json_file"]
  parts_dir=config["train_images"]["parts_dir"]
  fold_directory=config["train_images"]["fold_directory"]
    

  preprocess(train_csv,keypoint_json,images_folder,final_json_file,final_train_csv,parts_dir,fold_directory)