From 29df5943343a0f7d6fa4bff2918880501c97d07b Mon Sep 17 00:00:00 2001 From: Nikhil Mohan <47728081+nikmo33@users.noreply.github.com> Date: Fri, 16 Dec 2022 20:16:20 +0000 Subject: [PATCH] Minimal disk dataparser (#1128) * wip * merge upstream * clkeanup * cleanup * cleanup * cleanup * more cleanup * cleanup * cleanup docs * lint * style fix --- .../data/datamanagers/base_datamanager.py | 2 + .../data/datamanagers/semantic_datamanager.py | 1 + .../data/dataparsers/minimal_dataparser.py | 97 ++++++++++++++++++ tests/data/minimal_parser/train.npz | Bin 0 -> 8422 bytes tests/data/minimal_parser/val.npz | Bin 0 -> 8422 bytes tests/test_train.py | 10 ++ 6 files changed, 110 insertions(+) create mode 100644 nerfstudio/data/dataparsers/minimal_dataparser.py create mode 100644 tests/data/minimal_parser/train.npz create mode 100644 tests/data/minimal_parser/val.npz diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index 140d018ba..49b759a9d 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/nerfstudio/data/datamanagers/base_datamanager.py @@ -42,6 +42,7 @@ from nerfstudio.data.dataparsers.friends_dataparser import FriendsDataParserConf from nerfstudio.data.dataparsers.instant_ngp_dataparser import ( InstantNGPDataParserConfig, ) +from nerfstudio.data.dataparsers.minimal_dataparser import MinimalDataParserConfig from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig from nerfstudio.data.dataparsers.nuscenes_dataparser import NuScenesDataParserConfig from nerfstudio.data.dataparsers.phototourism_dataparser import ( @@ -65,6 +66,7 @@ AnnotatedDataParserUnion = tyro.conf.OmitSubcommandPrefixes[ # Omit prefixes of tyro.extras.subcommand_type_from_defaults( { "nerfstudio-data": NerfstudioDataParserConfig(), + "minimal-parser": MinimalDataParserConfig(), "blender-data": BlenderDataParserConfig(), "friends-data": FriendsDataParserConfig(), "instant-ngp-data": InstantNGPDataParserConfig(), diff --git a/nerfstudio/data/datamanagers/semantic_datamanager.py b/nerfstudio/data/datamanagers/semantic_datamanager.py index 188c7f6c3..322b1d3bd 100644 --- a/nerfstudio/data/datamanagers/semantic_datamanager.py +++ b/nerfstudio/data/datamanagers/semantic_datamanager.py @@ -41,6 +41,7 @@ class SemanticDataManager(VanillaDataManager): # pylint: disable=abstract-metho """ def create_train_dataset(self) -> SemanticDataset: + self.train_dataparser_outputs = self.dataparser.get_dataparser_outputs(split="train") return SemanticDataset( dataparser_outputs=self.train_dataparser_outputs, scale_factor=self.config.camera_res_scale_factor, diff --git a/nerfstudio/data/dataparsers/minimal_dataparser.py b/nerfstudio/data/dataparsers/minimal_dataparser.py new file mode 100644 index 000000000..b6ffd8518 --- /dev/null +++ b/nerfstudio/data/dataparsers/minimal_dataparser.py @@ -0,0 +1,97 @@ +# Copyright 2022 The Nerfstudio Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Data parser for pre-prepared datasets for all cameras, with no additional processing needed +Optional fields - semantics, mask_filenames, cameras.distortion_params, cameras.times +""" + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Type + +import numpy as np +import torch + +from nerfstudio.cameras.cameras import Cameras +from nerfstudio.data.dataparsers.base_dataparser import ( + DataParser, + DataParserConfig, + DataparserOutputs, + Semantics, +) +from nerfstudio.data.scene_box import SceneBox + + +@dataclass +class MinimalDataParserConfig(DataParserConfig): + """Minimal dataset config""" + + _target: Type = field(default_factory=lambda: MinimalDataParser) + """target class to instantiate""" + data: Path = Path("/home/nikhil/nerfstudio-main/tests/data/lego_test/minimal_parser") + + +@dataclass +class MinimalDataParser(DataParser): + """Minimal DatasetParser""" + + config: MinimalDataParserConfig + + def _generate_dataparser_outputs(self, split="train"): + filepath = self.config.data / f"{split}.npz" + data = np.load(filepath, allow_pickle=True) + + image_filenames = data["image_filenames"].tolist() + mask_filenames = data["mask_filenames"].tolist() if "mask_filenames" in data.keys() else None + + metadata = None + if "semantics" in data.keys(): + semantics = data["semantics"].item() + metadata = { + "semantics": Semantics( + filenames=semantics["filenames"].tolist(), + classes=semantics["classes"].tolist(), + colors=semantics["colors"].tolist(), + mask_classes=semantics["mask_classes"].tolist(), + ) + } + + scene_box_aabb = torch.from_numpy(data["scene_box"]) + scene_box = SceneBox(aabb=scene_box_aabb) + + camera_np = data["cameras"].item() + distortion_params = None + if "distortion_params" in camera_np.keys(): + distortion_params = torch.from_numpy(camera_np["distortion_params"]) + cameras = Cameras( + fx=torch.from_numpy(camera_np["fx"]), + fy=torch.from_numpy(camera_np["fy"]), + cx=torch.from_numpy(camera_np["cx"]), + cy=torch.from_numpy(camera_np["cy"]), + distortion_params=distortion_params, + height=torch.from_numpy(camera_np["height"]), + width=torch.from_numpy(camera_np["width"]), + camera_to_worlds=torch.from_numpy(camera_np["camera_to_worlds"])[:, :3, :4], + camera_type=torch.from_numpy(camera_np["camera_type"]), + times=torch.from_numpy(camera_np["times"]) if "times" in camera_np.keys() else None, + ) + + dataparser_outputs = DataparserOutputs( + image_filenames=image_filenames, + cameras=cameras, + scene_box=scene_box, + mask_filenames=mask_filenames, + metadata=metadata, + ) + return dataparser_outputs diff --git a/tests/data/minimal_parser/train.npz b/tests/data/minimal_parser/train.npz new file mode 100644 index 0000000000000000000000000000000000000000..9d8acbb15f7b0060ffaa046218d6aff7b79c0aff GIT binary patch literal 8422 zcmeHMdvH|M8Q(leAV3HZAPC5!LIQ;FNC*LvdoTg5ERVQ_&URWR>~1!>Ta#?g?hO#c zG2j5AqC8umh=72!+KPZRq9j@Yu>%OGQ&c)-%0oM&6cLr8)qcN|{npKrO>Tz%+0Xer z_nz~;zu!3fnm(~@fORyxZ)C1u_UPnEOM6UeHnu8pyz zsAZ@?477H}?CHfnkFm_NJeXeU3Ap{~6RhcrW~OIa)60B*!|$vr@%c+Vn7+qZ5%AD_ zKzCMq=sjah=Fm**L)KfNq`EkFRZV4eZI;{T_hePpR2W{T-|wtVD)D>VzN&!XuW=hm z($XZAR8`uvDr0Pxr82fC#wFvH%Xqu&kYS6l#V(ZzLBUd^w%U^<6CJUX62H$^DZ4v5(vs>#pd?53 zaHIt#OMW_f%#=~1OGYWm8<+K~>T=oBkg0arOHaM`Fd?8NABUi+TIpo+{tjG_wL`~X5 zh!4KlubEoHKI_3J@_GQ1CJ1q&9`FYnW{LelkMY9&w9*{vg!gh{otSMymg|Z7S&6X@ z43rXc?#P>U=rM#(#EL#xgLCX7TAd(^7-1JRfjndpmz5Z}KpgOe5pp2Z1@DkWJrScB zVqC)x@-a({>+qb-=)wnTf_&OT$cOgO)Gu}(c5pz~zt}S{34kwNg@@?f5uE#dv1A6?yIrxGu>oazv>0ll?z&>QH#Abi*QAfms-!m|Ln<2K3{sG@+VuW4J4>&bb zN9cf{$G=I47qu}%TzACCHq7YpUO--8#EP7)#Ed=E3i*fUj=r-jJ(?!i=JSO|V9iPl z`|t(Zk(kFC`VJjl3y+7UTO`Qf>~1f?2Vy}D*f*ZrBK~(3`hc^>Z12%$$g)0sn$cx^ z;ol3W8)}5OVGp&AM!*Q{BTulE82R!3Lcd$`kND67=<~kC83UjE8%d8h5#r)?*fvXC zd)P-kxh`$(3&=9Q+Vi`@eF*=1i_b!=Ar6iyJO-{k@@$Vk)DiWE9-htWqLy5fmik1? zEfVSiKd23!!SIb-@Enb}?!Xinc(W4oIe=Kg>p`!IBF{1_%|aithyk&;q{lTzY;&jpM`XM5DNR)YBA(k`Yd!TuDP?no&BZQwsmX47gwsae_g(;y|iak!QwiL#!^o` znJoDElS^gdgZJN8YZot+s(I+WbM&a1BRipI=&XF=iIAN)-S1X&}-ngN?T3@f7K6Sr(-|(A0;^6D~;?&9Fpl^!>xlkQxy-vJ+3wk&6~BC_vff*taDuv8}zSz84$Ft?c2Lo#j<=^f&jmJ;)J$l zb$79TU8-PN#PHgI3F64%DFQWt{tM3!SMrPN>Lz~t(MKB9N&E1OLp%BIO!4?*NlFGe zz;|OqgSL8AN&))#)ir}^y<2I@alN+%$6-w1H01Mw|5NnX216D>kCxQ zz&qij;D_(O*LFXbE5E56MuQ< zK7~)@{>|6(l{GAyZf-kER#KcJnws=fw;l? z@fx7_O{abzYG%i}%-G4-yw{6l8t$%|elKd}Z48or07-wl9H4Gk;N^C`SGd|hbx(od z+zX;yQN%)Ruu`MH!ivChh#}LPyt!*-274O{Z^P_zc*t93P%V+v5{u*rLuR=go#9O( zs03GwawO}Ig8pc`923&d4%Wn*Cv#Afv38jYP~TjQaa4_z7@&}dN112H@h(RSO|e5? z$$X!`f-V~;u-A$3ns1i{OBJHATp`x&jH=zE!PcsHwRUSpo;8ra;Owj^jyj? zIhY!c@*YFp>rxk5^Sz~pE`P!n?t_J?c4-S)n5Hb83tHe&PB-KXwe4p{ZhL)f-P{dS zJFvR@bbG2iB`)6rz5>Srh{qN;etua_OLyS2{526*;MVM0y#Gw1AI8QQ`f+SLsJQWG zQ_nXzrk>Y|Z85Q_aj8wgS|V@wgVQwBsoJz~-ydn9`K@rcfv&Wdn%xmcsy07M!&Q6q mW03=Ojps~1!>Ta#?g?hO#c zG2j5AqC8umh=72!+KPZRq9j@Yu>%OGQ&c)-%0oM&6cLr8)qcN|{npKrO>Tz%+0Xer z_nz~;zu!3fnm(~@fORyxZ)C1u_UPnEOM6UeHnu8pyz zsAZ@?477H}?CHfnkFm_NJeXeU3Ap{~6RhcrW~OIa)60B*!|$vr@%c+Vn7+qZ5%AD_ zKzCMq=sjah=Fm**L)KfNq`EkFRZV4eZI;{T_hePpR2W{T-|wtVD)D>VzN&!XuW=hm z($XZAR8`uvDr0Pxr82fC#wFvH%Xqu&kYS6l#V(ZzLBUd^w%U^<6CJUX62H$^DZ4v5(vs>#pd?53 zaHIt#OMW_f%#=~1OGYWm8<+K~>T=oBkg0arOHaM`Fd?8NABUi+TIpo+{tjG_wL`~X5 zh!4KlubEoHKI_3J@_GQ1CJ1q&9`FYnW{LelkMY9&w9*{vg!gh{otSMymg|Z7S&6X@ z43rXc?#P>U=rM#(#EL#xgLCX7TAd(^7-1JRfjndpmz5Z}KpgOe5pp2Z1@DkWJrScB zVqC)x@-a({>+qb-=)wnTf_&OT$cOgO)Gu}(c5pz~zt}S{34kwNg@@?f5uE#dv1A6?yIrxGu>oazv>0ll?z&>QH#Abi*QAfms-!m|Ln<2K3{sG@+VuW4J4>&bb zN9cf{$G=I47qu}%TzACCHq7YpUO--8#EP7)#Ed=E3i*fUj=r-jJ(?!i=JSO|V9iPl z`|t(Zk(kFC`VJjl3y+7UTO`Qf>~1f?2Vy}D*f*ZrBK~(3`hc^>Z12%$$g)0sn$cx^ z;ol3W8)}5OVGp&AM!*Q{BTulE82R!3Lcd$`kND67=<~kC83UjE8%d8h5#r)?*fvXC zd)P-kxh`$(3&=9Q+Vi`@eF*=1i_b!=Ar6iyJO-{k@@$Vk)DiWE9-htWqLy5fmik1? zEfVSiKd23!!SIb-@Enb}?!Xinc(W4oIe=Kg>p`!IBF{1_%|aithyk&;q{lTzY;&jpM`XM5DNR)YBA(k`Yd!TuDP?no&BZQwsmX47gwsae_g(;y|iak!QwiL#!^o` znJoDElS^gdgZJN8YZot+s(I+WbM&a1BRipI=&XF=iIAN)-S1X&}-ngN?T3@f7K6Sr(-|(A0;^6D~;?&9Fpl^!>xlkQxy-vJ+3wk&6~BC_vff*taDuv8}zSz84$Ft?c2Lo#j<=^f&jmJ;)J$l zb$79TU8-PN#PHgI3F64%DFQWt{tM3!SMrPN>Lz~t(MKB9N&E1OLp%BIO!4?*NlFGe zz;|OqgSL8AN&))#)ir}^y<2I@alN+%$6-w1H01Mw|5NnX216D>kCxQ zz&qij;D_(O*LFXbE5E56MuQ< zK7~)@{>|6(l{GAyZf-kER#KcJnws=fw;l? z@fx7_O{abzYG%i}%-G4-yw{6l8t$%|elKd}Z48or07-wl9H4Gk;N^C`SGd|hbx(od z+zX;yQN%)Ruu`MH!ivChh#}LPyt!*-274O{Z^P_zc*t93P%V+v5{u*rLuR=go#9O( zs03GwawO}Ig8pc`923&d4%Wn*Cv#Afv38jYP~TjQaa4_z7@&}dN112H@h(RSO|e5? z$$X!`f-V~;u-A$3ns1i{OBJHATp`x&jH=zE!PcsHwRUSpo;8ra;Owj^jyj? zIhY!c@*YFp>rxk5^Sz~pE`P!n?t_J?c4-S)n5Hb83tHe&PB-KXwe4p{ZhL)f-P{dS zJFvR@bbG2iB`)6rz5>Srh{qN;etua_OLyS2{526*;MVM0y#Gw1AI8QQ`f+SLsJQWG zQ_nXzrk>Y|Z85Q_aj8wgS|V@wgVQwBsoJz~-ydn9`K@rcfv&Wdn%xmcsy07M!&Q6q mW03=Ojps