import tensorflow as tf
print(tf.__version__)

import os
for dirname, _, filenames in os.walk('/kaggle/input'):
    for filename in filenames:
        print(os.path.join(dirname, filename))
"""
## **Download The Dataset**
"""
_URL = "https://storage.googleapis.com/ztm_tf_course/food_vision/pizza_steak.zip"

zip_file = tf.keras.utils.get_file(origin=_URL,
                                   fname="pizza_steak.zip",
                                   extract=True)

base_dir = os.path.join(os.path.dirname(zip_file), 'pizza_steak')
for dirpath, dirnames, filenames in os.walk(base_dir):
  print("Parent Directory Path : ",dirpath)
  print("Sub-Directories",f"({len(dirnames)}) : ",dirnames)
  print("Contained Files",f"({len(filenames)}) : ",filenames)
  print()
train_dir = os.path.join(base_dir,'train')
test_dir  = os.path.join(base_dir, 'test')
"""
## **Rename Image Directories within Train and Test directories**
"""
int2string = []
class_label = 0 
for class_name in os.listdir(train_dir):
  class_dir = train_dir + "/" + class_name
  if(os.path.isdir(class_dir)):
    os.rename(class_dir, train_dir + "/" + f"{class_label}")
    class_label = class_label + 1
    int2string.append(class_name)
int2string.index("pizza"), int2string.index("steak")
for class_name in os.listdir(test_dir):
  class_dir = test_dir + "/" + class_name
  if(os.path.isdir(class_dir)):
    os.rename(class_dir, test_dir + "/" + f"{int2string.index(class_name)}")
for dirpath, dirnames, filenames in os.walk(base_dir):
  print("Parent Directory Path : ",dirpath)
  print("Sub-Directories",f"({len(dirnames)}) : ",dirnames)
  print("Contained Files",f"({len(filenames)}) : ",filenames)
  print()
train_ds_files = tf.data.Dataset.list_files(train_dir + "/*/*", shuffle = False)
test_ds_files = tf.data.Dataset.list_files(test_dir + "/*/*", shuffle = False)
#Let's visualize our images
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import random
IMG_SHAPE = 224
"""
## **Mapping function that takes Image Path as an argument and returns Image Tensor + Label**
"""
def process_and_scale(filepath):
    label = tf.strings.split(filepath, os.sep)[-2]
    label = tf.strings.to_number(label, out_type=tf.dtypes.float32)
    img = tf.io.read_file(filepath)
    img = tf.image.decode_jpeg(img)
    img = tf.image.random_flip_left_right(img)
    img = tf.image.random_flip_up_down(img)
    img = tf.image.random_brightness(img,0.2)
    #img = tf.image.random_contrast(img,0.2,0.5)
    #img = tf.image.random_hue(img, 0.2)
    #img = tf.image.random_saturation(img, 5, 10) 
    img = tf.image.resize(img,[IMG_SHAPE,IMG_SHAPE])
    
    return img/255., label
"""
## **Remove Noise/Corrupt Data**
"""
for img_path in train_ds_files:
  img, label = process_and_scale(img_path)
  if img.shape == (224,224,1):
    plt.imshow(tf.squeeze(img).numpy())
    img_path = tf.compat.as_str_any(img_path.numpy())
    print(img_path)
    os.remove(img_path)       #remove noise/corrupt data
"""
## **Input Pipeline**:
* Shuffle the filepaths
* Repeat
* Map
* Take batches of data
* Prefetch for efficient data loading
"""
train_ds_files = tf.data.Dataset.list_files(train_dir + "/*/*", shuffle = False)
test_ds_files = tf.data.Dataset.list_files(test_dir + "/*/*", shuffle = False)
train_size = len(train_ds_files)
test_size = len(test_ds_files)
f"Train Size: {train_size}", f"Test Size: {test_size}"
num_epochs = 10
train_ds = train_ds_files.shuffle(train_size).repeat(num_epochs).take(train_size)
train_ds = train_ds.map(process_and_scale, num_parallel_calls = tf.data.AUTOTUNE).batch(16).prefetch(tf.data.AUTOTUNE)

