From adcd2064666ab528d10363c8fb223eb55a245b68 Mon Sep 17 00:00:00 2001 From: Chung Min Kim Date: Sun, 2 Apr 2023 18:35:21 +0000 Subject: [PATCH] WIP yaml integration --- projects-lerf/lerf/data/lerf_datamanager.py | 26 +++++++++++++++---- .../lerf/encoders/openclip_encoder.py | 2 +- projects-lerf/lerf/lerf.py | 2 +- 3 files changed, 23 insertions(+), 7 deletions(-) diff --git a/projects-lerf/lerf/data/lerf_datamanager.py b/projects-lerf/lerf/data/lerf_datamanager.py index 9a21fd8f5..43609e754 100644 --- a/projects-lerf/lerf/data/lerf_datamanager.py +++ b/projects-lerf/lerf/data/lerf_datamanager.py @@ -22,6 +22,7 @@ from dataclasses import dataclass, field from pathlib import Path from typing import Any, Dict, List, Optional, Tuple, Type, Union, Literal import os.path as osp +import yaml import torch from rich.progress import Console @@ -46,6 +47,9 @@ from lerf.encoders.image_encoder import BaseImageEncoder @dataclass class LERFDataManagerConfig(VanillaDataManagerConfig): _target: Type = field(default_factory=lambda: LERFDataManager) + patch_tile_size_range: Tuple[int, int] = (0.05, 0.5) + patch_tile_size_res: int = 7 + patch_stride_scaler: float = 0.5 class LERFDataManager(VanillaDataManager): # pylint: disable=abstract-method """Basic stored data manager implementation. @@ -79,26 +83,38 @@ class LERFDataManager(VanillaDataManager): # pylint: disable=abstract-method images = torch.cat(images) cache_dir = f"outputs/{self.config.dataparser.data.name}" + # clip_cache_path, dino_cache_path = self._search_cache(cache_dir) + clip_cache_path = osp.join(cache_dir, "clip") + dino_cache_path = osp.join(cache_dir, "dino.npy") self.clip_interpolator = PatchedDataloader( image_list=images, device=self.device, - cfg={ + cfg= + { "tile_size_range": (0.05, 0.5), "tile_size_res": 7, "stride_scaler": 0.5, "image_shape": tuple(images.shape[2:4]) }, - cache_path=Path(osp.join(cache_dir, "clip")), + cache_path=clip_cache_path, model=self.image_encoder, ) self.dino_dataloader = DinoDataloader( image_list=images, device=self.device, cfg={"image_shape": tuple(images.shape[2:4])}, - cache_path=Path(osp.join(cache_dir, "dino.npy")) + cache_path=dino_cache_path ) - def _search_cache(self): - pass + # def _search_cache(self, cache_path, prop_dict): + # if osp.exists(cache_path): + # # load yaml file + # with open(cache_path, "r") as f: + # cache = yaml.load(f, Loader=yaml.FullLoader) + # for k, v in cache.items(): + # if k in self.config: + # self.config[k] = v + # else: + # ... def next_train(self, step: int) -> Tuple[RayBundle, Dict]: """Returns the next batch of data from the train dataloader.""" diff --git a/projects-lerf/lerf/encoders/openclip_encoder.py b/projects-lerf/lerf/encoders/openclip_encoder.py index ad68a3897..559359be5 100644 --- a/projects-lerf/lerf/encoders/openclip_encoder.py +++ b/projects-lerf/lerf/encoders/openclip_encoder.py @@ -14,7 +14,7 @@ from lerf.encoders.image_encoder import BaseImageEncoder, BaseImageEncoderConfig @dataclass class OpenCLIPNetworkConfig(BaseImageEncoderConfig): _target: Type = field(default_factory=lambda: OpenCLIPNetwork) - name: str = "openclip" + name: str = "openclip_vit_b_16" description: str = "OpenCLIP network, images and text" """target class to instantiate""" clip_model_type: str = "ViT-B-16" diff --git a/projects-lerf/lerf/lerf.py b/projects-lerf/lerf/lerf.py index fe71aecc3..836e2a85b 100644 --- a/projects-lerf/lerf/lerf.py +++ b/projects-lerf/lerf/lerf.py @@ -176,7 +176,7 @@ class LERFModel(NerfactoModel): loss_dict["clip_loss"] = self.config.clip_loss_weight * torch.nn.functional.huber_loss( outputs["clip"], batch["clip"], delta=1.25 ) - # loss_dict["dino_loss"] = torch.nn.functional.mse_loss(outputs["dino"], batch["dino"]) + loss_dict["dino_loss"] = torch.nn.functional.mse_loss(outputs["dino"], batch["dino"]) return loss_dict def get_param_groups(self) -> Dict[str, List[Parameter]]: