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

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1I7cHYydsGwZjG_axddRZMqWvrZiapwyu

Código utilizado para gerar o data augmentation de imagens no geral
"""

#Imports necessários

import numpy as np
import keras,glob,os
import pandas as pd
import cv2
from shutil import copy2
from keras.preprocessing.image import ImageDataGenerator, array_to_img,img_to_array, load_img

img_path = r'C:\Users\Gabriel\Downloads\castasv3\normal'
outpath = r'C:\Users\Gabriel\Downloads\castasv3\normal-splited'

sep = os.path.sep
teste_path = os.path.join(outpath, 'test')
validacao_path = os.path.join(outpath, 'validation')
treino_path = os.path.join(outpath, 'train')

paths = [outpath, teste_path, validacao_path, treino_path]

#verifica se as paths de saida existem
for path in paths:
  if not os.path.lexists(path):
          os.mkdir(path)

#metodo que itera pelos arquivos do diretorio salvando a path da imagem e sua respectiva classe
def directory_from_df_with_class(df, directory):
  for r, d, f in os.walk(directory):
    for file in f:
      row = {}
      row['path'] = os.path.join(r, file)
      
      row['class'] = row['path'].split(os.path.sep)[-2]
      #print(row['class'])
      df= df.append(row, ignore_index=True)
  return df

columns_df = ['path', 'class']
df = pd.DataFrame(columns=columns_df)

df = directory_from_df_with_class(df, img_path)

"""Código que separa os datasets em treino, validação e teste."""

from random import randint

def separa_df_treino_test(df, coluna_classe, percentual_teste):
  #obtem as classes
  #classes = df[coluna_classe].unique()
  #contadores = df.groupby(coluna_classe).count()
  contadores = df[coluna_classe].value_counts()
  df_treino = pd.DataFrame(columns=df.columns)
  df_teste = pd.DataFrame(columns=df.columns)

  #define classe com valor minimo
  minimo = np.min(contadores.values)
  print("a classe com menor quantidade de amostras eh '{}' e possui {} amostras...".format(contadores.index[np.argmin(contadores.values)],minimo))

  #define quantidade de imagens que serao utilizadas
  tam_teste = int(minimo*percentual_teste)
  print("tamanho do conjunto de teste eh: {}...".format(tam_teste))
  print("o tamanho do dataframe original eh: {}...".format(len(df)))
  
  for coluna in contadores.index:
    #filtra o dataframe
    df_aux = df[df[coluna_classe] == coluna]

    #repete a quantidade de vezes que o tam teste e necessario
    for i in range(tam_teste):
      #obtem o indice
      indice = randint(0, len(df_aux)-1)
      
      #obtem valor  
      df_teste = df_teste.append(df_aux.iloc[[indice]], ignore_index=True)

      #remove do df
      df_aux.drop(df_aux.index[[indice]], inplace=True)
    
    df_treino = pd.concat([df_aux, df_treino], ignore_index=True)
  print("o tamanho do dataframe de teste eh: {}...".format(len(df_teste)))
  print("o tamanho do dataframe de treino eh: {}...".format(len(df_treino)))

  return df_treino, df_teste


def imagem_por_classe(df, coluna_classe, path_classe, quantidade=1):
  #captura as classes
  classes = df[coluna_classe].unique()

  lista_imagens = []
  #itera sobre as classes
  for classe in classes:
    df_classe = df[df[coluna_classe] == classe]
    #captura a path imagens
    for i in range(quantidade):
      indice = randint(0, len(df_classe)-1)
      lista_imagens.append(df_classe.iloc[[indice]][path_classe].values[0])
      df_classe.drop(df_classe.index[[indice]], inplace=True)

  return lista_imagens


df_treino, df_teste = separa_df_treino_test(df, 'class', 0.1)

imgs_visualizar = imagem_por_classe(df_teste, 'class', 'path', 1)

df_treino, df_validacao = separa_df_treino_test(df_treino, 'class', 0.2)

#imprimindo pesos
contadores = pd.DataFrame(df_treino['class'].value_counts())
contadores['peso'] = contadores.values/len(df_treino)

print('=================')
print(contadores)

#define funcao de salvar imagem
def salva_imagens_df(df, class_column, path_column, outpath):
  for i, row in df.iterrows():
    if not row[path_column] is None:
      nome = row[path_column].split(os.path.sep)[-1]
      path_destino = os.path.join(outpath, row[class_column])
      if not os.path.lexists(path_destino):
          os.mkdir(path_destino)
      arquivo_path = os.path.join(path_destino, nome)
      copy2(row[path_column], arquivo_path)

salva_imagens_df(df=df_treino, class_column='class', path_column='path', outpath=treino_path)
salva_imagens_df(df=df_teste, class_column='class', path_column='path', outpath=teste_path)
salva_imagens_df(df=df_validacao, class_column='class', path_column='path', outpath=validacao_path)

"""Faz o data augmentation:"""

filenames=df_treino['path'].values

for img in filenames:

    if "DS_Store" in img: continue
    src_fname, ext = os.path.splitext(img) 

    datagen = ImageDataGenerator(rotation_range=50,
            width_shift_range=0.2,
            height_shift_range=0.2,
            horizontal_flip=True,
            brightness_range=[0.6,1.0],
            vertical_flip=True)


    img = load_img(img)

    x = img_to_array(img)
    x = x.reshape((1,) + x.shape)

    img_name = src_fname.split(os.path.sep)[-1]
    class_name = src_fname.split(os.path.sep)[-2]
    new_dir = os.path.join(treino_path, class_name)
    if not os.path.lexists(new_dir):
        os.mkdir(new_dir)
    #save_fname = os.path.join(new_dir, os.path.basename(img_name))
    save_fname = new_dir

    i = 0
    for batch in datagen.flow (x, batch_size=1, save_to_dir = save_fname, 
                               save_prefix = img_name, save_format='jpg'):
        i+=1
        if i>5:
            break

#gerar um zip com o novo dataset
#!zip -r -q dataset_aumentado.zip dataset_aumentado/*

columns_df = ['path', 'class']
df = pd.DataFrame(columns=columns_df)
df = directory_from_df_with_class(df, treino_path)

#imprimindo pesos
contadores = pd.DataFrame(df['class'].value_counts())
contadores['peso'] = contadores.values/len(df)
contadores = contadores.sort_index()

print('=================')
print(contadores)
print(contadores['peso'].values.tolist())
print('tamanho df: ', len(df))