# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: 'Python 3.7.7 64-bit (''tf21'': conda)'
#     name: python_defaultSpec_1599298562404
# ---

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches


def add_layer(axes, rect_pos, filter_num, color, kernel_size):
    SCALE = 10
    kernel_size = SCALE * kernel_size

    axes.add_patch(
        patches.Rectangle(
            xy=rect_pos,
            width=filter_num,
            height=1,
            edgecolor = color, 
            fill=False))

    axes.add_patch(
        patches.Rectangle(
            xy=(rect_pos[0], 0.5-kernel_size/2),
            width=filter_num,
            height=kernel_size,
            edgecolor = color, 
            fill=True))


# + tags=[]
def model_layout(rect_pos_list, filter_num, filter_size):
    INPUT_DIM = 128
    filter_size = np.array(filter_size)/INPUT_DIM

    fig, ax = plt.subplots(figsize=(14,4))
    plot_extends = (rect_pos_list[0][0],
                    rect_pos_list[-1][0]+filter_num[-1])
    ax.plot(plot_extends,[0.5,0.5], color="red", linestyle="--")

    for rect_pos, filter_num, kernel_size in zip(rect_pos_list, filter_num, filter_size):
        add_layer(ax, rect_pos, filter_num, color= "blue", kernel_size=kernel_size)
    plt.xticks([])
    plt.yticks([])
    plt.xlim([-100, plot_extends[-1]+100])
    plt.show()


# -

# 6-Stack model
RECT_POS_LIST = [(0,0), 
(60,0), (80,0), (100,0), (120,0),  
(200,0), (240,0), (280,0), (320,0),
(420,0), (470,0), (520,0), (570,0), 
(700,0), (780,0), (860,0), (940,0),
(1100,0)]

# DenseNet21
filter_num = [8,8,8,8,8,16,16,16,16,32,32,32,32,64,64,64,64,20]
filter_size = [7,1,3,1,3,1,3,1,3,1,3,1,3,1,3,1,3,1,3,1,3,1]
model_layout(RECT_POS_LIST, filter_num, filter_size)

# DenseNet21_Small
filter_num = [8,16,16,16,16,16,16,16,16,16,16,16,16,32,32,32,32,20]
filter_size = [7,1,3,1,3,1,3,1,3,1,3,1,3,1,3,1,3,1,3,1,3,1]
model_layout(RECT_POS_LIST, filter_num, filter_size)

# 7-Stack model
RECT_POS_LIST = [
(60,0), (80,0), 
(200,0), (240,0),
(420,0), (470,0),
(600,0), (680,0),
(860,0), (940,0),
(1100,0)]

# VGG-16
filter_num = [8,8,16,16,32,32,64,64,64,64,20]
filter_size = [3,3,3,3,3,3,3,3,3,3,3]
model_layout(RECT_POS_LIST, filter_num, filter_size)



StarcoderdataJupyterScriptsDedupFiltered