save

2023-03-07 22:02:55 +00:00
parent cbecf668df
commit cb76d36c4f
20 changed files with 626 additions and 168 deletions
@@ -320,6 +320,87 @@ class PyramidNew(tf.keras.Model):

        return z_cond

+    def sample(self, inputs):
+        levels_input, conds = inputs
+        uv_static, uv_body, uv_curv = conds[0], conds[1], conds[2]
+        z_cond = self.encode_conds(uv_static, uv_body, uv_curv, training=False)
+    
+        pred = [None, None, None, None, None]
+        for i in range(len(self.levels)):
+            pred[i] = self.levels[i].sample(z_cond, levels_input[i][-2])
+        
+        pred_1 = pred[0]
+        pred_2 = pred[1]
+        pred_3 = pred[2]
+        pred_4 = pred[3]
+        pred_5 = pred[4]
+
+        pred_1 = self.levels[0].enhance(pred_1, levels_input[0][-2], training=False)
+        pred_1 = tf.image.resize(pred_1, (self.levels[1].size, self.levels[1].size//2), tf.image.ResizeMethod.BILINEAR)
+        pred_1 = pred_1 * levels_input[1][-2][:,:,:,tf.newaxis]
+        
+        pred_2 = pred_1 + pred_2
+        pred_2 = self.levels[1].enhance(pred_2, levels_input[1][-2], training=False)
+        pred_2 = tf.image.resize(pred_2, (self.levels[2].size, self.levels[2].size//2), tf.image.ResizeMethod.BILINEAR)
+        pred_2 = pred_2 * levels_input[2][-2][:,:,:,tf.newaxis]
+
+        pred_3 = pred_2 + pred_3
+        pred_3 = self.levels[2].enhance(pred_3, levels_input[2][-2], training=False)
+        pred_3 = tf.image.resize(pred_3, (self.levels[3].size, self.levels[3].size//2), tf.image.ResizeMethod.BILINEAR)
+        pred_3 = pred_3 * levels_input[3][-2][:,:,:,tf.newaxis]
+
+        pred_4 = pred_3 + pred_4
+        pred_4 = self.levels[3].enhance(pred_4, levels_input[3][-2], training=False)
+        pred_4 = tf.image.resize(pred_4, (self.levels[4].size, self.levels[4].size//2), tf.image.ResizeMethod.BILINEAR)
+        pred_4 = pred_4 * levels_input[4][-2][:,:,:,tf.newaxis]
+
+        pred_5 = pred_4 + pred_5
+        pred_5 = pred_5 * levels_input[4][-2][:,:,:,tf.newaxis]
+
+        return pred_5
+
+
+    def sample_seq(self, inputs):
+        levels_input, conds = inputs
+        uv_static, uv_body, uv_curv = conds[0], conds[1], conds[2]
+        z_cond = self.encode_conds(uv_static, uv_body, uv_curv, training=False)
+    
+        # level 1
+        pred_1 = self.levels[0].sample(z_cond, levels_input[0][-2])
+        pred_1 = self.levels[0].enhance(pred_1, levels_input[0][-2], training=False)
+        pred_1 = tf.image.resize(pred_1, (self.levels[1].size, self.levels[1].size//2), tf.image.ResizeMethod.BILINEAR)
+        pred_1 = pred_1 * levels_input[1][-2][:,:,:,tf.newaxis]
+
+        # level 2
+        pred_2 = self.levels[1].sample(z_cond, levels_input[1][-2])
+        pred_2 = pred_1 + pred_2
+        pred_2 = self.levels[1].enhance(pred_2, levels_input[1][-2], training=False)
+        pred_2 = tf.image.resize(pred_2, (self.levels[2].size, self.levels[2].size//2), tf.image.ResizeMethod.BILINEAR)
+        pred_2 = pred_2 * levels_input[2][-2][:,:,:,tf.newaxis]
+
+        # level 3
+        pred_3 = self.levels[2].sample(z_cond, levels_input[2][-2])
+        pred_3 = pred_2 + pred_3
+        pred_3 = self.levels[2].enhance(pred_3, levels_input[2][-2], training=False)
+        pred_3 = tf.image.resize(pred_3, (self.levels[3].size, self.levels[3].size//2), tf.image.ResizeMethod.BILINEAR)
+        pred_3 = pred_3 * levels_input[3][-2][:,:,:,tf.newaxis]
+
+        # level 4
+        pred_4 = self.levels[3].sample(z_cond, levels_input[3][-2])
+        pred_4 = pred_3 + pred_4
+        pred_4 = self.levels[3].enhance(pred_4, levels_input[3][-2], training=False)
+        pred_4 = tf.image.resize(pred_4, (self.levels[4].size, self.levels[4].size//2), tf.image.ResizeMethod.BILINEAR)
+        pred_4 = pred_4 * levels_input[4][-2][:,:,:,tf.newaxis]
+
+        # level 5
+        pred_5 = self.levels[4].sample(z_cond, levels_input[4][-2])
+        pred_5 = pred_5 * levels_input[4][-2][:,:,:,tf.newaxis]
+
+        pred_5 = pred_4 + pred_5
+
+        return pred_5
+    
+
    def call(self, inputs, training):
        levels_input, conds = inputs
        uv_static, uv_body, uv_curv = conds[0], conds[1], conds[2]