from kaffe.tensorflow import Network

class MobileNetYOLO(Network):
    def setup(self):
        (self.feed('data')
             .conv(3, 3, 32, 2, 2, biased=False, relu=False, name='conv0')
             .batch_normalization(relu=True, name='conv0_bn')
             .conv(3, 3, 32, 1, 1, biased=False, group=32, relu=False, name='conv1_dw')
             .batch_normalization(relu=True, name='conv1_dw_bn')
             .conv(1, 1, 64, 1, 1, biased=False, relu=False, name='conv1')
             .batch_normalization(relu=True, name='conv1_bn')
             .conv(3, 3, 64, 2, 2, biased=False, group=64, relu=False, name='conv2_dw')
             .batch_normalization(relu=True, name='conv2_dw_bn')
             .conv(1, 1, 128, 1, 1, biased=False, relu=False, name='conv2')
             .batch_normalization(relu=True, name='conv2_bn')
             .conv(3, 3, 128, 1, 1, biased=False, group=128, relu=False, name='conv3_dw')
             .batch_normalization(relu=True, name='conv3_dw_bn')
             .conv(1, 1, 128, 1, 1, biased=False, relu=False, name='conv3')
             .batch_normalization(relu=True, name='conv3_bn')
             .conv(3, 3, 128, 2, 2, biased=False, group=128, relu=False, name='conv4_dw')
             .batch_normalization(relu=True, name='conv4_dw_bn')
             .conv(1, 1, 256, 1, 1, biased=False, relu=False, name='conv4')
             .batch_normalization(relu=True, name='conv4_bn')
             .conv(3, 3, 256, 1, 1, biased=False, group=256, relu=False, name='conv5_dw')
             .batch_normalization(relu=True, name='conv5_dw_bn')
             .conv(1, 1, 256, 1, 1, biased=False, relu=False, name='conv5')
             .batch_normalization(relu=True, name='conv5_bn')
             .conv(3, 3, 256, 2, 2, biased=False, group=256, relu=False, name='conv6_dw')
             .batch_normalization(relu=True, name='conv6_dw_bn')
             .conv(1, 1, 512, 1, 1, biased=False, relu=False, name='conv6')
             .batch_normalization(relu=True, name='conv6_bn')
             .conv(3, 3, 512, 1, 1, biased=False, group=512, relu=False, name='conv7_dw')
             .batch_normalization(relu=True, name='conv7_dw_bn')
             .conv(1, 1, 512, 1, 1, biased=False, relu=False, name='conv7')
             .batch_normalization(relu=True, name='conv7_bn')
             .conv(3, 3, 512, 1, 1, biased=False, group=512, relu=False, name='conv8_dw')
             .batch_normalization(relu=True, name='conv8_dw_bn')
             .conv(1, 1, 512, 1, 1, biased=False, relu=False, name='conv8')
             .batch_normalization(relu=True, name='conv8_bn')
             .conv(3, 3, 512, 1, 1, biased=False, group=512, relu=False, name='conv9_dw')
             .batch_normalization(relu=True, name='conv9_dw_bn')
             .conv(1, 1, 512, 1, 1, biased=False, relu=False, name='conv9')
             .batch_normalization(relu=True, name='conv9_bn')
             .conv(3, 3, 512, 1, 1, biased=False, group=512, relu=False, name='conv10_dw')
             .batch_normalization(relu=True, name='conv10_dw_bn')
             .conv(1, 1, 512, 1, 1, biased=False, relu=False, name='conv10')
             .batch_normalization(relu=True, name='conv10_bn')
             .conv(3, 3, 512, 1, 1, biased=False, group=512, relu=False, name='conv11_dw')
             .batch_normalization(relu=True, name='conv11_dw_bn')
             .conv(1, 1, 512, 1, 1, biased=False, relu=False, name='conv11')
             .batch_normalization(relu=True, name='conv11_bn')
             .conv(3, 3, 512, 2, 2, biased=False, group=512, relu=False, name='conv12_dw')
             .batch_normalization(relu=True, name='conv12_dw_bn')
             .conv(1, 1, 1024, 1, 1, biased=False, relu=False, name='conv12')
             .batch_normalization(relu=True, name='conv12_bn')
             .conv(3, 3, 1024, 1, 1, biased=False, group=1024, relu=False, name='conv13_dw')
             .batch_normalization(relu=True, name='conv13_dw_bn')
             .conv(1, 1, 1024, 1, 1, biased=False, relu=False, name='conv13')
             .batch_normalization(relu=True, name='conv13_bn')
             .conv(3, 3, 1024, 1, 1, biased=False, group=1024, relu=False, name='conv15_dw')
             .batch_normalization(relu=True, name='conv15_dw_bn')
             .conv(1, 1, 512, 1, 1, biased=False, relu=False, name='conv15')
             .batch_normalization(relu=True, name='conv15_bn')
             .upsample(name='upsample'))

        (self.feed('conv11_bn')
             .conv(3, 3, 512, 1, 1, biased=False, group=512, relu=False, name='conv17_dw')
             .batch_normalization(relu=True, name='conv17_dw_bn')
             .conv(1, 1, 512, 1, 1, biased=False, relu=False, name='conv17')
             .batch_normalization(relu=True, name='conv17_bn'))

        (self.feed('upsample', 
                   'conv17_bn')
             .add(name='conv17_sum')
             .conv(3, 3, 512, 1, 1, biased=False, group=512, relu=False, name='conv18_dw')
             .batch_normalization(relu=True, name='conv18_dw_bn')
             .conv(1, 1, 512, 1, 1, biased=False, relu=False, name='conv18')
             .batch_normalization(relu=True, name='conv18_bn')
             .conv(1, 1, 256, 1, 1, biased=False, relu=False, name='conv19')
             .batch_normalization(relu=True, name='conv19_bn')
             .upsample(name='upsample3'))

        (self.feed('conv15_bn')
             .conv(1, 1, 45, 1, 1, relu=False, name='conv22'))

        (self.feed('conv18_bn')
             .conv(1, 1, 45, 1, 1, relu=False, name='conv23'))

        (self.feed('upsample3', 
                   'conv5_bn')
             .add(name='upsample3_sum')
             .conv(3, 3, 256, 1, 1, biased=False, group=256, relu=False, name='conv25_dw')
             .batch_normalization(relu=True, name='conv25_dw_bn')
             .conv(1, 1, 256, 1, 1, biased=False, relu=False, name='conv25')
             .batch_normalization(relu=True, name='conv25_bn')
             .conv(1, 1, 1, 1, 1, relu=False, name='conv26'))