text: # -*- coding: utf-8 -*-
  """Oxford_Pet_Localization+Classification.ipynb
  
  Automatically generated by Colaboratory.
  
  Original file is located at
      https://colab.research.google.com/drive/1HV5CFFetaKmH99VycgYJMsfK9-dek2yh
  """
  
  ## library import
  import tensorflow as tf
  from tensorflow import keras
  import numpy as np
  import os
  import re
  from PIL import Image
  import shutil
  import xml.etree.ElementTree as et
  import random
  import matplotlib.pyplot as plt
  from matplotlib.patches import Rectangle
  
  ## google drive에서 압축된 dataset download
  import gdown
  url = 'https://drive.google.com/uc?id=1dIR9ANjUsV9dWa0pS9J0c2KUGMfpIRG0'
  fname = 'oxford_pet.zip'
  gdown.download(url, fname, quiet=False)
  
  
  ## 압축풀기
  !unzip -q oxford_pet.zip -d oxford_pet
  
  ## directory 정보
  cur_dir = os.getcwd()
  data_dir = os.path.join(cur_dir, 'oxford_pet')
  image_dir = os.path.join(data_dir, 'images')
  bbox_dir = os.path.join(data_dir, 'annotations', 'xmls')
  seg_dir = os.path.join(data_dir, 'annotations', 'trimaps')
  
  ## image file 수 확인
  image_files = [fname for fname in os.listdir(image_dir) if os.path.splitext(fname)[-1] == '.jpg']
  print(len(image_files))
  
  ## localization을 위한 annotation이 되어 있는 file의 수 확인
  bbox_files = [fname for fname in os.listdir(bbox_dir) if os.path.splitext(fname)[-1] == '.xml']
  len(bbox_files)
  
  ## image file들을 읽어서 channel이 3이 아닌 image는 삭제, xml도 같이 삭제
  for image_file in image_files:
    image_path = os.path.join(image_dir, image_file)
    bbox_file = os.path.splitext(image_file)[0]+'.xml'
    bbox_path = os.path.join(bbox_dir, bbox_file)
    image = Image.open(image_path)
    image_mode = image.mode
    if image_mode != 'RGB':
      print(image_file, image_mode)
      image = np.asarray(image)
      print(image.shape)
      os.remove(image_path)
      try:
        os.remove(bbox_path)
      except FileNotFoundError:
        pass
  
  ## image file 수 확인
  image_files = [fname for fname in os.listdir(image_dir) if os.path.splitext(fname)[-1] == '.jpg']
  print(len(image_files))
  
  ## localization을 위한 annotation이 되어 있는 file의 수 확인
  bbox_files = [fname for fname in os.listdir(bbox_dir) if os.path.splitext(fname)[-1] == '.xml']
  len(bbox_files)
  
  class_list = set()
  for image_file in image_files:
      file_name = os.path.splitext(image_file)[0]
      class_name = re.sub('_\d+', '', file_name)
      class_list.add(class_name)
  class_list = list(class_list)
  print(len(class_list))
  
  class_list.sort()
  class_list
  
  class2idx = {cls:idx for idx, cls in enumerate(class_list)}
  class2idx
  
  IMG_SIZE = 224
  N_BBOX = len(bbox_files)
  N_TRAIN = 3000
  N_VAL = N_BBOX - N_TRAIN
  
  ## TFRecord 저장할 directory와 file 경로 설정
  tfr_dir = os.path.join(data_dir, 'tfrecord')
  os.makedirs(tfr_dir, exist_ok=True)
  
  tfr_train_dir = os.path.join(tfr_dir, 'loc_train.tfr')
  tfr_val_dir = os.path.join(tfr_dir, 'loc_val.tfr')
  
  ## TFRecord writer 생성
  writer_train = tf.io.TFRecordWriter(tfr_train_dir)
  writer_val = tf.io.TFRecordWriter(tfr_val_dir)
  
  # The following functions can be used to convert a value to a type compatible
  # with tf.Example.
  
  def _bytes_feature(value):
    """Returns a bytes_list from a string / byte."""
    if isinstance(value, type(tf.constant(0))):
      value = value.numpy() # BytesList won't unpack a string from an EagerTensor.
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))
  
  def _float_feature(value):
    """Returns a float_list from a float / double."""
    return tf.train.Feature(float_list=tf.train.FloatList(value=[value]))
  
  def _int64_feature(value):
    """Returns an int64_list from a bool / enum / int / uint."""
    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))
  
  shuffle_list = list(range(N_BBOX))
  random.shuffle(shuffle_list)
  
  train_idx_list = shuffle_list[:N_TRAIN]
  val_idx_list = shuffle_list[N_TRAIN:]
  
  shuffle_list
  
  for idx in train_idx_list:
      
    bbox_file = bbox_files[idx]
    bbox_path = os.path.join(bbox_dir, bbox_file)
  
    tree = et.parse(bbox_path)
    width = float(tree.find('./size/width').text)
    height = float(tree.find('.size/height').text)
    xmin = float(tree.find('./object/bndbox/xmin').text)
    ymin = float(tree.find('./object/bndbox/ymin').text)
    xmax = float(tree.find('./object/bndbox/xmax').text)
    ymax = float(tree.find('./object/bndbox/ymax').text)
    xc = (xmin + xmax) / 2.
    yc = (ymin + ymax) / 2.
    x = xc / width
    y = yc / height
    w = (xmax - xmin) / width
    h = (ymax - ymin) / height
  
    file_name = os.path.splitext(bbox_file)[0]
    image_file = file_name + '.jpg'
    image_path = os.path.join(image_dir, image_file)
    image = Image.open(image_path)
    image = image.resize((IMG_SIZE, IMG_SIZE))
    bimage = image.tobytes()
  
    class_name = re.sub('_\d+', '', file_name)
    class_num = class2idx[class_name]
  
    if file_name[0].islower(): # dog
      bi_cls_num = 0
    else: # cat
      bi_cls_num = 1
    
    example = tf.train.Example(features=tf.train.Features(feature={
        'image': _bytes_feature(bimage),
        'cls_num': _int64_feature(class_num),
        'bi_cls_num': _int64_feature(bi_cls_num),
        'x': _float_feature(x),
        'y': _float_feature(y),
        'w': _float_feature(w),
        'h': _float_feature(h)
    }))
    writer_train.write(example.SerializeToString())
  
  writer_train.close()
  
  for idx in val_idx_list:
    bbox_file = bbox_files[idx]
    bbox_path = os.path.join(bbox_dir, bbox_file)
  
    tree = et.parse(bbox_path)
    width = float(tree.find('./size/width').text)
    height = float(tree.find('.size/height').text)
    xmin = float(tree.find('./object/bndbox/xmin').text)
    ymin = float(tree.find('./object/bndbox/ymin').text)
    xmax = float(tree.find('./object/bndbox/xmax').text)
    ymax = float(tree.find('./object/bndbox/ymax').text)
    xc = (xmin + xmax) / 2.
    yc = (ymin + ymax) / 2.
    x = xc / width
    y = yc / height
    w = (xmax - xmin) / width
    h = (ymax - ymin) / height
  
    file_name = os.path.splitext(bbox_file)[0]
    image_file = file_name + '.jpg'
    image_path = os.path.join(image_dir, image_file)
    image = Image.open(image_path)
    image = image.resize((IMG_SIZE, IMG_SIZE))
    bimage = image.tobytes()
  
    class_name = re.sub('_\d+', '', file_name)
    class_num = class2idx[class_name]
  
    if file_name[0].islower(): # dog
      bi_cls_num = 0
    else: # cat
      bi_cls_num = 1
    
    example = tf.train.Example(features=tf.train.Features(feature={
        'image': _bytes_feature(bimage),
        'cls_num': _int64_feature(class_num),
        'bi_cls_num': _int64_feature(bi_cls_num),
        'x': _float_feature(x),
        'y': _float_feature(y),
        'w': _float_feature(w),
        'h': _float_feature(h)
    }))
    writer_val.write(example.SerializeToString())
  
  writer_val.close()
  
  ## Hyper Parameters
  N_CLASS = len(class_list)
  N_EPOCHS = 70
  N_BATCH = 40
  IMG_SIZE = 224
  learning_rate = 0.0001
  steps_per_epoch = N_TRAIN / N_BATCH
  validation_steps = int(np.ceil(N_VAL / N_BATCH))
  
  from tensorflow.keras import optimizers
  from tensorflow.keras.applications.resnet import ResNet152
  from tensorflow.keras.layers import Conv2D, ReLU, MaxPooling2D, Dense, BatchNormalization, GlobalAveragePooling2D, Concatenate
  
  resnet152=ResNet152(include_top=False, weights='imagenet', input_shape=(IMG_SIZE,IMG_SIZE,3))
  
  """Classification을 추가하여 Multi-task Learning으로 Localization 학습하기"""
  
  ## tfrecord file을 data로 parsing해주는 function
  def _parse_function2(tfrecord_serialized):
      features={'image': tf.io.FixedLenFeature([], tf.string),
                'cls_num': tf.io.FixedLenFeature([], tf.int64),
                'bi_cls_num': tf.io.FixedLenFeature([], tf.int64),
                'x': tf.io.FixedLenFeature([], tf.float32),
                'y': tf.io.FixedLenFeature([], tf.float32),
                'w': tf.io.FixedLenFeature([], tf.float32),
                'h': tf.io.FixedLenFeature([], tf.float32)              
               }
      parsed_features = tf.io.parse_single_example(tfrecord_serialized, features)
      
      image = tf.io.decode_raw(parsed_features['image'], tf.uint8)    
      image = tf.reshape(image, [IMG_SIZE, IMG_SIZE, 3])
      image = tf.cast(image, tf.float32)/255.
  
      cls_label = tf.cast(parsed_features['cls_num'], tf.float32)
      bi_cls_label = tf.cast(parsed_features['bi_cls_num'], tf.float32)
      
      x = tf.cast(parsed_features['x'], tf.float32)
      y = tf.cast(parsed_features['y'], tf.float32)
      w = tf.cast(parsed_features['w'], tf.float32)
      h = tf.cast(parsed_features['h'], tf.float32)
      gt = tf.stack([bi_cls_label, x, y, w, h], -1)
      
      return image, gt
  
  ## train dataset 만들기
  train_dataset2 = tf.data.TFRecordDataset(tfr_train_dir)
  train_dataset2 = train_dataset2.map(_parse_function2, num_parallel_calls=tf.data.experimental.AUTOTUNE)
  train_dataset2 = train_dataset2.shuffle(buffer_size=N_TRAIN).prefetch(
      tf.data.experimental.AUTOTUNE).batch(N_BATCH).repeat()
  
  ## validation dataset 만들기
  val_dataset2 = tf.data.TFRecordDataset(tfr_val_dir)
  val_dataset2 = val_dataset2.map(_parse_function2, num_parallel_calls=tf.data.experimental.AUTOTUNE)
  val_dataset2 = val_dataset2.batch(N_BATCH).repeat()
  
  ## Create model, compile & summary
  def model_loc_cls():
    resnet152=ResNet152(include_top=False, weights='imagenet', input_shape=(IMG_SIZE,IMG_SIZE,3))
    gap=GlobalAveragePooling2D()(resnet152.output)
  
    dense_b1_1=Dense(256)(gap)
    bn_b1_2=BatchNormalization()(dense_b1_1)
    relu_b1_3=ReLU()(bn_b1_2)
    dense_b1_4=Dense(64)(relu_b1_3)
    bn_b1_5=BatchNormalization()(dense_b1_4)
    relu_b1_6=ReLU()(bn_b1_5)
    output1=Dense(2,activation='softmax')(relu_b1_6)
  
    dense_b2_1=Dense(256)(gap)
    bn_b2_2=BatchNormalization()(dense_b2_1)
    relu_b2_3=ReLU()(bn_b2_2)
    dense_b2_4=Dense(64)(relu_b2_3)
    bn_b2_5=BatchNormalization()(dense_b2_4)
    relu_b2_6=ReLU()(bn_b2_5)
    output2=Dense(4,activation='sigmoid')(relu_b2_6)
  
    concat=Concatenate()([output1,output2])
  
    return keras.Model(inputs=resnet152.input,outputs=concat)
  
  loc_cls_model=model_loc_cls()
  loc_cls_model.summary()
  
  def new_loss_fn(y_true,y_pred):
    cls_labels=tf.cast(y_true[:,:1],tf.int64) # float32로 되어 있던 것을 바꿔줘야함
    loc_labels=y_true[:,1:]
    cls_preds=y_pred[:,:2]
    loc_preds=y_pred[:,2:]
    cls_loss=tf.keras.losses.SparseCategoricalCrossentropy()(cls_labels,cls_preds)
    loc_loss=tf.keras.losses.MeanSquaredError()(loc_labels,loc_preds)
    return cls_loss+5*loc_loss
  
  ## learning rate scheduing
  lr_schedule = keras.optimizers.schedules.ExponentialDecay(initial_learning_rate=learning_rate,
                                                            decay_steps=steps_per_epoch*10,
                                                            decay_rate=0.5,
                                                            staircase=True)
  
  loc_cls_model.compile(keras.optimizers.Adam(lr_schedule) ,loss=new_loss_fn)
  
  ## Train
  history_loc_cls=loc_cls_model.fit(
      train_dataset2,
      epochs=N_EPOCHS,
      steps_per_epoch=steps_per_epoch,
      validation_data=val_dataset2,
      validation_steps=validation_steps
  )
  
  # 
  avg_iou = 0
  num_imgs = validation_steps
  res = N_VAL % N_BATCH #5
  for i, (val_data, val_gt) in enumerate(val_dataset2.take(num_imgs)):    
      flag = (i == validation_steps-1)
      x = val_gt[:,1]
      y = val_gt[:,2]
      w = val_gt[:,3]
      h = val_gt[:,4]
      prediction = loc_cls_model.predict(val_data)
      pred_x = prediction[:,2]
      pred_y = prediction[:,3]
      pred_w = prediction[:,4]
      pred_h = prediction[:,5]
      for idx in range(N_BATCH):
          if(flag):
              if idx == res:
                  flag = False
                  break          
          xmin = int((x[idx].numpy() - w[idx].numpy()/2.)*IMG_SIZE)
          ymin = int((y[idx].numpy() - h[idx].numpy()/2.)*IMG_SIZE)
          xmax = int((x[idx].numpy() + w[idx].numpy()/2.)*IMG_SIZE)
          ymax = int((y[idx].numpy() + h[idx].numpy()/2.)*IMG_SIZE)
  
          pred_xmin = int((pred_x[idx] - pred_w[idx]/2.)*IMG_SIZE)
          pred_ymin = int((pred_y[idx] - pred_h[idx]/2.)*IMG_SIZE)
          pred_xmax = int((pred_x[idx] + pred_w[idx]/2.)*IMG_SIZE)
          pred_ymax = int((pred_y[idx] + pred_h[idx]/2.)*IMG_SIZE)
  
          if xmin > pred_xmax or xmax < pred_xmin:
              continue
          if ymin > pred_ymax or ymax < pred_ymin:
              continue
          w_union = np.max((xmax, pred_xmax)) - np.min((xmin, pred_xmin))
          h_union = np.max((ymax, pred_ymax)) - np.min((ymin, pred_ymin))
          w_inter = np.min((xmax, pred_xmax)) - np.max((xmin, pred_xmin))
          h_inter = np.min((ymax, pred_ymax)) - np.max((ymin, pred_ymin))
  
          w_sub1 = np.abs(xmax - pred_xmax)
          h_sub1 = np.abs(ymax - pred_ymax)
          w_sub2 = np.abs(xmin - pred_xmin)
          h_sub2 = np.abs(ymin - pred_ymin)
  
          iou = (w_inter * h_inter) / ((w_union * h_union) - (w_sub1 * h_sub1) - (w_sub2 * h_sub2))
          avg_iou += iou / N_VAL
  
  print(avg_iou)
  