Визуализируйте выпадающий слой в керасе

Question

Я нашел этот github репозиторий , который так хорошо рисует архитектуру глубокого обучения.

Теперь я хочу как-то нарисовать слой Dropout, который недоступен в этом репо. Я попытался создать 2 новых класса: DropoutConv2D для выпадающего слоя после слоя Conv и DropoutDense для выпадающего слоя после плотного наподобие этого.

class DropoutDense(Layer):
    def __init__(self):
        super(DropoutDense, self).__init__(filters = units)
    def get_description(self):
        return ["dropout"]
    def set_objects(self):
        x1 = self.prev_feature_map.right
        y11 = - math.pow(self.prev_feature_map.c, config.channel_scale) / 2
        y12 = math.pow(self.prev_feature_map.c, config.channel_scale) / 2
        x2 = self.next_feature_map.left
        y2 = - math.pow(self.next_feature_map.c, config.channel_scale) / 4
        line_color = config.line_color_layer
        self.objects.append(Line(x1, y11, x2, y2, color=line_color, dasharray=2))
        self.objects.append(Line(x1, y12, x2, y2, color=line_color, dasharray=2))

        x = (self.prev_feature_map.right + self.next_feature_map.left) / 2
        y = max(self.prev_feature_map.get_bottom(), self.next_feature_map.get_bottom()) + config.text_margin \
            + config.text_size

        for i, description in enumerate(self.get_description()):
            self.objects.append(Text(x, y + i * config.text_size, "{}".format(description),
                                     color=config.text_color_layer, size=config.text_size))

Кто-нибудь здесь работал с этим репо и давал мне предложения.