test_ds  = test_ds_files.shuffle(test_size).repeat(num_epochs).take(test_size)
test_ds  = test_ds.map(process_and_scale, num_parallel_calls = tf.data.AUTOTUNE).batch(16).prefetch(tf.data.AUTOTUNE)
from tensorflow.keras.layers import Conv2D, MaxPool2D, Flatten, Dense, Concatenate, Activation, Add
from tensorflow.keras import Input
from tensorflow.keras.models import Model
"""
## **Model Architecture - Inspired From Inception-Resnet V2**
"""
input_placeholder = Input(shape = (224,224,3))
conv1_out = Conv2D(filters = 32, 
           kernel_size = (3,3),
           strides=(2,2),
           activation = 'relu',
           padding = 'valid')(input_placeholder) #(111,111,32)
conv2_out = Conv2D(filters=32,
           kernel_size = (3,3),
           strides = (1,1),
           activation='relu',
           padding = 'valid')(conv1_out) #(109, 109, 32)
conv3_out = MaxPool2D(pool_size = (3,3),
                      strides=(2,2),
                      padding='valid')(conv2_out) #(54,54,32)
conv4_out = Conv2D(filters=64,
           kernel_size = (1,1),
           strides = (1,1),
           activation='relu',
           padding = 'same')(conv2_out) #(109,109,64)

conv5_out = Conv2D(filters=64,
           kernel_size = (3,1),
           strides = (1,1),
           activation='relu',
           padding = 'same')(conv4_out) #(109,109,64)
conv6_out = Conv2D(filters=64,
           kernel_size = (1,3),
           strides = (1,1),
           activation='relu',
           padding = 'same')(conv5_out) #(109,109,64)
conv7_out = Conv2D(filters=96,
           kernel_size = (3,3),
           strides = (2,2),
           activation='relu',
           padding = 'valid')(conv6_out) #(54,54,96)

concat1_out = Concatenate(axis = -1)([conv3_out, conv7_out]) #(54,54,128)
concat1_out = Conv2D(filters=256,
              kernel_size = (3,3),
              strides = (1,1),
              activation='relu',
              padding = 'same')(conv6_out) #(54,54,256)
concat1_relu =  Activation('relu')(concat1_out)

conv8_out = Conv2D(filters=32,
           kernel_size = (1,1),
           strides = (1,1),
           activation='relu',
           padding = 'same')(concat1_relu) #(54,54,32)
conv9_out = Conv2D(filters=32,
           kernel_size = (1,1),
           strides = (1,1),
           activation='relu',
           padding = 'same')(concat1_relu) #(54,54,32)

conv10_out = Conv2D(filters=32,
           kernel_size = (3,3),
           strides = (1,1),
           activation='relu',
           padding = 'same')(conv9_out) #(54,54,32)

conv11_out = Conv2D(filters=32,
           kernel_size = (1,1),
           strides = (1,1),
           activation='relu',
           padding = 'same')(concat1_relu ) #(54,54,32)
conv12_out = Conv2D(filters=48,
           kernel_size = (3,3),
           strides = (1,1),
           activation='relu',
           padding = 'same')(conv11_out) #(54,54,48)

conv13_out = Conv2D(filters=64,
           kernel_size = (3,3),
           strides = (1,1),
           activation='relu',
           padding = 'same')(conv12_out) #(54,54,64)

concat2_out = Concatenate(axis=-1)([conv8_out, conv10_out,conv13_out]) #(54,54,128)
conv14_out = Conv2D(filters=256,
           kernel_size = (1,1),
           strides = (1,1),
           activation='relu',
           padding = 'same')(concat2_out) #(54,54,256)

addedTensor = Add()([concat1_relu, tf.math.multiply(conv14_out,0.1)])
relu_out = Activation('relu')(addedTensor)

x = MaxPool2D(pool_size=(2,2))(relu_out)
x = Flatten()(x)
out = Dense(1, activation='sigmoid')(x)

model = Model(inputs=[input_placeholder],
                outputs=[out])
model.compile(loss="binary_crossentropy",
              optimizer=  tf.keras.optimizers.Adam(),
              metrics=["accuracy"])
model.summary()
"""
## **Model Training**
"""
history = model.fit(train_ds,
                    epochs=num_epochs, #10,
                    steps_per_epoch=len(train_ds), #2*len(train_ds),
                    validation_data=test_ds,
                    validation_steps=len(test_ds)) #2*len(test_ds),
# Plot loss per iteration
plt.plot(history.history['loss'], label='loss')
plt.plot(history.history['val_loss'], label='val_loss')
plt.legend()
# Plot accuracy per iteration
plt.plot(history.history['accuracy'], label='acc')
plt.plot(history.history['val_accuracy'], label='val_acc')
plt.legend()
"""
## **Save Model**
"""
# Save model
model.save('Food_Vision_Model')