diff --git a/.gitignore b/.gitignore
index f78ebcd..ee1e6da 100644
--- a/.gitignore
+++ b/.gitignore
@@ -132,4 +132,6 @@ dmypy.json
 .idea/
 
 # vscode
-.vscode/
\ No newline at end of file
+.vscode/
+
+wandb/
\ No newline at end of file
diff --git a/Dockerfile b/Dockerfile
new file mode 100644
index 0000000..717a469
--- /dev/null
+++ b/Dockerfile
@@ -0,0 +1,8 @@
+FROM tensorflow/tensorflow:latest-gpu
+
+RUN apt-get update
+RUN apt-get upgrade -y
+RUN apt-get install git ffmpeg libsm6 libxext6 libopenexr-dev zip unzip -y
+RUN pip install -r requierements.txt
+
+# RUN wandb login
\ No newline at end of file
diff --git a/convnext/LICENSE b/convnext/LICENSE
new file mode 100644
index 0000000..b93be90
--- /dev/null
+++ b/convnext/LICENSE
@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) Meta Platforms, Inc. and affiliates.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
diff --git a/convnext/README.md b/convnext/README.md
new file mode 100644
index 0000000..f0572ce
--- /dev/null
+++ b/convnext/README.md
@@ -0,0 +1,42 @@
+# ConvNeXt Tensorflow
+
+This is unofficial tensorflow keras implementation of ConvNeXt.  
+Its based on official [PyTorch implementation](https://github.com/facebookresearch/ConvNeXt).
+
+## Pre-trained Models
+
+| name | resolution | pretrain | acc@1 | #params | FLOPs | model |
+|:---:|:---:|:---:|:---:| :---:|:---:|:---:|
+| `convnext_tiny_224` | 224x224 |ImageNet-1K | 82.1 | 28M | 4.5G | [github](https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth) |
+| `convnext_small_224` | 224x224 |ImageNet-1K | 83.1 | 50M | 8.7G | [github](https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth) |
+| `convnext_base_224` | 224x224 |ImageNet-21K-1K | 85.8 | 89M | 15.4G | [github](https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth) |
+| `convnext_base_384` | 384x384 |ImageNet-21K-1K | 86.8 | 89M | 45.0G | [github](https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_384.pth) |
+| `convnext_large_224` | 224x224 |ImageNet-21K-1K | 86.6 | 198M | 34.4G | [github](https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth) |
+| `convnext_large_384` | 384x384 |ImageNet-21K-1K | 87.5 | 198M | 101.0G | [github](https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_384.pth) |
+| `convnext_xlarge_224` | 224x224 |ImageNet-21K-1K | 87.0 | 350M | 60.9G | [github](https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth) |
+| `convnext_xlarge_384` | 384x384 |ImageNet-21K-1K | 87.8 | 350M | 179.0G | [github](https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_384.pth) |
+
+## Note
+I've ported only ImageNet-21K-1K weights for base, large and xlarge models.  
+If you want to convert another pretrained weight in official repo, you can refer to [this script](https://github.com/bamps53/convnext-tf/blob/master/convert_weights.py) or just let me know.
+
+## Examples
+```python
+import tensorflow as tf
+from models.convnext_tf import create_model
+
+x = tf.zeros((1, 224, 224, 3), dtype=tf.float32)
+
+model = create_model('convnext_tiny_224', input_shape=(224, 224), pretrained=True)
+out = model(x) # (1, 1000)
+
+model = create_model('convnext_tiny_224', input_shape=(224, 224), num_classes=1, pretrained=True)
+out = model(x) # (1, 1)
+
+model = create_model('convnext_tiny_224', input_shape=(224, 224), include_top=False, pretrained=True)
+out = model(x) # (1, 16, 16, 768)
+```
+
+## Reference
+https://github.com/facebookresearch/ConvNeXt  
+https://github.com/rishigami/Swin-Transformer-TF  
\ No newline at end of file
diff --git a/convnext/convert_weights.py b/convnext/convert_weights.py
new file mode 100644
index 0000000..53028a0
--- /dev/null
+++ b/convnext/convert_weights.py
@@ -0,0 +1,118 @@
+import os
+
+import cv2
+import numpy as np
+import tensorflow as tf
+import timm
+import torch
+from models.convnext import *  # just to register model
+from models.convnext_tf import create_model
+
+model_urls = [
+    'https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth',
+    'https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth',
+    'https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_224.pth',
+    'https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_384.pth',
+    'https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_224.pth',
+    'https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_384.pth',
+    'https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_224_ema.pth',
+    'https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_384_ema.pth',
+]
+
+
+def load_test_image(image_size=(224, 224)):
+    from skimage.data import chelsea
+    img = chelsea()  # Chelsea the cat
+    img = cv2.resize(img, image_size)
+    img = img / 255
+    img = (img - np.array([0.485, 0.456, 0.406])) / \
+        np.array([0.229, 0.224, 0.225])
+    return img
+
+
+def load_from_url(model, url):
+    checkpoint = torch.hub.load_state_dict_from_url(
+        url=url, map_location="cpu", check_hash=True)
+    model.load_state_dict(checkpoint["model"])
+    return model
+
+
+def get_pt_state_dict(pt_model):
+    pt_state_dict = {}
+    for n, p in pt_model.named_parameters():
+        pt_state_dict[n] = p.detach().numpy()
+    return pt_state_dict
+
+
+def get_tf_state_dict(tf_model):
+    tf_state_dict = {}
+    for layer in tf_model.layers:
+        for var in layer.weights:
+            if ('dense' in var.name) and ('kernel' in var.name):
+                w = tf.transpose(var, [1, 0]).numpy()
+            elif ('depthwise_conv2d' in var.name) and ('kernel' in var.name):
+                w = tf.transpose(var, [2, 3, 0, 1]).numpy()
+            elif ('conv2d' in var.name) and ('kernel' in var.name):
+                w = tf.transpose(var, [3, 2, 0, 1]).numpy()
+            else:
+                w = var.numpy()
+            tf_state_dict[var.name] = w
+    return tf_state_dict
+
+
+def convert_weights(tf_model, pt_state_dict, key_map):
+    for layer in tf_model.layers:
+        for var in layer.weights:
+            pt_key = key_map[var.name]
+            pt_w = pt_state_dict[pt_key]
+            if ('dense' in var.name) and ('kernel' in var.name):
+                tf_w = np.transpose(pt_w, [1, 0])
+            elif ('depthwise_conv2d' in var.name) and ('kernel' in var.name):
+                tf_w = np.transpose(pt_w, [2, 3, 0, 1])
+            elif ('conv2d' in var.name) and ('kernel' in var.name):
+                tf_w = np.transpose(pt_w, [2, 3, 1, 0])
+            else:
+                tf_w = pt_w
+            var.assign(tf_w)
+
+
+def test_models(pt_model, tf_model):
+    img = load_test_image()
+    pt_x = torch.tensor(img[None, ]).permute(0, 3, 1, 2).to(torch.float32)
+    with torch.no_grad():
+        pt_y = torch.softmax(pt_model(pt_x), -1).numpy()
+
+    tf_x = tf.convert_to_tensor(img[None, ])
+    tf_y = tf.nn.softmax(tf_model(tf_x)).numpy()
+    np.testing.assert_allclose(pt_y, tf_y, rtol=1e-5)
+
+
+def main():
+    os.makedirs('weights', exist_ok=True)
+    for model_url in model_urls:
+        file_name = os.path.basename(model_url)
+        model_name = '_'.join(file_name.split('_')[:2])
+        save_name = file_name.replace("pth", "h5")
+        if os.path.exists(f'weights/{save_name}'):
+            continue
+
+        pt_model = timm.create_model(model_name)
+        pt_model = load_from_url(pt_model, model_url)
+        pt_model.eval()
+        pt_state_dict = get_pt_state_dict(pt_model)
+
+        tf_model = create_model(model_name)
+        tf_state_dict = get_tf_state_dict(tf_model)
+
+        key_map = {t: p for t, p in zip(
+            tf_state_dict.keys(), pt_state_dict.keys())}
+        convert_weights(tf_model, pt_state_dict, key_map)
+
+        test_models(pt_model, tf_model)
+        print(f'successfully converted {model_name}!')
+
+        tf_model.save_weights(f'weights/{save_name}')
+
+
+if __name__ == '__main__':
+    main()
diff --git a/convnext/models/convnext.py b/convnext/models/convnext.py
new file mode 100644
index 0000000..6de81df
--- /dev/null
+++ b/convnext/models/convnext.py
@@ -0,0 +1,201 @@
+# copied from https://github.com/facebookresearch/ConvNeXt/blob/main/models/convnext.py
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from timm.models.layers import trunc_normal_, DropPath
+from timm.models.registry import register_model
+
+class Block(nn.Module):
+    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, drop_path=0., layer_scale_init_value=1e-6):
+        super().__init__()
+        self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim) # depthwise conv
+        self.norm = LayerNorm(dim, eps=1e-6)
+        self.pwconv1 = nn.Linear(dim, 4 * dim) # pointwise/1x1 convs, implemented with linear layers
+        self.act = nn.GELU()
+        self.pwconv2 = nn.Linear(4 * dim, dim)
+        self.gamma = nn.Parameter(layer_scale_init_value * torch.ones((dim)), 
+                                    requires_grad=True) if layer_scale_init_value > 0 else None
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+    def forward(self, x):
+        input = x
+        x = self.dwconv(x)
+        x = x.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C)
+        x = self.norm(x)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.pwconv2(x)
+        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)
+        return x
+
+class ConvNeXt(nn.Module):
+    r""" ConvNeXt
+        A PyTorch impl of : `A ConvNet for the 2020s`  -
+          https://arxiv.org/pdf/2201.03545.pdf
+
+    Args:
+        in_chans (int): Number of input image channels. Default: 3
+        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]
+        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, in_chans=3, num_classes=1000, 
+                 depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], drop_path_rate=0., 
+                 layer_scale_init_value=1e-6, head_init_scale=1.,
+                 ):
+        super().__init__()
+
+        self.downsample_layers = nn.ModuleList() # stem and 3 intermediate downsampling conv layers
+        stem = nn.Sequential(
+            nn.Conv2d(in_chans, dims[0], kernel_size=4, stride=4),
+            LayerNorm(dims[0], eps=1e-6, data_format="channels_first")
+        )
+        self.downsample_layers.append(stem)
+        for i in range(3):
+            downsample_layer = nn.Sequential(
+                    LayerNorm(dims[i], eps=1e-6, data_format="channels_first"),
+                    nn.Conv2d(dims[i], dims[i+1], kernel_size=2, stride=2),
+            )
+            self.downsample_layers.append(downsample_layer)
+
+        self.stages = nn.ModuleList() # 4 feature resolution stages, each consisting of multiple residual blocks
+        dp_rates=[x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] 
+        cur = 0
+        for i in range(4):
+            stage = nn.Sequential(
+                *[Block(dim=dims[i], drop_path=dp_rates[cur + j], 
+                layer_scale_init_value=layer_scale_init_value) for j in range(depths[i])]
+            )
+            self.stages.append(stage)
+            cur += depths[i]
+
+        self.norm = nn.LayerNorm(dims[-1], eps=1e-6) # final norm layer
+        self.head = nn.Linear(dims[-1], num_classes)
+
+        self.apply(self._init_weights)
+        self.head.weight.data.mul_(head_init_scale)
+        self.head.bias.data.mul_(head_init_scale)
+
+    def _init_weights(self, m):
+        if isinstance(m, (nn.Conv2d, nn.Linear)):
+            trunc_normal_(m.weight, std=.02)
+            nn.init.constant_(m.bias, 0)
+
+    def forward_features(self, x):
+        for i in range(4):
+            x = self.downsample_layers[i](x)
+            x = self.stages[i](x)
+        return self.norm(x.mean([-2, -1])) # global average pooling, (N, C, H, W) -> (N, C)
+
+    def forward(self, x):
+        x = self.forward_features(x)
+        x = self.head(x)
+        return x
+
+class LayerNorm(nn.Module):
+    r""" LayerNorm that supports two data formats: channels_last (default) or channels_first. 
+    The ordering of the dimensions in the inputs. channels_last corresponds to inputs with 
+    shape (batch_size, height, width, channels) while channels_first corresponds to inputs 
+    with shape (batch_size, channels, height, width).
+    """
+    def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(normalized_shape))
+        self.bias = nn.Parameter(torch.zeros(normalized_shape))
+        self.eps = eps
+        self.data_format = data_format
+        if self.data_format not in ["channels_last", "channels_first"]:
+            raise NotImplementedError 
+        self.normalized_shape = (normalized_shape, )
+    
+    def forward(self, x):
+        if self.data_format == "channels_last":
+            return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+        elif self.data_format == "channels_first":
+            u = x.mean(1, keepdim=True)
+            s = (x - u).pow(2).mean(1, keepdim=True)
+            x = (x - u) / torch.sqrt(s + self.eps)
+            x = self.weight[:, None, None] * x + self.bias[:, None, None]
+            return x
+
+
+model_urls = {
+    "convnext_tiny_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth",
+    "convnext_small_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth",
+    "convnext_base_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth",
+    "convnext_large_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth",
+    "convnext_base_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth",
+    "convnext_large_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth",
+    "convnext_xlarge_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth",
+}
+
+@register_model
+def convnext_tiny(pretrained=False, **kwargs):
+    model = ConvNeXt(depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], **kwargs)
+    if pretrained:
+        url = model_urls['convnext_tiny_1k']
+        checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu", check_hash=True)
+        model.load_state_dict(checkpoint["model"])
+    return model
+
+@register_model
+def convnext_small(pretrained=False, **kwargs):
+    model = ConvNeXt(depths=[3, 3, 27, 3], dims=[96, 192, 384, 768], **kwargs)
+    if pretrained:
+        url = model_urls['convnext_small_1k']
+        checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu", check_hash=True)
+        model.load_state_dict(checkpoint["model"])
+    return model
+
+@register_model
+def convnext_base(pretrained=False, in_22k=False, **kwargs):
+    model = ConvNeXt(depths=[3, 3, 27, 3], dims=[128, 256, 512, 1024], **kwargs)
+    if pretrained:
+        url = model_urls['convnext_base_22k'] if in_22k else model_urls['convnext_base_1k']
+        checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu", check_hash=True)
+        model.load_state_dict(checkpoint["model"])
+    return model
+
+@register_model
+def convnext_large(pretrained=False, in_22k=False, **kwargs):
+    model = ConvNeXt(depths=[3, 3, 27, 3], dims=[192, 384, 768, 1536], **kwargs)
+    if pretrained:
+        url = model_urls['convnext_large_22k'] if in_22k else model_urls['convnext_large_1k']
+        checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu", check_hash=True)
+        model.load_state_dict(checkpoint["model"])
+    return model
+
+@register_model
+def convnext_xlarge(pretrained=False, in_22k=False, **kwargs):
+    model = ConvNeXt(depths=[3, 3, 27, 3], dims=[256, 512, 1024, 2048], **kwargs)
+    if pretrained:
+        url = model_urls['convnext_xlarge_22k'] if in_22k else model_urls['convnext_xlarge_1k']
+        checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu", check_hash=True)
+        model.load_state_dict(checkpoint["model"])
+    return model
diff --git a/convnext/models/convnext_tf.py b/convnext/models/convnext_tf.py
new file mode 100644
index 0000000..a2ee7c7
--- /dev/null
+++ b/convnext/models/convnext_tf.py
@@ -0,0 +1,207 @@
+import numpy as np
+import tensorflow as tf
+from tensorflow.keras.layers import (Conv2D, Dense, DepthwiseConv2D,
+                                     GlobalAveragePooling2D, Layer,
+                                     LayerNormalization)
+
+
+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
+
+    # Compute keep_prob
+    keep_prob = 1.0 - drop_prob
+
+    # 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
+
+
+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):
+        super().__init__()
+        self.drop_prob = drop_prob
+
+    def call(self, x, training=None):
+        return drop_path(x, self.drop_prob, training)
+
+
+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, drop_path=0., layer_scale_init_value=1e-6, prefix=''):
+        super().__init__()
+        self.dwconv = DepthwiseConv2D(
+            kernel_size=7, 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
+
+    def call(self, x):
+        input = x
+        x = self.dwconv(x)
+        # x = x.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C)
+        x = self.norm(x)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.pwconv2(x)
+        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)
+        return 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,
+                 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.,
+                 ):
+        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)]
+        )
+        self.downsample_layers.append(stem)
+        for i in range(3):
+            downsample_layer = tf.keras.Sequential([
+                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(4):
+            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]
+
+        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
+
+    def forward_features(self, x):
+        for i in range(4):
+            x = self.downsample_layers[i](x)
+            x = self.stages[i](x)
+        return x
+
+    def call(self, x):
+        x = self.forward_features(x)
+        if self.include_top:
+            x = self.avg(x)
+            x = self.norm(x)
+            x = self.head(x)
+        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]
+    ),
+)
+
+
+def create_model(model_name='convnext_tiny_1k', input_shape=(224, 224), num_classes=1000, include_top=True, pretrained=True, use_tpu=False, **kwargs):
+    cfg = model_configs['_'.join(model_name.split('_')[:2])]
+    net = ConvNeXt(num_classes, cfg['depths'],
+                   cfg['dims'], include_top, **kwargs)
+    net(tf.keras.Input(shape=(*input_shape, 3)))
+    if pretrained is True:
+        url = model_urls[model_name]
+        pretrained_ckpt = tf.keras.utils.get_file(
+            f'{model_name}.h5', url, untar=False)
+        if use_tpu:
+            load_locally = tf.saved_model.LoadOptions(
+                experimental_io_device='/job:localhost')
+            net.load_weights(
+                pretrained_ckpt, options=load_locally, skip_mismatch=True)
+        else:
+            net.load_weights(pretrained_ckpt, skip_mismatch=True, by_name=True)
+
+    return net
diff --git a/custom_utils/custom_utils.py b/custom_utils/custom_utils.py
new file mode 100644
index 0000000..b611790
--- /dev/null
+++ b/custom_utils/custom_utils.py
@@ -0,0 +1,65 @@
+import numpy as np
+import tensorflow as tf
+
+
+def m2uv_array(m2uv):
+    result = np.zeros((14475, 2), dtype=np.int32)
+    for k, v in m2uv.items():
+        if len(v) == 2:
+            result[k] = list(v)
+        elif len(v) == 1:
+            x = v.pop()
+            result[k] = [x, x]
+        else:
+            assert False
+
+    return result
+
+
+class UV2Mesh():
+    def __init__(self, res=256):
+        super(UV2Mesh, self).__init__()
+        self.res = res
+        uv_map = np.load('custom_utils/stretched_uv_map_3.npy')
+        self.uv_map_discreet = np.round(uv_map * self.res).astype(np.int32)
+        m2uv = np.load('custom_utils/m2uv.npy', allow_pickle=True).item()
+        m2uv = m2uv_array(m2uv)
+        self.uv_pixels = self.uv_map_discreet[m2uv]
+        # self.batch_size = batch_size
+
+        # self.x1 = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [self.batch_size, 1, 1]), tf.zeros((self.batch_size, 14475, 1), dtype=tf.int32)), axis=-1)
+        # self.y1 = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [self.batch_size, 1, 1]), tf.ones((self.batch_size, 14475, 1), dtype=tf.int32)), axis=-1)
+        # self.z1 = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [self.batch_size, 1, 1]), tf.ones((self.batch_size, 14475, 1), dtype=tf.int32) * 2), axis=-1)
+        # self.x2 = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [self.batch_size, 1, 1]), tf.zeros((self.batch_size, 14475, 1), dtype=tf.int32)), axis=-1)
+        # self.y2 = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [self.batch_size, 1, 1]), tf.ones((self.batch_size, 14475, 1), dtype=tf.int32)), axis=-1)
+        # self.z2 = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [self.batch_size, 1, 1]), tf.ones((self.batch_size, 14475, 1), dtype=tf.int32) * 2), axis=-1)
+
+    # input = channels_first
+    def uv_to_mesh(self, uv):
+        uv = tf.transpose(uv, [0,2,3,1])
+        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, 14475, 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, 14475, 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, 14475, 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, 14475, 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, 14475, 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, 14475, 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)
+
+        p1 = tf.stack((x1, y1, z1), axis=-1)
+        p2 = tf.stack((x2, y2, z2), axis=-1)
+
+        mesh = (p1 + p2) / 2
+        return mesh
+
+    def __call__(self, uv):
+        return self.uv_to_mesh(uv)
+        
\ No newline at end of file
diff --git a/custom_utils/m2uv.npy b/custom_utils/m2uv.npy
new file mode 100644
index 0000000..56708d7
Binary files /dev/null and b/custom_utils/m2uv.npy differ
diff --git a/custom_utils/smpl_faces.npy b/custom_utils/smpl_faces.npy
new file mode 100644
index 0000000..d1d74c9
Binary files /dev/null and b/custom_utils/smpl_faces.npy differ
diff --git a/custom_utils/ssmpl0.mat b/custom_utils/ssmpl0.mat
new file mode 100644
index 0000000..e3cea72
Binary files /dev/null and b/custom_utils/ssmpl0.mat differ
diff --git a/custom_utils/stretched_uv_map_3.npy b/custom_utils/stretched_uv_map_3.npy
new file mode 100644
index 0000000..5bd87e3
Binary files /dev/null and b/custom_utils/stretched_uv_map_3.npy differ
diff --git a/dataset.py b/dataset.py
index 7277d28..d3e6c9f 100644
--- a/dataset.py
+++ b/dataset.py
@@ -2,6 +2,7 @@ import os
 import sys
 import time
 import cv2
+import random
 import numpy as np
 import tensorflow as tf
 
@@ -10,32 +11,74 @@ class DatasetSimple(tf.data.Dataset):
     def _generator(folder):
         # Opening the file
         folder = folder.decode("utf-8")
-        uv_folder = os.path.join(folder, 'npz') 
+        uv_folder = os.path.join(folder, 'uv_posed') 
         pose_hape_folder = os.path.join(folder, 'poses_shapes') 
         frames = os.listdir(uv_folder)
+        random.shuffle(frames)
 
         for frame in frames:
-            uv = np.load(os.path.join(uv_folder, frame))
-            uv = uv['uv_gt']
-            # uv = cv2.resize(uv, [64, 128])
-            uv = np.transpose(uv, [2, 0, 1])
+            id, frame_nr = frame.split('.')[0].split('_')
+            uv_posed = np.load(os.path.join(folder, 'uv_posed', id + '_' + frame_nr + '.npz')); uv_posed = uv_posed['uv_gt']; uv_posed = np.transpose(uv_posed, [2, 0, 1])
+            pose_shape = np.load(os.path.join(folder, 'pose_shape',  id + '_' + frame_nr + '.npy'))
+            uv_static = np.load(os.path.join(folder, 'uv_static', id +  '.npz')); uv_static = uv_static['uv_static']; uv_static = np.transpose(uv_static, [2, 0, 1])
+            mesh_static = np.load(os.path.join(folder, 'mesh_static', id +  '.npz')); mesh_static = mesh_static['mesh_static']
+            mesh_body_posed = np.load(os.path.join(folder, 'mesh_body_posed', id + '_' + frame_nr + '.npy'))
+            mask = np.load(os.path.join(folder, 'mask', id + '.npy')); mask = mask.astype(np.float32)
 
-            pose_shape = np.load(os.path.join(pose_hape_folder, frame.split('.')[0] + '.npy'))
+            # uv_posed, (pose, shape, uv_static), mesh_static, mesh_body_posed, mask
+            yield uv_posed, (pose_shape[:72], pose_shape[72:], uv_static), mesh_static, mesh_body_posed, mask
 
-            yield uv, pose_shape
-
-    def __new__(cls, folder, res=64):
+    def __new__(cls, folder, res=256):
         return tf.data.Dataset.from_generator(
             cls._generator,
-            output_signature = (tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), tf.TensorSpec(shape = (83), dtype = tf.float32)),
+            output_signature = (tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), (tf.TensorSpec(shape = (72), dtype = tf.float32), tf.TensorSpec(shape = (11), dtype = tf.float32), tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32) ), tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32), tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32), tf.TensorSpec(shape = (14475,), dtype = tf.float32) ),
             args=(folder, )
         )
 
+os.environ["CUDA_VISIBLE_DEVICES"] = "1"
+
+
 if __name__ == '__main__':
     
-    dataset = DatasetSimple('/workspace/Dataset/stylegan/npz_1')
+    dataset = DatasetSimple('/home/workspace/data/dataset/stylegan_new/train').batch(2)
+
+    from losses import collision_loss as collision_penalty
+    from losses import collision_penalty as collision_penalty_2
+    from custom_utils.custom_utils import UV2Mesh
+    import matplotlib.pyplot as plt
+
+    def plot_3d(x, y=None):
+        fig = plt.figure(figsize=(15, 15))
+        ax = fig.add_subplot(111, projection='3d')
+        ax.set_xlim3d(-1, 1)
+        ax.set_ylim3d(-1, 1)
+        ax.set_zlim3d(-1, 1)
+        ax.set_xlabel("X")
+        ax.set_ylabel("Y")
+        ax.set_zlabel("Z")
+        ax.grid(False)
+        ax.scatter(x[:, 0], x[:, 1], x[:, 2], s=5)
+        if y is not None:
+            ax.scatter(y[:, 0], y[:, 1], y[:, 2], s=5)
+        plt.savefig('temp.jpg')
+
+    uv2mesh = UV2Mesh()
 
     for d in dataset:
         d
+
+        v = uv2mesh(d[0])
+        b = d[3]
+        m = d[4]
+
+        b2 = tf.concat((b[1][tf.newaxis,:,:], b[0][tf.newaxis,:,:]), axis=0)
+
+        c = collision_penalty(v, b, m)
+        c2 = collision_penalty(v, b2, m)
+
+        c_ = collision_penalty_2(v, b, m)
+        c2_ = collision_penalty_2(v, b2, m)
+
+        print()
         pass
     
\ No newline at end of file
diff --git a/inference_custom.py b/inference_custom.py
index 9515f39..9e48c15 100644
--- a/inference_custom.py
+++ b/inference_custom.py
@@ -44,8 +44,8 @@ def inference(ckpt_dir, use_custom_cuda, res, folder=None, samples=1000):
     featuremaps = [512, 512, 512, 512, 512, 256, 128, 64, 32]
     filter_index = resolutions.index(res)
     g_params = {
-        'z_dim': 512,
-        'w_dim': 512,
+        'z_dim': 256,
+        'w_dim': 256,
         'labels_dim': 0,
         'n_mapping': 8,
         'resolutions': resolutions[:filter_index + 1],
diff --git a/load_models.py b/load_models.py
index c10cb34..6537c29 100644
--- a/load_models.py
+++ b/load_models.py
@@ -18,11 +18,11 @@ def load_generator(g_params=None, is_g_clone=False, ckpt_dir=None, custom_cuda=T
         }
 
     test_latent = tf.ones((1, g_params['z_dim']), dtype=tf.float32)
-    test_labels = tf.ones((1, g_params['labels_dim']), dtype=tf.float32)
+    test_conditions = (tf.ones((1, 72), dtype=tf.float32), tf.ones((1, 11), dtype=tf.float32), tf.ones((1, 3, 512, 256), dtype=tf.float32))
 
     # build generator model
     generator = Generator(g_params)
-    _ = generator([test_latent, test_labels])
+    _ = generator([test_latent, test_conditions])
 
     if ckpt_dir is not None:
         if is_g_clone:
@@ -51,11 +51,10 @@ def load_discriminator(d_params=None, ckpt_dir=None, custom_cuda=True):
 
     res = d_params['resolutions'][-1]
     test_images = tf.ones((1, 3, res, int(res/2)), dtype=tf.float32)
-    test_labels = tf.ones((1, d_params['labels_dim']), dtype=tf.float32)
 
     # build discriminator model
     discriminator = Discriminator(d_params)
-    _ = discriminator([test_images, test_labels])
+    _ = discriminator(test_images)
 
     if ckpt_dir is not None:
         ckpt = tf.train.Checkpoint(discriminator=discriminator)
diff --git a/losses.py b/losses.py
index 32d2319..017f8f8 100644
--- a/losses.py
+++ b/losses.py
@@ -1,34 +1,37 @@
+import numpy as np
+import tensorflow_graphics as tfg
 import tensorflow as tf
+import scipy as sp
+import scipy.io as sio
 
+from scipy.spatial import cKDTree
+from tensorflow_graphics.geometry.representation.mesh.normals import vertex_normals
 
-def d_logistic(real_images, generator, discriminator, z_dim, labels):
+def d_logistic(real_images, generator, discriminator, z_dim, conditions):
     batch_size = tf.shape(real_images)[0]
     z = tf.random.normal(shape=[batch_size, z_dim], dtype=tf.float32)
 
     # forward pass
-    fake_images = generator([z, labels], training=True)
-    real_scores, real_cond = discriminator([real_images, labels], training=True)
-    fake_scores, fake_cond = discriminator([fake_images, labels], training=True)
+    fake_images = generator([z, conditions], training=True)
+    real_scores = discriminator(real_images, training=True)
+    fake_scores = discriminator(fake_images, training=True)
 
     # gan loss
     d_loss = tf.math.softplus(fake_scores)
     d_loss += tf.math.softplus(-real_scores)
 
-    cond_loss = tf.nn.l2_loss(labels - real_cond) + tf.nn.l2_loss(labels - fake_cond)
 
-    return d_loss, cond_loss
+    return fake_images, d_loss
 
 
-def d_logistic_r1_reg(real_images, generator, discriminator, z_dim, labels):
+def d_logistic_r1_reg(real_images, generator, discriminator, z_dim, conditions):
     batch_size = tf.shape(real_images)[0]
     z = tf.random.normal(shape=[batch_size, z_dim], dtype=tf.float32)
 
     # forward pass
-    fake_images = generator([z, labels], training=True)
-    real_scores, real_cond = discriminator([real_images, labels], training=True)
-    fake_scores, fake_cond = discriminator([fake_images, labels], training=True)
-
-    cond_loss = tf.nn.l2_loss(labels - real_cond) + tf.nn.l2_loss(labels - fake_cond)
+    fake_images = generator([z, conditions], training=True)
+    real_scores = discriminator(real_images, training=True)
+    fake_scores = discriminator(fake_images, training=True)
 
     # gan loss
     d_loss = tf.math.softplus(fake_scores)
@@ -36,50 +39,48 @@ def d_logistic_r1_reg(real_images, generator, discriminator, z_dim, labels):
 
     # gradient penalty
     with tf.GradientTape() as r1_tape:
-        r1_tape.watch([real_images, labels])
-        real_scores, _ = discriminator([real_images, labels], training=True)
+        r1_tape.watch([real_images, conditions])
+        real_scores = discriminator(real_images, training=True)
         real_loss = tf.reduce_sum(real_scores)
 
     real_grads = r1_tape.gradient(real_loss, real_images)
     r1_penalty = tf.reduce_sum(tf.math.square(real_grads), axis=[1, 2, 3])
     r1_penalty = tf.expand_dims(r1_penalty, axis=1)
-    return d_loss, cond_loss, r1_penalty
+    return fake_images, d_loss, r1_penalty
 
 
-def g_logistic_non_saturating(real_images, generator, discriminator, z_dim, labels):
+def g_logistic_non_saturating(real_images, generator, discriminator, z_dim, conditions):
     batch_size = tf.shape(real_images)[0]
     z = tf.random.normal(shape=[batch_size, z_dim], dtype=tf.float32)
 
     # forward pass
-    fake_images = generator([z, labels], training=True)
-    fake_scores, _ = discriminator([fake_images, labels], training=True)
-    # TODO: do something with cond output?
+    fake_images = generator([z, conditions], training=True)
+    fake_scores = discriminator(fake_images, training=True)
 
     # gan loss
     g_loss = tf.math.softplus(-fake_scores)
-    return g_loss
+    return fake_images, g_loss
 
 
 def g_logistic_ns_pathreg(real_images, generator, discriminator, z_dim,
                           pl_mean, pl_minibatch_shrink, pl_denorm, pl_decay,
-                          labels):
+                          conditions):
     batch_size = tf.shape(real_images)[0]
     z = tf.random.normal(shape=[batch_size, z_dim], dtype=tf.float32)
 
     pl_minibatch = tf.maximum(1, tf.math.floordiv(batch_size, pl_minibatch_shrink))
     pl_z = tf.random.normal(shape=[pl_minibatch, z_dim], dtype=tf.float32)
-    pl_labels = labels[:pl_minibatch]
+    pl_conditions = (conditions[0][:pl_minibatch], conditions[1][:pl_minibatch], conditions[2][:pl_minibatch])
 
     # forward pass
-    fake_images = generator([z, labels], ret_w_broadcasted=False, training=True)
-    fake_scores, _ = discriminator([fake_images, labels], training=True)
-    # TODO: do something with cond output?
+    fake_images = generator([z, conditions], ret_w_broadcasted=False, training=True)
+    fake_scores = discriminator(fake_images, training=True)
     g_loss = tf.math.softplus(-fake_scores)
 
     # Evaluate the regularization term using a smaller minibatch to conserve memory.
     with tf.GradientTape() as pl_tape:
-        pl_tape.watch([pl_z, pl_labels])
-        pl_fake_images, pl_w_broadcasted = generator([pl_z, pl_labels], ret_w_broadcasted=True, training=True)
+        pl_tape.watch([pl_z, pl_conditions])
+        pl_fake_images, pl_w_broadcasted = generator([pl_z, pl_conditions], ret_w_broadcasted=True, training=True)
 
         pl_noise = tf.random.normal(tf.shape(pl_fake_images)) * pl_denorm
         pl_noise_applied = tf.reduce_sum(pl_fake_images * pl_noise)
@@ -93,4 +94,104 @@ def g_logistic_ns_pathreg(real_images, generator, discriminator, z_dim,
 
     # Calculate (|J*y|-a)^2.
     pl_penalty = tf.square(pl_lengths - pl_mean)
-    return g_loss, pl_penalty
+    return fake_images, g_loss, pl_penalty
+
+
+# from hbertiche/PBNS
+""" COLLISION FUNCTIONS START """
+F = np.load("custom_utils/smpl_faces.npy")
+
+
+def _nearest_neighbour(V, B):
+	tree = cKDTree(B.numpy())
+	return tree.query(V.numpy(), n_jobs=-1)[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):
+    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.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.reduce_sum(tf.minimum(dot * mask - thr, 0) ** 2)
+
+    return loss
+
+""" COLLISION FUNCTIONS END """
+
+
+laplacian_loss = sio.loadmat("custom_utils/ssmpl0.mat")['L'].tocoo()
+edges = np.stack(np.nonzero(laplacian_loss.todense())).transpose()
+
+def get_edge_mask(masks):
+    return tf.reduce_min(tf.gather(masks, edges, axis=1), axis=-1)
+
+def edge_loss(x, y, mask):
+    e_mask = get_edge_mask(mask)
+    x_e = tf.gather(x, edges[:,0], axis=1) - tf.gather(x, edges[:,1], axis=1)
+    x_e = tf.reduce_sum(x_e ** 2, -1)
+    x_e = tf.sqrt(x_e + 1e-9)
+    y_e = tf.gather(y, edges[:,0], axis=1) - tf.gather(y, edges[:,1], axis=1)
+    y_e = tf.reduce_sum(y_e ** 2, -1)
+    y_e = tf.sqrt(y_e + 1e-9)
+    d_e = y_e - x_e
+    d_e = d_e * e_mask
+    loss = tf.reduce_sum(d_e ** 2)
+    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
+
+
+if __name__ == '__main__':
+    V = tf.ones((4, 14475, 3))
+    B = tf.zeros((4, 14475, 3))
+    B_vn = vertex_normals(B, tf.tile(F[None], [B.shape[0], 1, 1]))
+    loss_1 = collision_loss(V, B, thr=.004, stop_gradient=False)
+    loss_2 = collision_penalty(V, B, B_vn)
+    print()
diff --git a/requierements.txt b/requierements.txt
new file mode 100644
index 0000000..51dd24e
--- /dev/null
+++ b/requierements.txt
@@ -0,0 +1,7 @@
+numpy
+matplotlib
+opencv-python
+wandb
+OpenEXR
+tensorflow-graphics
+scipy
\ No newline at end of file
diff --git a/stylegan2/discriminator.py b/stylegan2/discriminator.py
index 5719c06..4544d2c 100644
--- a/stylegan2/discriminator.py
+++ b/stylegan2/discriminator.py
@@ -130,14 +130,9 @@ class Discriminator(tf.keras.Model):
         self.last_dense = Dense(max(self.labels_dim, 1), gain=1.0, lrmul=1.0, name='last_dense')
         self.last_bias = BiasAct(lrmul=1.0, act='linear', name='last_bias')
 
-        # conditional output
-        self.last_block_condition = DiscriminatorLastBlock(n_f0, n_f1, res, name='{:d}x{:d}'.format(res, res))
-        self.condition_dense = Dense(self.labels_dim, gain=1.0, lrmul=1.0, name='last_dense')
-        self.condition_bias = BiasAct(lrmul=1.0, act='linear', name='last_bias')
-
 
     def call(self, inputs, training=None, mask=None):
-        images, _ = inputs
+        images = inputs
 
         x = self.initial_fromrgb(images)
         for block in self.blocks:
@@ -150,13 +145,8 @@ class Discriminator(tf.keras.Model):
         x = self.last_dense(x)
         x = self.last_bias(x)
 
-        # condition output
-        cond = self.last_block_condition(features)
-        cond = self.condition_dense(cond)
-        cond = self.condition_bias(cond)
-
         scores_out = x
-        return scores_out, cond
+        return scores_out
 
     def compute_output_shape(self, input_shape):
         return input_shape[0][0], max(self.labels_dim, 1)
diff --git a/stylegan2/generator.py b/stylegan2/generator.py
index 1ccba03..bd28579 100644
--- a/stylegan2/generator.py
+++ b/stylegan2/generator.py
@@ -53,10 +53,10 @@ class Generator(tf.keras.Model):
         self.w_avg.assign(lerp(batch_avg, self.w_avg, self.w_ema_decay))
         return
 
-    def style_mixing_regularization(self, latents1, labels, w_broadcasted1):
+    def style_mixing_regularization(self, latents1, conditions, w_broadcasted1):
         # get another w and broadcast it
         latents2 = tf.random.normal(shape=tf.shape(latents1), dtype=tf.dtypes.float32)
-        dlatents2 = self.g_mapping([latents2, labels])
+        dlatents2 = self.g_mapping([latents2, conditions])
         w_broadcasted2 = self.broadcast(dlatents2)
 
         # find mixing limit index
@@ -89,14 +89,14 @@ class Generator(tf.keras.Model):
         return truncated_w_broadcasted
 
     def call(self, inputs, ret_w_broadcasted=False, truncation_psi=1.0, truncation_cutoff=None, training=None, mask=None):
-        latents, labels = inputs
+        latents, conditions = inputs
 
-        dlatents = self.g_mapping([latents, labels])
+        dlatents = self.g_mapping([latents, conditions])
         w_broadcasted = self.broadcast(dlatents)
 
         if training:
             self.update_moving_average_of_w(w_broadcasted)
-            w_broadcasted = self.style_mixing_regularization(latents, labels, w_broadcasted)
+            w_broadcasted = self.style_mixing_regularization(latents, conditions, w_broadcasted)
 
         if not training:
             w_broadcasted = self.truncation_trick(w_broadcasted, truncation_psi, truncation_cutoff)
diff --git a/stylegan2/layers/embedding.py b/stylegan2/layers/embedding.py
index c0a9fb1..90e29c8 100644
--- a/stylegan2/layers/embedding.py
+++ b/stylegan2/layers/embedding.py
@@ -1,24 +1,135 @@
 import tensorflow as tf
 
+from convnext.models.convnext_tf import ConvNeXt
 
-class LabelEmbedding(tf.keras.layers.Layer):
+def rodrigues(r):
+    theta = tf.norm(r, axis=(1, 2), keepdims=True)
+    theta = tf.maximum(theta, 1e-7)
+    r_hat = r / theta
+    z_stick = tf.zeros(theta.get_shape().as_list()[0], dtype=tf.float32)
+    m = tf.stack(
+        (z_stick, -r_hat[:, 0, 2], r_hat[:, 0, 1], r_hat[:, 0, 2], z_stick,
+        -r_hat[:, 0, 0], -r_hat[:, 0, 1], r_hat[:, 0, 0], z_stick), axis=1)
+    m = tf.reshape(m, (-1, 3, 3))
+    i_cube = tf.expand_dims(tf.eye(3, dtype=tf.float32), axis=0) + tf.zeros(
+        (theta.get_shape().as_list()[0], 3, 3), dtype=tf.float32)
+    # R = cos * i_cube + (1 - cos) * dot + tf.sin(theta) * m
+    R = i_cube + m * tf.sin(theta) + tf.matmul(m, m) * (1 - tf.cos(theta))
+    return R
+
+
+def run_rodrigues(pose):
+    batch_size = pose.shape[0]
+    R_out = rodrigues(tf.reshape(pose, (-1, 1, 3)))
+    R_out = tf.reshape(R_out, (batch_size, -1, 3, 3))
+    return R_out
+
+
+# if __name__ == '__main__':
+#     x = tf.ones((4, 72))
+#     y = run_rodrigues(x)
+#     print()
+#     pass
+
+
+class ParameterEmbedding(tf.keras.layers.Layer):
     def __init__(self, embed_dim, **kwargs):
-        super(LabelEmbedding, self).__init__(**kwargs)
+        super(ParameterEmbedding, self).__init__(**kwargs)
         self.embed_dim = embed_dim
 
-    def build(self, input_shape):
-        weight_shape = [input_shape[1], self.embed_dim]
-        # tf 1.15 mean(0.0), std(1.0) default value of tf.initializers.random_normal()
-        w_init = tf.random.normal(shape=weight_shape, mean=0.0, stddev=1.0)
-        self.w = tf.Variable(w_init, name='w', trainable=True)
+        self.dense_1 = tf.keras.layers.Dense(128, activation='tanh', name='dense_1')
+        self.bn_1 = tf.keras.layers.BatchNormalization()
+        self.dense_2 = tf.keras.layers.Dense(512, activation='tanh', name='dense_2')
+        self.bn_2 = tf.keras.layers.BatchNormalization()
+        self.dense_3 = tf.keras.layers.Dense(self.embed_dim, activation='tanh', name='dense_3')
 
     def call(self, inputs, training=None, mask=None):
-        x = tf.matmul(inputs, self.w)
+        x = inputs
+        x = self.dense_1(x)
+        x = self.bn_1(x)
+        x = self.dense_2(x)
+        x = self.bn_2(x)
+        x = self.dense_3(x)
         return x
 
     def get_config(self):
-        config = super(LabelEmbedding, self).get_config()
+        config = super(ParameterEmbedding, self).get_config()
         config.update({
             'embed_dim': self.embed_dim,
         })
         return config
+
+
+class PoseShapeEmbedding(tf.keras.layers.Layer):
+    def __init__(self, **kwargs):
+        super(PoseShapeEmbedding, self).__init__(**kwargs)
+
+        # self.pose_embedding = ParameterEmbedding(int(self.embed_dim/2))
+        # self.shape_embedding = ParameterEmbedding(int(self.embed_dim/2))
+        self.flatten = tf.keras.layers.Flatten()
+
+    def call(self, inputs, training=None, mask=None):
+        pose, shape = inputs
+        rotations = run_rodrigues(pose)
+        rotations = self.flatten(rotations)
+        # pose_latent = self.pose_embedding(rotations, training)
+        # shape_latent = self.shape_embedding(shape, training)
+        return tf.concat((rotations, shape), axis=-1)
+
+    def get_config(self):
+        config = super(PoseShapeEmbedding, self).get_config()
+        config.update({
+            'embed_dim': 216 + 11,
+        })
+        return config
+
+
+class GarmentEmbedding(tf.keras.layers.Layer):
+    def __init__(self, embed_dim, **kwargs):
+        super(GarmentEmbedding, self).__init__(**kwargs)
+        self.embed_dim = embed_dim
+
+        # TODO: channels_first?
+
+        self.encoder = ConvNeXt(depths=[3, 3, 27, 3], dims=[96, 192, 256, self.embed_dim], include_top=False)
+        # self.encoder(tf.keras.Input(shape=(*input_shape, 3)))
+        self.pooling = tf.keras.layers.GlobalAveragePooling2D()
+
+    def call(self, inputs, training=None, mask=None):
+        x = inputs
+
+        x = tf.transpose(x, [0,2,3,1])
+        x = self.encoder(x, training=training)
+        x = self.pooling(x)
+
+        return x
+
+    def get_config(self):
+        config = super(GarmentEmbedding, self).get_config()
+        config.update({
+            'embed_dim': self.embed_dim,
+        })
+        return config
+
+
+# class LabelEmbedding(tf.keras.layers.Layer):
+#     def __init__(self, embed_dim, **kwargs):
+#         super(LabelEmbedding, self).__init__(**kwargs)
+#         self.embed_dim = embed_dim
+
+#     def build(self, input_shape):
+#         weight_shape = [input_shape[1], self.embed_dim]
+#         # tf 1.15 mean(0.0), std(1.0) default value of tf.initializers.random_normal()
+#         w_init = tf.random.normal(shape=weight_shape, mean=0.0, stddev=1.0)
+#         self.w = tf.Variable(w_init, name='w', trainable=True)
+
+#     def call(self, inputs, training=None, mask=None):
+#         x = tf.matmul(inputs, self.w)
+#         return x
+
+#     def get_config(self):
+#         config = super(LabelEmbedding, self).get_config()
+#         config.update({
+#             'embed_dim': self.embed_dim,
+#         })
+#         return config
diff --git a/stylegan2/layers/mapping_block.py b/stylegan2/layers/mapping_block.py
index e7ad138..75eddc0 100644
--- a/stylegan2/layers/mapping_block.py
+++ b/stylegan2/layers/mapping_block.py
@@ -1,6 +1,6 @@
 import tensorflow as tf
 
-from stylegan2.layers.embedding import LabelEmbedding
+from stylegan2.layers.embedding import PoseShapeEmbedding, GarmentEmbedding
 from stylegan2.layers.dense import Dense
 from stylegan2.layers.bias_act import BiasAct
 
@@ -14,10 +14,8 @@ class Mapping(tf.keras.layers.Layer):
         self.gain = 1.0
         self.lrmul = 0.01
 
-        if self.labels_dim != self.w_dim:
-            self.labels_embedding = LabelEmbedding(embed_dim=self.w_dim, name='labels_embedding')
-        else:
-            self.labels_embedding = tf.keras.layers.Lambda(lambda x: x)
+        self.pose_shape_embedding = PoseShapeEmbedding()
+        self.garment_embedding = GarmentEmbedding((self.w_dim * 2) - (216 + 11))
 
         self.normalize = tf.keras.layers.Lambda(lambda x: x * tf.math.rsqrt(tf.reduce_mean(tf.square(x), axis=1, keepdims=True) + 1e-8))
 
@@ -28,13 +26,13 @@ class Mapping(tf.keras.layers.Layer):
             self.bias_act_layers.append(BiasAct(lrmul=self.lrmul, act='lrelu', name='bias_{:d}'.format(ii)))
 
     def call(self, inputs, training=None, mask=None):
-        latents, labels = inputs
+        latents, conditions = inputs
         x = latents
 
         # embed label if any
-        if self.labels_dim > 0:
-            y = self.labels_embedding(labels)
-            x = tf.concat([x, y], axis=1)
+        pose_shape_latent = self.pose_shape_embedding([conditions[0], conditions[1]])
+        garment_latent = self.garment_embedding(conditions[2])
+        x = tf.concat([x, pose_shape_latent, garment_latent], axis=1)
 
         # normalize inputs
         x = self.normalize(x)
diff --git a/stylegan2_ref/discriminator.py b/stylegan2_ref/discriminator.py
index 5045d99..a8de81c 100644
--- a/stylegan2_ref/discriminator.py
+++ b/stylegan2_ref/discriminator.py
@@ -119,10 +119,6 @@ class Discriminator(tf.keras.Model):
         self.last_dense = Dense(1, gain=1.0, lrmul=1.0, name='last_dense')
         self.last_bias = BiasAct(lrmul=1.0, act='linear', name='last_bias')
 
-        # conditional output
-        self.last_block_condition = DiscriminatorLastBlock(n_f0, n_f1, res, name='{:d}x{:d}'.format(res, res))
-        self.condition_dense = Dense(self.labels_dim, gain=1.0, lrmul=1.0, name='last_dense')
-        self.condition_bias = BiasAct(lrmul=1.0, act='linear', name='last_bias')
 
     def call(self, inputs, training=None, mask=None):
         images, _ = inputs
@@ -138,13 +134,8 @@ class Discriminator(tf.keras.Model):
         x = self.last_dense(x)
         x = self.last_bias(x)
 
-        # condition output
-        cond = self.last_block_condition(features)
-        cond = self.condition_dense(cond)
-        cond = self.condition_bias(cond)
-
         scores_out = x
-        return scores_out, cond
+        return scores_out
 
     def compute_output_shape(self, input_shape):
         return input_shape[0][0], max(self.labels_dim, 1)
\ No newline at end of file
diff --git a/temp.jpg b/temp.jpg
new file mode 100644
index 0000000..f8c4515
Binary files /dev/null and b/temp.jpg differ
diff --git a/train.py b/train.py
index 6219033..98d71f6 100644
--- a/train.py
+++ b/train.py
@@ -5,13 +5,16 @@ import wandb
 from cv2 import drawFrameAxes
 import numpy as np
 import tensorflow as tf
+from custom_utils.custom_utils import UV2Mesh
 
 from utils import str_to_bool
 from tf_utils import check_tf_version, allow_memory_growth, split_gpu_for_testing
 from load_models import load_generator, load_discriminator
 from dataset_ffhq import get_ffhq_dataset
 from dataset import DatasetSimple
-from losses import d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg
+from losses import d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg, edge_loss
+# from losses import collision_penalty as collision_penalty
+from losses import collision_loss as collision_penalty
 
 
 os.environ["CUDA_VISIBLE_DEVICES"] = "0"
@@ -28,6 +31,8 @@ def initiate_models(g_params, d_params, use_custom_cuda):
     return discriminator, generator, g_clone
 
 
+uv2mesh = UV2Mesh(res=256)
+
 class Trainer(object):
     def __init__(self, t_params, name):
         self.cur_tf_ver = t_params['cur_tf_ver']
@@ -41,12 +46,13 @@ class Trainer(object):
         self.train_res = t_params['train_res']
         self.print_step = 10
         self.save_step = 1000
-        self.image_summary_step = 1000
+        self.image_summary_step = 500
         self.reached_max_steps = False
         self.log_template = '{:s}, {:s}, {:s}'.format(
             'step {}: elapsed: {:.2f}s, d_loss: {:.3f}, g_loss: {:.3f}',
             'd_gan_loss: {:.3f}, g_gan_loss: {:.3f}',
-            'r1_penalty: {:.3f}, pl_penalty: {:.3f}')
+            'r1_penalty: {:.3f}, pl_penalty: {:.3f}',
+            'e_loss: {:.3f}, c_loss: {:.3f}')
 
         # copy network params
         self.g_params = t_params['g_params']
@@ -114,85 +120,114 @@ class Trainer(object):
         return params_copy
 
     def d_train_step(self, dist_inputs):
-        real_images, condition = dist_inputs[0][0], dist_inputs[0][1]
+        real_images, conditions, mesh_static, mesh_body_posed, mask = dist_inputs[0]
 
         with tf.GradientTape() as d_tape:
             # compute losses
-            d_loss, cond_loss = d_logistic(real_images, self.generator, self.discriminator, self.g_params['z_dim'], labels=condition)
+            fake_images, d_loss = d_logistic(real_images, self.generator, self.discriminator, self.g_params['z_dim'], conditions=conditions)
+
+            mesh_pred = uv2mesh(fake_images)
+
+            e_loss = edge_loss(mesh_pred, mesh_static, mask)
+            c_loss = collision_penalty(mesh_pred, mesh_body_posed)
 
             # scale loss
             d_loss = tf.reduce_sum(d_loss) * self.global_batch_scaler
-            cond_loss = tf.reduce_sum(cond_loss) * self.global_batch_scaler
+            e_loss = tf.reduce_sum(e_loss) * self.global_batch_scaler
+            c_loss = tf.reduce_sum(c_loss) * self.global_batch_scaler
 
-            d_loss = d_loss + cond_loss
+            d_loss = d_loss + e_loss + c_loss
 
         d_gradients = d_tape.gradient(d_loss, self.discriminator.trainable_variables)
         self.d_optimizer.apply_gradients(zip(d_gradients, self.discriminator.trainable_variables))
-        return d_loss, cond_loss
+        return d_loss, e_loss, c_loss
 
     def d_train_step_reg(self, dist_inputs):
-        real_images, condition = dist_inputs[0][0], dist_inputs[0][1]
+        real_images, conditions, mesh_static, mesh_body_posed, mask = dist_inputs[0]
 
         with tf.GradientTape() as d_tape:
             # compute losses
-            d_gan_loss, cond_loss, r1_penalty = d_logistic_r1_reg(real_images, self.generator, self.discriminator,
-                                                       self.g_params['z_dim'], labels=condition)
+            fake_images, d_gan_loss, r1_penalty = d_logistic_r1_reg(real_images, self.generator, self.discriminator,
+                                                       self.g_params['z_dim'], conditions=conditions)
             r1_penalty = r1_penalty * (0.5 * self.r1_gamma) * self.d_opt['reg_interval']
 
+            mesh_pred = uv2mesh(fake_images)
+
+            e_loss = edge_loss(mesh_pred, mesh_static, mask)
+            c_loss = collision_penalty(mesh_pred, mesh_body_posed)
+
             # scale losses
             r1_penalty = tf.reduce_sum(r1_penalty) * self.global_batch_scaler
             d_gan_loss = tf.reduce_sum(d_gan_loss) * self.global_batch_scaler
-            cond_loss = tf.reduce_sum(cond_loss) * self.global_batch_scaler
+            e_loss = tf.reduce_sum(e_loss) * self.global_batch_scaler
+            c_loss = tf.reduce_sum(c_loss) * self.global_batch_scaler
 
             # combine
-            d_loss = d_gan_loss + r1_penalty + cond_loss
+            d_loss = d_gan_loss + r1_penalty + e_loss + c_loss
 
         d_gradients = d_tape.gradient(d_loss, self.discriminator.trainable_variables)
         self.d_optimizer.apply_gradients(zip(d_gradients, self.discriminator.trainable_variables))
-        return d_loss, d_gan_loss, r1_penalty, cond_loss
+        return d_loss, d_gan_loss, r1_penalty, e_loss, c_loss
 
     def g_train_step(self, dist_inputs):
-        real_images, condition = dist_inputs[0][0], dist_inputs[0][1]
+        real_images, conditions, mesh_static, mesh_body_posed, mask = dist_inputs[0]
 
         with tf.GradientTape() as g_tape:
             # compute losses
-            g_loss = g_logistic_non_saturating(real_images, self.generator, self.discriminator, self.g_params['z_dim'], labels=condition)
+            fake_images, g_loss = g_logistic_non_saturating(real_images, self.generator, self.discriminator, self.g_params['z_dim'], conditions=conditions)
+
+            mesh_pred = uv2mesh(fake_images)
+
+            e_loss = edge_loss(mesh_pred, mesh_static, mask)
+            c_loss = collision_penalty(mesh_pred, mesh_body_posed)
 
             # scale loss
             g_loss = tf.reduce_sum(g_loss) * self.global_batch_scaler
+            e_loss = tf.reduce_sum(e_loss) * self.global_batch_scaler
+            c_loss = tf.reduce_sum(c_loss) * self.global_batch_scaler
 
+            g_loss = g_loss + e_loss + c_loss
+            
         g_gradients = g_tape.gradient(g_loss, self.generator.trainable_variables)
         self.g_optimizer.apply_gradients(zip(g_gradients, self.generator.trainable_variables))
-        return g_loss
+        return g_loss, e_loss, c_loss
 
     def g_train_step_reg(self, dist_inputs):
-        real_images, condition = dist_inputs[0][0], dist_inputs[0][1]
+        real_images, conditions, mesh_static, mesh_body_posed, mask = dist_inputs[0]
 
         with tf.GradientTape() as g_tape:
             # compute losses
-            g_gan_loss, pl_penalty = g_logistic_ns_pathreg(real_images, self.generator, self.discriminator,
+            fake_images, g_gan_loss, pl_penalty = g_logistic_ns_pathreg(real_images, self.generator, self.discriminator,
                                                            self.g_params['z_dim'], self.pl_mean,
-                                                           self.pl_minibatch_shrink, self.pl_denorm, self.pl_decay, labels=condition)
+                                                           self.pl_minibatch_shrink, self.pl_denorm, self.pl_decay, conditions=conditions)
             pl_penalty = pl_penalty * self.pl_weight * self.g_opt['reg_interval']
 
+            mesh_pred = uv2mesh(fake_images)
+
+            e_loss = edge_loss(mesh_pred, mesh_static, mask)
+            c_loss = collision_penalty(mesh_pred, mesh_body_posed)
+
             # scale loss
             pl_penalty = tf.reduce_sum(pl_penalty) * self.global_batch_scaler
             g_gan_loss = tf.reduce_sum(g_gan_loss) * self.global_batch_scaler
-
+            e_loss = tf.reduce_sum(e_loss) * self.global_batch_scaler
+            c_loss = tf.reduce_sum(c_loss) * self.global_batch_scaler
+            
             # combine
-            g_loss = g_gan_loss + pl_penalty
+            g_loss = g_gan_loss + pl_penalty + e_loss + c_loss
 
         g_gradients = g_tape.gradient(g_loss, self.generator.trainable_variables)
         self.g_optimizer.apply_gradients(zip(g_gradients, self.generator.trainable_variables))
-        return g_loss, g_gan_loss, pl_penalty
+        return g_loss, g_gan_loss, pl_penalty, e_loss, c_loss
 
     def train(self, dist_datasets, strategy):
         @tf.function
         def dist_d_train_step(inputs):
             per_replica_losses = strategy.run(fn=self.d_train_step, args=(inputs,))
             mean_d_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[0], axis=None)
-            mean_cond_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[1], axis=None)
-            return mean_d_loss, mean_cond_loss
+            mean_e_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[1], axis=None)
+            mean_c_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[2], axis=None)
+            return mean_d_loss, mean_e_loss, mean_c_loss
 
         @tf.function
         def dist_d_train_step_reg(inputs):
@@ -200,14 +235,17 @@ class Trainer(object):
             mean_d_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[0], axis=None)
             mean_d_gan_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[1], axis=None)
             mean_r1_penalty = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[2], axis=None)
-            mean_cond_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[3], axis=None)
-            return mean_d_loss, mean_d_gan_loss, mean_r1_penalty, mean_cond_loss
+            mean_e_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[3], axis=None)
+            mean_c_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[4], axis=None)
+            return mean_d_loss, mean_d_gan_loss, mean_r1_penalty, mean_e_loss, mean_c_loss
 
         @tf.function
         def dist_g_train_step(inputs):
             per_replica_losses = strategy.run(fn=self.g_train_step, args=(inputs,))
-            mean_g_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses, axis=None)
-            return mean_g_loss
+            mean_g_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[0], axis=None)
+            mean_e_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[1], axis=None)
+            mean_c_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[2], axis=None)
+            return mean_g_loss, mean_e_loss, mean_c_loss
 
         @tf.function
         def dist_g_train_step_reg(inputs):
@@ -215,7 +253,9 @@ class Trainer(object):
             mean_g_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[0], axis=None)
             mean_g_gan_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[1], axis=None)
             mean_pl_penalty = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[2], axis=None)
-            return mean_g_loss, mean_g_gan_loss, mean_pl_penalty
+            mean_e_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[3], axis=None)
+            mean_c_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[4], axis=None)
+            return mean_g_loss, mean_g_gan_loss, mean_pl_penalty, mean_e_loss, mean_c_loss
 
         @tf.function
         def dist_gen_samples(dist_inputs):
@@ -238,28 +278,29 @@ class Trainer(object):
         metric_g_gan_loss = tf.keras.metrics.Mean('g_gan_loss', dtype=tf.float32)
         metric_r1_penalty = tf.keras.metrics.Mean('r1_penalty', dtype=tf.float32)
         metric_pl_penalty = tf.keras.metrics.Mean('pl_penalty', dtype=tf.float32)
-        metric_cond_loss = tf.keras.metrics.Mean('cond_loss', dtype=tf.float32)
+        metric_e_loss = tf.keras.metrics.Mean('edge_loss', dtype=tf.float32)
+        metric_c_loss = tf.keras.metrics.Mean('collision_loss', dtype=tf.float32)
 
         # start training
         zero = tf.constant(0.0, dtype=tf.float32)
         print('max_steps: {}'.format(self.max_steps))
         t_start = time.time()
         step = self.g_optimizer.iterations.numpy()
-        for real_images in dist_datasets:
+        for batch in dist_datasets:
 
             # d train step
             if (step + 1) % self.d_opt['reg_interval'] == 0:
-                d_loss, d_gan_loss, r1_penalty, cond_loss = dist_d_train_step_reg((real_images, ))
+                d_loss, d_gan_loss, r1_penalty, e_loss, c_loss = dist_d_train_step_reg((batch, ))
             else:
-                d_loss, cond_loss = dist_d_train_step((real_images, ))
+                d_loss, e_loss, c_loss = dist_d_train_step((batch, ))
                 d_gan_loss = d_loss
                 r1_penalty = zero
 
             # g train step
             if (step + 1) % self.g_opt['reg_interval'] == 0:
-                g_loss, g_gan_loss, pl_penalty = dist_g_train_step_reg((real_images, ))
+                g_loss, g_gan_loss, pl_penalty, e_loss, c_loss = dist_g_train_step_reg((batch, ))
             else:
-                g_loss = dist_g_train_step((real_images,))
+                g_loss, e_loss, c_loss = dist_g_train_step((batch,))
                 g_gan_loss = g_loss
                 pl_penalty = zero
 
@@ -273,7 +314,8 @@ class Trainer(object):
             metric_g_gan_loss(g_gan_loss)
             metric_r1_penalty(r1_penalty)
             metric_pl_penalty(pl_penalty)
-            metric_cond_loss(cond_loss)
+            metric_e_loss(e_loss)
+            metric_c_loss(c_loss)
 
             # get current step
             step += 1
@@ -286,7 +328,8 @@ class Trainer(object):
                 tf.summary.scalar('g_gan_loss', metric_g_gan_loss.result(), step=step)
                 tf.summary.scalar('r1_penalty', metric_r1_penalty.result(), step=step)
                 tf.summary.scalar('pl_penalty', metric_pl_penalty.result(), step=step)
-                tf.summary.scalar('cond_loss', metric_cond_loss.result(), step=step)
+                tf.summary.scalar('edge_loss', metric_e_loss.result(), step=step)
+                tf.summary.scalar('collision_loss', metric_c_loss.result(), step=step)
 
             # save every self.save_step
             if step % self.save_step == 0:
@@ -296,8 +339,8 @@ class Trainer(object):
             if step % self.image_summary_step == 0:
                 # add summary image
                 test_z = tf.random.normal(shape=(self.n_samples, self.g_params['z_dim']), dtype=tf.dtypes.float32)
-                test_labels = tf.ones((self.n_samples, self.g_params['labels_dim']), dtype=tf.dtypes.float32)
-                summary_image = dist_gen_samples((test_z, test_labels))
+                test_conditions = batch[1] #tf.ones((self.n_samples, self.g_params['labels_dim']), dtype=tf.dtypes.float32)
+                summary_image = dist_gen_samples((test_z, test_conditions))
 
                 # convert to tensor image
                 summary_image = self.convert_per_replica_image(summary_image, strategy)
@@ -312,7 +355,8 @@ class Trainer(object):
                 elapsed = time.time() - t_start
                 print(self.log_template.format(step, elapsed, d_loss.numpy(), g_loss.numpy(),
                                                d_gan_loss.numpy(), g_gan_loss.numpy(),
-                                               r1_penalty.numpy(), pl_penalty.numpy()))
+                                               r1_penalty.numpy(), pl_penalty.numpy(),
+                                               e_loss.numpy(), c_loss.numpy()))
 
                 # reset timer
                 t_start = time.time()
@@ -364,7 +408,7 @@ def main():
     parser.add_argument('--tfrecord_dir', default='./tfrecords', type=str)
     parser.add_argument('--train_res', default=512, type=int)
     parser.add_argument('--shuffle_buffer_size', default=1000, type=int)
-    parser.add_argument('--batch_size_per_replica', default=12, type=int)
+    parser.add_argument('--batch_size_per_replica', default=4, type=int)
     args = vars(parser.parse_args())
 
     # check tensorflow version
@@ -382,15 +426,15 @@ def main():
     # featuremaps = [64, 64, 64, 64, 64, 32, 16, 8, 4]
     train_resolutions, train_featuremaps = filter_resolutions_featuremaps(resolutions, featuremaps, args['train_res'])
     g_params = {
-        'z_dim': 512,
-        'w_dim': 512,
-        'labels_dim': 83,
+        'z_dim': 256,
+        'w_dim': 256,
+        'labels_dim': -1,
         'n_mapping': 8,
         'resolutions': train_resolutions,
         'featuremaps': train_featuremaps,
     }
     d_params = {
-        'labels_dim': 83,
+        'labels_dim': -1,
         'resolutions': train_resolutions,
         'featuremaps': train_featuremaps,
     }
@@ -401,7 +445,14 @@ def main():
 
     # prepare dataset
     # dataset = get_ffhq_dataset(args['tfrecord_dir'], args['train_res'], batch_size=global_batch_size, epochs=None)
-    dataset = DatasetSimple('/workspace/Dataset/stylegan').batch(global_batch_size).shuffle(buffer_size=1000).prefetch(buffer_size=tf.data.experimental.AUTOTUNE).cache().repeat()
+    dataset = DatasetSimple('/home/workspace/data/dataset/stylegan_new/train').cache().batch(global_batch_size).repeat()
+    # limit = int(25594/global_batch_size)
+    # count = 0
+    # for d in dataset:
+    #     count += global_batch_size
+    #     if count > limit:
+    #         break
+
     options = tf.data.Options()
     options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.OFF
     dataset = dataset.with_options(options)