Add DSINE model prediction functionality and user agreement modal

.dockerignore

@@ -1 +1,2 @@
-.pixi/
+.pixi/
+checkpoints/

.gitignore

@@ -1,3 +1,165 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
 # pixi environments
 .pixi
+samples/*
+checkpoints/*
 

main.py

@@ -1,9 +1,80 @@
 import gradio as gr
+from gradio_modal import Modal
+from gradio_imageslider import ImageSlider
 
+import numpy as np
 
-def greet(name):
-    return "Hello " + name + "!"
+import torch
+import torch.nn.functional as F
+from torchvision import transforms
+from PIL import Image
+import utils.utils as utils
 
+from models.dsine import DSINE
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch(server_name="0.0.0.0", server_port=7860)
+device = torch.device("cpu")
+
+model = DSINE().to(device)
+model.pixel_coords = model.pixel_coords.to(device)
+model = utils.load_checkpoint("./checkpoints/dsine.pt", model)
+model.eval()
+
+
+def predict_normal(img_np: np.ndarray):
+    # normalize
+    normalize = transforms.Normalize(
+        mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
+    )
+
+    with torch.no_grad():
+        img = np.array(img_np).astype(np.float32) / 255.0
+        img = torch.from_numpy(img).permute(2, 0, 1).unsqueeze(0).to("cpu")
+        _, _, orig_H, orig_W = img.shape
+
+        # zero-pad the input image so that both the width and height are multiples of 32
+        l, r, t, b = utils.pad_input(orig_H, orig_W)
+        img = F.pad(img, (l, r, t, b), mode="constant", value=0.0)
+        img = normalize(img)
+
+        # NOTE: if intrins is not given, we just assume that the principal point is at the center
+        # and that the field-of-view is 60 degrees (feel free to modify this assumption)
+        intrins = utils.get_intrins_from_fov(
+            new_fov=60.0, H=orig_H, W=orig_W, device="cpu"
+        ).unsqueeze(0)
+
+        intrins[:, 0, 2] += l
+        intrins[:, 1, 2] += t
+
+        pred_norm = model(img, intrins=intrins)[-1]
+        pred_norm = pred_norm[:, :, t : t + orig_H, l : l + orig_W]
+
+        # save to output folder
+        # NOTE: by saving the prediction as uint8 png format, you lose a lot of precision
+        # if you want to use the predicted normals for downstream tasks, we recommend saving them as float32 NPY files
+        pred_norm_np = (
+            pred_norm.cpu().detach().numpy()[0, :, :, :].transpose(1, 2, 0)
+        )  # (H, W, 3)
+        pred_norm_np = ((pred_norm_np + 1.0) / 2.0 * 255.0).astype(np.uint8)
+
+    return (img_np, pred_norm_np)
+
+
+with gr.Blocks() as demo:
+    with gr.Group():
+        with gr.Row():
+            input_img = gr.Image(label="Input image", image_mode="RGB")
+            output_img = ImageSlider(label="Surface Normal", type="numpy")
+            # output_img = gr.Image(label="Normal")
+
+        btn = gr.Button("Predict")
+        btn.click(fn=predict_normal, inputs=[input_img], outputs=[output_img])
+
+    with Modal(visible=True, allow_user_close=False) as modal:
+        gr.Markdown(
+            "To use this space, you must agree to the terms and conditions. found [here](https://github.com/baegwangbin/DSINE/blob/main/LICENSE)."
+        )
+        btn = gr.Button("I agree")
+        btn.click(lambda: Modal(visible=False), None, modal)
+
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860)

models/dsine.py

@@ -0,0 +1,233 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from models.submodules import Encoder, ConvGRU, UpSampleBN, UpSampleGN, RayReLU, \
+                                convex_upsampling, get_unfold, get_prediction_head, \
+                                INPUT_CHANNELS_DICT
+from utils.rotation import axis_angle_to_matrix
+
+
+class Decoder(nn.Module):
+    def __init__(self, output_dims, B=5, NF=2048, BN=False, downsample_ratio=8):
+        super(Decoder, self).__init__()
+        input_channels = INPUT_CHANNELS_DICT[B]
+        output_dim, feature_dim, hidden_dim = output_dims
+        features = bottleneck_features = NF
+        self.downsample_ratio = downsample_ratio
+
+        UpSample = UpSampleBN if BN else UpSampleGN 
+        self.conv2 = nn.Conv2d(bottleneck_features + 2, features, kernel_size=1, stride=1, padding=0)
+        self.up1 = UpSample(skip_input=features // 1 + input_channels[1] + 2, output_features=features // 2, align_corners=False)
+        self.up2 = UpSample(skip_input=features // 2 + input_channels[2] + 2, output_features=features // 4, align_corners=False)
+
+        # prediction heads
+        i_dim = features // 4
+        h_dim = 128
+        self.normal_head = get_prediction_head(i_dim+2, h_dim, output_dim)
+        self.feature_head = get_prediction_head(i_dim+2, h_dim, feature_dim)
+        self.hidden_head = get_prediction_head(i_dim+2, h_dim, hidden_dim)
+
+    def forward(self, features, uvs):
+        _, _, x_block2, x_block3, x_block4 = features[4], features[5], features[6], features[8], features[11]
+        uv_32, uv_16, uv_8 = uvs
+
+        x_d0 = self.conv2(torch.cat([x_block4, uv_32], dim=1))
+        x_d1 = self.up1(x_d0, torch.cat([x_block3, uv_16], dim=1))
+        x_feat = self.up2(x_d1, torch.cat([x_block2, uv_8], dim=1))
+        x_feat = torch.cat([x_feat, uv_8], dim=1)
+
+        normal = self.normal_head(x_feat)
+        normal = F.normalize(normal, dim=1)
+        f = self.feature_head(x_feat)
+        h = self.hidden_head(x_feat)
+        return normal, f, h
+
+
+class DSINE(nn.Module):
+    def __init__(self):
+        super(DSINE, self).__init__()
+        self.downsample_ratio = 8
+        self.ps = 5           # patch size
+        self.num_iter = 5     # num iterations
+
+        # define encoder
+        self.encoder = Encoder(B=5, pretrained=True)
+
+        # define decoder
+        self.output_dim = output_dim = 3
+        self.feature_dim = feature_dim = 64
+        self.hidden_dim = hidden_dim = 64
+        self.decoder = Decoder([output_dim, feature_dim, hidden_dim], B=5, NF=2048, BN=False)
+
+        # ray direction-based ReLU
+        self.ray_relu = RayReLU(eps=1e-2)
+
+        # pixel_coords (1, 3, H, W)
+        # NOTE: this is set to some arbitrarily high number, 
+        # if your input is 2000+ pixels wide/tall, increase these values
+        h = 2000
+        w = 2000
+        pixel_coords = np.ones((3, h, w)).astype(np.float32)
+        x_range = np.concatenate([np.arange(w).reshape(1, w)] * h, axis=0)
+        y_range = np.concatenate([np.arange(h).reshape(h, 1)] * w, axis=1)
+        pixel_coords[0, :, :] = x_range + 0.5
+        pixel_coords[1, :, :] = y_range + 0.5
+        self.pixel_coords = torch.from_numpy(pixel_coords).unsqueeze(0)
+
+        # define ConvGRU cell
+        self.gru = ConvGRU(hidden_dim=hidden_dim, input_dim=feature_dim+2, ks=self.ps)
+
+        # padding used during NRN
+        self.pad = (self.ps - 1) // 2
+
+        # prediction heads
+        self.prob_head = get_prediction_head(self.hidden_dim+2, 64, self.ps*self.ps)   # weights assigned for each nghbr pixel 
+        self.xy_head = get_prediction_head(self.hidden_dim+2, 64, self.ps*self.ps*2)   # rotation axis for each nghbr pixel
+        self.angle_head = get_prediction_head(self.hidden_dim+2, 64, self.ps*self.ps)  # rotation angle for each nghbr pixel
+
+        # prediction heads - weights used for upsampling the coarse resolution output
+        self.up_prob_head = get_prediction_head(self.hidden_dim+2, 64, 9 * self.downsample_ratio * self.downsample_ratio)
+
+    def get_ray(self, intrins, H, W, orig_H, orig_W, return_uv=False):
+        B, _, _ = intrins.shape
+        fu = intrins[:, 0, 0][:,None,None] * (W / orig_W)
+        cu = intrins[:, 0, 2][:,None,None] * (W / orig_W)
+        fv = intrins[:, 1, 1][:,None,None] * (H / orig_H)
+        cv = intrins[:, 1, 2][:,None,None] * (H / orig_H)
+
+        # (B, 2, H, W)
+        ray = self.pixel_coords[:, :, :H, :W].repeat(B, 1, 1, 1)
+        ray[:, 0, :, :] = (ray[:, 0, :, :] - cu) / fu
+        ray[:, 1, :, :] = (ray[:, 1, :, :] - cv) / fv
+
+        if return_uv:
+            return ray[:, :2, :, :]
+        else:
+            return F.normalize(ray, dim=1)
+
+    def upsample(self, h, pred_norm, uv_8):
+        up_mask = self.up_prob_head(torch.cat([h, uv_8], dim=1))
+        up_pred_norm = convex_upsampling(pred_norm, up_mask, self.downsample_ratio)
+        up_pred_norm = F.normalize(up_pred_norm, dim=1)
+        return up_pred_norm
+
+    def refine(self, h, feat_map, pred_norm, intrins, orig_H, orig_W, uv_8, ray_8):
+        B, C, H, W = pred_norm.shape
+        fu = intrins[:, 0, 0][:,None,None,None] * (W / orig_W)  # (B, 1, 1, 1)
+        cu = intrins[:, 0, 2][:,None,None,None] * (W / orig_W)
+        fv = intrins[:, 1, 1][:,None,None,None] * (H / orig_H)
+        cv = intrins[:, 1, 2][:,None,None,None] * (H / orig_H)
+
+        h_new = self.gru(h, feat_map)
+
+        # get nghbr prob (B, 1, ps*ps, h, w)
+        nghbr_prob = self.prob_head(torch.cat([h_new, uv_8], dim=1)).unsqueeze(1)
+        nghbr_prob = torch.sigmoid(nghbr_prob)
+
+        # get nghbr normals (B, 3, ps*ps, h, w)
+        nghbr_normals = get_unfold(pred_norm, ps=self.ps, pad=self.pad)
+
+        # get nghbr xy (B, 2, ps*ps, h, w)
+        nghbr_xys = self.xy_head(torch.cat([h_new, uv_8], dim=1))
+        nghbr_xs, nghbr_ys = torch.split(nghbr_xys, [self.ps*self.ps, self.ps*self.ps], dim=1)
+        nghbr_xys = torch.cat([nghbr_xs.unsqueeze(1), nghbr_ys.unsqueeze(1)], dim=1)        
+        nghbr_xys = F.normalize(nghbr_xys, dim=1)
+
+        # get nghbr theta (B, 1, ps*ps, h, w)
+        nghbr_angle = self.angle_head(torch.cat([h_new, uv_8], dim=1)).unsqueeze(1)
+        nghbr_angle = torch.sigmoid(nghbr_angle) * np.pi
+
+        # get nghbr pixel coord (1, 3, ps*ps, h, w)
+        nghbr_pixel_coord = get_unfold(self.pixel_coords[:, :, :H, :W], ps=self.ps, pad=self.pad)
+
+        # nghbr axes (B, 3, ps*ps, h, w)
+        nghbr_axes = torch.zeros_like(nghbr_normals)
+
+        du_over_fu = nghbr_xys[:, 0, ...] / fu                                      # (B, ps*ps, h, w)
+        dv_over_fv = nghbr_xys[:, 1, ...] / fv                                      # (B, ps*ps, h, w)
+
+        term_u = (nghbr_pixel_coord[:, 0, ...] + nghbr_xys[:, 0, ...] - cu) / fu    # (B, ps*ps, h, w)
+        term_v = (nghbr_pixel_coord[:, 1, ...] + nghbr_xys[:, 1, ...] - cv) / fv    # (B, ps*ps, h, w)
+
+        nx = nghbr_normals[:, 0, ...]                                               # (B, ps*ps, h, w)
+        ny = nghbr_normals[:, 1, ...]                                               # (B, ps*ps, h, w)
+        nz = nghbr_normals[:, 2, ...]                                               # (B, ps*ps, h, w)
+
+        nghbr_delta_z_num = - (du_over_fu * nx + dv_over_fv * ny)
+        nghbr_delta_z_denom = (term_u * nx + term_v * ny + nz)
+        nghbr_delta_z_denom[torch.abs(nghbr_delta_z_denom) < 1e-8] = 1e-8 * torch.sign(nghbr_delta_z_denom[torch.abs(nghbr_delta_z_denom) < 1e-8])
+        nghbr_delta_z = nghbr_delta_z_num / nghbr_delta_z_denom
+
+        nghbr_axes[:, 0, ...] = du_over_fu + nghbr_delta_z * term_u
+        nghbr_axes[:, 1, ...] = dv_over_fv + nghbr_delta_z * term_v
+        nghbr_axes[:, 2, ...] = nghbr_delta_z
+        nghbr_axes = F.normalize(nghbr_axes, dim=1)                                 # (B, 3, ps*ps, h, w)
+
+        # make sure axes are all valid
+        invalid = torch.sum(torch.logical_or(torch.isnan(nghbr_axes), torch.isinf(nghbr_axes)).float(), dim=1) > 0.5    # (B, ps*ps, h, w)
+        nghbr_axes[:, 0, ...][invalid] = 0.0
+        nghbr_axes[:, 1, ...][invalid] = 0.0
+        nghbr_axes[:, 2, ...][invalid] = 0.0
+
+        # nghbr_axes_angle (B, 3, ps*ps, h, w)
+        nghbr_axes_angle = nghbr_axes * nghbr_angle
+        nghbr_axes_angle = nghbr_axes_angle.permute(0, 2, 3, 4, 1)  # (B, ps*ps, h, w, 3)
+        nghbr_R = axis_angle_to_matrix(nghbr_axes_angle)            # (B, ps*ps, h, w, 3, 3)
+
+        # (B, 3, ps*ps, h, w)
+        nghbr_normals_rot = torch.bmm(
+            nghbr_R.reshape(B * self.ps * self.ps * H * W, 3, 3),
+            nghbr_normals.permute(0, 2, 3, 4, 1).reshape(B * self.ps * self.ps * H * W, 3).unsqueeze(-1)
+        ).reshape(B, self.ps*self.ps, H, W, 3, 1).squeeze(-1).permute(0, 4, 1, 2, 3)        # (B, 3, ps*ps, h, w)
+        nghbr_normals_rot = F.normalize(nghbr_normals_rot, dim=1)
+
+        # ray ReLU
+        nghbr_normals_rot = torch.cat([
+            self.ray_relu(nghbr_normals_rot[:, :, i, :, :], ray_8).unsqueeze(2) 
+            for i in range(nghbr_normals_rot.size(2))
+            ], dim=2)
+
+        # (B, 1, ps*ps, h, w) * (B, 3, ps*ps, h, w)
+        pred_norm = torch.sum(nghbr_prob * nghbr_normals_rot, dim=2)    # (B, C, H, W)
+        pred_norm = F.normalize(pred_norm, dim=1)
+    
+        up_mask = self.up_prob_head(torch.cat([h_new, uv_8], dim=1))
+        up_pred_norm = convex_upsampling(pred_norm, up_mask, self.downsample_ratio)
+        up_pred_norm = F.normalize(up_pred_norm, dim=1)
+
+        return h_new, pred_norm, up_pred_norm
+
+    
+    def forward(self, img, intrins=None):
+        # Step 1. encoder
+        features = self.encoder(img)
+
+        # Step 2. get uv encoding
+        B, _, orig_H, orig_W = img.shape
+        intrins[:, 0, 2] += 0.5
+        intrins[:, 1, 2] += 0.5
+        uv_32 = self.get_ray(intrins, orig_H//32, orig_W//32, orig_H, orig_W, return_uv=True)
+        uv_16 = self.get_ray(intrins, orig_H//16, orig_W//16, orig_H, orig_W, return_uv=True)
+        uv_8 = self.get_ray(intrins, orig_H//8, orig_W//8, orig_H, orig_W, return_uv=True)
+        ray_8 = self.get_ray(intrins, orig_H//8, orig_W//8, orig_H, orig_W)
+
+        # Step 3. decoder - initial prediction
+        pred_norm, feat_map, h = self.decoder(features, uvs=(uv_32, uv_16, uv_8))
+        pred_norm = self.ray_relu(pred_norm, ray_8)
+
+        # Step 4. add ray direction encoding
+        feat_map = torch.cat([feat_map, uv_8], dim=1)
+
+        # iterative refinement
+        up_pred_norm = self.upsample(h, pred_norm, uv_8)
+        pred_list = [up_pred_norm]
+        for i in range(self.num_iter):
+            h, pred_norm, up_pred_norm = self.refine(h, feat_map, 
+                                                     pred_norm.detach(), 
+                                                     intrins, orig_H, orig_W, uv_8, ray_8)
+            pred_list.append(up_pred_norm)
+        return pred_list
+
+

models/submodules.py

@@ -0,0 +1,194 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import geffnet
+
+
+INPUT_CHANNELS_DICT = {
+    0: [1280, 112, 40, 24, 16],
+    1: [1280, 112, 40, 24, 16],
+    2: [1408, 120, 48, 24, 16],
+    3: [1536, 136, 48, 32, 24],
+    4: [1792, 160, 56, 32, 24],
+    5: [2048, 176, 64, 40, 24],
+    6: [2304, 200, 72, 40, 32],
+    7: [2560, 224, 80, 48, 32]
+}
+
+
+class Encoder(nn.Module):
+    def __init__(self, B=5, pretrained=True):
+        """ e.g. B=5 will return EfficientNet-B5
+        """
+        super(Encoder, self).__init__()
+        basemodel = geffnet.create_model('tf_efficientnet_b%s_ap' % B, pretrained=pretrained)
+        # Remove last layer
+        basemodel.global_pool = nn.Identity()
+        basemodel.classifier = nn.Identity()
+        self.original_model = basemodel
+
+    def forward(self, x):
+        features = [x]
+        for k, v in self.original_model._modules.items():
+            if (k == 'blocks'):
+                for ki, vi in v._modules.items():
+                    features.append(vi(features[-1]))
+            else:
+                features.append(v(features[-1]))
+        return features
+
+
+class ConvGRU(nn.Module):
+    def __init__(self, hidden_dim, input_dim, ks=3):
+        super(ConvGRU, self).__init__()
+        p = (ks - 1) // 2
+        self.convz = nn.Conv2d(hidden_dim+input_dim, hidden_dim, ks, padding=p)
+        self.convr = nn.Conv2d(hidden_dim+input_dim, hidden_dim, ks, padding=p)
+        self.convq = nn.Conv2d(hidden_dim+input_dim, hidden_dim, ks, padding=p)
+
+    def forward(self, h, x):
+        hx = torch.cat([h, x], dim=1)
+        z = torch.sigmoid(self.convz(hx))
+        r = torch.sigmoid(self.convr(hx))
+        q = torch.tanh(self.convq(torch.cat([r*h, x], dim=1)))
+        h = (1-z) * h + z * q
+        return h
+    
+
+class RayReLU(nn.Module):
+    def __init__(self, eps=1e-2):
+        super(RayReLU, self).__init__()
+        self.eps = eps
+
+    def forward(self, pred_norm, ray):
+        # angle between the predicted normal and ray direction
+        cos = torch.cosine_similarity(pred_norm, ray, dim=1).unsqueeze(1) # (B, 1, H, W)
+
+        # component of pred_norm along view
+        norm_along_view = ray * cos
+
+        # cos should be bigger than eps
+        norm_along_view_relu = ray * (torch.relu(cos - self.eps) + self.eps)
+
+        # difference        
+        diff = norm_along_view_relu - norm_along_view
+
+        # updated pred_norm
+        new_pred_norm = pred_norm + diff
+        new_pred_norm = F.normalize(new_pred_norm, dim=1)
+
+        return new_pred_norm    
+
+
+class UpSampleBN(nn.Module):
+    def __init__(self, skip_input, output_features, align_corners=True):
+        super(UpSampleBN, self).__init__()
+        self._net = nn.Sequential(nn.Conv2d(skip_input, output_features, kernel_size=3, stride=1, padding=1),
+                                  nn.BatchNorm2d(output_features),
+                                  nn.LeakyReLU(),
+                                  nn.Conv2d(output_features, output_features, kernel_size=3, stride=1, padding=1),
+                                  nn.BatchNorm2d(output_features),
+                                  nn.LeakyReLU())
+        self.align_corners = align_corners
+
+    def forward(self, x, concat_with):
+        up_x = F.interpolate(x, size=[concat_with.size(2), concat_with.size(3)], mode='bilinear', align_corners=self.align_corners)
+        f = torch.cat([up_x, concat_with], dim=1)
+        return self._net(f)
+
+
+class Conv2d_WS(nn.Conv2d):
+    """ weight standardization
+    """ 
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1,
+                 padding=0, dilation=1, groups=1, bias=True):
+        super(Conv2d_WS, self).__init__(in_channels, out_channels, kernel_size, stride,
+                 padding, dilation, groups, bias)
+
+    def forward(self, x):
+        weight = self.weight
+        weight_mean = weight.mean(dim=1, keepdim=True).mean(dim=2,
+                                  keepdim=True).mean(dim=3, keepdim=True)
+        weight = weight - weight_mean
+        std = weight.view(weight.size(0), -1).std(dim=1).view(-1, 1, 1, 1) + 1e-5
+        weight = weight / std.expand_as(weight)
+        return F.conv2d(x, weight, self.bias, self.stride, self.padding, self.dilation, self.groups)
+
+
+class UpSampleGN(nn.Module):
+    """ UpSample with GroupNorm
+    """
+    def __init__(self, skip_input, output_features, align_corners=True):
+        super(UpSampleGN, self).__init__()
+        self._net = nn.Sequential(Conv2d_WS(skip_input, output_features, kernel_size=3, stride=1, padding=1),
+                                  nn.GroupNorm(8, output_features),
+                                  nn.LeakyReLU(),
+                                  Conv2d_WS(output_features, output_features, kernel_size=3, stride=1, padding=1),
+                                  nn.GroupNorm(8, output_features),
+                                  nn.LeakyReLU())
+        self.align_corners = align_corners
+
+    def forward(self, x, concat_with):
+        up_x = F.interpolate(x, size=[concat_with.size(2), concat_with.size(3)], mode='bilinear', align_corners=self.align_corners)
+        f = torch.cat([up_x, concat_with], dim=1)
+        return self._net(f)
+
+
+def upsample_via_bilinear(out, up_mask, downsample_ratio):
+    """ bilinear upsampling (up_mask is a dummy variable)
+    """
+    return F.interpolate(out, scale_factor=downsample_ratio, mode='bilinear', align_corners=True)
+
+
+def upsample_via_mask(out, up_mask, downsample_ratio):
+    """ convex upsampling
+    """
+    # out: low-resolution output (B, o_dim, H, W)
+    # up_mask: (B, 9*k*k, H, W)
+    k = downsample_ratio
+
+    N, o_dim, H, W = out.shape
+    up_mask = up_mask.view(N, 1, 9, k, k, H, W)
+    up_mask = torch.softmax(up_mask, dim=2)             # (B, 1, 9, k, k, H, W)
+
+    up_out = F.unfold(out, [3, 3], padding=1)       # (B, 2, H, W) -> (B, 2 X 3*3, H*W)
+    up_out = up_out.view(N, o_dim, 9, 1, 1, H, W)   # (B, 2, 3*3, 1, 1, H, W)
+    up_out = torch.sum(up_mask * up_out, dim=2)     # (B, 2, k, k, H, W)
+
+    up_out = up_out.permute(0, 1, 4, 2, 5, 3)       # (B, 2, H, k, W, k)
+    return up_out.reshape(N, o_dim, k*H, k*W)   # (B, 2, kH, kW)
+
+
+def convex_upsampling(out, up_mask, k):
+    # out: low-resolution output    (B, C, H, W)
+    # up_mask:                      (B, 9*k*k, H, W)
+    B, C, H, W = out.shape
+    up_mask = up_mask.view(B, 1, 9, k, k, H, W)
+    up_mask = torch.softmax(up_mask, dim=2)             # (B, 1, 9, k, k, H, W)
+
+    out = F.pad(out, pad=(1,1,1,1), mode='replicate')
+    up_out = F.unfold(out, [3, 3], padding=0)           # (B, C, H, W) -> (B, C X 3*3, H*W)
+    up_out = up_out.view(B, C, 9, 1, 1, H, W)           # (B, C, 9, 1, 1, H, W)
+
+    up_out = torch.sum(up_mask * up_out, dim=2)         # (B, C, k, k, H, W)
+    up_out = up_out.permute(0, 1, 4, 2, 5, 3)           # (B, C, H, k, W, k)
+    return up_out.reshape(B, C, k*H, k*W)               # (B, C, kH, kW)
+
+
+def get_unfold(pred_norm, ps, pad):
+    B, C, H, W = pred_norm.shape
+    pred_norm = F.pad(pred_norm, pad=(pad,pad,pad,pad), mode='replicate')       # (B, C, h, w)
+    pred_norm_unfold = F.unfold(pred_norm, [ps, ps], padding=0)                 # (B, C X ps*ps, h*w)
+    pred_norm_unfold = pred_norm_unfold.view(B, C, ps*ps, H, W)                 # (B, C, ps*ps, h, w)
+    return pred_norm_unfold
+
+
+def get_prediction_head(input_dim, hidden_dim, output_dim):
+    return nn.Sequential(
+        nn.Conv2d(input_dim, hidden_dim, 3, padding=1), 
+        nn.ReLU(inplace=True),
+        nn.Conv2d(hidden_dim, hidden_dim, 1), 
+        nn.ReLU(inplace=True),
+        nn.Conv2d(hidden_dim, output_dim, 1),
+    )
+

pixi.lock

@@ -3,11 +3,13 @@ environments:
   default:
     channels:
     - url: https://conda.anaconda.org/conda-forge/
+    - url: https://conda.anaconda.org/pytorch/
     packages:
       linux-64:
       - conda: https://conda.anaconda.org/conda-forge/linux-64/_libgcc_mutex-0.1-conda_forge.tar.bz2
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/_openmp_mutex-4.5-2_gnu.tar.bz2
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/_openmp_mutex-4.5-2_kmp_llvm.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/linux-64/aom-3.8.1-h59595ed_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/blas-1.0-mkl.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-1.1.0-hd590300_1.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-bin-1.1.0-hd590300_1.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-python-1.1.0-py311hb755f60_1.conda
@@ -24,6 +26,7 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/linux-64/fribidi-1.0.10-h36c2ea0_0.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/linux-64/gettext-0.21.1-h27087fc_0.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/linux-64/gmp-6.3.0-h59595ed_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/gmpy2-2.1.2-py311h6a5fa03_1.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/linux-64/gnutls-3.7.9-hb077bed_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/graphite2-1.3.13-h58526e2_1001.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/linux-64/harfbuzz-8.3.0-h3d44ed6_0.conda
@@ -35,11 +38,11 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/linux-64/lerc-4.0.0-h27087fc_0.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/linux-64/libabseil-20240116.1-cxx17_h59595ed_2.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/libass-0.17.1-h8fe9dca_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/libblas-3.9.0-21_linux64_openblas.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/libblas-3.9.0-1_h86c2bf4_netlib.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/linux-64/libbrotlicommon-1.1.0-hd590300_1.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/libbrotlidec-1.1.0-hd590300_1.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/libbrotlienc-1.1.0-hd590300_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/libcblas-3.9.0-21_linux64_openblas.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/libcblas-3.9.0-5_h92ddd45_netlib.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/linux-64/libdeflate-1.19-hd590300_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/libdrm-2.4.120-hd590300_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/libexpat-2.5.0-hcb278e6_1.conda
@@ -48,14 +51,12 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/linux-64/libgfortran-ng-13.2.0-h69a702a_5.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/libgfortran5-13.2.0-ha4646dd_5.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/libglib-2.78.4-h783c2da_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/libgomp-13.2.0-h807b86a_5.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/libhwloc-2.9.3-default_h554bfaf_1009.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/libiconv-1.17-hd590300_2.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/libidn2-2.3.7-hd590300_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/libjpeg-turbo-3.0.0-hd590300_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/liblapack-3.9.0-21_linux64_openblas.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/liblapack-3.9.0-5_h92ddd45_netlib.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/linux-64/libnsl-2.0.1-hd590300_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/libopenblas-0.3.26-pthreads_h413a1c8_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/libopenvino-2023.3.0-h2e90f83_2.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/libopenvino-auto-batch-plugin-2023.3.0-hd5fc58b_2.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/libopenvino-auto-plugin-2023.3.0-hd5fc58b_2.conda
@@ -85,8 +86,12 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/linux-64/libxcrypt-4.4.36-hd590300_1.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/libxml2-2.12.5-h232c23b_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/libzlib-1.2.13-hd590300_5.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/llvm-openmp-15.0.7-h0cdce71_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/markupsafe-2.1.5-py311h459d7ec_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/matplotlib-base-3.8.3-py311h54ef318_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/mkl-2023.0.0-h84fe81f_26648.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/mpc-1.3.1-hfe3b2da_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/mpfr-4.2.1-h9458935_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/ncurses-6.4-h59595ed_2.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/nettle-3.9.1-h7ab15ed_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/numpy-1.26.4-py311h64a7726_0.conda
@@ -139,6 +144,7 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/annotated-types-0.6.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/anyio-4.3.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/attrs-23.2.0-pyh71513ae_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/beautifulsoup4-4.12.3-pyha770c72_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/certifi-2024.2.2-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/charset-normalizer-3.3.2-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/click-8.1.7-unix_pyh707e725_0.conda
@@ -155,6 +161,7 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/fonts-conda-ecosystem-1-0.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/fonts-conda-forge-1-0.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/fsspec-2024.2.0-pyhca7485f_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/gdown-5.1.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/gradio-4.19.2-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/gradio-client-0.10.1-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/h11-0.14.0-pyhd8ed1ab_0.tar.bz2
@@ -173,7 +180,9 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/jsonschema-specifications-2023.12.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/markdown-it-py-3.0.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/mdurl-0.1.2-pyhd8ed1ab_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/mpmath-1.3.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/munkres-1.1.4-pyh9f0ad1d_0.tar.bz2
+      - conda: https://conda.anaconda.org/conda-forge/noarch/networkx-3.2.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/packaging-23.2-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pkgutil-resolve-name-1.3.10-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pydantic-2.6.3-pyhd8ed1ab_0.conda
@@ -192,7 +201,9 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/shellingham-1.5.4-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/six-1.16.0-pyh6c4a22f_0.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/sniffio-1.3.1-pyhd8ed1ab_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/soupsieve-2.5-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/starlette-0.27.0-pyhd8ed1ab_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/sympy-1.12-pypyh9d50eac_103.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/tomlkit-0.12.0-pyha770c72_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/toolz-0.12.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/tqdm-4.66.2-pyhd8ed1ab_0.conda
@@ -202,6 +213,28 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/tzdata-2024a-h0c530f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/urllib3-2.2.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/zipp-3.17.0-pyhd8ed1ab_0.conda
+      - conda: https://conda.anaconda.org/pytorch/linux-64/pytorch-2.2.1-py3.11_cpu_0.tar.bz2
+      - conda: https://conda.anaconda.org/pytorch/linux-64/torchvision-0.17.1-py311_cpu.tar.bz2
+      - conda: https://conda.anaconda.org/pytorch/noarch/pytorch-mutex-1.0-cpu.tar.bz2
+      - pypi: https://files.pythonhosted.org/packages/01/37/e25f147a4219931ba1703f29e4d445b3b772017e9d0fa817cfe341f5fdf8/gradio_imageslider-0.0.18-py3-none-any.whl#sha256=69d3de62f86fdc6480354c7a751a9a26d16cadbf0ce792f5fb179082349686bf
+      - pypi: https://files.pythonhosted.org/packages/1c/0e/ff7a83382d42fdf6f3f5e1ea4b6761d45945d9901ba6900167833252b709/gradio_modal-0.0.2-py3-none-any.whl#sha256=7fb1b30dd967e5d64251b763afc922e189f24cadc1ae428126575c26cdd42b83
+      - pypi: https://files.pythonhosted.org/packages/1e/07/bf730d44c2fe1b676ad9cc2be5f5f861eb5d153fb6951987a2d6a96379a9/nvidia_nvjitlink_cu12-12.3.101-py3-none-manylinux1_x86_64.whl#sha256=64335a8088e2b9d196ae8665430bc6a2b7e6ef2eb877a9c735c804bd4ff6467c
+      - pypi: https://files.pythonhosted.org/packages/2c/df/5810707da6f2fd4be57f0cc417987c0fa16a2eecf0b1b71f82ea555dc619/torch-2.2.1-cp311-cp311-manylinux1_x86_64.whl#sha256=d86664ec85902967d902e78272e97d1aff1d331f7619d398d3ffab1c9b8e9157
+      - pypi: https://files.pythonhosted.org/packages/37/6d/121efd7382d5b0284239f4ab1fc1590d86d34ed4a4a2fdb13b30ca8e5740/nvidia_cublas_cu12-12.1.3.1-py3-none-manylinux1_x86_64.whl#sha256=ee53ccca76a6fc08fb9701aa95b6ceb242cdaab118c3bb152af4e579af792728
+      - pypi: https://files.pythonhosted.org/packages/38/00/d0d4e48aef772ad5aebcf70b73028f88db6e5640b36c38e90445b7a57c45/nvidia_nccl_cu12-2.19.3-py3-none-manylinux1_x86_64.whl#sha256=a9734707a2c96443331c1e48c717024aa6678a0e2a4cb66b2c364d18cee6b48d
+      - pypi: https://files.pythonhosted.org/packages/3a/49/12fc5188602c68a789a0fdaee63d176a71ad5c1e34d25aeb8554abe46089/torchvision-0.17.1-cp311-cp311-manylinux1_x86_64.whl#sha256=5966936c669a08870f6547cd0a90d08b157aeda03293f79e2adbb934687175ed
+      - pypi: https://files.pythonhosted.org/packages/44/31/4890b1c9abc496303412947fc7dcea3d14861720642b49e8ceed89636705/nvidia_curand_cu12-10.3.2.106-py3-none-manylinux1_x86_64.whl#sha256=9d264c5036dde4e64f1de8c50ae753237c12e0b1348738169cd0f8a536c0e1e0
+      - pypi: https://files.pythonhosted.org/packages/65/5b/cfaeebf25cd9fdec14338ccb16f6b2c4c7fa9163aefcf057d86b9cc248bb/nvidia_cusparse_cu12-12.1.0.106-py3-none-manylinux1_x86_64.whl#sha256=f3b50f42cf363f86ab21f720998517a659a48131e8d538dc02f8768237bd884c
+      - pypi: https://files.pythonhosted.org/packages/7e/00/6b218edd739ecfc60524e585ba8e6b00554dd908de2c9c66c1af3e44e18d/nvidia_cuda_cupti_cu12-12.1.105-py3-none-manylinux1_x86_64.whl#sha256=e54fde3983165c624cb79254ae9818a456eb6e87a7fd4d56a2352c24ee542d7e
+      - pypi: https://files.pythonhosted.org/packages/86/94/eb540db023ce1d162e7bea9f8f5aa781d57c65aed513c33ee9a5123ead4d/nvidia_cufft_cu12-11.0.2.54-py3-none-manylinux1_x86_64.whl#sha256=794e3948a1aa71fd817c3775866943936774d1c14e7628c74f6f7417224cdf56
+      - pypi: https://files.pythonhosted.org/packages/b6/9f/c64c03f49d6fbc56196664d05dba14e3a561038a81a638eeb47f4d4cfd48/nvidia_cuda_nvrtc_cu12-12.1.105-py3-none-manylinux1_x86_64.whl#sha256=339b385f50c309763ca65456ec75e17bbefcbbf2893f462cb8b90584cd27a1c2
+      - pypi: https://files.pythonhosted.org/packages/bc/1d/8de1e5c67099015c834315e333911273a8c6aaba78923dd1d1e25fc5f217/nvidia_cusolver_cu12-11.4.5.107-py3-none-manylinux1_x86_64.whl#sha256=8a7ec542f0412294b15072fa7dab71d31334014a69f953004ea7a118206fe0dd
+      - pypi: https://files.pythonhosted.org/packages/bd/ac/3974caaa459bf2c3a244a84be8d17561f631f7d42af370fc311defeca2fb/triton-2.2.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl#sha256=da58a152bddb62cafa9a857dd2bc1f886dbf9f9c90a2b5da82157cd2b34392b0
+      - pypi: https://files.pythonhosted.org/packages/d7/a5/bbbc3b74a94fbdbd7915e7ad030f16539bfdc1362f7e9003b594f0537950/glob2-0.7.tar.gz#sha256=85c3dbd07c8aa26d63d7aacee34fa86e9a91a3873bc30bf62ec46e531f92ab8c
+      - pypi: https://files.pythonhosted.org/packages/da/d3/8057f0587683ed2fcd4dbfbdfdfa807b9160b809976099d36b8f60d08f03/nvidia_nvtx_cu12-12.1.105-py3-none-manylinux1_x86_64.whl#sha256=dc21cf308ca5691e7c04d962e213f8a4aa9bbfa23d95412f452254c2caeb09e5
+      - pypi: https://files.pythonhosted.org/packages/eb/d5/c68b1d2cdfcc59e72e8a5949a37ddb22ae6cade80cd4a57a84d4c8b55472/nvidia_cuda_runtime_cu12-12.1.105-py3-none-manylinux1_x86_64.whl#sha256=6e258468ddf5796e25f1dc591a31029fa317d97a0a94ed93468fc86301d61e40
+      - pypi: https://files.pythonhosted.org/packages/ed/1f/95cb961341bf33e534a472672ef4916c90046dc3c57b1442f5586dda72f1/geffnet-1.0.2-py3-none-any.whl#sha256=259013030bf429bc7a4fc4e8dc3fef6364e9750a6cee24d28e1674f8c3048d3a
+      - pypi: https://files.pythonhosted.org/packages/ff/74/a2e2be7fb83aaedec84f391f082cf765dfb635e7caa9b49065f73e4835d8/nvidia_cudnn_cu12-8.9.2.26-py3-none-manylinux1_x86_64.whl#sha256=5ccb288774fdfb07a7e7025ffec286971c06d8d7b4fb162525334616d7629ff9
 packages:
 - kind: conda
   name: _libgcc_mutex
@@ -217,21 +250,19 @@ packages:
 - kind: conda
   name: _openmp_mutex
   version: '4.5'
-  build: 2_gnu
-  build_number: 16
+  build: 2_kmp_llvm
+  build_number: 2
   subdir: linux-64
-  url: https://conda.anaconda.org/conda-forge/linux-64/_openmp_mutex-4.5-2_gnu.tar.bz2
-  sha256: fbe2c5e56a653bebb982eda4876a9178aedfc2b545f25d0ce9c4c0b508253d22
-  md5: 73aaf86a425cc6e73fcf236a5a46396d
+  url: https://conda.anaconda.org/conda-forge/linux-64/_openmp_mutex-4.5-2_kmp_llvm.tar.bz2
+  sha256: 84a66275da3a66e3f3e70e9d8f10496d807d01a9e4ec16cd2274cc5e28c478fc
+  md5: 562b26ba2e19059551a811e72ab7f793
   depends:
   - _libgcc_mutex 0.1 conda_forge
-  - libgomp >=7.5.0
-  constrains:
-  - openmp_impl 9999
+  - llvm-openmp >=9.0.1
   license: BSD-3-Clause
   license_family: BSD
-  size: 23621
-  timestamp: 1650670423406
+  size: 5744
+  timestamp: 1650742457817
 - kind: conda
   name: aiofiles
   version: 23.2.1
@@ -245,6 +276,8 @@ packages:
   - python >=3.7
   license: Apache-2.0
   license_family: Apache
+  purls:
+  - pkg:pypi/aiofiles
   size: 19507
   timestamp: 1698946114329
 - kind: conda
@@ -268,6 +301,8 @@ packages:
   - typing-extensions >=4.0.1
   license: BSD-3-Clause
   license_family: BSD
+  purls:
+  - pkg:pypi/altair
   size: 457454
   timestamp: 1701735672437
 - kind: conda
@@ -284,6 +319,8 @@ packages:
   - typing-extensions >=4.0.0
   license: MIT
   license_family: MIT
+  purls:
+  - pkg:pypi/annotated-types
   size: 17026
   timestamp: 1696634393637
 - kind: conda
@@ -306,6 +343,8 @@ packages:
   - uvloop >=0.17
   license: MIT
   license_family: MIT
+  purls:
+  - pkg:pypi/anyio
   size: 102331
   timestamp: 1708355504396
 - kind: conda
@@ -336,8 +375,40 @@ packages:
   - python >=3.7
   license: MIT
   license_family: MIT
+  purls:
+  - pkg:pypi/attrs
   size: 54582
   timestamp: 1704011393776
+- kind: conda
+  name: beautifulsoup4
+  version: 4.12.3
+  build: pyha770c72_0
+  subdir: noarch
+  noarch: python
+  url: https://conda.anaconda.org/conda-forge/noarch/beautifulsoup4-4.12.3-pyha770c72_0.conda
+  sha256: 7b05b2d0669029326c623b9df7a29fa49d1982a9e7e31b2fea34b4c9a4a72317
+  md5: 332493000404d8411859539a5a630865
+  depends:
+  - python >=3.6
+  - soupsieve >=1.2
+  license: MIT
+  license_family: MIT
+  size: 118200
+  timestamp: 1705564819537
+- kind: conda
+  name: blas
+  version: '1.0'
+  build: mkl
+  subdir: linux-64
+  url: https://conda.anaconda.org/conda-forge/linux-64/blas-1.0-mkl.tar.bz2
+  sha256: a9a9125029a66905fc9e932dfd4f595be3a59a30db37fd7bf4a675a5c6151d62
+  md5: 349aef876b1d8c9dccae01de20d5b385
+  depends:
+  - mkl
+  track_features:
+  - blas_mkl
+  license: BSD 3-clause
+  size: 1381
 - kind: conda
   name: brotli
   version: 1.1.0
@@ -475,6 +546,8 @@ packages:
   - python >=3.7
   license: MIT
   license_family: MIT
+  purls:
+  - pkg:pypi/charset-normalizer
   size: 46597
   timestamp: 1698833765762
 - kind: conda
@@ -491,6 +564,8 @@ packages:
   - python >=3.8
   license: BSD-3-Clause
   license_family: BSD
+  purls:
+  - pkg:pypi/click
   size: 84437
   timestamp: 1692311973840
 - kind: conda
@@ -506,6 +581,8 @@ packages:
   - python >=3.7
   license: BSD-3-Clause
   license_family: BSD
+  purls:
+  - pkg:pypi/colorama
   size: 25170
   timestamp: 1666700778190
 - kind: conda
@@ -524,6 +601,8 @@ packages:
   - python_abi 3.11.* *_cp311
   license: BSD-3-Clause
   license_family: BSD
+  purls:
+  - pkg:pypi/contourpy
   size: 255843
   timestamp: 1699041590533
 - kind: conda
@@ -539,6 +618,8 @@ packages:
   - python >=3.8
   license: BSD-3-Clause
   license_family: BSD
+  purls:
+  - pkg:pypi/cycler
   size: 13458
   timestamp: 1696677888423
 - kind: conda
@@ -568,6 +649,8 @@ packages:
   depends:
   - python >=3.7
   license: MIT and PSF-2.0
+  purls:
+  - pkg:pypi/exceptiongroup
   size: 20551
   timestamp: 1704921321122
 - kind: conda
@@ -602,6 +685,8 @@ packages:
   - typing-extensions >=4.5.0
   license: MIT
   license_family: MIT
+  purls:
+  - pkg:pypi/fastapi
   size: 56205
   timestamp: 1693161441022
 - kind: conda
@@ -670,6 +755,8 @@ packages:
   - python >=3.6
   license: MIT
   license_family: MIT
+  purls:
+  - pkg:pypi/ffmpy
   size: 10415
   timestamp: 1659475099796
 - kind: conda
@@ -684,6 +771,8 @@ packages:
   depends:
   - python >=3.7
   license: Unlicense
+  purls:
+  - pkg:pypi/filelock
   size: 15605
   timestamp: 1698715139726
 - kind: conda
@@ -806,6 +895,8 @@ packages:
   - python_abi 3.11.* *_cp311
   license: MIT
   license_family: MIT
+  purls:
+  - pkg:pypi/fonttools
   size: 2817685
   timestamp: 1708049363863
 - kind: conda
@@ -850,8 +941,40 @@ packages:
   - python >=3.8
   license: BSD-3-Clause
   license_family: BSD
+  purls:
+  - pkg:pypi/fsspec
   size: 128662
   timestamp: 1707102616711
+- kind: conda
+  name: gdown
+  version: 5.1.0
+  build: pyhd8ed1ab_0
+  subdir: noarch
+  noarch: python
+  url: https://conda.anaconda.org/conda-forge/noarch/gdown-5.1.0-pyhd8ed1ab_0.conda
+  sha256: 1ab1e5cf5c851f91abebfc6a6c094bc6e2afa3639e6586f6ff890acc8551a63d
+  md5: 6f880647c0270648f710f334c60bc76c
+  depends:
+  - beautifulsoup4
+  - filelock
+  - python >=3.8
+  - requests
+  - tqdm
+  license: MIT
+  license_family: MIT
+  purls:
+  - pkg:pypi/gdown
+  size: 21479
+  timestamp: 1706981727598
+- kind: pypi
+  name: geffnet
+  version: 1.0.2
+  url: https://files.pythonhosted.org/packages/ed/1f/95cb961341bf33e534a472672ef4916c90046dc3c57b1442f5586dda72f1/geffnet-1.0.2-py3-none-any.whl#sha256=259013030bf429bc7a4fc4e8dc3fef6364e9750a6cee24d28e1674f8c3048d3a
+  sha256: 259013030bf429bc7a4fc4e8dc3fef6364e9750a6cee24d28e1674f8c3048d3a
+  requires_dist:
+  - torch >=1.4
+  - torchvision
+  requires_python: '>=3.6'
 - kind: conda
   name: gettext
   version: 0.21.1
@@ -865,6 +988,11 @@ packages:
   license: LGPL-2.1-or-later AND GPL-3.0-or-later
   size: 4320628
   timestamp: 1665673494324
+- kind: pypi
+  name: glob2
+  version: '0.7'
+  url: https://files.pythonhosted.org/packages/d7/a5/bbbc3b74a94fbdbd7915e7ad030f16539bfdc1362f7e9003b594f0537950/glob2-0.7.tar.gz#sha256=85c3dbd07c8aa26d63d7aacee34fa86e9a91a3873bc30bf62ec46e531f92ab8c
+  sha256: 85c3dbd07c8aa26d63d7aacee34fa86e9a91a3873bc30bf62ec46e531f92ab8c
 - kind: conda
   name: gmp
   version: 6.3.0
@@ -879,6 +1007,28 @@ packages:
   license: GPL-2.0-or-later AND LGPL-3.0-or-later
   size: 563123
   timestamp: 1699629991732
+- kind: conda
+  name: gmpy2
+  version: 2.1.2
+  build: py311h6a5fa03_1
+  build_number: 1
+  subdir: linux-64
+  url: https://conda.anaconda.org/conda-forge/linux-64/gmpy2-2.1.2-py311h6a5fa03_1.tar.bz2
+  sha256: 20862200f4d07ba583ab6ae9b56d7de2462474240872100973711dfa20d562d7
+  md5: 3515bd4a3d92bbd3cc2d25aac335e34d
+  depends:
+  - gmp >=6.2.1,<7.0a0
+  - libgcc-ng >=12
+  - mpc >=1.2.1,<2.0a0
+  - mpfr >=4.1.0,<5.0a0
+  - python >=3.11,<3.12.0a0
+  - python_abi 3.11.* *_cp311
+  license: LGPL-3.0-or-later
+  license_family: LGPL
+  purls:
+  - pkg:pypi/gmpy2
+  size: 219751
+  timestamp: 1666808798297
 - kind: conda
   name: gnutls
   version: 3.7.9
@@ -936,6 +1086,8 @@ packages:
   - typing-extensions >=4.0,<5.dev0
   - uvicorn >=0.14.0
   license: MIT AND Apache-2.0
+  purls:
+  - pkg:pypi/gradio
   size: 12582364
   timestamp: 1708676706303
 - kind: conda
@@ -958,8 +1110,26 @@ packages:
   - websockets >=10.0,<12.0
   license: Apache-2.0
   license_family: APACHE
+  purls:
+  - pkg:pypi/gradio-client
   size: 299846
   timestamp: 1708677035723
+- kind: pypi
+  name: gradio-imageslider
+  version: 0.0.18
+  url: https://files.pythonhosted.org/packages/01/37/e25f147a4219931ba1703f29e4d445b3b772017e9d0fa817cfe341f5fdf8/gradio_imageslider-0.0.18-py3-none-any.whl#sha256=69d3de62f86fdc6480354c7a751a9a26d16cadbf0ce792f5fb179082349686bf
+  sha256: 69d3de62f86fdc6480354c7a751a9a26d16cadbf0ce792f5fb179082349686bf
+  requires_python: '>=3.8'
+- kind: pypi
+  name: gradio-modal
+  version: 0.0.2
+  url: https://files.pythonhosted.org/packages/1c/0e/ff7a83382d42fdf6f3f5e1ea4b6761d45945d9901ba6900167833252b709/gradio_modal-0.0.2-py3-none-any.whl#sha256=7fb1b30dd967e5d64251b763afc922e189f24cadc1ae428126575c26cdd42b83
+  sha256: 7fb1b30dd967e5d64251b763afc922e189f24cadc1ae428126575c26cdd42b83
+  requires_dist:
+  - gradio <5.0, >=4.0
+  - build ; extra == 'dev'
+  - twine ; extra == 'dev'
+  requires_python: '>=3.8'
 - kind: conda
   name: graphite2
   version: 1.3.13
@@ -989,6 +1159,8 @@ packages:
   - typing_extensions
   license: MIT
   license_family: MIT
+  purls:
+  - pkg:pypi/h11
   size: 48251
   timestamp: 1664132995560
 - kind: conda
@@ -1006,6 +1178,8 @@ packages:
   - python >=3.6.1
   license: MIT
   license_family: MIT
+  purls:
+  - pkg:pypi/h2
   size: 46754
   timestamp: 1634280590080
 - kind: conda
@@ -1041,6 +1215,8 @@ packages:
   - python
   license: MIT
   license_family: MIT
+  purls:
+  - pkg:pypi/hpack
   size: 25341
   timestamp: 1598856368685
 - kind: conda
@@ -1061,6 +1237,8 @@ packages:
   - sniffio 1.*
   license: BSD-3-Clause
   license_family: BSD
+  purls:
+  - pkg:pypi/httpcore
   size: 45727
   timestamp: 1708529429006
 - kind: conda
@@ -1081,6 +1259,8 @@ packages:
   - sniffio
   license: BSD-3-Clause
   license_family: BSD
+  purls:
+  - pkg:pypi/httpx
   size: 64651
   timestamp: 1708531043505
 - kind: conda
@@ -1103,6 +1283,8 @@ packages:
   - typing-extensions >=3.7.4.3
   license: Apache-2.0
   license_family: APACHE
+  purls:
+  - pkg:pypi/huggingface-hub
   size: 225834
   timestamp: 1709252375076
 - kind: conda
@@ -1118,6 +1300,8 @@ packages:
   - python >=3.6
   license: MIT
   license_family: MIT
+  purls:
+  - pkg:pypi/hyperframe
   size: 14646
   timestamp: 1619110249723
 - kind: conda
@@ -1148,6 +1332,8 @@ packages:
   - python >=3.6
   license: BSD-3-Clause
   license_family: BSD
+  purls:
+  - pkg:pypi/idna
   size: 50124
   timestamp: 1701027126206
 - kind: conda
@@ -1198,6 +1384,8 @@ packages:
   - importlib-resources >=6.1.2,<6.1.3.0a0
   license: Apache-2.0
   license_family: APACHE
+  purls:
+  - pkg:pypi/importlib-resources
   size: 30448
   timestamp: 1708948623752
 - kind: conda
@@ -1214,6 +1402,8 @@ packages:
   - python >=3.7
   license: BSD-3-Clause
   license_family: BSD
+  purls:
+  - pkg:pypi/jinja2
   size: 111589
   timestamp: 1704967140287
 - kind: conda
@@ -1252,6 +1442,8 @@ packages:
   - referencing >=0.31.0
   license: MIT
   license_family: MIT
+  purls:
+  - pkg:pypi/jsonschema-specifications
   size: 16431
   timestamp: 1703778502971
 - kind: conda
@@ -1270,6 +1462,8 @@ packages:
   - python_abi 3.11.* *_cp311
   license: BSD-3-Clause
   license_family: BSD
+  purls:
+  - pkg:pypi/kiwisolver
   size: 73273
   timestamp: 1695380140676
 - kind: conda
@@ -1376,24 +1570,22 @@ packages:
 - kind: conda
   name: libblas
   version: 3.9.0
-  build: 21_linux64_openblas
-  build_number: 21
+  build: 1_h86c2bf4_netlib
+  build_number: 1
   subdir: linux-64
-  url: https://conda.anaconda.org/conda-forge/linux-64/libblas-3.9.0-21_linux64_openblas.conda
-  sha256: ebd5c91f029f779fb88a1fcbd1e499559a9c258e3674ff58a2fbb4e375ae56d9
-  md5: 0ac9f44fc096772b0aa092119b00c3ca
+  url: https://conda.anaconda.org/conda-forge/linux-64/libblas-3.9.0-1_h86c2bf4_netlib.tar.bz2
+  sha256: c95bf82fbf258350d5f61c653a1a14d2e1cca7302bb348cc8347caa84e7ac9b1
+  md5: 91273b91f1023d5c72ffc79e02e7812e
   depends:
-  - libopenblas >=0.3.26,<0.3.27.0a0
-  - libopenblas >=0.3.26,<1.0a0
-  constrains:
-  - liblapacke 3.9.0 21_linux64_openblas
-  - blas * openblas
-  - libcblas 3.9.0 21_linux64_openblas
-  - liblapack 3.9.0 21_linux64_openblas
+  - libgcc-ng >=9.3.0
+  - libgfortran-ng
+  - libgfortran5 >=9.3.0
+  track_features:
+  - blas_netlib
   license: BSD-3-Clause
   license_family: BSD
-  size: 14691
-  timestamp: 1705979549006
+  size: 203305
+  timestamp: 1603052017387
 - kind: conda
   name: libbrotlicommon
   version: 1.1.0
@@ -1444,22 +1636,23 @@ packages:
 - kind: conda
   name: libcblas
   version: 3.9.0
-  build: 21_linux64_openblas
-  build_number: 21
+  build: 5_h92ddd45_netlib
+  build_number: 5
   subdir: linux-64
-  url: https://conda.anaconda.org/conda-forge/linux-64/libcblas-3.9.0-21_linux64_openblas.conda
-  sha256: 467bbfbfe1a1aeb8b1f9f6485eedd8ed1b6318941bf3702da72336ccf4dc25a6
-  md5: 4a3816d06451c4946e2db26b86472cb6
+  url: https://conda.anaconda.org/conda-forge/linux-64/libcblas-3.9.0-5_h92ddd45_netlib.tar.bz2
+  sha256: 93086c3586c5b1f961cb468995dc35d9b99ca8f10d0b76d594a554029f60670c
+  md5: 6a3f536ec30f6e6948211a07b1d04ced
   depends:
-  - libblas 3.9.0 21_linux64_openblas
-  constrains:
-  - liblapacke 3.9.0 21_linux64_openblas
-  - blas * openblas
-  - liblapack 3.9.0 21_linux64_openblas
+  - libblas 3.9.0.*
+  - libgcc-ng >=9.3.0
+  - libgfortran-ng
+  - libgfortran5 >=9.3.0
+  track_features:
+  - blas_netlib
   license: BSD-3-Clause
   license_family: BSD
-  size: 14614
-  timestamp: 1705979564122
+  size: 55516
+  timestamp: 1618011709500
 - kind: conda
   name: libdeflate
   version: '1.19'
@@ -1592,21 +1785,6 @@ packages:
   license: LGPL-2.1-or-later
   size: 2692079
   timestamp: 1708284870228
-- kind: conda
-  name: libgomp
-  version: 13.2.0
-  build: h807b86a_5
-  build_number: 5
-  subdir: linux-64
-  url: https://conda.anaconda.org/conda-forge/linux-64/libgomp-13.2.0-h807b86a_5.conda
-  sha256: 0d3d4b1b0134283ea02d58e8eb5accf3655464cf7159abf098cc694002f8d34e
-  md5: d211c42b9ce49aee3734fdc828731689
-  depends:
-  - _libgcc_mutex 0.1 conda_forge
-  license: GPL-3.0-only WITH GCC-exception-3.1
-  license_family: GPL
-  size: 419751
-  timestamp: 1706819107383
 - kind: conda
   name: libhwloc
   version: 2.9.3
@@ -1672,22 +1850,23 @@ packages:
 - kind: conda
   name: liblapack
   version: 3.9.0
-  build: 21_linux64_openblas
-  build_number: 21
+  build: 5_h92ddd45_netlib
+  build_number: 5
   subdir: linux-64
-  url: https://conda.anaconda.org/conda-forge/linux-64/liblapack-3.9.0-21_linux64_openblas.conda
-  sha256: 64b5c35dce00dd6f9f53178b2fe87116282e00967970bd6551a5a42923806ded
-  md5: 1a42f305615c3867684e049e85927531
+  url: https://conda.anaconda.org/conda-forge/linux-64/liblapack-3.9.0-5_h92ddd45_netlib.tar.bz2
+  sha256: 41eb83ea6ce409fd8c81fc38e029422840f581022e6f33998a680c0b23884cd6
+  md5: ffb80081cf8f43903b07045630188851
   depends:
-  - libblas 3.9.0 21_linux64_openblas
-  constrains:
-  - liblapacke 3.9.0 21_linux64_openblas
-  - libcblas 3.9.0 21_linux64_openblas
-  - blas * openblas
+  - libblas 3.9.0.*
+  - libgcc-ng >=9.3.0
+  - libgfortran-ng
+  - libgfortran5 >=9.3.0
+  track_features:
+  - blas_netlib
   license: BSD-3-Clause
   license_family: BSD
-  size: 14599
-  timestamp: 1705979579648
+  size: 3111106
+  timestamp: 1618011727809
 - kind: conda
   name: libnsl
   version: 2.0.1
@@ -1702,24 +1881,6 @@ packages:
   license_family: GPL
   size: 33408
   timestamp: 1697359010159
-- kind: conda
-  name: libopenblas
-  version: 0.3.26
-  build: pthreads_h413a1c8_0
-  subdir: linux-64
-  url: https://conda.anaconda.org/conda-forge/linux-64/libopenblas-0.3.26-pthreads_h413a1c8_0.conda
-  sha256: b626954b5a1113dafec8df89fa8bf18ce9b4701464d9f084ddd7fc9fac404bbd
-  md5: 760ae35415f5ba8b15d09df5afe8b23a
-  depends:
-  - libgcc-ng >=12
-  - libgfortran-ng
-  - libgfortran5 >=12.3.0
-  constrains:
-  - openblas >=0.3.26,<0.3.27.0a0
-  license: BSD-3-Clause
-  license_family: BSD
-  size: 5578031
-  timestamp: 1704950143521
 - kind: conda
   name: libopenvino
   version: 2023.3.0
@@ -2183,6 +2344,22 @@ packages:
   license_family: Other
   size: 61588
   timestamp: 1686575217516
+- kind: conda
+  name: llvm-openmp
+  version: 15.0.7
+  build: h0cdce71_0
+  subdir: linux-64
+  url: https://conda.anaconda.org/conda-forge/linux-64/llvm-openmp-15.0.7-h0cdce71_0.conda
+  sha256: 7c67d383a8b1f3e7bf9e046e785325c481f6868194edcfb9d78d261da4ad65d4
+  md5: 589c9a3575a050b583241c3d688ad9aa
+  depends:
+  - libzlib >=1.2.13,<1.3.0a0
+  constrains:
+  - openmp 15.0.7|15.0.7.*
+  license: Apache-2.0 WITH LLVM-exception
+  license_family: APACHE
+  size: 3268766
+  timestamp: 1673584331056
 - kind: conda
   name: markdown-it-py
   version: 3.0.0
@@ -2197,6 +2374,8 @@ packages:
   - python >=3.8
   license: MIT
   license_family: MIT
+  purls:
+  - pkg:pypi/markdown-it-py
   size: 64356
   timestamp: 1686175179621
 - kind: conda
@@ -2215,6 +2394,8 @@ packages:
   - jinja2 >=3.0.0
   license: BSD-3-Clause
   license_family: BSD
+  purls:
+  - pkg:pypi/markupsafe
   size: 27502
   timestamp: 1706900084436
 - kind: conda
@@ -2245,6 +2426,8 @@ packages:
   - tk >=8.6.13,<8.7.0a0
   license: PSF-2.0
   license_family: PSF
+  purls:
+  - pkg:pypi/matplotlib
   size: 7927557
   timestamp: 1708026755428
 - kind: conda
@@ -2260,8 +2443,76 @@ packages:
   - python >=3.6
   license: MIT
   license_family: MIT
+  purls:
+  - pkg:pypi/mdurl
   size: 14680
   timestamp: 1704317789138
+- kind: conda
+  name: mkl
+  version: 2023.0.0
+  build: h84fe81f_26648
+  build_number: 26648
+  subdir: linux-64
+  url: https://conda.anaconda.org/conda-forge/linux-64/mkl-2023.0.0-h84fe81f_26648.conda
+  sha256: 989c614ec8417a22f61e9d4ec920c88910bbc3239885a6f20ad4eacad09f277c
+  md5: 2a37425a272410241c7613c38e3955b2
+  depends:
+  - _openmp_mutex * *_llvm
+  - _openmp_mutex >=4.5
+  - llvm-openmp >=15.0.7
+  - tbb 2021.*
+  license: LicenseRef-ProprietaryIntel
+  license_family: Proprietary
+  size: 161440193
+  timestamp: 1674673370494
+- kind: conda
+  name: mpc
+  version: 1.3.1
+  build: hfe3b2da_0
+  subdir: linux-64
+  url: https://conda.anaconda.org/conda-forge/linux-64/mpc-1.3.1-hfe3b2da_0.conda
+  sha256: 2f88965949ba7b4b21e7e5facd62285f7c6efdb17359d1b365c3bb4ecc968d29
+  md5: 289c71e83dc0daa7d4c81f04180778ca
+  depends:
+  - gmp >=6.2.1,<7.0a0
+  - libgcc-ng >=12
+  - mpfr >=4.1.0,<5.0a0
+  license: LGPL-3.0-or-later
+  license_family: LGPL
+  size: 116276
+  timestamp: 1674263855481
+- kind: conda
+  name: mpfr
+  version: 4.2.1
+  build: h9458935_0
+  subdir: linux-64
+  url: https://conda.anaconda.org/conda-forge/linux-64/mpfr-4.2.1-h9458935_0.conda
+  sha256: 008230a53ff15cf61966476b44f7ba2c779826825b9ca639a0a2b44d8f7aa6cb
+  md5: 4c28f3210b30250037a4a627eeee9e0f
+  depends:
+  - gmp >=6.2.1,<7.0a0
+  - libgcc-ng >=12
+  license: LGPL-3.0-only
+  license_family: LGPL
+  size: 641512
+  timestamp: 1698004072487
+- kind: conda
+  name: mpmath
+  version: 1.3.0
+  build: pyhd8ed1ab_0
+  subdir: noarch
+  noarch: python
+  url: https://conda.anaconda.org/conda-forge/noarch/mpmath-1.3.0-pyhd8ed1ab_0.conda
+  sha256: a4f025c712ec1502a55c471b56a640eaeebfce38dd497d5a1a33729014cac47a
+  md5: dbf6e2d89137da32fa6670f3bffc024e
+  depends:
+  - python >=3.6
+  license: BSD-3-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/mpmath
+  size: 438339
+  timestamp: 1678228210181
 - kind: conda
   name: munkres
   version: 1.1.4
@@ -2305,6 +2556,28 @@ packages:
   license_family: GPL
   size: 1011638
   timestamp: 1686309814836
+- kind: conda
+  name: networkx
+  version: 3.2.1
+  build: pyhd8ed1ab_0
+  subdir: noarch
+  noarch: python
+  url: https://conda.anaconda.org/conda-forge/noarch/networkx-3.2.1-pyhd8ed1ab_0.conda
+  sha256: 7629aa4f9f8cdff45ea7a4701fe58dccce5bf2faa01c26eb44cbb27b7e15ca9d
+  md5: 425fce3b531bed6ec3c74fab3e5f0a1c
+  depends:
+  - python >=3.9
+  constrains:
+  - matplotlib >=3.5
+  - scipy >=1.9,!=1.11.0,!=1.11.1
+  - numpy >=1.22
+  - pandas >=1.4
+  license: BSD-3-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/networkx
+  size: 1149552
+  timestamp: 1698504905258
 - kind: conda
   name: numpy
   version: 1.26.4
@@ -2325,8 +2598,90 @@ packages:
   - numpy-base <0a0
   license: BSD-3-Clause
   license_family: BSD
+  purls:
+  - pkg:pypi/numpy
   size: 8065890
   timestamp: 1707225944355
+- kind: pypi
+  name: nvidia-cublas-cu12
+  version: 12.1.3.1
+  url: https://files.pythonhosted.org/packages/37/6d/121efd7382d5b0284239f4ab1fc1590d86d34ed4a4a2fdb13b30ca8e5740/nvidia_cublas_cu12-12.1.3.1-py3-none-manylinux1_x86_64.whl#sha256=ee53ccca76a6fc08fb9701aa95b6ceb242cdaab118c3bb152af4e579af792728
+  sha256: ee53ccca76a6fc08fb9701aa95b6ceb242cdaab118c3bb152af4e579af792728
+  requires_python: '>=3'
+- kind: pypi
+  name: nvidia-cuda-cupti-cu12
+  version: 12.1.105
+  url: https://files.pythonhosted.org/packages/7e/00/6b218edd739ecfc60524e585ba8e6b00554dd908de2c9c66c1af3e44e18d/nvidia_cuda_cupti_cu12-12.1.105-py3-none-manylinux1_x86_64.whl#sha256=e54fde3983165c624cb79254ae9818a456eb6e87a7fd4d56a2352c24ee542d7e
+  sha256: e54fde3983165c624cb79254ae9818a456eb6e87a7fd4d56a2352c24ee542d7e
+  requires_python: '>=3'
+- kind: pypi
+  name: nvidia-cuda-nvrtc-cu12
+  version: 12.1.105
+  url: https://files.pythonhosted.org/packages/b6/9f/c64c03f49d6fbc56196664d05dba14e3a561038a81a638eeb47f4d4cfd48/nvidia_cuda_nvrtc_cu12-12.1.105-py3-none-manylinux1_x86_64.whl#sha256=339b385f50c309763ca65456ec75e17bbefcbbf2893f462cb8b90584cd27a1c2
+  sha256: 339b385f50c309763ca65456ec75e17bbefcbbf2893f462cb8b90584cd27a1c2
+  requires_python: '>=3'
+- kind: pypi
+  name: nvidia-cuda-runtime-cu12
+  version: 12.1.105
+  url: https://files.pythonhosted.org/packages/eb/d5/c68b1d2cdfcc59e72e8a5949a37ddb22ae6cade80cd4a57a84d4c8b55472/nvidia_cuda_runtime_cu12-12.1.105-py3-none-manylinux1_x86_64.whl#sha256=6e258468ddf5796e25f1dc591a31029fa317d97a0a94ed93468fc86301d61e40
+  sha256: 6e258468ddf5796e25f1dc591a31029fa317d97a0a94ed93468fc86301d61e40
+  requires_python: '>=3'
+- kind: pypi
+  name: nvidia-cudnn-cu12
+  version: 8.9.2.26
+  url: https://files.pythonhosted.org/packages/ff/74/a2e2be7fb83aaedec84f391f082cf765dfb635e7caa9b49065f73e4835d8/nvidia_cudnn_cu12-8.9.2.26-py3-none-manylinux1_x86_64.whl#sha256=5ccb288774fdfb07a7e7025ffec286971c06d8d7b4fb162525334616d7629ff9
+  sha256: 5ccb288774fdfb07a7e7025ffec286971c06d8d7b4fb162525334616d7629ff9
+  requires_dist:
+  - nvidia-cublas-cu12
+  requires_python: '>=3'
+- kind: pypi
+  name: nvidia-cufft-cu12
+  version: 11.0.2.54
+  url: https://files.pythonhosted.org/packages/86/94/eb540db023ce1d162e7bea9f8f5aa781d57c65aed513c33ee9a5123ead4d/nvidia_cufft_cu12-11.0.2.54-py3-none-manylinux1_x86_64.whl#sha256=794e3948a1aa71fd817c3775866943936774d1c14e7628c74f6f7417224cdf56
+  sha256: 794e3948a1aa71fd817c3775866943936774d1c14e7628c74f6f7417224cdf56
+  requires_python: '>=3'
+- kind: pypi
+  name: nvidia-curand-cu12
+  version: 10.3.2.106
+  url: https://files.pythonhosted.org/packages/44/31/4890b1c9abc496303412947fc7dcea3d14861720642b49e8ceed89636705/nvidia_curand_cu12-10.3.2.106-py3-none-manylinux1_x86_64.whl#sha256=9d264c5036dde4e64f1de8c50ae753237c12e0b1348738169cd0f8a536c0e1e0
+  sha256: 9d264c5036dde4e64f1de8c50ae753237c12e0b1348738169cd0f8a536c0e1e0
+  requires_python: '>=3'
+- kind: pypi
+  name: nvidia-cusolver-cu12
+  version: 11.4.5.107
+  url: https://files.pythonhosted.org/packages/bc/1d/8de1e5c67099015c834315e333911273a8c6aaba78923dd1d1e25fc5f217/nvidia_cusolver_cu12-11.4.5.107-py3-none-manylinux1_x86_64.whl#sha256=8a7ec542f0412294b15072fa7dab71d31334014a69f953004ea7a118206fe0dd
+  sha256: 8a7ec542f0412294b15072fa7dab71d31334014a69f953004ea7a118206fe0dd
+  requires_dist:
+  - nvidia-cublas-cu12
+  - nvidia-nvjitlink-cu12
+  - nvidia-cusparse-cu12
+  requires_python: '>=3'
+- kind: pypi
+  name: nvidia-cusparse-cu12
+  version: 12.1.0.106
+  url: https://files.pythonhosted.org/packages/65/5b/cfaeebf25cd9fdec14338ccb16f6b2c4c7fa9163aefcf057d86b9cc248bb/nvidia_cusparse_cu12-12.1.0.106-py3-none-manylinux1_x86_64.whl#sha256=f3b50f42cf363f86ab21f720998517a659a48131e8d538dc02f8768237bd884c
+  sha256: f3b50f42cf363f86ab21f720998517a659a48131e8d538dc02f8768237bd884c
+  requires_dist:
+  - nvidia-nvjitlink-cu12
+  requires_python: '>=3'
+- kind: pypi
+  name: nvidia-nccl-cu12
+  version: 2.19.3
+  url: https://files.pythonhosted.org/packages/38/00/d0d4e48aef772ad5aebcf70b73028f88db6e5640b36c38e90445b7a57c45/nvidia_nccl_cu12-2.19.3-py3-none-manylinux1_x86_64.whl#sha256=a9734707a2c96443331c1e48c717024aa6678a0e2a4cb66b2c364d18cee6b48d
+  sha256: a9734707a2c96443331c1e48c717024aa6678a0e2a4cb66b2c364d18cee6b48d
+  requires_python: '>=3'
+- kind: pypi
+  name: nvidia-nvjitlink-cu12
+  version: 12.3.101
+  url: https://files.pythonhosted.org/packages/1e/07/bf730d44c2fe1b676ad9cc2be5f5f861eb5d153fb6951987a2d6a96379a9/nvidia_nvjitlink_cu12-12.3.101-py3-none-manylinux1_x86_64.whl#sha256=64335a8088e2b9d196ae8665430bc6a2b7e6ef2eb877a9c735c804bd4ff6467c
+  sha256: 64335a8088e2b9d196ae8665430bc6a2b7e6ef2eb877a9c735c804bd4ff6467c
+  requires_python: '>=3'
+- kind: pypi
+  name: nvidia-nvtx-cu12
+  version: 12.1.105
+  url: https://files.pythonhosted.org/packages/da/d3/8057f0587683ed2fcd4dbfbdfdfa807b9160b809976099d36b8f60d08f03/nvidia_nvtx_cu12-12.1.105-py3-none-manylinux1_x86_64.whl#sha256=dc21cf308ca5691e7c04d962e213f8a4aa9bbfa23d95412f452254c2caeb09e5
+  sha256: dc21cf308ca5691e7c04d962e213f8a4aa9bbfa23d95412f452254c2caeb09e5
+  requires_python: '>=3'
 - kind: conda
   name: ocl-icd
   version: 2.3.2
@@ -2405,6 +2760,8 @@ packages:
   - python_abi 3.11.* *_cp311
   license: Apache-2.0
   license_family: Apache
+  purls:
+  - pkg:pypi/orjson
   size: 274831
   timestamp: 1708717354229
 - kind: conda
@@ -2436,6 +2793,8 @@ packages:
   - python >=3.7
   license: Apache-2.0
   license_family: APACHE
+  purls:
+  - pkg:pypi/packaging
   size: 49452
   timestamp: 1696202521121
 - kind: conda
@@ -2527,6 +2886,8 @@ packages:
   depends:
   - python >=3.6
   license: MIT AND PSF-2.0
+  purls:
+  - pkg:pypi/pkgutil-resolve-name
   size: 10778
   timestamp: 1694617398467
 - kind: conda
@@ -2575,6 +2936,8 @@ packages:
   - typing-extensions >=4.6.1
   license: MIT
   license_family: MIT
+  purls:
+  - pkg:pypi/pydantic
   size: 271902
   timestamp: 1709075341323
 - kind: conda
@@ -2592,6 +2955,8 @@ packages:
   - typing-extensions >=4.6.0,!=4.7.0
   license: MIT
   license_family: MIT
+  purls:
+  - pkg:pypi/pydantic-core
   size: 1659194
   timestamp: 1708700970198
 - kind: conda
@@ -2607,6 +2972,8 @@ packages:
   - python >=3.6
   license: MIT
   license_family: MIT
+  purls:
+  - pkg:pypi/pydub
   size: 31190
   timestamp: 1615612572693
 - kind: conda
@@ -2622,6 +2989,8 @@ packages:
   - python >=3.7
   license: BSD-2-Clause
   license_family: BSD
+  purls:
+  - pkg:pypi/pygments
   size: 860425
   timestamp: 1700608076927
 - kind: conda
@@ -2637,6 +3006,8 @@ packages:
   - python >=3.6
   license: MIT
   license_family: MIT
+  purls:
+  - pkg:pypi/pyparsing
   size: 89521
   timestamp: 1690737983548
 - kind: conda
@@ -2654,6 +3025,8 @@ packages:
   - python >=3.8
   license: BSD-3-Clause
   license_family: BSD
+  purls:
+  - pkg:pypi/pysocks
   size: 18981
   timestamp: 1661604969727
 - kind: conda
@@ -2700,6 +3073,8 @@ packages:
   - six >=1.5
   license: Apache-2.0
   license_family: APACHE
+  purls:
+  - pkg:pypi/python-dateutil
   size: 222742
   timestamp: 1709299922152
 - kind: conda
@@ -2715,6 +3090,8 @@ packages:
   - python >=3.7
   license: Apache-2.0
   license_family: Apache
+  purls:
+  - pkg:pypi/python-multipart
   size: 26439
   timestamp: 1707760278735
 - kind: conda
@@ -2730,6 +3107,8 @@ packages:
   - python >=3.6
   license: Apache-2.0
   license_family: APACHE
+  purls:
+  - pkg:pypi/tzdata
   size: 144024
   timestamp: 1707747742930
 - kind: conda
@@ -2747,6 +3126,45 @@ packages:
   license_family: BSD
   size: 6385
   timestamp: 1695147338551
+- kind: conda
+  name: pytorch
+  version: 2.2.1
+  build: py3.11_cpu_0
+  subdir: linux-64
+  url: https://conda.anaconda.org/pytorch/linux-64/pytorch-2.2.1-py3.11_cpu_0.tar.bz2
+  sha256: a7397aaeb0a67d73bfca7f3e6c1a29286ff57fc7e6c3b510fe9271ab906f0848
+  md5: 00b065484baa6166c50a5f067980edca
+  depends:
+  - blas * mkl
+  - filelock
+  - jinja2
+  - llvm-openmp <16
+  - mkl >=2018
+  - networkx
+  - python >=3.11,<3.12.0a0
+  - pytorch-mutex 1.0 cpu
+  - pyyaml
+  - sympy
+  - typing_extensions
+  constrains:
+  - cpuonly
+  license: BSD 3-Clause
+  license_family: BSD
+  size: 86421712
+  timestamp: 1708027733815
+- kind: conda
+  name: pytorch-mutex
+  version: '1.0'
+  build: cpu
+  subdir: noarch
+  noarch: generic
+  url: https://conda.anaconda.org/pytorch/noarch/pytorch-mutex-1.0-cpu.tar.bz2
+  sha256: d48c964188ca49660d750cffd73698d217cf94e694cd51987f9f186425435e76
+  md5: 49565ed726991fd28d08a39885caa88d
+  track_features:
+  - pytorch-mutex
+  size: 2897
+  timestamp: 1628062935164
 - kind: conda
   name: pytz
   version: '2024.1'
@@ -2760,6 +3178,8 @@ packages:
   - python >=3.7
   license: MIT
   license_family: MIT
+  purls:
+  - pkg:pypi/pytz
   size: 188538
   timestamp: 1706886944988
 - kind: conda
@@ -2811,6 +3231,8 @@ packages:
   - rpds-py >=0.7.0
   license: MIT
   license_family: MIT
+  purls:
+  - pkg:pypi/referencing
   size: 39055
   timestamp: 1706711589688
 - kind: conda
@@ -2832,6 +3254,8 @@ packages:
   - chardet >=3.0.2,<6
   license: Apache-2.0
   license_family: APACHE
+  purls:
+  - pkg:pypi/requests
   size: 56690
   timestamp: 1684774408600
 - kind: conda
@@ -2866,6 +3290,8 @@ packages:
   - python_abi 3.11.* *_cp311
   license: MIT
   license_family: MIT
+  purls:
+  - pkg:pypi/rpds-py
   size: 915849
   timestamp: 1707923007711
 - kind: conda
@@ -2883,6 +3309,8 @@ packages:
   - python_abi 3.11.* *_cp311
   license: MIT
   license_family: MIT
+  purls:
+  - pkg:pypi/ruff
   size: 5738078
   timestamp: 1709229162825
 - kind: conda
@@ -2898,6 +3326,8 @@ packages:
   - python >=2.7
   license: BSD-2-Clause
   license_family: BSD
+  purls:
+  - pkg:pypi/semantic-version
   size: 17923
   timestamp: 1653579450384
 - kind: conda
@@ -2913,6 +3343,8 @@ packages:
   - python >=3.7
   license: MIT
   license_family: MIT
+  purls:
+  - pkg:pypi/shellingham
   size: 14568
   timestamp: 1698144516278
 - kind: conda
@@ -2928,6 +3360,8 @@ packages:
   - python
   license: MIT
   license_family: MIT
+  purls:
+  - pkg:pypi/six
   size: 14259
   timestamp: 1620240338595
 - kind: conda
@@ -2958,8 +3392,28 @@ packages:
   - python >=3.7
   license: Apache-2.0
   license_family: Apache
+  purls:
+  - pkg:pypi/sniffio
   size: 15064
   timestamp: 1708953086199
+- kind: conda
+  name: soupsieve
+  version: '2.5'
+  build: pyhd8ed1ab_1
+  build_number: 1
+  subdir: noarch
+  noarch: python
+  url: https://conda.anaconda.org/conda-forge/noarch/soupsieve-2.5-pyhd8ed1ab_1.conda
+  sha256: 54ae221033db8fbcd4998ccb07f3c3828b4d77e73b0c72b18c1d6a507059059c
+  md5: 3f144b2c34f8cb5a9abd9ed23a39c561
+  depends:
+  - python >=3.8
+  license: MIT
+  license_family: MIT
+  purls:
+  - pkg:pypi/soupsieve
+  size: 36754
+  timestamp: 1693929424267
 - kind: conda
   name: starlette
   version: 0.27.0
@@ -2992,6 +3446,28 @@ packages:
   license_family: BSD
   size: 2644060
   timestamp: 1702359203835
+- kind: conda
+  name: sympy
+  version: '1.12'
+  build: pypyh9d50eac_103
+  build_number: 103
+  subdir: noarch
+  noarch: python
+  url: https://conda.anaconda.org/conda-forge/noarch/sympy-1.12-pypyh9d50eac_103.conda
+  sha256: 0025dd4e6411423903bf478d1b9fbff0cbbbe546f51c9375dfd6729ef2e1a1ac
+  md5: 2f7d6347d7acf6edf1ac7f2189f44c8f
+  depends:
+  - __unix
+  - gmpy2 >=2.0.8
+  - mpmath >=0.19
+  - python * *_cpython
+  - python >=3.8
+  license: BSD-3-Clause
+  license_family: BSD
+  purls:
+  - pkg:pypi/sympy
+  size: 4256289
+  timestamp: 1684180689319
 - kind: conda
   name: tbb
   version: 2021.11.0
@@ -3038,6 +3514,8 @@ packages:
   - python >=3.7
   license: MIT
   license_family: MIT
+  purls:
+  - pkg:pypi/tomlkit
   size: 37086
   timestamp: 1690458478787
 - kind: conda
@@ -3055,6 +3533,67 @@ packages:
   license_family: BSD
   size: 52358
   timestamp: 1706112720607
+- kind: pypi
+  name: torch
+  version: 2.2.1
+  url: https://files.pythonhosted.org/packages/2c/df/5810707da6f2fd4be57f0cc417987c0fa16a2eecf0b1b71f82ea555dc619/torch-2.2.1-cp311-cp311-manylinux1_x86_64.whl#sha256=d86664ec85902967d902e78272e97d1aff1d331f7619d398d3ffab1c9b8e9157
+  sha256: d86664ec85902967d902e78272e97d1aff1d331f7619d398d3ffab1c9b8e9157
+  requires_dist:
+  - filelock
+  - typing-extensions >=4.8.0
+  - sympy
+  - networkx
+  - jinja2
+  - fsspec
+  - nvidia-cuda-nvrtc-cu12 ==12.1.105 ; platform_system == 'Linux' and platform_machine == 'x86_64'
+  - nvidia-cuda-runtime-cu12 ==12.1.105 ; platform_system == 'Linux' and platform_machine == 'x86_64'
+  - nvidia-cuda-cupti-cu12 ==12.1.105 ; platform_system == 'Linux' and platform_machine == 'x86_64'
+  - nvidia-cudnn-cu12 ==8.9.2.26 ; platform_system == 'Linux' and platform_machine == 'x86_64'
+  - nvidia-cublas-cu12 ==12.1.3.1 ; platform_system == 'Linux' and platform_machine == 'x86_64'
+  - nvidia-cufft-cu12 ==11.0.2.54 ; platform_system == 'Linux' and platform_machine == 'x86_64'
+  - nvidia-curand-cu12 ==10.3.2.106 ; platform_system == 'Linux' and platform_machine == 'x86_64'
+  - nvidia-cusolver-cu12 ==11.4.5.107 ; platform_system == 'Linux' and platform_machine == 'x86_64'
+  - nvidia-cusparse-cu12 ==12.1.0.106 ; platform_system == 'Linux' and platform_machine == 'x86_64'
+  - nvidia-nccl-cu12 ==2.19.3 ; platform_system == 'Linux' and platform_machine == 'x86_64'
+  - nvidia-nvtx-cu12 ==12.1.105 ; platform_system == 'Linux' and platform_machine == 'x86_64'
+  - triton ==2.2.0 ; platform_system == 'Linux' and platform_machine == 'x86_64' and python_version < '3.12'
+  - opt-einsum >=3.3 ; extra == 'opt-einsum'
+  - optree >=0.9.1 ; extra == 'optree'
+  requires_python: '>=3.8.0'
+- kind: pypi
+  name: torchvision
+  version: 0.17.1
+  url: https://files.pythonhosted.org/packages/3a/49/12fc5188602c68a789a0fdaee63d176a71ad5c1e34d25aeb8554abe46089/torchvision-0.17.1-cp311-cp311-manylinux1_x86_64.whl#sha256=5966936c669a08870f6547cd0a90d08b157aeda03293f79e2adbb934687175ed
+  sha256: 5966936c669a08870f6547cd0a90d08b157aeda03293f79e2adbb934687175ed
+  requires_dist:
+  - numpy
+  - torch ==2.2.1
+  - pillow !=8.3.*, >=5.3.0
+  - scipy ; extra == 'scipy'
+  requires_python: '>=3.8'
+- kind: conda
+  name: torchvision
+  version: 0.17.1
+  build: py311_cpu
+  subdir: linux-64
+  url: https://conda.anaconda.org/pytorch/linux-64/torchvision-0.17.1-py311_cpu.tar.bz2
+  sha256: 4ec2bd12953402d2304d3ae5f752d3597a3d3336ad523dff4f6c4107e27a80ab
+  md5: c30958c550b69bba88314247da8a6a5d
+  depends:
+  - ffmpeg >=4.2
+  - libjpeg-turbo
+  - libpng
+  - numpy >=1.23.5
+  - pillow >=5.3.0,!=8.3.*
+  - python >=3.11,<3.12.0a0
+  - pytorch 2.2.1
+  - pytorch-mutex 1.0 cpu
+  - requests
+  constrains:
+  - cpuonly
+  license: BSD
+  size: 7351821
+  timestamp: 1708371226538
 - kind: conda
   name: tqdm
   version: 4.66.2
@@ -3068,8 +3607,30 @@ packages:
   - colorama
   - python >=3.7
   license: MPL-2.0 or MIT
+  purls:
+  - pkg:pypi/tqdm
   size: 89567
   timestamp: 1707598746354
+- kind: pypi
+  name: triton
+  version: 2.2.0
+  url: https://files.pythonhosted.org/packages/bd/ac/3974caaa459bf2c3a244a84be8d17561f631f7d42af370fc311defeca2fb/triton-2.2.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl#sha256=da58a152bddb62cafa9a857dd2bc1f886dbf9f9c90a2b5da82157cd2b34392b0
+  sha256: da58a152bddb62cafa9a857dd2bc1f886dbf9f9c90a2b5da82157cd2b34392b0
+  requires_dist:
+  - filelock
+  - cmake >=3.20 ; extra == 'build'
+  - lit ; extra == 'build'
+  - autopep8 ; extra == 'tests'
+  - flake8 ; extra == 'tests'
+  - isort ; extra == 'tests'
+  - numpy ; extra == 'tests'
+  - pytest ; extra == 'tests'
+  - scipy >=1.7.1 ; extra == 'tests'
+  - torch ; extra == 'tests'
+  - matplotlib ; extra == 'tutorials'
+  - pandas ; extra == 'tutorials'
+  - tabulate ; extra == 'tutorials'
+  - torch ; extra == 'tutorials'
 - kind: conda
   name: typer
   version: 0.9.0
@@ -3088,6 +3649,8 @@ packages:
   - typing-extensions >=3.7.4.3
   license: MIT
   license_family: MIT
+  purls:
+  - pkg:pypi/typer
   size: 79064
   timestamp: 1683029411797
 - kind: conda
@@ -3118,6 +3681,8 @@ packages:
   - python >=3.8
   license: PSF-2.0
   license_family: PSF
+  purls:
+  - pkg:pypi/typing-extensions
   size: 37018
   timestamp: 1708904796013
 - kind: conda
@@ -3147,6 +3712,8 @@ packages:
   - python >=3.7
   license: MIT
   license_family: MIT
+  purls:
+  - pkg:pypi/urllib3
   size: 94669
   timestamp: 1708239595549
 - kind: conda
@@ -3181,6 +3748,8 @@ packages:
   - python_abi 3.11.* *_cp311
   license: BSD-3-Clause
   license_family: BSD
+  purls:
+  - pkg:pypi/websockets
   size: 200745
   timestamp: 1695410107238
 - kind: conda
@@ -3457,6 +4026,8 @@ packages:
   - python >=3.8
   license: MIT
   license_family: MIT
+  purls:
+  - pkg:pypi/zipp
   size: 18954
   timestamp: 1695255262261
 - kind: conda

pixi.toml

@@ -3,12 +3,22 @@ name = "DSINE-space"
 version = "0.1.0"
 description = "Add a short description here"
 authors = ["pablovela5620 <pablovela5620@gmail.com>"]
-channels = ["conda-forge"]
+channels = ["conda-forge", "pytorch"]
 platforms = ["linux-64"]
 
 [tasks]
-start="python main.py"
+download-ckpt = "gdown 1B60LahYi3IuHu0nP2EySEWH29guWavZI --output checkpoints/"
+start={cmd="python main.py", depends_on=["download-ckpt"]}
 
 [dependencies]
 python = "3.11.*"
+pytorch = {version = ">=2.2.0,<2.3", channel="pytorch"}
+torchvision = {version = ">=0.17.0,<0.18", channel="pytorch"}
 gradio = ">=4.19.2,<4.20"
+gdown = ">=5.1.0,<5.2"
+
+[pypi-dependencies]
+gradio-modal = "*"
+gradio-imageslider = "*"
+geffnet = "*"
+glob2 = "*"

test.py

@@ -0,0 +1,86 @@
+import os
+import sys
+import glob
+import argparse
+import numpy as np
+
+import torch
+import torch.nn.functional as F
+from torchvision import transforms
+from PIL import Image
+import utils.utils as utils
+
+
+def test_samples(args, model, intrins=None, device="cpu"):
+    img_paths = glob.glob("./samples/img/*.png") + glob.glob("./samples/img/*.jpg")
+    img_paths.sort()
+
+    # normalize
+    normalize = transforms.Normalize(
+        mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
+    )
+
+    with torch.no_grad():
+        for img_path in img_paths:
+            print(img_path)
+            ext = os.path.splitext(img_path)[1]
+            img = Image.open(img_path).convert("RGB")
+            img = np.array(img).astype(np.float32) / 255.0
+            img = torch.from_numpy(img).permute(2, 0, 1).unsqueeze(0).to(device)
+            _, _, orig_H, orig_W = img.shape
+
+            # zero-pad the input image so that both the width and height are multiples of 32
+            l, r, t, b = utils.pad_input(orig_H, orig_W)
+            img = F.pad(img, (l, r, t, b), mode="constant", value=0.0)
+            img = normalize(img)
+
+            intrins_path = img_path.replace(ext, ".txt")
+            if os.path.exists(intrins_path):
+                # NOTE: camera intrinsics should be given as a txt file
+                # it should contain the values of fx, fy, cx, cy
+                intrins = utils.get_intrins_from_txt(
+                    intrins_path, device=device
+                ).unsqueeze(0)
+            else:
+                # NOTE: if intrins is not given, we just assume that the principal point is at the center
+                # and that the field-of-view is 60 degrees (feel free to modify this assumption)
+                intrins = utils.get_intrins_from_fov(
+                    new_fov=60.0, H=orig_H, W=orig_W, device=device
+                ).unsqueeze(0)
+
+            intrins[:, 0, 2] += l
+            intrins[:, 1, 2] += t
+
+            pred_norm = model(img, intrins=intrins)[-1]
+            pred_norm = pred_norm[:, :, t : t + orig_H, l : l + orig_W]
+
+            # save to output folder
+            # NOTE: by saving the prediction as uint8 png format, you lose a lot of precision
+            # if you want to use the predicted normals for downstream tasks, we recommend saving them as float32 NPY files
+            pred_norm_np = (
+                pred_norm.cpu().detach().numpy()[0, :, :, :].transpose(1, 2, 0)
+            )  # (H, W, 3)
+            pred_norm_np = ((pred_norm_np + 1.0) / 2.0 * 255.0).astype(np.uint8)
+            target_path = img_path.replace("/img/", "/output/").replace(ext, ".png")
+            im = Image.fromarray(pred_norm_np)
+            im.save(target_path)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--ckpt", default="dsine", type=str, help="model checkpoint")
+    parser.add_argument("--mode", default="samples", type=str, help="{samples}")
+    args = parser.parse_args()
+
+    # define model
+    device = torch.device("cpu")
+
+    from models.dsine import DSINE
+
+    model = DSINE().to(device)
+    model.pixel_coords = model.pixel_coords.to(device)
+    model = utils.load_checkpoint("./checkpoints/%s.pt" % args.ckpt, model)
+    model.eval()
+
+    if args.mode == "samples":
+        test_samples(args, model, intrins=None, device=device)

utils/rotation.py

@@ -0,0 +1,85 @@
+import torch
+import numpy as np
+
+
+# NOTE: from PyTorch3D
+def axis_angle_to_quaternion(axis_angle: torch.Tensor) -> torch.Tensor:
+    """
+    Convert rotations given as axis/angle to quaternions.
+
+    Args:
+        axis_angle: Rotations given as a vector in axis angle form,
+            as a tensor of shape (..., 3), where the magnitude is
+            the angle turned anticlockwise in radians around the
+            vector's direction.
+
+    Returns:
+        quaternions with real part first, as tensor of shape (..., 4).
+    """
+    angles = torch.norm(axis_angle, p=2, dim=-1, keepdim=True)
+    half_angles = angles * 0.5
+    eps = 1e-6
+    small_angles = angles.abs() < eps
+    sin_half_angles_over_angles = torch.empty_like(angles)
+    sin_half_angles_over_angles[~small_angles] = (
+        torch.sin(half_angles[~small_angles]) / angles[~small_angles]
+    )
+    # for x small, sin(x/2) is about x/2 - (x/2)^3/6
+    # so sin(x/2)/x is about 1/2 - (x*x)/48
+    sin_half_angles_over_angles[small_angles] = (
+        0.5 - (angles[small_angles] * angles[small_angles]) / 48
+    )
+    quaternions = torch.cat(
+        [torch.cos(half_angles), axis_angle * sin_half_angles_over_angles], dim=-1
+    )
+    return quaternions
+
+
+# NOTE: from PyTorch3D
+def quaternion_to_matrix(quaternions: torch.Tensor) -> torch.Tensor:
+    """
+    Convert rotations given as quaternions to rotation matrices.
+
+    Args:
+        quaternions: quaternions with real part first,
+            as tensor of shape (..., 4).
+
+    Returns:
+        Rotation matrices as tensor of shape (..., 3, 3).
+    """
+    r, i, j, k = torch.unbind(quaternions, -1)
+    # pyre-fixme[58]: `/` is not supported for operand types `float` and `Tensor`.
+    two_s = 2.0 / (quaternions * quaternions).sum(-1)
+
+    o = torch.stack(
+        (
+            1 - two_s * (j * j + k * k),
+            two_s * (i * j - k * r),
+            two_s * (i * k + j * r),
+            two_s * (i * j + k * r),
+            1 - two_s * (i * i + k * k),
+            two_s * (j * k - i * r),
+            two_s * (i * k - j * r),
+            two_s * (j * k + i * r),
+            1 - two_s * (i * i + j * j),
+        ),
+        -1,
+    )
+    return o.reshape(quaternions.shape[:-1] + (3, 3))
+
+
+# NOTE: from PyTorch3D
+def axis_angle_to_matrix(axis_angle: torch.Tensor) -> torch.Tensor:
+    """
+    Convert rotations given as axis/angle to rotation matrices.
+
+    Args:
+        axis_angle: Rotations given as a vector in axis angle form,
+            as a tensor of shape (..., 3), where the magnitude is
+            the angle turned anticlockwise in radians around the
+            vector's direction.
+
+    Returns:
+        Rotation matrices as tensor of shape (..., 3, 3).
+    """
+    return quaternion_to_matrix(axis_angle_to_quaternion(axis_angle))

utils/utils.py

@@ -0,0 +1,105 @@
+""" utils
+"""
+import os
+import torch
+import numpy as np
+
+
+def load_checkpoint(fpath, model):
+    print('loading checkpoint... {}'.format(fpath))
+
+    ckpt = torch.load(fpath, map_location='cpu')['model']
+
+    load_dict = {}
+    for k, v in ckpt.items():
+        if k.startswith('module.'):
+            k_ = k.replace('module.', '')
+            load_dict[k_] = v
+        else:
+            load_dict[k] = v
+
+    model.load_state_dict(load_dict)
+    print('loading checkpoint... / done')
+    return model
+
+
+def compute_normal_error(pred_norm, gt_norm):
+    pred_error = torch.cosine_similarity(pred_norm, gt_norm, dim=1)
+    pred_error = torch.clamp(pred_error, min=-1.0, max=1.0)
+    pred_error = torch.acos(pred_error) * 180.0 / np.pi
+    pred_error = pred_error.unsqueeze(1)    # (B, 1, H, W)
+    return pred_error
+
+
+def compute_normal_metrics(total_normal_errors):
+    total_normal_errors = total_normal_errors.detach().cpu().numpy()
+    num_pixels = total_normal_errors.shape[0]
+
+    metrics = {
+        'mean': np.average(total_normal_errors),
+        'median': np.median(total_normal_errors),
+        'rmse': np.sqrt(np.sum(total_normal_errors * total_normal_errors) / num_pixels),
+        'a1': 100.0 * (np.sum(total_normal_errors < 5) / num_pixels),
+        'a2': 100.0 * (np.sum(total_normal_errors < 7.5) / num_pixels),
+        'a3': 100.0 * (np.sum(total_normal_errors < 11.25) / num_pixels),
+        'a4': 100.0 * (np.sum(total_normal_errors < 22.5) / num_pixels),
+        'a5': 100.0 * (np.sum(total_normal_errors < 30) / num_pixels)
+    }
+
+    return metrics
+
+
+def pad_input(orig_H, orig_W):
+    if orig_W % 32 == 0:
+        l = 0
+        r = 0
+    else:
+        new_W = 32 * ((orig_W // 32) + 1)
+        l = (new_W - orig_W) // 2
+        r = (new_W - orig_W) - l
+
+    if orig_H % 32 == 0:
+        t = 0
+        b = 0
+    else:
+        new_H = 32 * ((orig_H // 32) + 1)
+        t = (new_H - orig_H) // 2
+        b = (new_H - orig_H) - t
+    return l, r, t, b
+
+
+def get_intrins_from_fov(new_fov, H, W, device):
+    # NOTE: top-left pixel should be (0,0)
+    if W >= H:
+        new_fu = (W / 2.0) / np.tan(np.deg2rad(new_fov / 2.0))
+        new_fv = (W / 2.0) / np.tan(np.deg2rad(new_fov / 2.0))
+    else:
+        new_fu = (H / 2.0) / np.tan(np.deg2rad(new_fov / 2.0))
+        new_fv = (H / 2.0) / np.tan(np.deg2rad(new_fov / 2.0))
+
+    new_cu = (W / 2.0) - 0.5
+    new_cv = (H / 2.0) - 0.5
+
+    new_intrins = torch.tensor([
+        [new_fu,    0,          new_cu  ],
+        [0,         new_fv,     new_cv  ],
+        [0,         0,          1       ]
+    ], dtype=torch.float32, device=device)
+
+    return new_intrins
+
+
+def get_intrins_from_txt(intrins_path, device):
+    # NOTE: top-left pixel should be (0,0)
+    with open(intrins_path, 'r') as f:
+        intrins_ = f.readlines()[0].split()[0].split(',')
+        intrins_ = [float(i) for i in intrins_]
+        fx, fy, cx, cy = intrins_
+
+    intrins = torch.tensor([
+        [fx, 0,cx],
+        [ 0,fy,cy],
+        [ 0, 0, 1]
+    ], dtype=torch.float32, device=device)
+
+    return intrins