diff --git a/configs/ae_static.py b/configs/ae_static.py new file mode 100644 index 0000000..d5a9873 --- /dev/null +++ b/configs/ae_static.py @@ -0,0 +1,71 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '5' +config['name'] = 'static_torch_1' +config['notes'] = ' ' +config['tags'] = ['torch', 'static'] + +config['batch_size'] = 4 +config['train_dir'] = '/data-local/127/train' +config['valid_dir'] = '/data-local/127/valid' +config['test_dir'] = '/data-net/dataset/test_ready' +config['train_frames'] = None +config['valid_frames'] = None + +config['resolution'] = 512 +config['is_skirt'] = False +config['dataset_in_memory'] = True +config['is_cape'] = False +config['remove_global_orientation'] = True +config['device'] = 'cuda' + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['latent_dim'] = 128 + +config['bottleneck_scale'] = 2 + +config['cond_ae'] = dict() +config['cond_ae']['device'] = config['device'] +config['cond_ae']['is_label'] = False +config['cond_ae']['latent_dim'] = config['latent_dim'] +config['cond_ae']['kernel_size'] = 7 +config['cond_ae']['bottleneck_size'] = 8 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(8, 16, 3), + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, config['cond_ae']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae']['bottleneck_size'], 64, 3), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 8, 3), + CNBlockConfig(8, 3, 3), + CNBlockConfig(3, 3, 3), +] + +config['cond_ae']['encoder_block_settings'] = block_setting +config['cond_ae']['decoder_block_settings'] = block_setting_reverse +config['cond_ae']['activation'] = 'tanh' + + +config['lr'] = 0.0005 +config['cond_ae']['upscale_output'] = False +config['pretrained_model'] = '' + +config['signature'] = { 'id', + 'frame_nr', + 'uv_static', + 'mesh_static', + 'mesh_mask', + 'uv_mask', + } \ No newline at end of file diff --git a/custom_utils/custom_utils.py b/custom_utils/custom_utils.py index e8b48f8..097742c 100644 --- a/custom_utils/custom_utils.py +++ b/custom_utils/custom_utils.py @@ -1,12 +1,16 @@ import json import numpy as np -import tensorflow as tf +# import tensorflow as tf +import torch +import torch.nn as nn import scipy as sp import plotly.graph_objects as go from plotly.subplots import make_subplots +from torchvision.transforms import Resize +from torchvision.transforms.functional import InterpolationMode -from tensorflow_graphics.geometry.convolution.graph_pooling import pool, unpool -from tensorflow_graphics.geometry.convolution.graph_convolution import edge_convolution_template +# from tensorflow_graphics.geometry.convolution.graph_pooling import pool, unpool +# from tensorflow_graphics.geometry.convolution.graph_convolution import edge_convolution_template def m2uv_array(m2uv, is_smpl_orig=False): @@ -23,8 +27,8 @@ def m2uv_array(m2uv, is_smpl_orig=False): return result -class UV2Mesh(tf.keras.Model): - def __init__(self, res=256, is_skirt=False, is_smpl_orig=False, is_adaptive=False): +class UV2Mesh(nn.Module): + def __init__(self, res=256, is_skirt=False, is_smpl_orig=False): super(UV2Mesh, self).__init__() self.res = res uv_map = np.load('custom_utils/stretched_uv_map_' + ('skirt' if is_skirt else '3') + ('_orig' if is_smpl_orig else '') + '.npy') @@ -33,48 +37,40 @@ class UV2Mesh(tf.keras.Model): m2uv = m2uv_array(m2uv, is_smpl_orig=is_smpl_orig) self.uv_pixels = self.uv_map_discreet[m2uv] self.smpl_res = 14475 if not is_smpl_orig else 3638 - self.is_adaptive = is_adaptive - - if self.is_adaptive: - laplacian = sp.sparse.load_npz("custom_utils/laplacian/laplacian0" + ("_skirt" if is_skirt else "") + ".npz") - self.gc = GraphConvolution((3, 3), L=laplacian, batchSize=4) - @tf.function def uv_to_mesh(self, uv): if uv.shape[1] == 3: - uv = tf.transpose(uv, [0,2,3,1]) + uv = torch.permute(uv, [0,2,3,1]) if uv.shape[1] != 512: - uv = tf.image.resize(uv, (512, 256), tf.image.ResizeMethod.LANCZOS5) + uv = Resize(uv[:,:,None], size=(512, 256), interpolation=InterpolationMode.BILINEAR) - batch_size = uv.shape[0] + # batch_size = uv.shape[0] - x1_ind = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.zeros((batch_size, self.smpl_res, 1), dtype=tf.int32)), axis=-1) - y1_ind = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, self.smpl_res, 1), dtype=tf.int32)), axis=-1) - z1_ind = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, self.smpl_res, 1), dtype=tf.int32) * 2), axis=-1) - x2_ind = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.zeros((batch_size, self.smpl_res, 1), dtype=tf.int32)), axis=-1) - y2_ind = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, self.smpl_res, 1), dtype=tf.int32)), axis=-1) - z2_ind = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, self.smpl_res, 1), dtype=tf.int32) * 2), axis=-1) + # x1_ind = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.zeros((batch_size, self.smpl_res, 1), dtype=tf.int32)), axis=-1) + # y1_ind = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, self.smpl_res, 1), dtype=tf.int32)), axis=-1) + # z1_ind = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, self.smpl_res, 1), dtype=tf.int32) * 2), axis=-1) + # x2_ind = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.zeros((batch_size, self.smpl_res, 1), dtype=tf.int32)), axis=-1) + # y2_ind = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, self.smpl_res, 1), dtype=tf.int32)), axis=-1) + # z2_ind = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, self.smpl_res, 1), dtype=tf.int32) * 2), axis=-1) - x1 = tf.gather_nd(uv, x1_ind, batch_dims=1) - y1 = tf.gather_nd(uv, y1_ind, batch_dims=1) - z1 = tf.gather_nd(uv, z1_ind, batch_dims=1) - x2 = tf.gather_nd(uv, x2_ind, batch_dims=1) - y2 = tf.gather_nd(uv, y2_ind, batch_dims=1) - z2 = tf.gather_nd(uv, z2_ind, batch_dims=1) + # x1 = tf.gather_nd(uv, x1_ind, batch_dims=1) + # y1 = tf.gather_nd(uv, y1_ind, batch_dims=1) + # z1 = tf.gather_nd(uv, z1_ind, batch_dims=1) + # x2 = tf.gather_nd(uv, x2_ind, batch_dims=1) + # y2 = tf.gather_nd(uv, y2_ind, batch_dims=1) + # z2 = tf.gather_nd(uv, z2_ind, batch_dims=1) - p1 = tf.stack((x1, y1, z1), axis=-1) - p2 = tf.stack((x2, y2, z2), axis=-1) + # p1 = tf.stack((x1, y1, z1), axis=-1) + # p2 = tf.stack((x2, y2, z2), axis=-1) - mesh = (p1 + p2) / 2 + mesh = torch.mean(uv[:, self.uv_pixels[:,:,0], self.uv_pixels[:,:,1]], dim=2) return mesh def __call__(self, uv): x = self.uv_to_mesh(uv) - if self.is_adaptive: - x = self.gc(x) return x @@ -313,10 +309,16 @@ def load_model(model, checkpoint, disable_assert=False, assert_consumed_on=False def get_summary_str(model): - lines = [] - model.summary(print_fn=lines.append, line_length=150,expand_nested=True,show_trainable=True) + # from torchview import draw_graph + # draw_graph(model, input_size=((4,3,512,256), (4,512,256)), expand_nested=True) + + # lines = [] + # model.summary(print_fn=lines.append, line_length=150,expand_nested=True,show_trainable=True) # Add initial spaces to avoid markdown formatting in TensorBoard - return lines + trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad) + total_params = sum(param.numel() for param in model.parameters()) + + return [f'all_params: {total_params}, trainable: {trainable_params}'] def get_summary_html(model): lines = get_summary_str(model) @@ -341,7 +343,7 @@ def merge_verts(Vs, batchDim): return V -class GraphConvolution(tf.keras.layers.Layer): +class GraphConvolution(nn.Module): def __init__(self, shape, activation=None, name='gcn', L=None, batchSize=None, normalization=None): super(GraphConvolution, self).__init__() self.shape = shape @@ -405,4 +407,46 @@ def group_inputs_for_levels(batch, config): uv_orig = batch['uv_posed' + res_modifier] levels_input.append([uv_orig, batch['uv_static' + res_modifier], batch['uv_body_posed' + res_modifier], batch['uv_curvature' + res_modifier], batch['uv_mask' + res_modifier], batch['uv_mask']]) conds = [batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature']] - return levels_input, conds \ No newline at end of file + return levels_input, conds + + +class Accumulator(): + def __init__(self): + self.dict = dict() + + def update(self, dict): + for key in dict.keys(): + if key in self.dict.keys(): + self.dict[key].append(dict[key]) + else: + self.dict[key] = [dict[key]] + + def get(self): + result = dict() + for key in self.dict.keys(): + result[key] = np.mean(self.dict[key]) + return result + + def reset(self): + self.dict = dict() + + +from stl import mesh as mesh_stl + +def save_stl(file, mesh, faces): + obj = mesh_stl.Mesh(np.zeros(faces.shape[0], dtype=mesh_stl.Mesh.dtype)) + for i, f in enumerate(faces): + for j in range(3): + obj.vectors[i][j] = mesh[f[j],:] + + obj.save(file) + + +def write_obj(file, vertices, faces, colors): + with open(file, 'w') as file: + for v, c in zip(vertices, colors): + line = 'v ' + ' '.join([str(_) for _ in v]) + ' ' + ' '.join([str(_) for _ in c]) + '\n' + file.write(line) + for f in faces: + line = 'f ' + ' '.join([str(_) for _ in f]) + '\n' + file.write(line) \ No newline at end of file diff --git a/data/dataset.py b/data/dataset.py index c667724..a516f41 100644 --- a/data/dataset.py +++ b/data/dataset.py @@ -37,6 +37,8 @@ def read_frames(folder, shuffle=True, is_cape=False, testing=False, frames_file= frames = sorted(os.listdir(os.path.join(folder, 'mesh_static') )) frames = [frame.split('.')[0] + '_4' for frame in frames] + frames = [Frame(frame.split('_')[0], frame.split('_')[1]) for frame in frames] + if frames_file is not None: frames = [] with open(os.path.join(folder, frames_file), 'r') as f: @@ -122,7 +124,7 @@ class StaticCache(): class Dataset(Dataset): - def __init__(self, folder, signature, resolutions, frames_file='frames_all.txt', device='cuda', use_cache=False, is_cape=False): + def __init__(self, folder, signature, resolutions=[512], frames_file='frames_all.txt', device='cuda', use_cache=False, is_cape=False): self.folder = folder self.device = device self.is_cape = is_cape @@ -210,6 +212,10 @@ class Dataset(Dataset): data = { key: (torch.from_numpy(value).to(self.device) if not (type(value) == str or type(value) == torch.Tensor) else value) for key, value in data.items()} + for key, value in data.items(): + if key.startswith('uv_') and value.shape[-1] == 3: + data[key] = torch.permute(value, [2,0,1]) + if self.cache is not None: self.cache[frame] = data return data diff --git a/losses/losses.py b/losses/losses.py new file mode 100644 index 0000000..18c948e --- /dev/null +++ b/losses/losses.py @@ -0,0 +1,244 @@ +import torch + +import numpy as np +# import tensorflow_graphics as tfg +import scipy as sp +import scipy.io as sio + +from pytorch3d.structures import Meshes +from scipy.spatial import cKDTree +# from tensorflow_graphics.geometry.representation.mesh.normals import vertex_normals +# from tensorflow_graphics.geometry.convolution.graph_convolution import edge_convolution_template + + + +# from hbertiche/PBNS +""" COLLISION FUNCTIONS START """ +F_coll = np.load("smpl/smpl_faces.npy") +F_coll_skirt = np.load("smpl/smpl_faces_skirt.npy") +F_coll_orig = np.load("smpl/smpl_faces_orig.npy") +F_norm = np.tile(F_coll[np.newaxis,:,:], [4,1,1]) +F_norm_skirt = np.tile(F_coll_skirt[np.newaxis,:,:], [4,1,1]) +F_norm_orig = np.tile(F_coll_orig[np.newaxis,:,:], [4,1,1]) + +def _nearest_neighbour(V, B): + tree = cKDTree(B.numpy()) + return tree.query(V.numpy(), workers=-1, distance_upper_bound=0.2)[1] + +# @tf.function +# def _tf_nn_parallel(V, B): +# return tf.map_fn(fn=lambda elem: tf.py_function(func=_nearest_neighbour, inp=elem, Tout=tf.int32), elems=(V,B), fn_output_signature=tf.int32, parallel_iterations=V.shape[0]) + +# @tf.function +def collision_loss(V, B, mask, thr=.004): + mesh = Meshes(B, torch.tile(F_coll[None], [B.shape[0], 1, 1])) + B_vn = torch.stack(mesh.verts_normals_list()) + + # compute correspondences + idx = torch.stack([_nearest_neighbour(V[i], B[i]) for i in range(V.shape[0])]) + N, M = V.shape[:2] + # _idx = tf.tile(np.array(list(range(N)), np.int32)[:,None], [1, M]) + _idx = torch.ones((N, M), dtype=torch.int32, device=V.get_device()) * np.array(list(range(N)), np.int32)[:,None] + idx = torch.stack((_idx, idx), axis=-1) + # compute loss + D = V - B[idx] + b_vn = B_vn[idx] + dot = torch.einsum('abc,abc->ab', D, b_vn) + thresholded = torch.minimum(dot - thr, 0) * mask + loss = torch.sum(thresholded ** 2, axis=-1) + ratio = torch.sum((thresholded < 0).to(torch.float32)) / torch.sum(mask) + + return loss, ratio + + +# def collision_ratio(V, B, mask, thr=.004): + # B_vn = vertex_normals(B, tf.tile(F[None], [B.shape[0], 1, 1])) + + # # compute correspondences + # idx = _tf_nn_parallel(V, B) + # N, M = V.shape[:2] + # # _idx = tf.tile(np.array(list(range(N)), np.int32)[:,None], [1, M]) + # _idx = tf.ones((N, M), dtype=tf.int32) * np.array(list(range(N)), np.int32)[:,None] + # idx = tf.stack((_idx, idx), axis=-1) + # # compute loss + # D = V - tf.gather_nd(B, idx) + # b_vn = tf.gather_nd(B_vn, idx) + # dot = tf.einsum('abc,abc->ab', D, b_vn) + # loss = tf.minimum(dot - thr, 0) * mask # / tf.reduce_sum(mask, axis=-1) + # ratio = tf.reduce_sum(tf.cast(loss < 0, dtype=tf.float32)) / tf.reduce_sum(mask) + + # return ratio + +""" COLLISION FUNCTIONS END """ + + +laplacian_loss = sio.loadmat("custom_utils/ssmpl0.mat")['L'].tocoo() +edges: torch.Tensor = torch.tensor(np.stack(np.nonzero(laplacian_loss.todense())).transpose(), device='cuda') + +laplacian_loss_skirt = sio.loadmat("custom_utils/ssmpl_skirt0.mat")['L'].tocoo() +edges_skirt: torch.Tensor = torch.tensor(np.stack(np.nonzero(laplacian_loss.todense())).transpose(), device='cuda') + + +def get_edge_mask(masks: torch.Tensor): + if edges.get_device() != masks.get_device(): + edges.to(masks.get_device()) + return torch.min(masks[:,edges], dim=-1)[0] + +def edge_loss(x, y, mask): + e_mask = get_edge_mask(mask) + x_e = x[:,edges[:,0],:] - x[:,edges[:,1],:] + x_e = torch.sum(x_e ** 2, -1) + x_e = torch.sqrt(x_e + 1e-9) + y_e = y[:,edges[:,0],:] - y[:,edges[:,1],:] + y_e = torch.sum(y_e ** 2, -1) + y_e = torch.sqrt(y_e + 1e-9) + d_e = y_e - x_e + d_e = d_e * e_mask + loss = torch.sum(d_e ** 2, dim=-1)# / tf.reduce_sum(e_mask, axis=-1) + return loss + +def get_edge_mask_skirt(masks: torch.Tensor): + if edges.get_device() != masks.get_device(): + edges.to(masks.get_device()) + return torch.min(masks[:,edges_skirt], dim=-1)[0] + +def edge_loss_skirt(x, y, mask): + e_mask = get_edge_mask_skirt(mask) + x_e = x[:,edges_skirt[:,0],:] - x[:,edges_skirt[:,1],:] + x_e = torch.sum(x_e ** 2, -1) + x_e = torch.sqrt(x_e + 1e-9) + y_e = y[:,edges_skirt[:,0],:] - y[:,edges_skirt[:,1],:] + y_e = torch.sum(y_e ** 2, -1) + y_e = torch.sqrt(y_e + 1e-9) + d_e = y_e - x_e + d_e = d_e * e_mask + loss = torch.sum(d_e ** 2, dim=-1)# / tf.reduce_sum(e_mask, axis=-1) + return loss + + +# class PairwiseDistance(tf.keras.layers.Layer): +# def __init__(self, **kwargs): +# super(PairwiseDistance, self).__init__(**kwargs) + +# def call(self, A, B): +# rA = tf.reduce_sum(tf.square(A), axis=-1) +# rB = tf.reduce_sum(tf.square(B), axis=-1) +# transpose_axes = [0, 2, 1] +# distances = - 2*tf.matmul(A, tf.transpose(B, transpose_axes)) + rA[:, :, tf.newaxis] + rB[:, tf.newaxis, :] +# return distances + + +# class NearestNeighbour(tf.keras.layers.Layer): +# def __init__(self, **kwargs): +# super(NearestNeighbour, self).__init__(**kwargs) + +# def call(self, A, B): +# distances = PairwiseDistance(dtype=self.dtype)(A, B) +# nearest_neighbour = tf.argmin(distances, axis=-1) +# return tf.cast(nearest_neighbour, dtype=tf.int32) + +# +# @tf.function +def collision_penalty(va, vb, mask, eps=1e-3): + nb = vertex_normals(vb, tf.tile(F[None], [vb.shape[0], 1, 1])) + batch_size = tf.cast(tf.shape(va)[0], va.dtype) + + closest_vertices = NearestNeighbour(dtype=va.dtype)(va, vb) + vb = tf.gather(vb, closest_vertices, batch_dims=1) + nb = tf.gather(nb, closest_vertices, batch_dims=1) + + distance = tf.reduce_sum(nb*(va - vb), axis=-1) + interpenetration = tf.maximum(eps - distance * mask, 0) + + return tf.reduce_sum(interpenetration**3) / batch_size + +whole_garment_uv_mask = np.load('custom_utils/mask_256x512.npy') + +def uv_mask_loss(pred, uv_mask): + inv_mask = (1 - uv_mask) * whole_garment_uv_mask + pred = tf.reduce_sum(pred, axis=1) + invalid = pred * inv_mask + invalid = invalid > 0.01 + invalid_count = tf.math.count_nonzero(invalid, dtype=tf.float32) + total_count = tf.math.count_nonzero(inv_mask, dtype=tf.float32) + return invalid_count / total_count + + + + +def merge_verts(Vs, batchDim): + V = tf.concat(tf.unstack(Vs, num=batchDim), axis=0) + return V + + +def _bend_loss_single(x, F, L, mask): + verts = tf.boolean_mask(x, mask, axis=0) + VN = vertex_normals(verts, F) + bend = edge_convolution_template( + VN, L, sizes=None, + edge_function=lambda x,y: (x - y)**2, + reduction='weighted', + edge_function_kwargs={} + ) + loss = tf.reduce_sum(bend) + return loss + +def bend_loss(x, F, L, mask): + # x: predicted outfits + # F: faces + # L: laplacian + loss = tf.stack([_bend_loss_single(x[i], F[i], L[i], mask[i]) for i in range(x.shape[0])]) + return loss + + +# def bend_loss(x, F, L, mask): +# # x: predicted outfits +# # F: faces +# # L: laplacian +# # x = x * mask[:,:,tf.newaxis] +# batch_size = x.shape[0] +# verts, faces = maskFacesBatch(x, F, mask) +# nr_valid_verts = [tf.reduce_sum(mask[i]) for i in range(batch_size)] +# x = tf.concat(verts, axis=0) +# F = tf.concat([faces[i] + tf.reduce_sum(nr_valid_verts[:i]) for i in range(1, batch_size)]) +# F = tf.concat((faces[0], F), axis=0) +# VN = vertex_normals(x, F) +# bend = edge_convolution_template( +# VN, L, sizes=None, +# edge_function=lambda x,y: (x - y)**2, +# reduction='weighted', +# edge_function_kwargs={} +# ) +# loss = tf.reduce_sum(bend, axis=[1]) +# loss = tf.reshape(loss, (batch_size, 14475)) +# loss = tf.reduce_mean(loss, axis=1) +# return loss + + +def log_normal_pdf(sample, mean, logvar, raxis=1): + log2pi = tf.math.log(2. * np.pi) + res = tf.reduce_sum(-.5 * ((sample - mean) ** 2. * tf.exp(-logvar) + logvar + log2pi), axis=raxis) + return res + + +def kl_loss(mean, logvar): + kl_loss = 1 + logvar - tf.square(mean) - tf.exp(logvar) + kl_loss = tf.reduce_mean(kl_loss) + kl_loss *= -0.5 + return kl_loss + + +def normal_loss(pred, gt, mask, is_skirt=False, is_smpl_orig=False): + f_norm = None + if not is_skirt and not is_smpl_orig: + f_norm = F_norm + elif is_skirt: + f_norm = F_norm_skirt + elif is_smpl_orig: + f_norm = F_norm_orig + else: + raise NotImplementedError + + pred_norm = vertex_normals(pred, f_norm) + gt_norm = vertex_normals(gt, f_norm) + return tf.reduce_mean(tf.reduce_mean((tf.reduce_sum(tf.math.abs(pred_norm - gt_norm), axis=-1) * mask), axis=-1)) diff --git a/models/base.py b/models/base.py index e9cfd67..f918b3e 100644 --- a/models/base.py +++ b/models/base.py @@ -1,134 +1,21 @@ -import tensorflow as tf -from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block, ConvNeXtReverseLabels +import torch +from torch import nn +from models.convnext import ConvNeXt, CNBlock, CNBlockConfig +from typing import Any, Callable, List, Optional, Sequence -class Bottleneck(tf.keras.Model): + +class VariationalEncoder(nn.Module): - def __init__(self, dim, latent_dim, kernel_size=7): - super(Bottleneck, self).__init__() - - self.bottleneck = tf.keras.Sequential( - [Block(dim=dim, drop_path=0, kernel_size=kernel_size, - layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)] - ) - - self.flatten = tf.keras.layers.Flatten() - self.dense = tf.keras.layers.Dense(latent_dim) - - def call(self, inputs, training): - x = inputs - - x = self.bottleneck(x, training=training) - x = self.flatten(x) - x = self.dense(x, training=training) - - return x - - -class VariationalEncoder(tf.keras.Model): - - def __init__(self, latent_dim, depths, dims, kernel_size=7): + def __init__(self, block_setting: CNBlockConfig, kernel_size: int = 7): super(VariationalEncoder, self).__init__() - self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size) + self.encoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size) - self.flatten = tf.keras.layers.Flatten() - self.mean = tf.keras.layers.Dense(latent_dim) - self.log_var = tf.keras.layers.Dense(latent_dim) - - def call(self, inputs, training): - x = inputs - - x = self.encoder(x, training=training) - x = self.flatten(x) - mean = self.mean(x) - log_var = self.log_var(x) - - return mean, log_var - - -class AEEncoder(tf.keras.Model): - - def __init__(self, latent_dim, depths, dims, bottleneck_size=32, kernel_size=7): - super(AEEncoder, self).__init__() - - self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size) - self.bottleneck = Bottleneck(bottleneck_size, latent_dim, kernel_size) - - - def call(self, inputs, training): - if len(inputs) == 2: - x, _ = inputs - else: - x = inputs - - x = self.encoder(x, training=training) - x = self.bottleneck(x, training=training) - - if len(inputs) == 2: - return x, inputs[1] - else: - return x - - -class Decoder(tf.keras.Model): - - def __init__(self, input_size, depths, dims, kernel_size=7, is_label=False, upscale=False): - super(Decoder, self).__init__() - - self.dense = tf.keras.layers.Dense(tf.reduce_prod(input_size)) - self.reshape = tf.keras.layers.Reshape(input_size) - self.upscale = upscale - - if not is_label: - self.decoder = ConvNeXtReverse(depths=list(reversed(depths)), dims=list(reversed(dims)), kernel_size=kernel_size) - else: - self.decoder = ConvNeXtReverseLabels(depths=list(reversed(depths)), dims=list(reversed(dims)), nr_of_classes=7) - - - def call(self, inputs, training): - if len(inputs) == 2: - x, mask = inputs - else: - x = inputs - - x = self.dense(x) - x = self.reshape(x) - x = self.decoder(x, training=training) - - if self.upscale: - x = tf.image.resize(x, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5) - - - if len(inputs) == 2: - if len(mask.shape) == 2: - x = x * mask[tf.newaxis,:,:,tf.newaxis] - else: - x = x * mask[:,:,:,tf.newaxis] - - return x - - - -class VariationalEncoderNew(tf.keras.Model): - - def __init__(self, depths, dims, kernel_size=7): - super(VariationalEncoderNew, self).__init__() - - self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size) - - self.mean_branch = tf.keras.Sequential( - [Block(dim=64, drop_path=0, kernel_size=kernel_size, - layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)] - ) - - self.var_branch = tf.keras.Sequential( - [Block(dim=64, drop_path=0, kernel_size=kernel_size, - layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)] - ) - - self.flatten_mean = tf.keras.layers.Flatten() - self.flatten_var = tf.keras.layers.Flatten() + self.mean_branch = nn.Sequential([CNBlock(block_setting[-1].out_channels, kernel_size=kernel_size) for _ in range(5)]) + self.var_branch = nn.Sequential([CNBlock(block_setting[-1].out_channels, kernel_size=kernel_size) for _ in range(5)]) + self.flatten_mean = nn.Flatten() + self.flatten_var = nn.Flatten() def call(self, inputs, training): @@ -144,23 +31,27 @@ class VariationalEncoderNew(tf.keras.Model): return mean, log_var - -class AEEncoderNew(tf.keras.Model): +class AEEncoder(nn.Module): - def __init__(self, depths, dims, kernel_size=7, activation='tanh'): - super(AEEncoderNew, self).__init__() + def __init__(self, block_setting: CNBlockConfig, kernel_size: int = 7, activation='tanh', device='cuda'): + super(AEEncoder, self).__init__() - self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size) - self.act = tf.keras.layers.Activation(activation) + self.encoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size, device=device) + if activation == 'tanh': + self.act = nn.Tanh() + elif activation == 'sigmoid': + self.act = nn.Sigmoid() + else: + raise ValueError('Activation function not supported') - def call(self, inputs, training): + def forward(self, inputs): if len(inputs) == 2: x, _ = inputs else: x = inputs - x = self.encoder(x, training=training) + x = self.encoder(x) x = self.act(x) if len(inputs) == 2: @@ -169,39 +60,36 @@ class AEEncoderNew(tf.keras.Model): return x -class DecoderNew(tf.keras.Model): +class Decoder(nn.Module): - def __init__(self, depths, dims, kernel_size=7, is_label=False): - super(DecoderNew, self).__init__() + def __init__(self, block_setting: CNBlockConfig, kernel_size=7, is_label=False, nr_of_classes=None, device='cuda'): + super(Decoder, self).__init__() self.is_label = is_label - if not is_label: - self.decoder = ConvNeXtReverse(depths=depths, dims=dims, kernel_size=kernel_size) - else: - self.decoder = ConvNeXtReverseLabels(depths=depths, dims=dims, nr_of_classes=7) + self.decoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size, is_label=is_label, nr_of_labels=nr_of_classes, device=device, reversed=True) + self.softmax = nn.Softmax(dim=1) - def call(self, inputs, training): + def forward(self, inputs): if len(inputs) == 2: x, mask = inputs else: x = inputs - x = self.decoder(x, training=training) + x = self.decoder(x) if self.is_label: - x = x * mask[:,:,:,tf.newaxis] + x = x * mask[:,:,:,None] - background = tf.ones((mask.shape[0],512,256,1)) * (1 - mask[:,:,:,tf.newaxis]) - x = tf.concat((background, x), axis=-1) - - x = tf.math.softmax(x) + background = torch.ones((mask.shape[0],512,256,1), device=mask.get_device()) * (1 - mask[:,:,:,None]) + x = torch.concat((background, x), axis=-1) + x = self.softmax(x) if len(inputs) == 2: if len(mask.shape) == 2: - x = x * mask[tf.newaxis,:,:,tf.newaxis] + x = x * mask[None,None,:,:] else: - x = x * mask[:,:,:,tf.newaxis] + x = x * mask[:,None,:,:] return x \ No newline at end of file diff --git a/models/cond_encoder.py b/models/cond_encoder.py index c86f54c..34ab6ee 100644 --- a/models/cond_encoder.py +++ b/models/cond_encoder.py @@ -1,38 +1,21 @@ -import tensorflow as tf -from models.base import AEEncoder, Decoder, AEEncoderNew, DecoderNew +import torch +from torch import nn +from models.base import AEEncoder, Decoder -class AutoEncoder(tf.keras.Model): +class AutoEncoder(nn.Module): def __init__(self, config): super(AutoEncoder, self).__init__() - self.encoder : AEEncoder = AEEncoder(config['latent_dim'], config['encoder_depths'], config['encoder_dims'], config['bottleneck_size'], config['kernel_size']) - self.decoder : Decoder = Decoder(config['decoder_input'], config['decoder_depths'], config['decoder_dims'], kernel_size=config['kernel_size'], is_label=config['is_label'], upscale=config['upscale_output']) + self.encoder : AEEncoder = AEEncoder(config['encoder_block_settings'], config['kernel_size'], config['activation'], device=config['device']) + self.decoder : Decoder = Decoder(config['decoder_block_settings'], kernel_size=config['kernel_size'], is_label=config['is_label'], device=config['device']) - def call(self, inputs, training): - x = inputs + def forward(self, input): + x = input - x = self.encoder(x, training=training) - x = self.decoder(x, training=training) - - return x - - -class AutoEncoderNew(tf.keras.Model): - - def __init__(self, config): - super(AutoEncoderNew, self).__init__() - - self.encoder : AEEncoderNew = AEEncoderNew(config['encoder_depths'], config['encoder_dims'], config['kernel_size'], config['activation']) - self.decoder : DecoderNew = DecoderNew(config['decoder_depths'], config['decoder_dims'], kernel_size=config['kernel_size'], is_label=config['is_label']) - - - def call(self, inputs, training): - x = inputs - - x = self.encoder(x, training=training) - x = self.decoder(x, training=training) + x = self.encoder(x) + x = self.decoder(x) return x \ No newline at end of file diff --git a/models/convnext.py b/models/convnext.py index 55bad56..e055370 100644 --- a/models/convnext.py +++ b/models/convnext.py @@ -1,306 +1,424 @@ -from unittest import skip -import numpy as np -import tensorflow as tf -from tensorflow.keras.layers import (Conv2D, Conv2DTranspose, Dense, DepthwiseConv2D, - GlobalAveragePooling2D, Layer, - LayerNormalization, Flatten, Reshape, UpSampling2D) +# taken from torchvision.models.convnext -def drop_path(inputs, drop_prob, is_training): - # borrowed from https://github.com/rishigami/Swin-Transformer-TF/blob/main/swintransformer/model.py - if (not is_training) or (drop_prob == 0.): - return inputs +from functools import partial +from typing import Any, Callable, List, Optional, Sequence - # Compute keep_prob - keep_prob = 1.0 - drop_prob +import torch +from torch import nn, Tensor +from torch.nn import functional as F - # Compute drop_connect tensor - random_tensor = keep_prob - shape = (tf.shape(inputs)[0],) + (1,) * \ - (len(tf.shape(inputs)) - 1) - random_tensor += tf.random.uniform(shape, dtype=inputs.dtype) - binary_tensor = tf.floor(random_tensor) - output = tf.math.divide(inputs, keep_prob) * binary_tensor - return output +from torchvision.ops.misc import Conv2dNormActivation, Permute +# from torchvision.ops.misc import ConvNormActivation as Conv2dNormActivation +# from torchvision.ops import Permute +from torchvision.ops.stochastic_depth import StochasticDepth +from utils.convnext_utils import CNBlockConfig -class DropPath(tf.keras.layers.Layer): - # borrowed from https://github.com/rishigami/Swin-Transformer-TF/blob/main/swintransformer/model.py - def __init__(self, drop_prob=None): +__all__ = [ + "ConvNeXt", + "ConvNeXt_Tiny_Weights", + "ConvNeXt_Small_Weights", + "ConvNeXt_Base_Weights", + "ConvNeXt_Large_Weights", + "convnext_tiny", + "convnext_small", + "convnext_base", + "convnext_large", +] + + +class LayerNorm2d(nn.LayerNorm): + def forward(self, x: Tensor) -> Tensor: + x = x.permute(0, 2, 3, 1) + x = F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps) + x = x.permute(0, 3, 1, 2) + return x + + +class CNBlock(nn.Module): + def __init__( + self, + dim, + layer_scale: float = 1e-6, + stochastic_depth_prob: float = 0.0, + kernel_size: int = 7, + norm_layer: Optional[Callable[..., nn.Module]] = None, + ) -> None: super().__init__() - self.drop_prob = drop_prob + if norm_layer is None: + norm_layer = partial(nn.LayerNorm, eps=1e-6) - def call(self, x, training): - return drop_path(x, self.drop_prob, training) + self.block = nn.Sequential( + nn.Conv2d(dim, dim, kernel_size=kernel_size, padding='same', groups=dim, bias=True), + Permute([0, 2, 3, 1]), + norm_layer(dim), + nn.Linear(in_features=dim, out_features=4 * dim, bias=True), + nn.GELU(), + nn.Linear(in_features=4 * dim, out_features=dim, bias=True), + Permute([0, 3, 1, 2]), + ) + self.layer_scale = nn.Parameter(torch.ones(dim, 1, 1) * layer_scale) + self.stochastic_depth = StochasticDepth(stochastic_depth_prob, "row") + + def forward(self, input: Tensor) -> Tensor: + result = self.layer_scale * self.block(input) + result = self.stochastic_depth(result) + result += input + return result -class Block(Layer): - r""" ConvNeXt Block. There are two equivalent implementations: - (1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W) - (2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back - We use (2) as we find it slightly faster in PyTorch - - Args: - dim (int): Number of input channels. - drop_path (float): Stochastic depth rate. Default: 0.0 - layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6. - """ - - def __init__(self, dim, kernel_size=7, drop_path=0., layer_scale_init_value=1e-6, prefix=''): +class ConvNeXtBase(nn.Module): + def __init__( + self, + block_setting: List[CNBlockConfig], + stochastic_depth_prob: float = 0.0, + layer_scale: float = 1e-6, + block: Optional[Callable[..., nn.Module]] = None, + norm_layer: Optional[Callable[..., nn.Module]] = None, + device: str = 'cuda', + kernel_size: int = 7, + ) -> None: super().__init__() - self.dwconv = DepthwiseConv2D( - kernel_size=kernel_size, padding='same') # depthwise conv - self.norm = LayerNormalization(epsilon=1e-6) - # pointwise/1x1 convs, implemented with linear layers - self.pwconv1 = Dense(4 * dim) - self.act = tf.keras.activations.gelu - self.pwconv2 = Dense(dim) - self.drop_path = DropPath(drop_path) - self.dim = dim - self.layer_scale_init_value = layer_scale_init_value - self.prefix = prefix - def build(self, input_shape): - self.gamma = tf.Variable( - initial_value=self.layer_scale_init_value * tf.ones((self.dim)), - trainable=True, - name=f'{self.prefix}/gamma') - self.built = True + if not block_setting: + raise ValueError("The block_setting should not be empty") + elif not (isinstance(block_setting, Sequence) and all([isinstance(s, CNBlockConfig) for s in block_setting])): + raise TypeError("The block_setting should be List[CNBlockConfig]") - def call(self, x, training): - input = x - x = self.dwconv(x, training=training) - # x = x.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C) - x = self.norm(x, training=training) - x = self.pwconv1(x, training=training) + if block is None: + block = CNBlock + + if norm_layer is None: + norm_layer = partial(LayerNorm2d, eps=1e-6) + + # self.permute_cf = Permute([0, 3, 1, 2]) + # self.permute_cl = Permute([0, 2, 3, 1]) + + self.stages: List[nn.Module] = [] + self.downsamples: List[nn.Module] = [] + + # Stem + firstconv_output_channels = block_setting[0].input_channels + self.stem = Conv2dNormActivation( + 3, + firstconv_output_channels, + kernel_size=4, + stride=4, + padding=0, + norm_layer=norm_layer, + activation_layer=None, + bias=True, + ).to(device) + + + total_stage_blocks = sum(cnf.num_layers for cnf in block_setting) + stage_block_id = 0 + for cnf in block_setting: + # Bottlenecks + stage: List[nn.Module] = [] + for _ in range(cnf.num_layers): + # adjust stochastic depth probability based on the depth of the stage block + sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0) + stage.append(block(cnf.input_channels, layer_scale, sd_prob, kernel_size)) + stage_block_id += 1 + self.stages.append(nn.Sequential(*stage).to(device)) + if cnf.out_channels is not None: + # Downsampling + self.downsamples.append( + nn.Sequential( + norm_layer(cnf.input_channels), + nn.Conv2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2), + ).to(device) + ) + + self.act = nn.Tanh() + + for m in self.modules(): + if isinstance(m, (nn.Conv2d, nn.Linear)): + nn.init.trunc_normal_(m.weight, std=0.02) + if m.bias is not None: + nn.init.zeros_(m.bias) + + def _forward_impl(self, x: Tensor) -> Tensor: + # x = self.permute_cf(x) + x = self.stem(x) + # features: List[Tensor] = [] + for i in range(len(self.stages) - 1): + x = self.stages[i](x) + # features.append(x) + x = self.downsamples[i](x) + x = self.stages[-1](x) x = self.act(x) - x = self.pwconv2(x, training=training) - if self.gamma is not None: - x = self.gamma * x - # x = x.permute(0, 3, 1, 2) # (N, H, W, C) -> (N, C, H, W) - x = input + self.drop_path(x, training=training) + # x = self.permute_cl(x) return x + def forward(self, x: Tensor) -> Tensor: + return self._forward_impl(x) -class ConvNeXt(tf.keras.Model): - r""" ConvNeXt - A Tensorflow keras impl of : `A ConvNet for the 2020s` - - https://arxiv.org/pdf/2201.03545.pdf - Args: - num_classes (int): Number of classes for classification head. Default: 1000 - depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3] - dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768] - include_top (bool): whether to add head or just use it as feature extractor. Default: True - drop_path_rate (float): Stochastic depth rate. Default: 0. - layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6. - head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1. - """ - def __init__(self, num_classes=1000, kernel_size=7, - depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], include_top=True, - drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1., - ): +class ConvNeXtReverseBase2(nn.Module): + def __init__( + self, + block_setting: List[CNBlockConfig], + stochastic_depth_prob: float = 0.0, + layer_scale: float = 1e-6, + block: Optional[Callable[..., nn.Module]] = None, + norm_layer: Optional[Callable[..., nn.Module]] = None, + device: str = 'cuda', + use_skip: bool = False, + is_label: bool = False, + nr_of_labels: Optional[int] = None, + kernel_size: int = 7, + ) -> None: super().__init__() - self.include_top = include_top - self.downsample_layers = [] # stem and 3 intermediate downsampling conv layers - # stem = tf.keras.Sequential([ - # Conv2D(dims[0], kernel_size=4, strides=4, padding='same'), - # LayerNormalization(epsilon=1e-6)] - # ) - stem = [Conv2D(dims[0], kernel_size=4, strides=4, padding='same'), LayerNormalization(epsilon=1e-6)] - self.downsample_layers.append(stem) - for i in range(len(depths) - 1): - # downsample_layer = tf.keras.Sequential([ - # LayerNormalization(epsilon=1e-6), - # Conv2D(dims[i+1], kernel_size=2, strides=2, padding='same')] - # ) - downsample_layer = [LayerNormalization(epsilon=1e-6), Conv2D(dims[i+1], kernel_size=2, strides=2, padding='same')] - self.downsample_layers.append(downsample_layer) - self.stages = [] # 4 feature resolution stages, each consisting of multiple residual blocks - dp_rates = [x for x in np.linspace(0, drop_path_rate, sum(depths))] - cur = 0 - for i in range(len(depths)): - stage = tf.keras.Sequential( - [Block(dim=dims[i], drop_path=dp_rates[cur + j], kernel_size=kernel_size, - layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])] - ) - self.stages.append(stage) - cur += depths[i] + self.use_skip = use_skip - if self.include_top: - self.avg = GlobalAveragePooling2D() - self.norm = LayerNormalization(epsilon=1e-6) # final norm layer - self.head = Dense(num_classes) - else: - self.avg = None - self.norm = None - self.head = None + if not block_setting: + raise ValueError("The block_setting should not be empty") + elif not (isinstance(block_setting, Sequence) and all([isinstance(s, CNBlockConfig) for s in block_setting])): + raise TypeError("The block_setting should be List[CNBlockConfig]") - def forward_features(self, x, training): - for i in range(len(self.downsample_layers)): - # x = self.downsample_layers[i](x) - x = self.downsample_layers[i][0](x, training=training) - x = self.downsample_layers[i][1](x, training=training) - x = self.stages[i](x, training=training) + if block is None: + block = CNBlock + + if norm_layer is None: + norm_layer = partial(LayerNorm2d, eps=1e-6) + + # self.permute_cf = Permute([0, 3, 1, 2]) + # self.permute_cl = Permute([0, 2, 3, 1]) + + # Stem + # firstconv_output_channels = block_setting[0].input_channels + # self.first = nn.Sequential( + # nn.ConvTranspose2d( + # firstconv_output_channels, + # firstconv_output_channels, + # kernel_size=2, + # stride=2, + # padding=0, + # ), + # norm_layer(firstconv_output_channels), + # ).to(device) + self.first = nn.Sequential( + nn.ConvTranspose2d(block_setting[0].input_channels, block_setting[0].out_channels, kernel_size=2, stride=2), + norm_layer(block_setting[0].out_channels), + ).to(device) + + + total_stage_blocks = sum(cnf.num_layers for cnf in block_setting) + stage_block_id = 0 + self.stages: List[nn.Module] = [] + self.upsamples: List[nn.Module] = [] + for cnf in block_setting[1:]: + if cnf.input_channels is not None: + # Bottlenecks + stage: List[nn.Module] = [] + for _ in range(cnf.num_layers): + # adjust stochastic depth probability based on the depth of the stage block + sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0) + stage.append(block(cnf.input_channels, layer_scale, sd_prob, kernel_size)) + stage_block_id += 1 + self.stages.append(nn.Sequential(*stage).to(device)) + else: + self.stages.append(nn.Identity().to(device)) + + if cnf.out_channels is not None: + # Upsampling + self.upsamples.append( + nn.Sequential( + norm_layer(cnf.input_channels), + nn.ConvTranspose2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2), + ).to(device) + ) + + self.last = nn.Conv2d( + block_setting[-1].input_channels, + 3 if not is_label else nr_of_labels, + kernel_size=3, + padding='same', + ).to(device) + self.activation = nn.Tanh() if not is_label else nn.Softmax(dim=1) + + # self.last = nn.Sequential( + # Permute([0, 2, 3, 1]), + # nn.Linear(in_features=block_setting[-1].input_channels, out_features=1, bias=True), + # Permute([0, 3, 1, 2]), + # ).to(device) + + + for m in self.modules(): + if isinstance(m, (nn.Conv2d, nn.Linear)): + nn.init.trunc_normal_(m.weight, std=0.02) + if m.bias is not None: + nn.init.zeros_(m.bias) + + + def _forward_impl(self, x) -> Tensor: + # input, features = x + # x = self.permute_cf(x) + + x = self.first(x) + for i in range(len(self.stages)): + # if self.use_skip: + # x = self.stages[i](torch.concat((x, features[2-i]), dim=1)) + # else: + x = self.stages[i](x) + x = self.upsamples[i](x) + + # x = self.stages[-2](x) + # x = self.upsamples[-1](x) + + # x = self.stages[-1](x) + x = self.last(x) + x = self.activation(x) + # x = self.permute_cl(x) + return x - def call(self, x, training): - x = self.forward_features(x, training=training) + def forward(self, x) -> Tensor: + return self._forward_impl(x) + + +class ConvNeXtReverseBase(nn.Module): + def __init__( + self, + block_setting: List[CNBlockConfig], + stochastic_depth_prob: float = 0.0, + layer_scale: float = 1e-6, + block: Optional[Callable[..., nn.Module]] = None, + norm_layer: Optional[Callable[..., nn.Module]] = None + ) -> None: + super().__init__() + + if not block_setting: + raise ValueError("The block_setting should not be empty") + elif not (isinstance(block_setting, Sequence) and all([isinstance(s, CNBlockConfig) for s in block_setting])): + raise TypeError("The block_setting should be List[CNBlockConfig]") + + if block is None: + block = CNBlock + + if norm_layer is None: + norm_layer = partial(LayerNorm2d, eps=1e-6) + + layers: List[nn.Module] = [] + + + # add bottleneck which will be executed at the end, to process after the final larger upsampling + + + + # Stem + firstconv_output_channels = block_setting[0].input_channels + + layers.append( + nn.Sequential( + nn.Conv2d(16, 1, kernel_size=7, padding=3, bias=True), + + ) + ) + + stage: List[nn.Module] = [] + for _ in range(3): + stage.append(block(16, layer_scale, 0.0)) + layers.append(nn.Sequential(*stage)) + + layers.append( + # Convtr NormActivation( + # firstconv_output_channels, + # 3, + # kernel_size=4, + # stride=4, + # padding=0, + # norm_layer=None, + # activation_layer=None, + # bias=True, + # ) + nn.Sequential( + nn.ConvTranspose2d( + firstconv_output_channels, + 16, + kernel_size=4, + stride=4, + padding=0, + ), + ) + ) + + total_stage_blocks = sum(cnf.num_layers for cnf in block_setting) + stage_block_id = 0 + for cnf in block_setting: + # Bottlenecks + stage: List[nn.Module] = [] + for _ in range(cnf.num_layers): + # adjust stochastic depth probability based on the depth of the stage block + sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0) + stage.append(block(cnf.input_channels, layer_scale, sd_prob)) + stage_block_id += 1 + layers.append(nn.Sequential(*stage)) + if cnf.out_channels is not None: + # Upsampling + layers.append( + nn.Sequential( + norm_layer(cnf.out_channels), + nn.ConvTranspose2d(cnf.out_channels, cnf.input_channels, kernel_size=2, stride=2), + ) + ) + + self.features = nn.Sequential(*reversed(layers)) + + for m in self.modules(): + if isinstance(m, (nn.Conv2d, nn.Linear)): + nn.init.trunc_normal_(m.weight, std=0.02) + if m.bias is not None: + nn.init.zeros_(m.bias) + + def _forward_impl(self, x: Tensor) -> Tensor: + x = self.features(x) + # x = torch.tanh(x) + return x + + def forward(self, x: Tensor) -> Tensor: + return self._forward_impl(x) + + +def _convnext( + block_setting: List[CNBlockConfig], + stochastic_depth_prob: float, + reversed: bool, + is_label: bool, + nr_of_labels: Optional[int], + kernel_size: int, + device='cuda', + **kwargs: Any, +) -> ConvNeXtBase: + + if reversed: + model = ConvNeXtReverseBase2(block_setting, stochastic_depth_prob=stochastic_depth_prob, is_label=is_label, nr_of_labels=nr_of_labels, kernel_size=kernel_size, device=device, **kwargs) + else: + model = ConvNeXtBase(block_setting, stochastic_depth_prob=stochastic_depth_prob, kernel_size=kernel_size, device=device, **kwargs) + + + return model + + +class ConvNeXt(nn.Module): + def __init__(self, block_setting, stochastic_depth_prob=0.0, reversed=False, is_label=False, nr_of_labels=None, kernel_size=7, device='cuda', **kwargs) -> None: + super().__init__() + + self.model = _convnext(block_setting, stochastic_depth_prob, reversed, is_label, nr_of_labels, kernel_size, device=device, **kwargs) + + def forward(self, x: Tensor) -> Tensor: + x = self.model(x) return x -class ConvNeXtReverse(tf.keras.Model): - r""" ConvNeXt - A Tensorflow keras impl of : `A ConvNet for the 2020s` - - https://arxiv.org/pdf/2201.03545.pdf - Args: - num_classes (int): Number of classes for classification head. Default: 1000 - depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3] - dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768] - include_top (bool): whether to add head or just use it as feature extractor. Default: True - drop_path_rate (float): Stochastic depth rate. Default: 0. - layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6. - head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1. - """ - - def __init__(self, - depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], kernel_size=7, - drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1., - ): - super().__init__() - self.upsample_layers = [] # stem and 3 intermediate downsampling conv layers - stem = tf.keras.Sequential([ - Conv2DTranspose(dims[0], kernel_size=4, strides=4, padding='same'), - # UpSampling2D((4, 4), interpolation='nearest'), - # Conv2D(dims[0], kernel_size=4, strides=1, padding='same'), - # Dense(dims[0]), - LayerNormalization(epsilon=1e-6)] - ) - self.upsample_layers.append(stem) - for i in range(len(depths) - 1): - downsample_layer = tf.keras.Sequential([ - LayerNormalization(epsilon=1e-6), - Conv2DTranspose(dims[i+1], kernel_size=2, strides=2, padding='same') - # UpSampling2D((2,2), interpolation='nearest'), - # Conv2D(dims[i+1], kernel_size=2, strides=1, padding='same'), - # Dense(dims[i+1]), - ]) - self.upsample_layers.append(downsample_layer) - - self.stages = [] # 4 feature resolution stages, each consisting of multiple residual blocks - dp_rates = [x for x in np.linspace(0, drop_path_rate, sum(depths))] - cur = 0 - for i in range(len(depths)): - stage = tf.keras.Sequential( - [Block(dim=dims[i], drop_path=dp_rates[cur + j], kernel_size=kernel_size, - layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])] - ) - self.stages.append(stage) - cur += depths[i] - - self.last = Dense(3, activation='tanh') - - def call(self, x, training): - for i in reversed(range(len(self.upsample_layers))): - x = self.upsample_layers[i](x, training=training) - x = self.stages[i](x, training=training) - x = self.last(x, training=training) - return x - - -class ConvNeXtReverseLabels(tf.keras.Model): - r""" ConvNeXt - A Tensorflow keras impl of : `A ConvNet for the 2020s` - - https://arxiv.org/pdf/2201.03545.pdf - - Args: - num_classes (int): Number of classes for classification head. Default: 1000 - depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3] - dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768] - include_top (bool): whether to add head or just use it as feature extractor. Default: True - drop_path_rate (float): Stochastic depth rate. Default: 0. - layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6. - head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1. - """ - - def __init__(self, - depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], - drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1., - nr_of_classes=7): - super().__init__() - self.downsample_layers = [] # stem and 3 intermediate downsampling conv layers - stem = tf.keras.Sequential([ - Conv2DTranspose(dims[0], kernel_size=4, strides=4, padding='same'), - LayerNormalization(epsilon=1e-6)] - ) - self.downsample_layers.append(stem) - for i in range(len(depths) - 1): - downsample_layer = tf.keras.Sequential([ - LayerNormalization(epsilon=1e-6), - Conv2DTranspose(dims[i+1], kernel_size=2, strides=2, padding='same')] - ) - self.downsample_layers.append(downsample_layer) - - self.stages = [] # 4 feature resolution stages, each consisting of multiple residual blocks - dp_rates = [x for x in np.linspace(0, drop_path_rate, sum(depths))] - cur = 0 - for i in range(len(depths)): - stage = tf.keras.Sequential( - [Block(dim=dims[i], drop_path=dp_rates[cur + j], - layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])] - ) - self.stages.append(stage) - cur += depths[i] - - # self.last = Dense(nr_of_classes, activation='softmax') - self.last = Dense(nr_of_classes - 1) - - def call(self, x, training): - for i in reversed(range(len(self.downsample_layers))): - x = self.downsample_layers[i](x, training=training) - x = self.stages[i](x, training=training) - x = self.last(x, training=training) - return x - - - - - -model_urls = { - "convnext_tiny_224": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_tiny_1k_224_ema.h5", - "convnext_small_224": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_small_1k_224_ema.h5", - "convnext_base_224": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_base_22k_1k_224.h5", - "convnext_base_384": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_base_22k_1k_384.h5", - "convnext_large_224": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_large_22k_1k_224.h5", - "convnext_large_384": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_large_22k_1k_384.h5", - "convnext_xlarge_224": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_xlarge_22k_1k_224_ema.h5", - "convnext_xlarge_384": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_xlarge_22k_1k_384_ema.h5", -} - - -model_configs = dict( - convnext_tiny=dict( - depths=[3, 3, 9, 3], - dims=[8, 16, 32, 512] - ), - convnext_small=dict( - depths=[3, 3, 27, 3], - dims=[96, 192, 384, 768] - ), - convnext_base=dict( - depths=[3, 3, 27, 3], - dims=[128, 256, 512, 1024] - ), - convnext_large=dict( - depths=[3, 3, 27, 3], - dims=[192, 384, 768, 1536] - ), - convnext_xlarge=dict( - depths=[3, 3, 27, 3], - dims=[256, 512, 1024, 2048] - ), -) +if __name__ == '__main__': + model = ConvNeXt().to('cuda') + test = model(torch.zeros((4, 3, 512, 256)).to('cuda')) + pass \ No newline at end of file diff --git a/run.py b/run.py new file mode 100644 index 0000000..e8a040f --- /dev/null +++ b/run.py @@ -0,0 +1,9 @@ +import os +from configs.ae_static import config + + +if __name__ == '__main__': + + name = config['name'] + current_file = '/'.join(__file__.split('/')[-2:]) + os.system(f'nohup python scripts/ae_static_train.py --name lhunor_{name} > nohup/torch/{name}.out 2>&1 &') \ No newline at end of file diff --git a/scripts/ae_static_train.py b/scripts/ae_static_train.py new file mode 100644 index 0000000..705a2ba --- /dev/null +++ b/scripts/ae_static_train.py @@ -0,0 +1,197 @@ +import os +import sys +sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') + +from configs.ae_static import config +os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" +os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu'] + + +import time +import argparse +import wandb +import random +import numpy as np +import torch +from torch import nn +from torchmetrics import MultiScaleStructuralSimilarityIndexMeasure as MSSSIM +from custom_utils.custom_utils import UV2Mesh, get_summary_html, load_model, maskFacesBatch, get_summary_str +import plotly.graph_objects as go +import plotly.express as ex +import torch.optim as optim + +if config['is_skirt']: + from losses.losses import edge_loss_skirt as edge_loss +else: + from losses.losses import edge_loss + +from data.dataset import Dataset +from torch.utils.data import DataLoader +from custom_utils.custom_utils import batch_to_dict, plot_3d +from models.cond_encoder import AutoEncoder +from utils.checkpoints import CheckpointSaver +from custom_utils.custom_utils import Accumulator + + +def init_wandb(): + run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/data-local/127/temp/wandb', tags=config['tags']) + artifact = wandb.Artifact('code', type='code') + artifact.add_file('scripts/ae_static_train.py') + artifact.add_file('losses/losses.py') + artifact.add_file('models/base.py') + artifact.add_file('models/cond_encoder.py') + artifact.add_file('models/convnext.py') + artifact.add_file('data/dataset.py') + config_file_names = [key for key in sys.modules.keys() if key.startswith('configs.')] + for config_file_name in config_file_names: + file_name = '/'.join(config_file_name.split('.')) + '.py' + if os.path.exists(file_name): + artifact.add_file(file_name) + run.log_artifact(artifact) + +uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape']) + + + +def train_step(batch): + model.train() + pred = model([batch['uv_static'], batch['uv_mask']]) + + mesh_pred = uv2mesh(pred) + mesh_orig = batch['mesh_static'] + + e_loss = 1e-2 * torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask'])) + l1_loss = torch.mean(torch.sum((torch.abs(batch['uv_static'] - pred)), axis=1)) + ms_ssim = torch.mean(_MSSSIM(pred + 1, batch['uv_static'] + 1)) + loss = 1e3 * (1 - ms_ssim) + + optimizer.zero_grad() + loss.backward() + optimizer.step() + + metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['mesh_mask']) + + return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim, } + + +def valid_step(batch): + model.eval() + + with torch.inference_mode(): + pred = model([batch['uv_static'], batch['uv_mask']]) + + mesh_pred = uv2mesh(pred) + mesh_orig = batch['mesh_static'] + + e_loss = 1e-2 * torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask'])) + l1_loss = torch.mean(torch.sum((torch.abs(batch['uv_static'] - pred)), axis=1)) + ms_ssim = torch.mean(_MSSSIM(pred + 1, batch['uv_static'] + 1)) + loss = 1e3 * (1 - ms_ssim) + + metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / torch.sum(batch['mesh_mask']) + + return { 'valid/loss': loss, 'valid/edge_loss': e_loss, 'valid/l1_loss': l1_loss, 'valid/euclid_error': metric, 'valid/ms_ssim': ms_ssim, } + + +def plot_everything(batch: dict, type: str): + pred = model([batch['uv_static'], batch['uv_mask']]).detach() + + mesh_pred = uv2mesh(pred) + mesh_orig = batch['mesh_static'] + + orig_img = torch.permute(batch['uv_static'], [0,2,3,1]) + pred_img = torch.permute(pred, [0,2,3,1]) + img = torch.concat((orig_img, pred_img), axis=2) + imgs = [wandb.Image((img[i].cpu().numpy() + 1) / 2.0) for i in range(4)] + + F = np.load("smpl/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy") + mask = batch['mesh_mask'].cpu().numpy().astype(bool) + garment_static, faces = maskFacesBatch(mesh_pred.detach().cpu().numpy(), F, mask) + gt = [mesh_orig[i].cpu().numpy()[mask[i], :] for i in range(pred.shape[0])] + + body = gt + garment_posed = garment_static + collision = [np.concatenate((body[i], garment_posed[i]), axis=0) for i in range(len(garment_posed))] + collision_faces = [np.concatenate((faces[i], faces[i] + np.sum(mask[i])), axis=0) for i in range(len(garment_posed))] + collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(garment_posed[i], dtype=np.int32)*230), axis=0) for i in range(len(garment_posed))] + + figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_static[i], faces[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) } + + return { **figs, **{ type + '/images': imgs } } + + +def train(): + valid_iterator = iter(valid_dataloader) + + step = 0 + for batch in train_dataloader: + train_loss = train_step(batch) + train_accumulator.update(train_loss) + + if step % 25 == 0: + data_valid = next(valid_iterator) + valid_loss = valid_step(data_valid) + valid_accumulator.update(valid_loss) + + if step % 1000 == 0 or (step < 1000 and step % 100 == 0): + to_log = dict({ **train_loss, **valid_loss, **{ 'current_step': step } }) + if step % 5000 == 0 or (step < 5000 and step % 1000 == 0): + train_figs = plot_everything(batch, 'train') + valid_figs = plot_everything(data_valid, 'valid') + to_log = { **to_log, **train_figs, **valid_figs } + if step == 0: + to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } } + + wandb.log(to_log, step=step) + train_accumulator.reset() + valid_accumulator.reset() + print(step) + + if step % 10000 == 0: + ckpt_saver_enc(model.encoder, step, optimizer) + ckpt_saver_dec(model.decoder, step, optimizer) + + print(step) + + step += 1 + + +if __name__ == '__main__': + init_wandb() + # device = "cuda" if torch.cuda.is_available() else "cpu" + + parser = argparse.ArgumentParser(description='') + parser.add_argument('--name', type=str) + args = vars(parser.parse_args()) + + resolution_modifier = '' + if config['resolution'] != 512: + resolution_modifier = '_' + str(config['resolution']) + skirt_modifier = '_skirt' if config['is_skirt'] else '' + + # dataset = Dataset('/data-local/127/train', config['signature'], resolutions=[512], device=device, frames_file='frames_jumpsuit.txt') + dataset_train = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file=config['train_frames']) + dataset_valid = Dataset(config['valid_dir'], config['signature'], device=config['device'], frames_file=config['valid_frames']) + train_dataloader = DataLoader(dataset_train, batch_size=4, drop_last=True, shuffle=False) + valid_dataloader = DataLoader(dataset_valid, batch_size=4, drop_last=True, shuffle=False) + train_accumulator = Accumulator() + valid_accumulator = Accumulator() + + model : AutoEncoder = AutoEncoder(config['cond_ae']) + optimizer = optim.Adam(model.parameters(), lr=config['lr']) + + _MSSSIM = MSSSIM(data_range=2).to(config['device']) + + if config['pretrained_model'] != '': + load_model(model.encoder, config['pretrained_model'] + '_enc') + load_model(model.decoder, config['pretrained_model'] + '_dec') + + ckpt_saver_enc = CheckpointSaver('checkpoints/iccv/' + config['name'] + '_enc', top_n=20) + ckpt_saver_dec = CheckpointSaver('checkpoints/iccv/' + config['name'] + '_dec', top_n=20) + + run_name = config['name'] + print(f'started training: {run_name}') + train() + + + diff --git a/utils/convnext_utils.py b/utils/convnext_utils.py new file mode 100644 index 0000000..c3f629d --- /dev/null +++ b/utils/convnext_utils.py @@ -0,0 +1,21 @@ +from typing import Optional + +class CNBlockConfig: + # Stores information listed at Section 3 of the ConvNeXt paper + def __init__( + self, + input_channels: int, + out_channels: Optional[int], + num_layers: int, + ) -> None: + self.input_channels = input_channels + self.out_channels = out_channels + self.num_layers = num_layers + + def __repr__(self) -> str: + s = self.__class__.__name__ + "(" + s += "input_channels={input_channels}" + s += ", out_channels={out_channels}" + s += ", num_layers={num_layers}" + s += ")" + return s.format(**self.__dict__) \ No newline at end of file