We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
why are there weird red margins of my loaded images?
**Describe the bug**
When I load the images, they have red margins.

**To Reproduce**
Steps to reproduce the behavior:
I just install the latest version of nerfstudio (commitid 9b6010ea), and pip install -e . to install it.

**Expected behavior**
A clear and concise description of what you expected to happen.

**Screenshots**
If applicable, add screenshots to help explain your problem.
<img width="508" alt="image" src="https://github.com/nerfstudio-project/nerfstudio/assets/44845151/5ae36028-17e7-4683-b1a4-cec33f24dcf6">


**Additional context**
Add any other context about the problem here.


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `nerfstudio/data/dataparsers/instant_ngp_dataparser.py`
Content:
```
1 # Copyright 2022 the Regents of the University of California, Nerfstudio Team and contributors. All rights reserved.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 
15 """Data parser for instant ngp data"""
16 
17 from __future__ import annotations
18 
19 from dataclasses import dataclass, field
20 from pathlib import Path
21 from typing import Dict, Tuple, Type
22 
23 import imageio
24 import numpy as np
25 import torch
26 
27 from nerfstudio.cameras import camera_utils
28 from nerfstudio.cameras.cameras import Cameras, CameraType
29 from nerfstudio.data.dataparsers.base_dataparser import (
30     DataParser,
31     DataParserConfig,
32     DataparserOutputs,
33 )
34 from nerfstudio.data.scene_box import SceneBox
35 from nerfstudio.utils.io import load_from_json
36 from nerfstudio.utils.rich_utils import CONSOLE
37 
38 
39 @dataclass
40 class InstantNGPDataParserConfig(DataParserConfig):
41     """Instant-NGP dataset parser config"""
42 
43     _target: Type = field(default_factory=lambda: InstantNGP)
44     """target class to instantiate"""
45     data: Path = Path("data/ours/posterv2")
46     """Directory or explicit json file path specifying location of data."""
47     scene_scale: float = 0.3333
48     """How much to scale the scene."""
49 
50 
51 @dataclass
52 class InstantNGP(DataParser):
53     """Instant NGP Dataset"""
54 
55     config: InstantNGPDataParserConfig
56 
57     def _generate_dataparser_outputs(self, split="train"):
58         if self.config.data.suffix == ".json":
59             meta = load_from_json(self.config.data)
60             data_dir = self.config.data.parent
61         else:
62             meta = load_from_json(self.config.data / "transforms.json")
63             data_dir = self.config.data
64 
65         image_filenames = []
66         mask_filenames = []
67         poses = []
68         num_skipped_image_filenames = 0
69         for frame in meta["frames"]:
70             fname = data_dir / Path(frame["file_path"])
71             # search for png file
72             if not fname.exists():
73                 fname = data_dir / Path(frame["file_path"] + ".png")
74             if not fname.exists():
75                 CONSOLE.log(f"couldn't find {fname} image")
76                 num_skipped_image_filenames += 1
77             else:
78                 if "w" not in meta:
79                     img_0 = imageio.imread(fname)
80                     h, w = img_0.shape[:2]
81                     meta["w"] = w
82                     if "h" in meta:
83                         meta_h = meta["h"]
84                         assert meta_h == h, f"height of image dont not correspond metadata {h} != {meta_h}"
85                     else:
86                         meta["h"] = h
87                 image_filenames.append(fname)
88                 poses.append(np.array(frame["transform_matrix"]))
89                 if "mask_path" in frame:
90                     mask_fname = data_dir / Path(frame["mask_path"])
91                     mask_filenames.append(mask_fname)
92         if num_skipped_image_filenames >= 0:
93             CONSOLE.print(f"Skipping {num_skipped_image_filenames} files in dataset split {split}.")
94         assert (
95             len(image_filenames) != 0
96         ), """
97         No image files found. 
98         You should check the file_paths in the transforms.json file to make sure they are correct.
99         """
100         poses = np.array(poses).astype(np.float32)
101         poses[:, :3, 3] *= self.config.scene_scale
102 
103         camera_to_world = torch.from_numpy(poses[:, :3])  # camera to world transform
104 
105         distortion_params = camera_utils.get_distortion_params(
106             k1=float(meta.get("k1", 0)),
107             k2=float(meta.get("k2", 0)),
108             k3=float(meta.get("k3", 0)),
109             k4=float(meta.get("k4", 0)),
110             p1=float(meta.get("p1", 0)),
111             p2=float(meta.get("p2", 0)),
112         )
113 
114         # in x,y,z order
115         # assumes that the scene is centered at the origin
116         aabb_scale = 0.5 * meta.get("aabb_scale", 1)
117 
118         scene_box = SceneBox(
119             aabb=torch.tensor(
120                 [[-aabb_scale, -aabb_scale, -aabb_scale], [aabb_scale, aabb_scale, aabb_scale]], dtype=torch.float32
121             )
122         )
123 
124         fl_x, fl_y = InstantNGP.get_focal_lengths(meta)
125 
126         w, h = meta["w"], meta["h"]
127 
128         camera_type = CameraType.PERSPECTIVE
129         if meta.get("is_fisheye", False):
130             camera_type = CameraType.FISHEYE
131 
132         cameras = Cameras(
133             fx=float(fl_x),
134             fy=float(fl_y),
135             cx=float(meta.get("cx", 0.5 * w)),
136             cy=float(meta.get("cy", 0.5 * h)),
137             distortion_params=distortion_params,
138             height=int(h),
139             width=int(w),
140             camera_to_worlds=camera_to_world,
141             camera_type=camera_type,
142         )
143 
144         # TODO(ethan): add alpha background color
145         dataparser_outputs = DataparserOutputs(
146             image_filenames=image_filenames,
147             cameras=cameras,
148             scene_box=scene_box,
149             mask_filenames=mask_filenames if len(mask_filenames) > 0 else None,
150             dataparser_scale=self.config.scene_scale,
151         )
152 
153         return dataparser_outputs
154 
155     @classmethod
156     def get_focal_lengths(cls, meta: Dict) -> Tuple[float, float]:
157         """Reads or computes the focal length from transforms dict.
158         Args:
159             meta: metadata from transforms.json file.
160         Returns:
161             Focal lengths in the x and y directions. Error is raised if these cannot be calculated.
162         """
163         fl_x, fl_y = 0, 0
164 
165         def fov_to_focal_length(rad, res):
166             return 0.5 * res / np.tan(0.5 * rad)
167 
168         if "fl_x" in meta:
169             fl_x = meta["fl_x"]
170         elif "x_fov" in meta:
171             fl_x = fov_to_focal_length(np.deg2rad(meta["x_fov"]), meta["w"])
172         elif "camera_angle_x" in meta:
173             fl_x = fov_to_focal_length(meta["camera_angle_x"], meta["w"])
174 
175         if "camera_angle_y" not in meta or "y_fov" not in meta:
176             fl_y = fl_x
177         else:
178             if "fl_y" in meta:
179                 fl_y = meta["fl_y"]
180             elif "y_fov" in meta:
181                 fl_y = fov_to_focal_length(np.deg2rad(meta["y_fov"]), meta["h"])
182             elif "camera_angle_y" in meta:
183                 fl_y = fov_to_focal_length(meta["camera_angle_y"], meta["h"])
184 
185         if fl_x == 0 or fl_y == 0:
186             raise AttributeError("Focal length cannot be calculated from transforms.json (missing fields).")
187 
188         return (fl_x, fl_y)
189
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/nerfstudio/data/dataparsers/instant_ngp_dataparser.py b/nerfstudio/data/dataparsers/instant_ngp_dataparser.py
--- a/nerfstudio/data/dataparsers/instant_ngp_dataparser.py
+++ b/nerfstudio/data/dataparsers/instant_ngp_dataparser.py
@@ -18,7 +18,7 @@
 
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Dict, Tuple, Type
+from typing import Dict, Tuple, Type, Literal
 
 import imageio
 import numpy as np
@@ -32,6 +32,12 @@
     DataparserOutputs,
 )
 from nerfstudio.data.scene_box import SceneBox
+from nerfstudio.data.utils.dataparsers_utils import (
+    get_train_eval_split_filename,
+    get_train_eval_split_fraction,
+    get_train_eval_split_interval,
+    get_train_eval_split_all,
+)
 from nerfstudio.utils.io import load_from_json
 from nerfstudio.utils.rich_utils import CONSOLE
 
@@ -46,6 +52,18 @@
     """Directory or explicit json file path specifying location of data."""
     scene_scale: float = 0.3333
     """How much to scale the scene."""
+    eval_mode: Literal["fraction", "filename", "interval", "all"] = "fraction"
+    """
+    The method to use for splitting the dataset into train and eval.
+    Fraction splits based on a percentage for train and the remaining for eval.
+    Filename splits based on filenames containing train/eval.
+    Interval uses every nth frame for eval.
+    All uses all the images for any split.
+    """
+    train_split_fraction: float = 0.9
+    """The percentage of the dataset to use for training. Only used when eval_mode is train-split-fraction."""
+    eval_interval: int = 8
+    """The interval between frames to use for eval. Only used when eval_mode is eval-interval."""
 
 
 @dataclass
@@ -94,12 +112,40 @@
         assert (
             len(image_filenames) != 0
         ), """
-        No image files found. 
+        No image files found.
         You should check the file_paths in the transforms.json file to make sure they are correct.
         """
         poses = np.array(poses).astype(np.float32)
         poses[:, :3, 3] *= self.config.scene_scale
 
+        # find train and eval indices based on the eval_mode specified
+        if self.config.eval_mode == "fraction":
+            i_train, i_eval = get_train_eval_split_fraction(image_filenames, self.config.train_split_fraction)
+        elif self.config.eval_mode == "filename":
+            i_train, i_eval = get_train_eval_split_filename(image_filenames)
+        elif self.config.eval_mode == "interval":
+            i_train, i_eval = get_train_eval_split_interval(image_filenames, self.config.eval_interval)
+        elif self.config.eval_mode == "all":
+            CONSOLE.log(
+                "[yellow] Be careful with '--eval-mode=all'. If using camera optimization, the cameras may diverge in the current implementation, giving unpredictable results."
+            )
+            i_train, i_eval = get_train_eval_split_all(image_filenames)
+        else:
+            raise ValueError(f"Unknown eval mode {self.config.eval_mode}")
+
+        if split == "train":
+            indices = i_train
+        elif split in ["val", "test"]:
+            indices = i_eval
+        else:
+            raise ValueError(f"Unknown dataparser split {split}")
+        # Choose image_filenames and poses based on split, but after auto orient and scaling the poses.
+        image_filenames = [image_filenames[i] for i in indices]
+        mask_filenames = [mask_filenames[i] for i in indices] if len(mask_filenames) > 0 else []
+
+        idx_tensor = torch.tensor(indices, dtype=torch.long)
+        poses = poses[idx_tensor]
+
         camera_to_world = torch.from_numpy(poses[:, :3])  # camera to world transform
 
         distortion_params = camera_utils.get_distortion_params(

{"golden_diff": "diff --git a/nerfstudio/data/dataparsers/instant_ngp_dataparser.py b/nerfstudio/data/dataparsers/instant_ngp_dataparser.py\n--- a/nerfstudio/data/dataparsers/instant_ngp_dataparser.py\n+++ b/nerfstudio/data/dataparsers/instant_ngp_dataparser.py\n@@ -18,7 +18,7 @@\n \n from dataclasses import dataclass, field\n from pathlib import Path\n-from typing import Dict, Tuple, Type\n+from typing import Dict, Tuple, Type, Literal\n \n import imageio\n import numpy as np\n@@ -32,6 +32,12 @@\n     DataparserOutputs,\n )\n from nerfstudio.data.scene_box import SceneBox\n+from nerfstudio.data.utils.dataparsers_utils import (\n+    get_train_eval_split_filename,\n+    get_train_eval_split_fraction,\n+    get_train_eval_split_interval,\n+    get_train_eval_split_all,\n+)\n from nerfstudio.utils.io import load_from_json\n from nerfstudio.utils.rich_utils import CONSOLE\n \n@@ -46,6 +52,18 @@\n     \"\"\"Directory or explicit json file path specifying location of data.\"\"\"\n     scene_scale: float = 0.3333\n     \"\"\"How much to scale the scene.\"\"\"\n+    eval_mode: Literal[\"fraction\", \"filename\", \"interval\", \"all\"] = \"fraction\"\n+    \"\"\"\n+    The method to use for splitting the dataset into train and eval.\n+    Fraction splits based on a percentage for train and the remaining for eval.\n+    Filename splits based on filenames containing train/eval.\n+    Interval uses every nth frame for eval.\n+    All uses all the images for any split.\n+    \"\"\"\n+    train_split_fraction: float = 0.9\n+    \"\"\"The percentage of the dataset to use for training. Only used when eval_mode is train-split-fraction.\"\"\"\n+    eval_interval: int = 8\n+    \"\"\"The interval between frames to use for eval. Only used when eval_mode is eval-interval.\"\"\"\n \n \n @dataclass\n@@ -94,12 +112,40 @@\n         assert (\n             len(image_filenames) != 0\n         ), \"\"\"\n-        No image files found. \n+        No image files found.\n         You should check the file_paths in the transforms.json file to make sure they are correct.\n         \"\"\"\n         poses = np.array(poses).astype(np.float32)\n         poses[:, :3, 3] *= self.config.scene_scale\n \n+        # find train and eval indices based on the eval_mode specified\n+        if self.config.eval_mode == \"fraction\":\n+            i_train, i_eval = get_train_eval_split_fraction(image_filenames, self.config.train_split_fraction)\n+        elif self.config.eval_mode == \"filename\":\n+            i_train, i_eval = get_train_eval_split_filename(image_filenames)\n+        elif self.config.eval_mode == \"interval\":\n+            i_train, i_eval = get_train_eval_split_interval(image_filenames, self.config.eval_interval)\n+        elif self.config.eval_mode == \"all\":\n+            CONSOLE.log(\n+                \"[yellow] Be careful with '--eval-mode=all'. If using camera optimization, the cameras may diverge in the current implementation, giving unpredictable results.\"\n+            )\n+            i_train, i_eval = get_train_eval_split_all(image_filenames)\n+        else:\n+            raise ValueError(f\"Unknown eval mode {self.config.eval_mode}\")\n+\n+        if split == \"train\":\n+            indices = i_train\n+        elif split in [\"val\", \"test\"]:\n+            indices = i_eval\n+        else:\n+            raise ValueError(f\"Unknown dataparser split {split}\")\n+        # Choose image_filenames and poses based on split, but after auto orient and scaling the poses.\n+        image_filenames = [image_filenames[i] for i in indices]\n+        mask_filenames = [mask_filenames[i] for i in indices] if len(mask_filenames) > 0 else []\n+\n+        idx_tensor = torch.tensor(indices, dtype=torch.long)\n+        poses = poses[idx_tensor]\n+\n         camera_to_world = torch.from_numpy(poses[:, :3])  # camera to world transform\n \n         distortion_params = camera_utils.get_distortion_params(\n", "issue": "why are there weird red margins of my loaded images?\n**Describe the bug**\r\nWhen I load the images, they have red margins.\r\n\r\n**To Reproduce**\r\nSteps to reproduce the behavior:\r\nI just install the latest version of nerfstudio (commitid 9b6010ea), and pip install -e . to install it.\r\n\r\n**Expected behavior**\r\nA clear and concise description of what you expected to happen.\r\n\r\n**Screenshots**\r\nIf applicable, add screenshots to help explain your problem.\r\n<img width=\"508\" alt=\"image\" src=\"https://github.com/nerfstudio-project/nerfstudio/assets/44845151/5ae36028-17e7-4683-b1a4-cec33f24dcf6\">\r\n\r\n\r\n**Additional context**\r\nAdd any other context about the problem here.\r\n\n", "before_files": [{"content": "# Copyright 2022 the Regents of the University of California, Nerfstudio Team and contributors. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Data parser for instant ngp data\"\"\"\n\nfrom __future__ import annotations\n\nfrom dataclasses import dataclass, field\nfrom pathlib import Path\nfrom typing import Dict, Tuple, Type\n\nimport imageio\nimport numpy as np\nimport torch\n\nfrom nerfstudio.cameras import camera_utils\nfrom nerfstudio.cameras.cameras import Cameras, CameraType\nfrom nerfstudio.data.dataparsers.base_dataparser import (\n    DataParser,\n    DataParserConfig,\n    DataparserOutputs,\n)\nfrom nerfstudio.data.scene_box import SceneBox\nfrom nerfstudio.utils.io import load_from_json\nfrom nerfstudio.utils.rich_utils import CONSOLE\n\n\n@dataclass\nclass InstantNGPDataParserConfig(DataParserConfig):\n    \"\"\"Instant-NGP dataset parser config\"\"\"\n\n    _target: Type = field(default_factory=lambda: InstantNGP)\n    \"\"\"target class to instantiate\"\"\"\n    data: Path = Path(\"data/ours/posterv2\")\n    \"\"\"Directory or explicit json file path specifying location of data.\"\"\"\n    scene_scale: float = 0.3333\n    \"\"\"How much to scale the scene.\"\"\"\n\n\n@dataclass\nclass InstantNGP(DataParser):\n    \"\"\"Instant NGP Dataset\"\"\"\n\n    config: InstantNGPDataParserConfig\n\n    def _generate_dataparser_outputs(self, split=\"train\"):\n        if self.config.data.suffix == \".json\":\n            meta = load_from_json(self.config.data)\n            data_dir = self.config.data.parent\n        else:\n            meta = load_from_json(self.config.data / \"transforms.json\")\n            data_dir = self.config.data\n\n        image_filenames = []\n        mask_filenames = []\n        poses = []\n        num_skipped_image_filenames = 0\n        for frame in meta[\"frames\"]:\n            fname = data_dir / Path(frame[\"file_path\"])\n            # search for png file\n            if not fname.exists():\n                fname = data_dir / Path(frame[\"file_path\"] + \".png\")\n            if not fname.exists():\n                CONSOLE.log(f\"couldn't find {fname} image\")\n                num_skipped_image_filenames += 1\n            else:\n                if \"w\" not in meta:\n                    img_0 = imageio.imread(fname)\n                    h, w = img_0.shape[:2]\n                    meta[\"w\"] = w\n                    if \"h\" in meta:\n                        meta_h = meta[\"h\"]\n                        assert meta_h == h, f\"height of image dont not correspond metadata {h} != {meta_h}\"\n                    else:\n                        meta[\"h\"] = h\n                image_filenames.append(fname)\n                poses.append(np.array(frame[\"transform_matrix\"]))\n                if \"mask_path\" in frame:\n                    mask_fname = data_dir / Path(frame[\"mask_path\"])\n                    mask_filenames.append(mask_fname)\n        if num_skipped_image_filenames >= 0:\n            CONSOLE.print(f\"Skipping {num_skipped_image_filenames} files in dataset split {split}.\")\n        assert (\n            len(image_filenames) != 0\n        ), \"\"\"\n        No image files found. \n        You should check the file_paths in the transforms.json file to make sure they are correct.\n        \"\"\"\n        poses = np.array(poses).astype(np.float32)\n        poses[:, :3, 3] *= self.config.scene_scale\n\n        camera_to_world = torch.from_numpy(poses[:, :3])  # camera to world transform\n\n        distortion_params = camera_utils.get_distortion_params(\n            k1=float(meta.get(\"k1\", 0)),\n            k2=float(meta.get(\"k2\", 0)),\n            k3=float(meta.get(\"k3\", 0)),\n            k4=float(meta.get(\"k4\", 0)),\n            p1=float(meta.get(\"p1\", 0)),\n            p2=float(meta.get(\"p2\", 0)),\n        )\n\n        # in x,y,z order\n        # assumes that the scene is centered at the origin\n        aabb_scale = 0.5 * meta.get(\"aabb_scale\", 1)\n\n        scene_box = SceneBox(\n            aabb=torch.tensor(\n                [[-aabb_scale, -aabb_scale, -aabb_scale], [aabb_scale, aabb_scale, aabb_scale]], dtype=torch.float32\n            )\n        )\n\n        fl_x, fl_y = InstantNGP.get_focal_lengths(meta)\n\n        w, h = meta[\"w\"], meta[\"h\"]\n\n        camera_type = CameraType.PERSPECTIVE\n        if meta.get(\"is_fisheye\", False):\n            camera_type = CameraType.FISHEYE\n\n        cameras = Cameras(\n            fx=float(fl_x),\n            fy=float(fl_y),\n            cx=float(meta.get(\"cx\", 0.5 * w)),\n            cy=float(meta.get(\"cy\", 0.5 * h)),\n            distortion_params=distortion_params,\n            height=int(h),\n            width=int(w),\n            camera_to_worlds=camera_to_world,\n            camera_type=camera_type,\n        )\n\n        # TODO(ethan): add alpha background color\n        dataparser_outputs = DataparserOutputs(\n            image_filenames=image_filenames,\n            cameras=cameras,\n            scene_box=scene_box,\n            mask_filenames=mask_filenames if len(mask_filenames) > 0 else None,\n            dataparser_scale=self.config.scene_scale,\n        )\n\n        return dataparser_outputs\n\n    @classmethod\n    def get_focal_lengths(cls, meta: Dict) -> Tuple[float, float]:\n        \"\"\"Reads or computes the focal length from transforms dict.\n        Args:\n            meta: metadata from transforms.json file.\n        Returns:\n            Focal lengths in the x and y directions. Error is raised if these cannot be calculated.\n        \"\"\"\n        fl_x, fl_y = 0, 0\n\n        def fov_to_focal_length(rad, res):\n            return 0.5 * res / np.tan(0.5 * rad)\n\n        if \"fl_x\" in meta:\n            fl_x = meta[\"fl_x\"]\n        elif \"x_fov\" in meta:\n            fl_x = fov_to_focal_length(np.deg2rad(meta[\"x_fov\"]), meta[\"w\"])\n        elif \"camera_angle_x\" in meta:\n            fl_x = fov_to_focal_length(meta[\"camera_angle_x\"], meta[\"w\"])\n\n        if \"camera_angle_y\" not in meta or \"y_fov\" not in meta:\n            fl_y = fl_x\n        else:\n            if \"fl_y\" in meta:\n                fl_y = meta[\"fl_y\"]\n            elif \"y_fov\" in meta:\n                fl_y = fov_to_focal_length(np.deg2rad(meta[\"y_fov\"]), meta[\"h\"])\n            elif \"camera_angle_y\" in meta:\n                fl_y = fov_to_focal_length(meta[\"camera_angle_y\"], meta[\"h\"])\n\n        if fl_x == 0 or fl_y == 0:\n            raise AttributeError(\"Focal length cannot be calculated from transforms.json (missing fields).\")\n\n        return (fl_x, fl_y)\n", "path": "nerfstudio/data/dataparsers/instant_ngp_dataparser.py"}], "after_files": [{"content": "# Copyright 2022 the Regents of the University of California, Nerfstudio Team and contributors. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Data parser for instant ngp data\"\"\"\n\nfrom __future__ import annotations\n\nfrom dataclasses import dataclass, field\nfrom pathlib import Path\nfrom typing import Dict, Tuple, Type, Literal\n\nimport imageio\nimport numpy as np\nimport torch\n\nfrom nerfstudio.cameras import camera_utils\nfrom nerfstudio.cameras.cameras import Cameras, CameraType\nfrom nerfstudio.data.dataparsers.base_dataparser import (\n    DataParser,\n    DataParserConfig,\n    DataparserOutputs,\n)\nfrom nerfstudio.data.scene_box import SceneBox\nfrom nerfstudio.data.utils.dataparsers_utils import (\n    get_train_eval_split_filename,\n    get_train_eval_split_fraction,\n    get_train_eval_split_interval,\n    get_train_eval_split_all,\n)\nfrom nerfstudio.utils.io import load_from_json\nfrom nerfstudio.utils.rich_utils import CONSOLE\n\n\n@dataclass\nclass InstantNGPDataParserConfig(DataParserConfig):\n    \"\"\"Instant-NGP dataset parser config\"\"\"\n\n    _target: Type = field(default_factory=lambda: InstantNGP)\n    \"\"\"target class to instantiate\"\"\"\n    data: Path = Path(\"data/ours/posterv2\")\n    \"\"\"Directory or explicit json file path specifying location of data.\"\"\"\n    scene_scale: float = 0.3333\n    \"\"\"How much to scale the scene.\"\"\"\n    eval_mode: Literal[\"fraction\", \"filename\", \"interval\", \"all\"] = \"fraction\"\n    \"\"\"\n    The method to use for splitting the dataset into train and eval.\n    Fraction splits based on a percentage for train and the remaining for eval.\n    Filename splits based on filenames containing train/eval.\n    Interval uses every nth frame for eval.\n    All uses all the images for any split.\n    \"\"\"\n    train_split_fraction: float = 0.9\n    \"\"\"The percentage of the dataset to use for training. Only used when eval_mode is train-split-fraction.\"\"\"\n    eval_interval: int = 8\n    \"\"\"The interval between frames to use for eval. Only used when eval_mode is eval-interval.\"\"\"\n\n\n@dataclass\nclass InstantNGP(DataParser):\n    \"\"\"Instant NGP Dataset\"\"\"\n\n    config: InstantNGPDataParserConfig\n\n    def _generate_dataparser_outputs(self, split=\"train\"):\n        if self.config.data.suffix == \".json\":\n            meta = load_from_json(self.config.data)\n            data_dir = self.config.data.parent\n        else:\n            meta = load_from_json(self.config.data / \"transforms.json\")\n            data_dir = self.config.data\n\n        image_filenames = []\n        mask_filenames = []\n        poses = []\n        num_skipped_image_filenames = 0\n        for frame in meta[\"frames\"]:\n            fname = data_dir / Path(frame[\"file_path\"])\n            # search for png file\n            if not fname.exists():\n                fname = data_dir / Path(frame[\"file_path\"] + \".png\")\n            if not fname.exists():\n                CONSOLE.log(f\"couldn't find {fname} image\")\n                num_skipped_image_filenames += 1\n            else:\n                if \"w\" not in meta:\n                    img_0 = imageio.imread(fname)\n                    h, w = img_0.shape[:2]\n                    meta[\"w\"] = w\n                    if \"h\" in meta:\n                        meta_h = meta[\"h\"]\n                        assert meta_h == h, f\"height of image dont not correspond metadata {h} != {meta_h}\"\n                    else:\n                        meta[\"h\"] = h\n                image_filenames.append(fname)\n                poses.append(np.array(frame[\"transform_matrix\"]))\n                if \"mask_path\" in frame:\n                    mask_fname = data_dir / Path(frame[\"mask_path\"])\n                    mask_filenames.append(mask_fname)\n        if num_skipped_image_filenames >= 0:\n            CONSOLE.print(f\"Skipping {num_skipped_image_filenames} files in dataset split {split}.\")\n        assert (\n            len(image_filenames) != 0\n        ), \"\"\"\n        No image files found.\n        You should check the file_paths in the transforms.json file to make sure they are correct.\n        \"\"\"\n        poses = np.array(poses).astype(np.float32)\n        poses[:, :3, 3] *= self.config.scene_scale\n\n        # find train and eval indices based on the eval_mode specified\n        if self.config.eval_mode == \"fraction\":\n            i_train, i_eval = get_train_eval_split_fraction(image_filenames, self.config.train_split_fraction)\n        elif self.config.eval_mode == \"filename\":\n            i_train, i_eval = get_train_eval_split_filename(image_filenames)\n        elif self.config.eval_mode == \"interval\":\n            i_train, i_eval = get_train_eval_split_interval(image_filenames, self.config.eval_interval)\n        elif self.config.eval_mode == \"all\":\n            CONSOLE.log(\n                \"[yellow] Be careful with '--eval-mode=all'. If using camera optimization, the cameras may diverge in the current implementation, giving unpredictable results.\"\n            )\n            i_train, i_eval = get_train_eval_split_all(image_filenames)\n        else:\n            raise ValueError(f\"Unknown eval mode {self.config.eval_mode}\")\n\n        if split == \"train\":\n            indices = i_train\n        elif split in [\"val\", \"test\"]:\n            indices = i_eval\n        else:\n            raise ValueError(f\"Unknown dataparser split {split}\")\n        # Choose image_filenames and poses based on split, but after auto orient and scaling the poses.\n        image_filenames = [image_filenames[i] for i in indices]\n        mask_filenames = [mask_filenames[i] for i in indices] if len(mask_filenames) > 0 else []\n\n        idx_tensor = torch.tensor(indices, dtype=torch.long)\n        poses = poses[idx_tensor]\n\n        camera_to_world = torch.from_numpy(poses[:, :3])  # camera to world transform\n\n        distortion_params = camera_utils.get_distortion_params(\n            k1=float(meta.get(\"k1\", 0)),\n            k2=float(meta.get(\"k2\", 0)),\n            k3=float(meta.get(\"k3\", 0)),\n            k4=float(meta.get(\"k4\", 0)),\n            p1=float(meta.get(\"p1\", 0)),\n            p2=float(meta.get(\"p2\", 0)),\n        )\n\n        # in x,y,z order\n        # assumes that the scene is centered at the origin\n        aabb_scale = 0.5 * meta.get(\"aabb_scale\", 1)\n\n        scene_box = SceneBox(\n            aabb=torch.tensor(\n                [[-aabb_scale, -aabb_scale, -aabb_scale], [aabb_scale, aabb_scale, aabb_scale]], dtype=torch.float32\n            )\n        )\n\n        fl_x, fl_y = InstantNGP.get_focal_lengths(meta)\n\n        w, h = meta[\"w\"], meta[\"h\"]\n\n        camera_type = CameraType.PERSPECTIVE\n        if meta.get(\"is_fisheye\", False):\n            camera_type = CameraType.FISHEYE\n\n        cameras = Cameras(\n            fx=float(fl_x),\n            fy=float(fl_y),\n            cx=float(meta.get(\"cx\", 0.5 * w)),\n            cy=float(meta.get(\"cy\", 0.5 * h)),\n            distortion_params=distortion_params,\n            height=int(h),\n            width=int(w),\n            camera_to_worlds=camera_to_world,\n            camera_type=camera_type,\n        )\n\n        # TODO(ethan): add alpha background color\n        dataparser_outputs = DataparserOutputs(\n            image_filenames=image_filenames,\n            cameras=cameras,\n            scene_box=scene_box,\n            mask_filenames=mask_filenames if len(mask_filenames) > 0 else None,\n            dataparser_scale=self.config.scene_scale,\n        )\n\n        return dataparser_outputs\n\n    @classmethod\n    def get_focal_lengths(cls, meta: Dict) -> Tuple[float, float]:\n        \"\"\"Reads or computes the focal length from transforms dict.\n        Args:\n            meta: metadata from transforms.json file.\n        Returns:\n            Focal lengths in the x and y directions. Error is raised if these cannot be calculated.\n        \"\"\"\n        fl_x, fl_y = 0, 0\n\n        def fov_to_focal_length(rad, res):\n            return 0.5 * res / np.tan(0.5 * rad)\n\n        if \"fl_x\" in meta:\n            fl_x = meta[\"fl_x\"]\n        elif \"x_fov\" in meta:\n            fl_x = fov_to_focal_length(np.deg2rad(meta[\"x_fov\"]), meta[\"w\"])\n        elif \"camera_angle_x\" in meta:\n            fl_x = fov_to_focal_length(meta[\"camera_angle_x\"], meta[\"w\"])\n\n        if \"camera_angle_y\" not in meta or \"y_fov\" not in meta:\n            fl_y = fl_x\n        else:\n            if \"fl_y\" in meta:\n                fl_y = meta[\"fl_y\"]\n            elif \"y_fov\" in meta:\n                fl_y = fov_to_focal_length(np.deg2rad(meta[\"y_fov\"]), meta[\"h\"])\n            elif \"camera_angle_y\" in meta:\n                fl_y = fov_to_focal_length(meta[\"camera_angle_y\"], meta[\"h\"])\n\n        if fl_x == 0 or fl_y == 0:\n            raise AttributeError(\"Focal length cannot be calculated from transforms.json (missing fields).\")\n\n        return (fl_x, fl_y)\n", "path": "nerfstudio/data/dataparsers/instant_ngp_dataparser.py"}]}