working pyramid pipeline

2022-09-27 21:26:15 +00:00
parent 1da83294ee
commit fa8f24f0c5
10 changed files with 997 additions and 213 deletions
@@ -0,0 +1,70 @@
+import os
+import time
+import random
+import tensorflow as tf
+import numpy as np
+# import plotly.graph_objects as go
+# import plotly.express as ex
+
+from custom_utils.custom_utils import load_model
+from models.base import AEEncoder
+from models.cvae_model import CVAE
+from models.pyramid_level import PyramidLevel
+
+
+class Pyramid(tf.keras.Model):
+    
+    def __init__(self, config):
+        super(Pyramid, self).__init__()
+
+        self.static_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size'])
+        if config['pyramid']['static_model'] != '':
+            load_model(self.static_enc, config['pyramid']['static_model'] + '_enc')
+
+        self.body_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size'])
+        if config['pyramid']['body_model'] != '':
+            load_model(self.body_enc, config['pyramid']['body_model'] + '_enc')
+
+        self.curv_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size'])
+        if config['pyramid']['curvature_model'] != '':
+            load_model(self.curv_enc, config['pyramid']['curvature_model'] + '_enc')
+
+        if config['pyramid']['freeze_encoders']:
+            self.static_enc.trainable = False
+            self.body_enc.trainable = False
+            self.curv_enc.trainable = False
+
+        self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(512, activation='relu'),
+                                                 tf.keras.layers.Dense(256, activation='relu'),
+                                                 tf.keras.layers.Dense(128, activation='relu'),
+                                                 tf.keras.layers.Dense(config['pyramid']['latent_dim']),])
+
+        self.levels = [PyramidLevel({ **config['pyramid']['levels'][i], **{ 'cond_ae': config['cond_ae']} }) for i in range(len(config['pyramid']['levels']))]
+
+
+    def encode_conds(self, uv_template, uv_body_posed, uv_curv):
+        z_static = self.static_enc(uv_template)
+        z_body = self.body_enc(uv_body_posed)
+        z_curv = self.curv_enc(uv_curv)
+
+        z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
+        z_cond = self.cond_encoder(z_cond)
+
+        return z_cond
+
+    def call(self, inputs, training=None):
+        uv_static, uv_body, uv_curv = inputs[-1][1], inputs[-1][2], inputs[-1][3]
+        z_cond = self.encode_conds(uv_static, uv_body, uv_curv)
+
+        out_levels = []
+        pred, mean, logvar = self.levels[0](inputs[0], z_cond, training=training)
+        out_levels.append([pred, mean, logvar])
+
+        for i in range(1, len(self.levels)):
+            pred_level, mean, logvar = self.levels[i](inputs[i], z_cond, training=training)
+            pred = pred + pred_level
+            out_levels.append([pred, mean, logvar])
+        
+        return out_levels
+
+