[DONOTMERGE] one-off builds for pytorch 1.13.1

do not merge

.circleci/config.in.yml

@@ -64,7 +64,7 @@ jobs:
       CUDA_VERSION: "11.3"
     resource_class: gpu.nvidia.small.multi
     machine:
-      image: linux-cuda-11:default
+      image: ubuntu-2004:202101-01
     steps:
       - checkout
       - <<: *setupcuda
@@ -116,7 +116,7 @@ jobs:
       # so we aren't running the tests.
       - run:
           name: build
-          no_output_timeout: 40m
+          no_output_timeout: 20m
           command: MAX_JOBS=15 TEST_FLAG=--no-test python3 packaging/build_conda.py
       - store_artifacts:
           path: /opt/conda/conda-bld/linux-64
@@ -128,7 +128,7 @@ jobs:
   binary_linux_conda_cuda:
     <<: *binary_common
     machine:
-      image: linux-cuda-11:default
+      image: ubuntu-1604-cuda-10.2:202012-01
     resource_class: gpu.nvidia.small.multi
     steps:
     - checkout
@@ -145,7 +145,7 @@ jobs:
           docker pull $TESTRUN_DOCKER_IMAGE
     - run:
         name: Build and run tests
-        no_output_timeout: 40m
+        no_output_timeout: 20m
         command: |
           set -e
 
@@ -156,6 +156,24 @@ jobs:
 
           docker run --gpus all  --ipc=host -v $(pwd):/remote -w /remote ${VARS_TO_PASS} ${TESTRUN_DOCKER_IMAGE} python3 ./packaging/build_conda.py
 
+  binary_macos_wheel:
+    <<: *binary_common
+    macos:
+      xcode: "13.4.1"
+    steps:
+      - checkout
+      - run:
+          # Cannot easily deduplicate this as source'ing activate
+          # will set environment variables which we need to propagate
+          # to build_wheel.sh
+          command: |
+            curl -o conda.sh https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-x86_64.sh
+            sh conda.sh -b
+            source $HOME/miniconda3/bin/activate
+            packaging/build_wheel.sh
+      - store_artifacts:
+          path: dist
+
 workflows:
   version: 2
   build_and_test:
@@ -164,8 +182,23 @@ workflows:
       #     context: DOCKERHUB_TOKEN
       {{workflows()}}
       - binary_linux_conda_cuda:
-          name: testrun_conda_cuda_py310_cu117_pyt201
+          name: testrun_conda_cuda_py38_cu102_pyt190
           context: DOCKERHUB_TOKEN
-          python_version: "3.10"
-          pytorch_version: '2.0.1'
-          cu_version: "cu117"
+          python_version: "3.8"
+          pytorch_version: '1.9.0'
+          cu_version: "cu102"
+      - binary_macos_wheel:
+          cu_version: cpu
+          name: macos_wheel_py3.8_cpu
+          python_version: '3.8'
+          pytorch_version: '1.13.0'
+      - binary_macos_wheel:
+          cu_version: cpu
+          name: macos_wheel_py3.9_cpu
+          python_version: '3.9'
+          pytorch_version: '1.13.0'
+      - binary_macos_wheel:
+          cu_version: cpu
+          name: macos_wheel_py3.10_cpu
+          python_version: '3.10'
+          pytorch_version: '1.13.0'

.circleci/config.yml

@@ -64,7 +64,7 @@ jobs:
       CUDA_VERSION: "11.3"
     resource_class: gpu.nvidia.small.multi
     machine:
-      image: linux-cuda-11:default
+      image: ubuntu-2004:202101-01
     steps:
       - checkout
       - <<: *setupcuda
@@ -116,7 +116,7 @@ jobs:
       # so we aren't running the tests.
       - run:
           name: build
-          no_output_timeout: 40m
+          no_output_timeout: 20m
           command: MAX_JOBS=15 TEST_FLAG=--no-test python3 packaging/build_conda.py
       - store_artifacts:
           path: /opt/conda/conda-bld/linux-64
@@ -128,7 +128,7 @@ jobs:
   binary_linux_conda_cuda:
     <<: *binary_common
     machine:
-      image: linux-cuda-11:default
+      image: ubuntu-1604-cuda-10.2:202012-01
     resource_class: gpu.nvidia.small.multi
     steps:
     - checkout
@@ -145,7 +145,7 @@ jobs:
           docker pull $TESTRUN_DOCKER_IMAGE
     - run:
         name: Build and run tests
-        no_output_timeout: 40m
+        no_output_timeout: 20m
         command: |
           set -e
 
@@ -156,54 +156,28 @@ jobs:
 
           docker run --gpus all  --ipc=host -v $(pwd):/remote -w /remote ${VARS_TO_PASS} ${TESTRUN_DOCKER_IMAGE} python3 ./packaging/build_conda.py
 
+  binary_macos_wheel:
+    <<: *binary_common
+    macos:
+      xcode: "13.4.1"
+    steps:
+      - checkout
+      - run:
+          # Cannot easily deduplicate this as source'ing activate
+          # will set environment variables which we need to propagate
+          # to build_wheel.sh
+          command: |
+            curl -o conda.sh https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-x86_64.sh
+            sh conda.sh -b
+            source $HOME/miniconda3/bin/activate
+            packaging/build_wheel.sh
+      - store_artifacts:
+          path: dist
+
 workflows:
   version: 2
   build_and_test:
     jobs:
-      # - main:
-      #     context: DOCKERHUB_TOKEN
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda113
-          context: DOCKERHUB_TOKEN
-          cu_version: cu113
-          name: linux_conda_py38_cu113_pyt1120
-          python_version: '3.8'
-          pytorch_version: 1.12.0
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda116
-          context: DOCKERHUB_TOKEN
-          cu_version: cu116
-          name: linux_conda_py38_cu116_pyt1120
-          python_version: '3.8'
-          pytorch_version: 1.12.0
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda113
-          context: DOCKERHUB_TOKEN
-          cu_version: cu113
-          name: linux_conda_py38_cu113_pyt1121
-          python_version: '3.8'
-          pytorch_version: 1.12.1
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda116
-          context: DOCKERHUB_TOKEN
-          cu_version: cu116
-          name: linux_conda_py38_cu116_pyt1121
-          python_version: '3.8'
-          pytorch_version: 1.12.1
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda116
-          context: DOCKERHUB_TOKEN
-          cu_version: cu116
-          name: linux_conda_py38_cu116_pyt1130
-          python_version: '3.8'
-          pytorch_version: 1.13.0
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda117
-          context: DOCKERHUB_TOKEN
-          cu_version: cu117
-          name: linux_conda_py38_cu117_pyt1130
-          python_version: '3.8'
-          pytorch_version: 1.13.0
       - binary_linux_conda:
           conda_docker_image: pytorch/conda-builder:cuda116
           context: DOCKERHUB_TOKEN
@@ -218,132 +192,6 @@ workflows:
           name: linux_conda_py38_cu117_pyt1131
           python_version: '3.8'
           pytorch_version: 1.13.1
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda117
-          context: DOCKERHUB_TOKEN
-          cu_version: cu117
-          name: linux_conda_py38_cu117_pyt200
-          python_version: '3.8'
-          pytorch_version: 2.0.0
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda118
-          context: DOCKERHUB_TOKEN
-          cu_version: cu118
-          name: linux_conda_py38_cu118_pyt200
-          python_version: '3.8'
-          pytorch_version: 2.0.0
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda117
-          context: DOCKERHUB_TOKEN
-          cu_version: cu117
-          name: linux_conda_py38_cu117_pyt201
-          python_version: '3.8'
-          pytorch_version: 2.0.1
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda118
-          context: DOCKERHUB_TOKEN
-          cu_version: cu118
-          name: linux_conda_py38_cu118_pyt201
-          python_version: '3.8'
-          pytorch_version: 2.0.1
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda118
-          context: DOCKERHUB_TOKEN
-          cu_version: cu118
-          name: linux_conda_py38_cu118_pyt210
-          python_version: '3.8'
-          pytorch_version: 2.1.0
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda121
-          context: DOCKERHUB_TOKEN
-          cu_version: cu121
-          name: linux_conda_py38_cu121_pyt210
-          python_version: '3.8'
-          pytorch_version: 2.1.0
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda118
-          context: DOCKERHUB_TOKEN
-          cu_version: cu118
-          name: linux_conda_py38_cu118_pyt211
-          python_version: '3.8'
-          pytorch_version: 2.1.1
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda121
-          context: DOCKERHUB_TOKEN
-          cu_version: cu121
-          name: linux_conda_py38_cu121_pyt211
-          python_version: '3.8'
-          pytorch_version: 2.1.1
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda118
-          context: DOCKERHUB_TOKEN
-          cu_version: cu118
-          name: linux_conda_py38_cu118_pyt212
-          python_version: '3.8'
-          pytorch_version: 2.1.2
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda121
-          context: DOCKERHUB_TOKEN
-          cu_version: cu121
-          name: linux_conda_py38_cu121_pyt212
-          python_version: '3.8'
-          pytorch_version: 2.1.2
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda118
-          context: DOCKERHUB_TOKEN
-          cu_version: cu118
-          name: linux_conda_py38_cu118_pyt220
-          python_version: '3.8'
-          pytorch_version: 2.2.0
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda121
-          context: DOCKERHUB_TOKEN
-          cu_version: cu121
-          name: linux_conda_py38_cu121_pyt220
-          python_version: '3.8'
-          pytorch_version: 2.2.0
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda113
-          context: DOCKERHUB_TOKEN
-          cu_version: cu113
-          name: linux_conda_py39_cu113_pyt1120
-          python_version: '3.9'
-          pytorch_version: 1.12.0
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda116
-          context: DOCKERHUB_TOKEN
-          cu_version: cu116
-          name: linux_conda_py39_cu116_pyt1120
-          python_version: '3.9'
-          pytorch_version: 1.12.0
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda113
-          context: DOCKERHUB_TOKEN
-          cu_version: cu113
-          name: linux_conda_py39_cu113_pyt1121
-          python_version: '3.9'
-          pytorch_version: 1.12.1
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda116
-          context: DOCKERHUB_TOKEN
-          cu_version: cu116
-          name: linux_conda_py39_cu116_pyt1121
-          python_version: '3.9'
-          pytorch_version: 1.12.1
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda116
-          context: DOCKERHUB_TOKEN
-          cu_version: cu116
-          name: linux_conda_py39_cu116_pyt1130
-          python_version: '3.9'
-          pytorch_version: 1.13.0
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda117
-          context: DOCKERHUB_TOKEN
-          cu_version: cu117
-          name: linux_conda_py39_cu117_pyt1130
-          python_version: '3.9'
-          pytorch_version: 1.13.0
       - binary_linux_conda:
           conda_docker_image: pytorch/conda-builder:cuda116
           context: DOCKERHUB_TOKEN
@@ -358,132 +206,6 @@ workflows:
           name: linux_conda_py39_cu117_pyt1131
           python_version: '3.9'
           pytorch_version: 1.13.1
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda117
-          context: DOCKERHUB_TOKEN
-          cu_version: cu117
-          name: linux_conda_py39_cu117_pyt200
-          python_version: '3.9'
-          pytorch_version: 2.0.0
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda118
-          context: DOCKERHUB_TOKEN
-          cu_version: cu118
-          name: linux_conda_py39_cu118_pyt200
-          python_version: '3.9'
-          pytorch_version: 2.0.0
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda117
-          context: DOCKERHUB_TOKEN
-          cu_version: cu117
-          name: linux_conda_py39_cu117_pyt201
-          python_version: '3.9'
-          pytorch_version: 2.0.1
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda118
-          context: DOCKERHUB_TOKEN
-          cu_version: cu118
-          name: linux_conda_py39_cu118_pyt201
-          python_version: '3.9'
-          pytorch_version: 2.0.1
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda118
-          context: DOCKERHUB_TOKEN
-          cu_version: cu118
-          name: linux_conda_py39_cu118_pyt210
-          python_version: '3.9'
-          pytorch_version: 2.1.0
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda121
-          context: DOCKERHUB_TOKEN
-          cu_version: cu121
-          name: linux_conda_py39_cu121_pyt210
-          python_version: '3.9'
-          pytorch_version: 2.1.0
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda118
-          context: DOCKERHUB_TOKEN
-          cu_version: cu118
-          name: linux_conda_py39_cu118_pyt211
-          python_version: '3.9'
-          pytorch_version: 2.1.1
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda121
-          context: DOCKERHUB_TOKEN
-          cu_version: cu121
-          name: linux_conda_py39_cu121_pyt211
-          python_version: '3.9'
-          pytorch_version: 2.1.1
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda118
-          context: DOCKERHUB_TOKEN
-          cu_version: cu118
-          name: linux_conda_py39_cu118_pyt212
-          python_version: '3.9'
-          pytorch_version: 2.1.2
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda121
-          context: DOCKERHUB_TOKEN
-          cu_version: cu121
-          name: linux_conda_py39_cu121_pyt212
-          python_version: '3.9'
-          pytorch_version: 2.1.2
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda118
-          context: DOCKERHUB_TOKEN
-          cu_version: cu118
-          name: linux_conda_py39_cu118_pyt220
-          python_version: '3.9'
-          pytorch_version: 2.2.0
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda121
-          context: DOCKERHUB_TOKEN
-          cu_version: cu121
-          name: linux_conda_py39_cu121_pyt220
-          python_version: '3.9'
-          pytorch_version: 2.2.0
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda113
-          context: DOCKERHUB_TOKEN
-          cu_version: cu113
-          name: linux_conda_py310_cu113_pyt1120
-          python_version: '3.10'
-          pytorch_version: 1.12.0
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda116
-          context: DOCKERHUB_TOKEN
-          cu_version: cu116
-          name: linux_conda_py310_cu116_pyt1120
-          python_version: '3.10'
-          pytorch_version: 1.12.0
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda113
-          context: DOCKERHUB_TOKEN
-          cu_version: cu113
-          name: linux_conda_py310_cu113_pyt1121
-          python_version: '3.10'
-          pytorch_version: 1.12.1
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda116
-          context: DOCKERHUB_TOKEN
-          cu_version: cu116
-          name: linux_conda_py310_cu116_pyt1121
-          python_version: '3.10'
-          pytorch_version: 1.12.1
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda116
-          context: DOCKERHUB_TOKEN
-          cu_version: cu116
-          name: linux_conda_py310_cu116_pyt1130
-          python_version: '3.10'
-          pytorch_version: 1.13.0
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda117
-          context: DOCKERHUB_TOKEN
-          cu_version: cu117
-          name: linux_conda_py310_cu117_pyt1130
-          python_version: '3.10'
-          pytorch_version: 1.13.0
       - binary_linux_conda:
           conda_docker_image: pytorch/conda-builder:cuda116
           context: DOCKERHUB_TOKEN
@@ -498,163 +220,3 @@ workflows:
           name: linux_conda_py310_cu117_pyt1131
           python_version: '3.10'
           pytorch_version: 1.13.1
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda117
-          context: DOCKERHUB_TOKEN
-          cu_version: cu117
-          name: linux_conda_py310_cu117_pyt200
-          python_version: '3.10'
-          pytorch_version: 2.0.0
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda118
-          context: DOCKERHUB_TOKEN
-          cu_version: cu118
-          name: linux_conda_py310_cu118_pyt200
-          python_version: '3.10'
-          pytorch_version: 2.0.0
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda117
-          context: DOCKERHUB_TOKEN
-          cu_version: cu117
-          name: linux_conda_py310_cu117_pyt201
-          python_version: '3.10'
-          pytorch_version: 2.0.1
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda118
-          context: DOCKERHUB_TOKEN
-          cu_version: cu118
-          name: linux_conda_py310_cu118_pyt201
-          python_version: '3.10'
-          pytorch_version: 2.0.1
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda118
-          context: DOCKERHUB_TOKEN
-          cu_version: cu118
-          name: linux_conda_py310_cu118_pyt210
-          python_version: '3.10'
-          pytorch_version: 2.1.0
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda121
-          context: DOCKERHUB_TOKEN
-          cu_version: cu121
-          name: linux_conda_py310_cu121_pyt210
-          python_version: '3.10'
-          pytorch_version: 2.1.0
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda118
-          context: DOCKERHUB_TOKEN
-          cu_version: cu118
-          name: linux_conda_py310_cu118_pyt211
-          python_version: '3.10'
-          pytorch_version: 2.1.1
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda121
-          context: DOCKERHUB_TOKEN
-          cu_version: cu121
-          name: linux_conda_py310_cu121_pyt211
-          python_version: '3.10'
-          pytorch_version: 2.1.1
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda118
-          context: DOCKERHUB_TOKEN
-          cu_version: cu118
-          name: linux_conda_py310_cu118_pyt212
-          python_version: '3.10'
-          pytorch_version: 2.1.2
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda121
-          context: DOCKERHUB_TOKEN
-          cu_version: cu121
-          name: linux_conda_py310_cu121_pyt212
-          python_version: '3.10'
-          pytorch_version: 2.1.2
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda118
-          context: DOCKERHUB_TOKEN
-          cu_version: cu118
-          name: linux_conda_py310_cu118_pyt220
-          python_version: '3.10'
-          pytorch_version: 2.2.0
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda121
-          context: DOCKERHUB_TOKEN
-          cu_version: cu121
-          name: linux_conda_py310_cu121_pyt220
-          python_version: '3.10'
-          pytorch_version: 2.2.0
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda118
-          context: DOCKERHUB_TOKEN
-          cu_version: cu118
-          name: linux_conda_py311_cu118_pyt210
-          python_version: '3.11'
-          pytorch_version: 2.1.0
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda121
-          context: DOCKERHUB_TOKEN
-          cu_version: cu121
-          name: linux_conda_py311_cu121_pyt210
-          python_version: '3.11'
-          pytorch_version: 2.1.0
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda118
-          context: DOCKERHUB_TOKEN
-          cu_version: cu118
-          name: linux_conda_py311_cu118_pyt211
-          python_version: '3.11'
-          pytorch_version: 2.1.1
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda121
-          context: DOCKERHUB_TOKEN
-          cu_version: cu121
-          name: linux_conda_py311_cu121_pyt211
-          python_version: '3.11'
-          pytorch_version: 2.1.1
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda118
-          context: DOCKERHUB_TOKEN
-          cu_version: cu118
-          name: linux_conda_py311_cu118_pyt212
-          python_version: '3.11'
-          pytorch_version: 2.1.2
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda121
-          context: DOCKERHUB_TOKEN
-          cu_version: cu121
-          name: linux_conda_py311_cu121_pyt212
-          python_version: '3.11'
-          pytorch_version: 2.1.2
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda118
-          context: DOCKERHUB_TOKEN
-          cu_version: cu118
-          name: linux_conda_py311_cu118_pyt220
-          python_version: '3.11'
-          pytorch_version: 2.2.0
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda121
-          context: DOCKERHUB_TOKEN
-          cu_version: cu121
-          name: linux_conda_py311_cu121_pyt220
-          python_version: '3.11'
-          pytorch_version: 2.2.0
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda118
-          context: DOCKERHUB_TOKEN
-          cu_version: cu118
-          name: linux_conda_py312_cu118_pyt220
-          python_version: '3.12'
-          pytorch_version: 2.2.0
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda121
-          context: DOCKERHUB_TOKEN
-          cu_version: cu121
-          name: linux_conda_py312_cu121_pyt220
-          python_version: '3.12'
-          pytorch_version: 2.2.0
-      - binary_linux_conda_cuda:
-          name: testrun_conda_cuda_py310_cu117_pyt201
-          context: DOCKERHUB_TOKEN
-          python_version: "3.10"
-          pytorch_version: '2.0.1'
-          cu_version: "cu117"

.circleci/regenerate.py

@@ -18,53 +18,55 @@ from packaging import version
 
 # The CUDA versions which have pytorch conda packages available for linux for each
 # version of pytorch.
+# Pytorch 1.4 also supports cuda 10.0 but we no longer build for cuda 10.0 at all.
 CONDA_CUDA_VERSIONS = {
-    "1.12.0": ["cu113", "cu116"],
-    "1.12.1": ["cu113", "cu116"],
+    "1.9.0": ["cu102", "cu111"],
+    "1.9.1": ["cu102", "cu111"],
+    "1.10.0": ["cu102", "cu111", "cu113"],
+    "1.10.1": ["cu102", "cu111", "cu113"],
+    "1.10.2": ["cu102", "cu111", "cu113"],
+    "1.11.0": ["cu102", "cu111", "cu113", "cu115"],
+    "1.12.0": ["cu102", "cu113", "cu116"],
+    "1.12.1": ["cu102", "cu113", "cu116"],
     "1.13.0": ["cu116", "cu117"],
-    "1.13.1": ["cu116", "cu117"],
-    "2.0.0": ["cu117", "cu118"],
-    "2.0.1": ["cu117", "cu118"],
-    "2.1.0": ["cu118", "cu121"],
-    "2.1.1": ["cu118", "cu121"],
-    "2.1.2": ["cu118", "cu121"],
-    "2.2.0": ["cu118", "cu121"],
 }
 
 
 def conda_docker_image_for_cuda(cuda_version):
-    if len(cuda_version) != 5:
-        raise ValueError("Unknown cuda version")
-    return "pytorch/conda-builder:cuda" + cuda_version[2:]
+    if cuda_version in ("cu101", "cu102", "cu111"):
+        return None
+    if cuda_version == "cu113":
+        return "pytorch/conda-builder:cuda113"
+    if cuda_version == "cu115":
+        return "pytorch/conda-builder:cuda115"
+    if cuda_version == "cu116":
+        return "pytorch/conda-builder:cuda116"
+    if cuda_version == "cu117":
+        return "pytorch/conda-builder:cuda117"
+    raise ValueError("Unknown cuda version")
 
 
 def pytorch_versions_for_python(python_version):
-    if python_version in ["3.8", "3.9"]:
+    if python_version in ["3.7", "3.8"]:
         return list(CONDA_CUDA_VERSIONS)
+    if python_version == "3.9":
+        return [
+            i
+            for i in CONDA_CUDA_VERSIONS
+            if version.Version(i) > version.Version("1.7.0")
+        ]
     if python_version == "3.10":
         return [
             i
             for i in CONDA_CUDA_VERSIONS
             if version.Version(i) >= version.Version("1.11.0")
         ]
-    if python_version == "3.11":
-        return [
-            i
-            for i in CONDA_CUDA_VERSIONS
-            if version.Version(i) >= version.Version("2.1.0")
-        ]
-    if python_version == "3.12":
-        return [
-            i
-            for i in CONDA_CUDA_VERSIONS
-            if version.Version(i) >= version.Version("2.2.0")
-        ]
 
 
 def workflows(prefix="", filter_branch=None, upload=False, indentation=6):
     w = []
     for btype in ["conda"]:
-        for python_version in ["3.8", "3.9", "3.10", "3.11", "3.12"]:
+        for python_version in ["3.8", "3.9", "3.10"]:
             for pytorch_version in pytorch_versions_for_python(python_version):
                 for cu_version in CONDA_CUDA_VERSIONS[pytorch_version]:
                     w += workflow_pair(

.flake8

@@ -1,8 +1,5 @@
 [flake8]
-# B028 No explicit stacklevel argument found.
-# B907 'foo' is manually surrounded by quotes, consider using the `!r` conversion flag.
-# B905 `zip()` without an explicit `strict=` parameter.
-ignore = E203, E266, E501, W503, E221, B028, B905, B907
+ignore = E203, E266, E501, W503, E221
 max-line-length = 88
 max-complexity = 18
 select = B,C,E,F,W,T4,B9

INSTALL.md

@@ -9,19 +9,19 @@ The core library is written in PyTorch. Several components have underlying imple
 
 - Linux or macOS or Windows
 - Python 3.8, 3.9 or 3.10
-- PyTorch 1.12.0, 1.12.1, 1.13.0, 2.0.0, 2.0.1, 2.1.0, 2.1.1, 2.1.2 or 2.2.0.
+- PyTorch 1.9.0, 1.9.1, 1.10.0, 1.10.1, 1.10.2, 1.11.0, 1.12.0, 1.12.1 or 1.13.0.
 - torchvision that matches the PyTorch installation. You can install them together as explained at pytorch.org to make sure of this.
 - gcc & g++ ≥ 4.9
 - [fvcore](https://github.com/facebookresearch/fvcore)
 - [ioPath](https://github.com/facebookresearch/iopath)
 - If CUDA is to be used, use a version which is supported by the corresponding pytorch version and at least version 9.2.
-- If CUDA older than 11.7 is to be used and you are building from source, the CUB library must be available. We recommend version 1.10.0.
+- If CUDA is to be used and you are building from source, the CUB library must be available. We recommend version 1.10.0.
 
 The runtime dependencies can be installed by running:
 ```
 conda create -n pytorch3d python=3.9
 conda activate pytorch3d
-conda install pytorch=1.13.0 torchvision pytorch-cuda=11.6 -c pytorch -c nvidia
+conda install -c pytorch pytorch=1.9.1 torchvision cudatoolkit=11.6
 conda install -c fvcore -c iopath -c conda-forge fvcore iopath
 ```
 
@@ -77,8 +77,13 @@ Or, to install a nightly (non-official, alpha) build:
 # Anaconda Cloud
 conda install pytorch3d -c pytorch3d-nightly
 ```
+### 2. Install from PyPI, on Mac only.
+This works with pytorch 1.13.0 only. The build is CPU only.
+```
+pip install pytorch3d
+```
 
-### 2. Install wheels for Linux
+### 3. Install wheels for Linux
 We have prebuilt wheels with CUDA for Linux for PyTorch 1.11.0, for each of the supported CUDA versions,
 for Python 3.8 and 3.9. This is for ease of use on Google Colab.
 These are installed in a special way.
@@ -97,7 +102,6 @@ version_str="".join([
     torch.version.cuda.replace(".",""),
     f"_pyt{pyt_version_str}"
 ])
-!pip install fvcore iopath
 !pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html
 ```
 

README.md

@@ -12,7 +12,7 @@ Key features include:
 - Data structure for storing and manipulating triangle meshes
 - Efficient operations on triangle meshes (projective transformations, graph convolution, sampling, loss functions)
 - A differentiable mesh renderer
-- Implicitron, see [its README](projects/implicitron_trainer), a framework for new-view synthesis via implicit representations. ([blog post](https://ai.facebook.com/blog/implicitron-a-new-modular-extensible-framework-for-neural-implicit-representations-in-pytorch3d/))
+- Implicitron, see [its README](projects/implicitron_trainer), a framework for new-view synthesis via implicit representations.
 
 PyTorch3D is designed to integrate smoothly with deep learning methods for predicting and manipulating 3D data.
 For this reason, all operators in PyTorch3D:
@@ -24,8 +24,6 @@ For this reason, all operators in PyTorch3D:
 
 Within FAIR, PyTorch3D has been used to power research projects such as [Mesh R-CNN](https://arxiv.org/abs/1906.02739).
 
-See our [blog post](https://ai.facebook.com/blog/-introducing-pytorch3d-an-open-source-library-for-3d-deep-learning/) to see more demos and learn about PyTorch3D.
-
 ## Installation
 
 For detailed instructions refer to [INSTALL.md](INSTALL.md).
@@ -146,14 +144,6 @@ If you are using the pulsar backend for sphere-rendering (the `PulsarPointRender
 
 Please see below for a timeline of the codebase updates in reverse chronological order. We are sharing updates on the releases as well as research projects which are built with PyTorch3D. The changelogs for the releases are available under [`Releases`](https://github.com/facebookresearch/pytorch3d/releases),  and the builds can be installed using `conda` as per the instructions in [INSTALL.md](INSTALL.md).
 
-**[Oct 31st 2023]:**   PyTorch3D [v0.7.5](https://github.com/facebookresearch/pytorch3d/releases/tag/v0.7.5) released.
-
-**[May 10th 2023]:**   PyTorch3D [v0.7.4](https://github.com/facebookresearch/pytorch3d/releases/tag/v0.7.4) released.
-
-**[Apr 5th 2023]:**   PyTorch3D [v0.7.3](https://github.com/facebookresearch/pytorch3d/releases/tag/v0.7.3) released.
-
-**[Dec 19th 2022]:**   PyTorch3D [v0.7.2](https://github.com/facebookresearch/pytorch3d/releases/tag/v0.7.2) released.
-
 **[Oct 23rd 2022]:**   PyTorch3D [v0.7.1](https://github.com/facebookresearch/pytorch3d/releases/tag/v0.7.1) released.
 
 **[Aug 10th 2022]:**   PyTorch3D [v0.7.0](https://github.com/facebookresearch/pytorch3d/releases/tag/v0.7.0) released with Implicitron and MeshRasterizerOpenGL.

docs/.readthedocs.yaml

@@ -1,27 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the BSD-style license found in the
-# LICENSE file in the root directory of this source tree.
-
-# Read the Docs configuration file
-# See https://docs.readthedocs.io/en/stable/config-file/v2.html for details
-
-# Required
-version: 2
-
-# Set the version of Python and other tools you might need
-build:
-  os: ubuntu-22.04
-  tools:
-    python: "3.11"
-
-# Build documentation in the docs/ directory with Sphinx
-sphinx:
-  configuration: docs/conf.py
-
-# We recommend specifying your dependencies to enable reproducible builds:
-# https://docs.readthedocs.io/en/stable/guides/reproducible-builds.html
-python:
-  install:
-  - requirements: docs/requirements.txt

docs/README.md

@@ -3,7 +3,7 @@
 ### Install dependencies
 
 ```
-pip install -U recommonmark sphinx sphinx_rtd_theme sphinx_markdown_tables
+pip install -U recommonmark mock sphinx sphinx_rtd_theme sphinx_markdown_tables
 ```
 
 ### Add symlink to the root README.md

docs/conf.py

@@ -20,8 +20,7 @@
 import os
 import sys
 
-import unittest.mock as mock
-
+import mock
 from recommonmark.parser import CommonMarkParser
 from recommonmark.states import DummyStateMachine
 from sphinx.builders.html import StandaloneHTMLBuilder

docs/notes/cameras.md

@@ -85,7 +85,7 @@ cameras_ndc = PerspectiveCameras(focal_length=fcl_ndc, principal_point=prp_ndc)
 # Screen space camera
 image_size = ((128, 256),)    # (h, w)
 fcl_screen = (76.8,)          # fcl_ndc * min(image_size) / 2
-prp_screen = ((115.2, 32), )  # w / 2 - px_ndc * min(image_size) / 2, h / 2 - py_ndc * min(image_size) / 2
+prp_screen = ((115.2, 48), )  # w / 2 - px_ndc * min(image_size) / 2, h / 2 - py_ndc * min(image_size) / 2
 cameras_screen = PerspectiveCameras(focal_length=fcl_screen, principal_point=prp_screen, in_ndc=False, image_size=image_size)
 ```
 

docs/notes/io.md

@@ -7,20 +7,20 @@ sidebar_label: File IO
 There is a flexible interface for loading and saving point clouds and meshes from different formats.
 
 The main usage is via the `pytorch3d.io.IO` object, and its methods
-`load_mesh`, `save_mesh`, `load_pointcloud` and `save_pointcloud`.
+`load_mesh`, `save_mesh`, `load_point_cloud` and `save_point_cloud`.
 
 For example, to load a mesh you might do
 ```
 from pytorch3d.io import IO
 
 device=torch.device("cuda:0")
-mesh = IO().load_mesh("mymesh.obj", device=device)
+mesh = IO().load_mesh("mymesh.ply", device=device)
 ```
 
 and to save a pointcloud you might do
 ```
 pcl = Pointclouds(...)
-IO().save_pointcloud(pcl, "output_pointcloud.ply")
+IO().save_point_cloud(pcl, "output_pointcloud.obj")
 ```
 
 For meshes, this supports OBJ, PLY and OFF files.
@@ -31,4 +31,4 @@ In addition, there is experimental support for loading meshes from
 [glTF 2 assets](https://github.com/KhronosGroup/glTF/tree/master/specification/2.0)
 stored either in a GLB container file or a glTF JSON file with embedded binary data.
 This must be enabled explicitly, as described in
-`pytorch3d/io/experimental_gltf_io.py`.
+`pytorch3d/io/experimental_gltf_io.ply`.

docs/requirements.txt

@@ -1,11 +1,12 @@
 docutils>=0.14
 Sphinx>=1.7
-recommonmark
+recommonmark==0.4.0
 sphinx_rtd_theme
 sphinx_markdown_tables
+mock
 numpy
 iopath
 fvcore
-https://download.pytorch.org/whl/cpu/torchvision-0.15.2%2Bcpu-cp311-cp311-linux_x86_64.whl
-https://download.pytorch.org/whl/cpu/torch-2.0.1%2Bcpu-cp311-cp311-linux_x86_64.whl
+https://download.pytorch.org/whl/cpu/torchvision-0.8.2%2Bcpu-cp37-cp37m-linux_x86_64.whl
+https://download.pytorch.org/whl/cpu/torch-1.7.1%2Bcpu-cp37-cp37m-linux_x86_64.whl
 omegaconf

docs/tutorials/bundle_adjustment.ipynb

@@ -89,7 +89,7 @@
     "except ModuleNotFoundError:\n",
     "    need_pytorch3d=True\n",
     "if need_pytorch3d:\n",
-    "    if torch.__version__.startswith(\"2.2.\") and sys.platform.startswith(\"linux\"):\n",
+    "    if torch.__version__.startswith(\"1.13.\") and sys.platform.startswith(\"linux\"):\n",
     "        # We try to install PyTorch3D via a released wheel.\n",
     "        pyt_version_str=torch.__version__.split(\"+\")[0].replace(\".\", \"\")\n",
     "        version_str=\"\".join([\n",
@@ -101,6 +101,9 @@
     "        !pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html\n",
     "    else:\n",
     "        # We try to install PyTorch3D from source.\n",
+    "        !curl -LO https://github.com/NVIDIA/cub/archive/1.10.0.tar.gz\n",
+    "        !tar xzf 1.10.0.tar.gz\n",
+    "        os.environ[\"CUB_HOME\"] = os.getcwd() + \"/cub-1.10.0\"\n",
     "        !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'"
    ]
   },

docs/tutorials/camera_position_optimization_with_differentiable_rendering.ipynb

@@ -76,7 +76,7 @@
     "except ModuleNotFoundError:\n",
     "    need_pytorch3d=True\n",
     "if need_pytorch3d:\n",
-    "    if torch.__version__.startswith(\"2.2.\") and sys.platform.startswith(\"linux\"):\n",
+    "    if torch.__version__.startswith(\"1.13.\") and sys.platform.startswith(\"linux\"):\n",
     "        # We try to install PyTorch3D via a released wheel.\n",
     "        pyt_version_str=torch.__version__.split(\"+\")[0].replace(\".\", \"\")\n",
     "        version_str=\"\".join([\n",
@@ -88,6 +88,9 @@
     "        !pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html\n",
     "    else:\n",
     "        # We try to install PyTorch3D from source.\n",
+    "        !curl -LO https://github.com/NVIDIA/cub/archive/1.10.0.tar.gz\n",
+    "        !tar xzf 1.10.0.tar.gz\n",
+    "        os.environ[\"CUB_HOME\"] = os.getcwd() + \"/cub-1.10.0\"\n",
     "        !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'"
    ]
   },

docs/tutorials/dataloaders_ShapeNetCore_R2N2.ipynb

@@ -51,7 +51,7 @@
     "except ModuleNotFoundError:\n",
     "    need_pytorch3d=True\n",
     "if need_pytorch3d:\n",
-    "    if torch.__version__.startswith(\"2.2.\") and sys.platform.startswith(\"linux\"):\n",
+    "    if torch.__version__.startswith(\"1.13.\") and sys.platform.startswith(\"linux\"):\n",
     "        # We try to install PyTorch3D via a released wheel.\n",
     "        pyt_version_str=torch.__version__.split(\"+\")[0].replace(\".\", \"\")\n",
     "        version_str=\"\".join([\n",
@@ -63,6 +63,9 @@
     "        !pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html\n",
     "    else:\n",
     "        # We try to install PyTorch3D from source.\n",
+    "        !curl -LO https://github.com/NVIDIA/cub/archive/1.10.0.tar.gz\n",
+    "        !tar xzf 1.10.0.tar.gz\n",
+    "        os.environ[\"CUB_HOME\"] = os.getcwd() + \"/cub-1.10.0\"\n",
     "        !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'"
    ]
   },

docs/tutorials/deform_source_mesh_to_target_mesh.ipynb

@@ -90,7 +90,7 @@
     "except ModuleNotFoundError:\n",
     "    need_pytorch3d=True\n",
     "if need_pytorch3d:\n",
-    "    if torch.__version__.startswith(\"2.2.\") and sys.platform.startswith(\"linux\"):\n",
+    "    if torch.__version__.startswith(\"1.13.\") and sys.platform.startswith(\"linux\"):\n",
     "        # We try to install PyTorch3D via a released wheel.\n",
     "        pyt_version_str=torch.__version__.split(\"+\")[0].replace(\".\", \"\")\n",
     "        version_str=\"\".join([\n",
@@ -102,6 +102,9 @@
     "        !pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html\n",
     "    else:\n",
     "        # We try to install PyTorch3D from source.\n",
+    "        !curl -LO https://github.com/NVIDIA/cub/archive/1.10.0.tar.gz\n",
+    "        !tar xzf 1.10.0.tar.gz\n",
+    "        os.environ[\"CUB_HOME\"] = os.getcwd() + \"/cub-1.10.0\"\n",
     "        !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'"
    ]
   },
@@ -189,7 +192,7 @@
    "outputs": [],
    "source": [
     "# Load the dolphin mesh.\n",
-    "trg_obj = 'dolphin.obj'"
+    "trg_obj = os.path.join('dolphin.obj')"
    ]
   },
   {
@@ -244,7 +247,7 @@
     "id": "dYWDl4VGWHRK"
    },
    "source": [
-    "## 2. Visualize the source and target meshes"
+    "###  Visualize the source and target meshes"
    ]
   },
   {
@@ -262,7 +265,7 @@
     "    points = sample_points_from_meshes(mesh, 5000)\n",
     "    x, y, z = points.clone().detach().cpu().squeeze().unbind(1)    \n",
     "    fig = plt.figure(figsize=(5, 5))\n",
-    "    ax = fig.add_subplot(111, projection='3d')\n",
+    "    ax = Axes3D(fig)\n",
     "    ax.scatter3D(x, z, -y)\n",
     "    ax.set_xlabel('x')\n",
     "    ax.set_ylabel('z')\n",
@@ -482,7 +485,7 @@
     "final_verts = final_verts * scale + center\n",
     "\n",
     "# Store the predicted mesh using save_obj\n",
-    "final_obj = 'final_model.obj'\n",
+    "final_obj = os.path.join('./', 'final_model.obj')\n",
     "save_obj(final_obj, final_verts, final_faces)"
    ]
   },

docs/tutorials/fit_simple_neural_radiance_field.ipynb

@@ -56,7 +56,7 @@
     "except ModuleNotFoundError:\n",
     "    need_pytorch3d=True\n",
     "if need_pytorch3d:\n",
-    "    if torch.__version__.startswith(\"2.2.\") and sys.platform.startswith(\"linux\"):\n",
+    "    if torch.__version__.startswith(\"1.13.\") and sys.platform.startswith(\"linux\"):\n",
     "        # We try to install PyTorch3D via a released wheel.\n",
     "        pyt_version_str=torch.__version__.split(\"+\")[0].replace(\".\", \"\")\n",
     "        version_str=\"\".join([\n",
@@ -68,6 +68,9 @@
     "        !pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html\n",
     "    else:\n",
     "        # We try to install PyTorch3D from source.\n",
+    "        !curl -LO https://github.com/NVIDIA/cub/archive/1.10.0.tar.gz\n",
+    "        !tar xzf 1.10.0.tar.gz\n",
+    "        os.environ[\"CUB_HOME\"] = os.getcwd() + \"/cub-1.10.0\"\n",
     "        !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'"
    ]
   },

docs/tutorials/fit_textured_mesh.ipynb

@@ -68,7 +68,7 @@
     "except ModuleNotFoundError:\n",
     "    need_pytorch3d=True\n",
     "if need_pytorch3d:\n",
-    "    if torch.__version__.startswith(\"2.2.\") and sys.platform.startswith(\"linux\"):\n",
+    "    if torch.__version__.startswith(\"1.13.\") and sys.platform.startswith(\"linux\"):\n",
     "        # We try to install PyTorch3D via a released wheel.\n",
     "        pyt_version_str=torch.__version__.split(\"+\")[0].replace(\".\", \"\")\n",
     "        version_str=\"\".join([\n",
@@ -80,6 +80,9 @@
     "        !pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html\n",
     "    else:\n",
     "        # We try to install PyTorch3D from source.\n",
+    "        !curl -LO https://github.com/NVIDIA/cub/archive/1.10.0.tar.gz\n",
+    "        !tar xzf 1.10.0.tar.gz\n",
+    "        os.environ[\"CUB_HOME\"] = os.getcwd() + \"/cub-1.10.0\"\n",
     "        !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'"
    ]
   },

docs/tutorials/fit_textured_volume.ipynb

@@ -47,7 +47,7 @@
     "except ModuleNotFoundError:\n",
     "    need_pytorch3d=True\n",
     "if need_pytorch3d:\n",
-    "    if torch.__version__.startswith(\"2.2.\") and sys.platform.startswith(\"linux\"):\n",
+    "    if torch.__version__.startswith(\"1.13.\") and sys.platform.startswith(\"linux\"):\n",
     "        # We try to install PyTorch3D via a released wheel.\n",
     "        pyt_version_str=torch.__version__.split(\"+\")[0].replace(\".\", \"\")\n",
     "        version_str=\"\".join([\n",
@@ -59,6 +59,9 @@
     "        !pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html\n",
     "    else:\n",
     "        # We try to install PyTorch3D from source.\n",
+    "        !curl -LO https://github.com/NVIDIA/cub/archive/1.10.0.tar.gz\n",
+    "        !tar xzf 1.10.0.tar.gz\n",
+    "        os.environ[\"CUB_HOME\"] = os.getcwd() + \"/cub-1.10.0\"\n",
     "        !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'"
    ]
   },

docs/tutorials/implicitron_config_system.ipynb

@@ -78,7 +78,7 @@
         "except ModuleNotFoundError:\n",
         "    need_pytorch3d=True\n",
         "if need_pytorch3d:\n",
-        "    if torch.__version__.startswith(\"2.2.\") and sys.platform.startswith(\"linux\"):\n",
+        "    if torch.__version__.startswith(\"1.13.\") and sys.platform.startswith(\"linux\"):\n",
         "        # We try to install PyTorch3D via a released wheel.\n",
         "        pyt_version_str=torch.__version__.split(\"+\")[0].replace(\".\", \"\")\n",
         "        version_str=\"\".join([\n",
@@ -90,6 +90,9 @@
         "        !pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html\n",
         "    else:\n",
         "        # We try to install PyTorch3D from source.\n",
+        "        !curl -LO https://github.com/NVIDIA/cub/archive/1.10.0.tar.gz\n",
+        "        !tar xzf 1.10.0.tar.gz\n",
+        "        os.environ[\"CUB_HOME\"] = os.getcwd() + \"/cub-1.10.0\"\n",
         "        !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'"
       ]
     },

docs/tutorials/implicitron_volumes.ipynb

@@ -72,7 +72,7 @@
         "except ModuleNotFoundError:\n",
         "    need_pytorch3d=True\n",
         "if need_pytorch3d:\n",
-        "    if torch.__version__.startswith(\"2.2.\") and sys.platform.startswith(\"linux\"):\n",
+        "    if torch.__version__.startswith(\"1.13.\") and sys.platform.startswith(\"linux\"):\n",
         "        # We try to install PyTorch3D via a released wheel.\n",
         "        pyt_version_str=torch.__version__.split(\"+\")[0].replace(\".\", \"\")\n",
         "        version_str=\"\".join([\n",
@@ -84,6 +84,9 @@
         "        !pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html\n",
         "    else:\n",
         "        # We try to install PyTorch3D from source.\n",
+        "        !curl -LO https://github.com/NVIDIA/cub/archive/1.10.0.tar.gz\n",
+        "        !tar xzf 1.10.0.tar.gz\n",
+        "        os.environ[\"CUB_HOME\"] = os.getcwd() + \"/cub-1.10.0\"\n",
         "        !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'"
       ]
     },

docs/tutorials/render_colored_points.ipynb

@@ -50,7 +50,7 @@
     "except ModuleNotFoundError:\n",
     "    need_pytorch3d=True\n",
     "if need_pytorch3d:\n",
-    "    if torch.__version__.startswith(\"2.2.\") and sys.platform.startswith(\"linux\"):\n",
+    "    if torch.__version__.startswith(\"1.13.\") and sys.platform.startswith(\"linux\"):\n",
     "        # We try to install PyTorch3D via a released wheel.\n",
     "        pyt_version_str=torch.__version__.split(\"+\")[0].replace(\".\", \"\")\n",
     "        version_str=\"\".join([\n",
@@ -62,6 +62,9 @@
     "        !pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html\n",
     "    else:\n",
     "        # We try to install PyTorch3D from source.\n",
+    "        !curl -LO https://github.com/NVIDIA/cub/archive/1.10.0.tar.gz\n",
+    "        !tar xzf 1.10.0.tar.gz\n",
+    "        os.environ[\"CUB_HOME\"] = os.getcwd() + \"/cub-1.10.0\"\n",
     "        !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'"
    ]
   },

docs/tutorials/render_densepose.ipynb

@@ -57,7 +57,7 @@
     "except ModuleNotFoundError:\n",
     "    need_pytorch3d=True\n",
     "if need_pytorch3d:\n",
-    "    if torch.__version__.startswith(\"2.2.\") and sys.platform.startswith(\"linux\"):\n",
+    "    if torch.__version__.startswith(\"1.13.\") and sys.platform.startswith(\"linux\"):\n",
     "        # We try to install PyTorch3D via a released wheel.\n",
     "        pyt_version_str=torch.__version__.split(\"+\")[0].replace(\".\", \"\")\n",
     "        version_str=\"\".join([\n",
@@ -69,6 +69,9 @@
     "        !pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html\n",
     "    else:\n",
     "        # We try to install PyTorch3D from source.\n",
+    "        !curl -LO https://github.com/NVIDIA/cub/archive/1.10.0.tar.gz\n",
+    "        !tar xzf 1.10.0.tar.gz\n",
+    "        os.environ[\"CUB_HOME\"] = os.getcwd() + \"/cub-1.10.0\"\n",
     "        !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'"
    ]
   },

docs/tutorials/render_textured_meshes.ipynb

@@ -73,7 +73,7 @@
     "except ModuleNotFoundError:\n",
     "    need_pytorch3d=True\n",
     "if need_pytorch3d:\n",
-    "    if torch.__version__.startswith(\"2.2.\") and sys.platform.startswith(\"linux\"):\n",
+    "    if torch.__version__.startswith(\"1.13.\") and sys.platform.startswith(\"linux\"):\n",
     "        # We try to install PyTorch3D via a released wheel.\n",
     "        pyt_version_str=torch.__version__.split(\"+\")[0].replace(\".\", \"\")\n",
     "        version_str=\"\".join([\n",
@@ -85,6 +85,9 @@
     "        !pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html\n",
     "    else:\n",
     "        # We try to install PyTorch3D from source.\n",
+    "        !curl -LO https://github.com/NVIDIA/cub/archive/1.10.0.tar.gz\n",
+    "        !tar xzf 1.10.0.tar.gz\n",
+    "        os.environ[\"CUB_HOME\"] = os.getcwd() + \"/cub-1.10.0\"\n",
     "        !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'"
    ]
   },

docs/tutorials/utils/camera_visualization.py

@@ -33,7 +33,7 @@ def plot_camera_scene(cameras, cameras_gt, status: str):
     a string passed inside the `status` argument.
     """
     fig = plt.figure()
-    ax = fig.add_subplot(projection="3d")
+    ax = fig.gca(projection="3d")
     ax.clear()
     ax.set_title(status)
     handle_cam = plot_cameras(ax, cameras, color="#FF7D1E")

packaging/build_conda.py

@@ -50,6 +50,7 @@ def setup_cuda():
     os.environ["FORCE_CUDA"] = "1"
 
     basic_nvcc_flags = (
+        "-gencode=arch=compute_35,code=sm_35 "
         "-gencode=arch=compute_50,code=sm_50 "
         "-gencode=arch=compute_60,code=sm_60 "
         "-gencode=arch=compute_70,code=sm_70 "
@@ -57,19 +58,13 @@ def setup_cuda():
         "-gencode=arch=compute_50,code=compute_50"
     )
     if CU_VERSION == "cu102":
-        nvcc_flags = "-gencode=arch=compute_35,code=sm_35 " + basic_nvcc_flags
-    elif CU_VERSION < ("cu118"):
-        nvcc_flags = (
-            "-gencode=arch=compute_35,code=sm_35 "
-            + "-gencode=arch=compute_80,code=sm_80 "
-            + "-gencode=arch=compute_86,code=sm_86 "
-            + basic_nvcc_flags
-        )
+        nvcc_flags = basic_nvcc_flags
+    elif CU_VERSION == "cu110":
+        nvcc_flags = "-gencode=arch=compute_80,code=sm_80 " + basic_nvcc_flags
     else:
         nvcc_flags = (
             "-gencode=arch=compute_80,code=sm_80 "
             + "-gencode=arch=compute_86,code=sm_86 "
-            + "-gencode=arch=compute_90,code=sm_90 "
             + basic_nvcc_flags
         )
 

packaging/linux_wheels/README.md

@@ -26,6 +26,5 @@ version_str="".join([
     torch.version.cuda.replace(".",""),
     f"_pyt{pyt_version_str}"
 ])
-!pip install fvcore iopath
 !pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html
 ```

packaging/linux_wheels/go.sh

@@ -5,13 +5,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-# Some directory to persist downloaded conda packages
-conda_cache=/raid/$USER/building_conda_cache
-
-mkdir -p "$conda_cache"
-
-sudo docker run --rm -v "$conda_cache:/conda_cache" -v "$PWD/../../:/inside" -e SELECTED_CUDA=cu113 pytorch/conda-builder:cuda113 bash inside/packaging/linux_wheels/inside.sh
-sudo docker run --rm -v "$conda_cache:/conda_cache" -v "$PWD/../../:/inside" -e SELECTED_CUDA=cu115 pytorch/conda-builder:cuda115 bash inside/packaging/linux_wheels/inside.sh
-sudo docker run --rm -v "$conda_cache:/conda_cache" -v "$PWD/../../:/inside" -e SELECTED_CUDA=cu116 pytorch/conda-builder:cuda116 bash inside/packaging/linux_wheels/inside.sh
-sudo docker run --rm -v "$conda_cache:/conda_cache" -v "$PWD/../../:/inside" -e SELECTED_CUDA=cu117 pytorch/conda-builder:cuda117 bash inside/packaging/linux_wheels/inside.sh
-sudo docker run --rm -v "$conda_cache:/conda_cache" -v "$PWD/../../:/inside" -e SELECTED_CUDA=cu118 pytorch/conda-builder:cuda118 bash inside/packaging/linux_wheels/inside.sh
+sudo docker run --rm  -v "$PWD/../../:/inside" pytorch/conda-cuda bash inside/packaging/linux_wheels/inside.sh
+sudo docker run --rm  -v "$PWD/../../:/inside" -e SELECTED_CUDA=cu113 pytorch/conda-builder:cuda113 bash inside/packaging/linux_wheels/inside.sh
+sudo docker run --rm  -v "$PWD/../../:/inside" -e SELECTED_CUDA=cu115 pytorch/conda-builder:cuda115 bash inside/packaging/linux_wheels/inside.sh
+sudo docker run --rm  -v "$PWD/../../:/inside" -e SELECTED_CUDA=cu116 pytorch/conda-builder:cuda116 bash inside/packaging/linux_wheels/inside.sh

packaging/linux_wheels/inside.sh

@@ -16,32 +16,23 @@ VERSION=$(python -c "exec(open('pytorch3d/__init__.py').read()); print(__version
 
 export BUILD_VERSION=$VERSION
 export FORCE_CUDA=1
-export MAX_JOBS=8
-export CONDA_PKGS_DIRS=/conda_cache
 
-if false
-then
-    # We used to have to do this for old versions of CUDA
-    wget --no-verbose https://github.com/NVIDIA/cub/archive/1.10.0.tar.gz
-    tar xzf 1.10.0.tar.gz
-    CUB_HOME=$(realpath ./cub-1.10.0)
-    export CUB_HOME
-    echo "CUB_HOME is now $CUB_HOME"
-fi
+wget --no-verbose https://github.com/NVIDIA/cub/archive/1.10.0.tar.gz
+tar xzf 1.10.0.tar.gz
+CUB_HOME=$(realpath ./cub-1.10.0)
+export CUB_HOME
+echo "CUB_HOME is now $CUB_HOME"
 
 # As a rule, we want to build for any combination of dependencies which is supported by
 # PyTorch3D and not older than the current Google Colab set up.
 
-PYTHON_VERSIONS="3.8 3.9 3.10"
+PYTHON_VERSIONS="3.7 3.8 3.9 3.10"
 # the keys are pytorch versions
 declare -A CONDA_CUDA_VERSIONS=(
-#    ["1.11.0"]="cu113"
-#    ["1.12.0"]="cu113"
-#    ["1.12.1"]="cu113"
-#    ["1.13.0"]="cu116"
-#    ["1.13.1"]="cu116 cu117"
-#    ["2.0.0"]="cu117 cu118"
-    ["2.0.1"]="cu117 cu118"
+    ["1.10.1"]="cu111 cu113"
+    ["1.10.2"]="cu111 cu113"
+    ["1.10.0"]="cu111 cu113"
+    ["1.11.0"]="cu111 cu113 cu115"
 )
 
 
@@ -50,43 +41,39 @@ for python_version in $PYTHON_VERSIONS
 do
     for pytorch_version in "${!CONDA_CUDA_VERSIONS[@]}"
     do
+        if [[ "3.7 3.8" != *$python_version* ]] && [[ "1.7.0" == *$pytorch_version* ]]
+        then
+            #python 3.9 and later not supported by pytorch 1.7.0 and before
+            continue
+        fi
         if [[ "3.7 3.8 3.9" != *$python_version* ]] && [[ "1.7.0 1.7.1 1.8.0 1.8.1 1.9.0 1.9.1 1.10.0 1.10.1 1.10.2" == *$pytorch_version* ]]
         then
             #python 3.10 and later not supported by pytorch 1.10.2 and before
             continue
         fi
 
-        extra_channel="-c nvidia"
-        cudatools="pytorch-cuda"
+        extra_channel="-c conda-forge"
         if [[ "1.11.0" == "$pytorch_version" ]]
         then
             extra_channel=""
-            cudatools="cudatoolkit"
-        fi
-        if [[ "1.12.0" == "$pytorch_version" ]] || [[ "1.12.1" == "$pytorch_version" ]]
-        then
-            extra_channel="-c conda-forge"
-            cudatools="cudatoolkit"
         fi
 
         for cu_version in ${CONDA_CUDA_VERSIONS[$pytorch_version]}
         do
+            if [[ "cu113 cu115 cu116" == *$cu_version* ]]
+            #       ^^^ CUDA versions listed here have to be built
+            # in their own containers.
+            then
             if [[ $SELECTED_CUDA != "$cu_version" ]]
+                then
+                    continue
+                fi
+            elif [[ $SELECTED_CUDA != "" ]]
             then
                 continue
             fi
 
             case "$cu_version" in
-                cu118)
-                    export CUDA_HOME=/usr/local/cuda-11.8/
-                    export CUDA_TAG=11.8
-                    export NVCC_FLAGS="-gencode=arch=compute_35,code=sm_35 -gencode=arch=compute_50,code=sm_50 -gencode=arch=compute_60,code=sm_60 -gencode=arch=compute_70,code=sm_70 -gencode=arch=compute_75,code=sm_75 -gencode=arch=compute_80,code=sm_80 -gencode=arch=compute_86,code=sm_86 -gencode=arch=compute_50,code=compute_50"
-                ;;
-                cu117)
-                    export CUDA_HOME=/usr/local/cuda-11.7/
-                    export CUDA_TAG=11.7
-                    export NVCC_FLAGS="-gencode=arch=compute_35,code=sm_35 -gencode=arch=compute_50,code=sm_50 -gencode=arch=compute_60,code=sm_60 -gencode=arch=compute_70,code=sm_70 -gencode=arch=compute_75,code=sm_75 -gencode=arch=compute_80,code=sm_80 -gencode=arch=compute_86,code=sm_86 -gencode=arch=compute_50,code=compute_50"
-                ;;
                 cu116)
                     export CUDA_HOME=/usr/local/cuda-11.6/
                     export CUDA_TAG=11.6
@@ -143,7 +130,7 @@ do
             conda create -y -n "$tag" "python=$python_version"
             conda activate "$tag"
             # shellcheck disable=SC2086
-            conda install -y -c pytorch $extra_channel "pytorch=$pytorch_version" "$cudatools=$CUDA_TAG"
+            conda install -y -c pytorch $extra_channel "pytorch=$pytorch_version" "cudatoolkit=$CUDA_TAG" torchvision
             pip install fvcore iopath
             echo "python version" "$python_version" "pytorch version" "$pytorch_version" "cuda version" "$cu_version" "tag" "$tag"
 

projects/implicitron_trainer/README.md

@@ -212,7 +212,9 @@ from pytorch3d.implicitron.tools.config import registry
 class XRayRenderer(BaseRenderer, torch.nn.Module):
     n_pts_per_ray: int = 64
 
+    # if there are other base classes, make sure to call `super().__init__()` explicitly
     def __post_init__(self):
+        super().__init__()
         # custom initialization
 
     def forward(
@@ -248,7 +250,7 @@ The main object for this trainer loop is `Experiment`. It has four top-level rep
 * `data_source`: This is a `DataSourceBase` which defaults to `ImplicitronDataSource`.
 It constructs the data sets and dataloaders.
 * `model_factory`: This is a `ModelFactoryBase` which defaults to `ImplicitronModelFactory`.
-It constructs the model, which is usually an instance of `OverfitModel` (for NeRF-style training with overfitting to one scene) or `GenericModel` (that is able to generalize to multiple scenes by NeRFormer-style conditioning on other scene views), and can load its weights from a checkpoint.
+It constructs the model, which is usually an instance of implicitron's main `GenericModel` class, and can load its weights from a checkpoint.
 * `optimizer_factory`: This is an `OptimizerFactoryBase` which defaults to `ImplicitronOptimizerFactory`.
 It constructs the optimizer and can load its weights from a checkpoint.
 * `training_loop`: This is a `TrainingLoopBase` which defaults to `ImplicitronTrainingLoop` and defines the main training loop.
@@ -292,43 +294,6 @@ model_GenericModel_args: GenericModel
         ╘== ReductionFeatureAggregator
 ```
 
-Here is the class structure of OverfitModel:
-
-```
-model_OverfitModel_args: OverfitModel
-└-- raysampler_*_args: RaySampler
-    ╘== AdaptiveRaysampler
-    ╘== NearFarRaysampler
-└-- renderer_*_args: BaseRenderer
-    ╘== MultiPassEmissionAbsorptionRenderer
-    ╘== LSTMRenderer
-    ╘== SignedDistanceFunctionRenderer
-        └-- ray_tracer_args: RayTracing
-        └-- ray_normal_coloring_network_args: RayNormalColoringNetwork
-└-- implicit_function_*_args: ImplicitFunctionBase
-    ╘== NeuralRadianceFieldImplicitFunction
-    ╘== SRNImplicitFunction
-        └-- raymarch_function_args: SRNRaymarchFunction
-        └-- pixel_generator_args: SRNPixelGenerator
-    ╘== SRNHyperNetImplicitFunction
-        └-- hypernet_args: SRNRaymarchHyperNet
-        └-- pixel_generator_args: SRNPixelGenerator
-    ╘== IdrFeatureField
-└-- coarse_implicit_function_*_args: ImplicitFunctionBase
-    ╘== NeuralRadianceFieldImplicitFunction
-    ╘== SRNImplicitFunction
-        └-- raymarch_function_args: SRNRaymarchFunction
-        └-- pixel_generator_args: SRNPixelGenerator
-    ╘== SRNHyperNetImplicitFunction
-        └-- hypernet_args: SRNRaymarchHyperNet
-        └-- pixel_generator_args: SRNPixelGenerator
-    ╘== IdrFeatureField
-```
-
-OverfitModel has been introduced to create a simple class to disantagle Nerfs which the overfit pattern
-from the GenericModel.
-
-
 Please look at the annotations of the respective classes or functions for the lists of hyperparameters.
 `tests/experiment.yaml` shows every possible option if you have no user-defined classes.
 

projects/implicitron_trainer/configs/overfit_base.yaml

@@ -1,79 +0,0 @@
-defaults:
-- default_config
-- _self_
-exp_dir: ./data/exps/overfit_base/
-training_loop_ImplicitronTrainingLoop_args:
-  visdom_port: 8097
-  visualize_interval: 0
-  max_epochs: 1000
-data_source_ImplicitronDataSource_args:
-  data_loader_map_provider_class_type: SequenceDataLoaderMapProvider
-  dataset_map_provider_class_type: JsonIndexDatasetMapProvider
-  data_loader_map_provider_SequenceDataLoaderMapProvider_args:
-    dataset_length_train: 1000
-    dataset_length_val: 1
-    num_workers: 8
-  dataset_map_provider_JsonIndexDatasetMapProvider_args:
-    dataset_root: ${oc.env:CO3D_DATASET_ROOT}
-    n_frames_per_sequence: -1
-    test_on_train: true
-    test_restrict_sequence_id: 0
-    dataset_JsonIndexDataset_args:
-      load_point_clouds: false
-      mask_depths: false
-      mask_images: false
-model_factory_ImplicitronModelFactory_args:
-  model_class_type: "OverfitModel"
-  model_OverfitModel_args:
-    loss_weights:
-      loss_mask_bce: 1.0
-      loss_prev_stage_mask_bce: 1.0
-      loss_autodecoder_norm: 0.01
-      loss_rgb_mse: 1.0
-      loss_prev_stage_rgb_mse: 1.0
-    output_rasterized_mc: false
-    chunk_size_grid: 102400
-    render_image_height: 400
-    render_image_width: 400
-    share_implicit_function_across_passes: false
-    implicit_function_class_type: "NeuralRadianceFieldImplicitFunction"
-    implicit_function_NeuralRadianceFieldImplicitFunction_args:
-      n_harmonic_functions_xyz: 10
-      n_harmonic_functions_dir: 4
-      n_hidden_neurons_xyz: 256
-      n_hidden_neurons_dir: 128
-      n_layers_xyz: 8
-      append_xyz:
-      - 5
-    coarse_implicit_function_class_type: "NeuralRadianceFieldImplicitFunction"
-    coarse_implicit_function_NeuralRadianceFieldImplicitFunction_args:
-      n_harmonic_functions_xyz: 10
-      n_harmonic_functions_dir: 4
-      n_hidden_neurons_xyz: 256
-      n_hidden_neurons_dir: 128
-      n_layers_xyz: 8
-      append_xyz:
-      - 5
-    raysampler_AdaptiveRaySampler_args:
-      n_rays_per_image_sampled_from_mask: 1024
-      scene_extent: 8.0
-      n_pts_per_ray_training: 64
-      n_pts_per_ray_evaluation: 64
-      stratified_point_sampling_training: true
-      stratified_point_sampling_evaluation: false
-    renderer_MultiPassEmissionAbsorptionRenderer_args:
-      n_pts_per_ray_fine_training: 64
-      n_pts_per_ray_fine_evaluation: 64
-      append_coarse_samples_to_fine: true
-      density_noise_std_train: 1.0
-optimizer_factory_ImplicitronOptimizerFactory_args:
-  breed: Adam
-  weight_decay: 0.0
-  lr_policy: MultiStepLR
-  multistep_lr_milestones: []
-  lr: 0.0005
-  gamma: 0.1
-  momentum: 0.9
-  betas:
-  - 0.9
-  - 0.999

projects/implicitron_trainer/configs/overfit_singleseq_base.yaml

@@ -1,42 +0,0 @@
-defaults:
-- overfit_base
-- _self_
-data_source_ImplicitronDataSource_args:
-  data_loader_map_provider_SequenceDataLoaderMapProvider_args:
-    batch_size: 1
-    dataset_length_train: 1000
-    dataset_length_val: 1
-    num_workers: 8
-  dataset_map_provider_JsonIndexDatasetMapProvider_args:
-    assert_single_seq: true
-    n_frames_per_sequence: -1
-    test_restrict_sequence_id: 0
-    test_on_train: false
-model_factory_ImplicitronModelFactory_args:
-  model_class_type: "OverfitModel"
-  model_OverfitModel_args:
-    render_image_height: 800
-    render_image_width: 800
-    log_vars:
-    - loss_rgb_psnr_fg
-    - loss_rgb_psnr
-    - loss_eikonal
-    - loss_prev_stage_rgb_psnr
-    - loss_mask_bce
-    - loss_prev_stage_mask_bce
-    - loss_rgb_mse
-    - loss_prev_stage_rgb_mse
-    - loss_depth_abs
-    - loss_depth_abs_fg
-    - loss_kl
-    - loss_mask_neg_iou
-    - objective
-    - epoch
-    - sec/it
-optimizer_factory_ImplicitronOptimizerFactory_args:
-  lr: 0.0005
-  multistep_lr_milestones:
-  - 200
-  - 300
-training_loop_ImplicitronTrainingLoop_args:
-  max_epochs: 400

projects/implicitron_trainer/configs/overfit_singleseq_nerf_blender.yaml

@@ -1,56 +0,0 @@
-defaults:
-- overfit_singleseq_base
-- _self_
-exp_dir: "./data/overfit_nerf_blender_repro/${oc.env:BLENDER_SINGLESEQ_CLASS}"
-data_source_ImplicitronDataSource_args:
-  data_loader_map_provider_SequenceDataLoaderMapProvider_args:
-    dataset_length_train: 100
-  dataset_map_provider_class_type: BlenderDatasetMapProvider
-  dataset_map_provider_BlenderDatasetMapProvider_args:
-    base_dir: ${oc.env:BLENDER_DATASET_ROOT}/${oc.env:BLENDER_SINGLESEQ_CLASS}
-    n_known_frames_for_test: null
-    object_name: ${oc.env:BLENDER_SINGLESEQ_CLASS}
-    path_manager_factory_class_type: PathManagerFactory
-    path_manager_factory_PathManagerFactory_args:
-      silence_logs: true
-
-model_factory_ImplicitronModelFactory_args:
-  model_class_type: "OverfitModel"
-  model_OverfitModel_args:
-    mask_images: false
-    raysampler_class_type: AdaptiveRaySampler
-    raysampler_AdaptiveRaySampler_args:
-      n_pts_per_ray_training: 64
-      n_pts_per_ray_evaluation: 64
-      n_rays_per_image_sampled_from_mask: 4096
-      stratified_point_sampling_training: true
-      stratified_point_sampling_evaluation: false
-      scene_extent: 2.0
-      scene_center:
-      - 0.0
-      - 0.0
-      - 0.0
-    renderer_MultiPassEmissionAbsorptionRenderer_args:
-      density_noise_std_train: 0.0
-      n_pts_per_ray_fine_training: 128
-      n_pts_per_ray_fine_evaluation: 128
-      raymarcher_EmissionAbsorptionRaymarcher_args:
-        blend_output: false
-    loss_weights:
-      loss_rgb_mse: 1.0
-      loss_prev_stage_rgb_mse: 1.0
-      loss_mask_bce: 0.0
-      loss_prev_stage_mask_bce: 0.0
-      loss_autodecoder_norm: 0.00
-
-optimizer_factory_ImplicitronOptimizerFactory_args:
-  exponential_lr_step_size: 3001
-  lr_policy: LinearExponential
-  linear_exponential_lr_milestone: 200
-
-training_loop_ImplicitronTrainingLoop_args:
-  max_epochs: 6000
-  metric_print_interval: 10
-  store_checkpoints_purge: 3
-  test_when_finished: true
-  validation_interval: 100

projects/implicitron_trainer/experiment.py

@@ -59,7 +59,7 @@ from pytorch3d.implicitron.dataset.data_source import (
     DataSourceBase,
     ImplicitronDataSource,
 )
-from pytorch3d.implicitron.models.base_model import ImplicitronModelBase
+from pytorch3d.implicitron.models.generic_model import ImplicitronModelBase
 
 from pytorch3d.implicitron.models.renderer.multipass_ea import (
     MultiPassEmissionAbsorptionRenderer,
@@ -207,6 +207,12 @@ class Experiment(Configurable):  # pyre-ignore: 13
                 val_loader,
             ) = accelerator.prepare(model, optimizer, train_loader, val_loader)
 
+        # pyre-fixme[16]: Optional type has no attribute `is_multisequence`.
+        if not self.training_loop.evaluator.is_multisequence:
+            all_train_cameras = self.data_source.all_train_cameras
+        else:
+            all_train_cameras = None
+
         # Enter the main training loop.
         self.training_loop.run(
             train_loader=train_loader,
@@ -217,6 +223,7 @@ class Experiment(Configurable):  # pyre-ignore: 13
             model=model,
             optimizer=optimizer,
             scheduler=scheduler,
+            all_train_cameras=all_train_cameras,
             accelerator=accelerator,
             device=device,
             exp_dir=self.exp_dir,

projects/implicitron_trainer/impl/optimizer_factory.py

@@ -121,6 +121,7 @@ class ImplicitronOptimizerFactory(OptimizerFactoryBase):
         """
         # Get the parameters to optimize
         if hasattr(model, "_get_param_groups"):  # use the model function
+            # pyre-ignore[29]
             p_groups = model._get_param_groups(self.lr, wd=self.weight_decay)
         else:
             p_groups = [

projects/implicitron_trainer/impl/training_loop.py

@@ -21,6 +21,7 @@ from pytorch3d.implicitron.tools.config import (
     run_auto_creation,
 )
 from pytorch3d.implicitron.tools.stats import Stats
+from pytorch3d.renderer.cameras import CamerasBase
 from torch.utils.data import DataLoader, Dataset
 
 from .utils import seed_all_random_engines
@@ -110,8 +111,6 @@ class ImplicitronTrainingLoop(TrainingLoopBase):
     def __post_init__(self):
         run_auto_creation(self)
 
-    # pyre-fixme[14]: `run` overrides method defined in `TrainingLoopBase`
-    #  inconsistently.
     def run(
         self,
         *,
@@ -123,6 +122,7 @@ class ImplicitronTrainingLoop(TrainingLoopBase):
         optimizer: torch.optim.Optimizer,
         scheduler: Any,
         accelerator: Optional[Accelerator],
+        all_train_cameras: Optional[CamerasBase],
         device: torch.device,
         exp_dir: str,
         stats: Stats,
@@ -142,6 +142,7 @@ class ImplicitronTrainingLoop(TrainingLoopBase):
             if test_loader is not None:
                 # pyre-fixme[16]: `Optional` has no attribute `run`.
                 self.evaluator.run(
+                    all_train_cameras=all_train_cameras,
                     dataloader=test_loader,
                     device=device,
                     dump_to_json=True,
@@ -199,6 +200,7 @@ class ImplicitronTrainingLoop(TrainingLoopBase):
                     and epoch % self.test_interval == 0
                 ):
                     self.evaluator.run(
+                        all_train_cameras=all_train_cameras,
                         device=device,
                         dataloader=test_loader,
                         model=model,
@@ -215,6 +217,7 @@ class ImplicitronTrainingLoop(TrainingLoopBase):
         if self.test_when_finished:
             if test_loader is not None:
                 self.evaluator.run(
+                    all_train_cameras=all_train_cameras,
                     device=device,
                     dump_to_json=True,
                     epoch=stats.epoch,
@@ -257,6 +260,7 @@ class ImplicitronTrainingLoop(TrainingLoopBase):
             list(log_vars),
             plot_file=os.path.join(exp_dir, "train_stats.pdf"),
             visdom_env=visdom_env_charts,
+            verbose=False,
             visdom_server=self.visdom_server,
             visdom_port=self.visdom_port,
         )
@@ -382,8 +386,7 @@ class ImplicitronTrainingLoop(TrainingLoopBase):
 
             # print textual status update
             if it % self.metric_print_interval == 0 or last_iter:
-                std_out = stats.get_status_string(stat_set=trainmode, max_it=n_batches)
-                logger.info(std_out)
+                stats.print(stat_set=trainmode, max_it=n_batches)
 
             # visualize results
             if (
@@ -393,6 +396,7 @@ class ImplicitronTrainingLoop(TrainingLoopBase):
             ):
                 prefix = f"e{stats.epoch}_it{stats.it[trainmode]}"
                 if hasattr(model, "visualize"):
+                    # pyre-ignore [29]
                     model.visualize(
                         viz,
                         visdom_env_imgs,

projects/implicitron_trainer/tests/experiment.yaml

@@ -103,10 +103,8 @@ data_source_ImplicitronDataSource_args:
     num_views: 40
     data_file: null
     azimuth_range: 180.0
-    distance: 2.7
     resolution: 128
     use_point_light: true
-    gpu_idx: 0
     path_manager_factory_class_type: PathManagerFactory
     path_manager_factory_PathManagerFactory_args:
       silence_logs: true
@@ -129,19 +127,6 @@ data_source_ImplicitronDataSource_args:
     dataset_length_train: 0
     dataset_length_val: 0
     dataset_length_test: 0
-  data_loader_map_provider_TrainEvalDataLoaderMapProvider_args:
-    batch_size: 1
-    num_workers: 0
-    dataset_length_train: 0
-    dataset_length_val: 0
-    dataset_length_test: 0
-    train_conditioning_type: SAME
-    val_conditioning_type: SAME
-    test_conditioning_type: KNOWN
-    images_per_seq_options: []
-    sample_consecutive_frames: false
-    consecutive_frames_max_gap: 0
-    consecutive_frames_max_gap_seconds: 0.1
 model_factory_ImplicitronModelFactory_args:
   resume: true
   model_class_type: GenericModel
@@ -216,7 +201,6 @@ model_factory_ImplicitronModelFactory_args:
       n_rays_total_training: null
       stratified_point_sampling_training: true
       stratified_point_sampling_evaluation: false
-      cast_ray_bundle_as_cone: false
       scene_extent: 8.0
       scene_center:
       - 0.0
@@ -229,7 +213,6 @@ model_factory_ImplicitronModelFactory_args:
       n_rays_total_training: null
       stratified_point_sampling_training: true
       stratified_point_sampling_evaluation: false
-      cast_ray_bundle_as_cone: false
       min_depth: 0.1
       max_depth: 8.0
     renderer_LSTMRenderer_args:
@@ -249,8 +232,6 @@ model_factory_ImplicitronModelFactory_args:
       append_coarse_samples_to_fine: true
       density_noise_std_train: 0.0
       return_weights: false
-      blurpool_weights: false
-      sample_pdf_eps: 1.0e-05
       raymarcher_CumsumRaymarcher_args:
         surface_thickness: 1
         bg_color:
@@ -363,7 +344,6 @@ model_factory_ImplicitronModelFactory_args:
       n_hidden_neurons_dir: 128
       input_xyz: true
       xyz_ray_dir_in_camera_coords: false
-      use_integrated_positional_encoding: false
       transformer_dim_down_factor: 2.0
       n_hidden_neurons_xyz: 80
       n_layers_xyz: 2
@@ -375,7 +355,6 @@ model_factory_ImplicitronModelFactory_args:
       n_hidden_neurons_dir: 128
       input_xyz: true
       xyz_ray_dir_in_camera_coords: false
-      use_integrated_positional_encoding: false
       transformer_dim_down_factor: 1.0
       n_hidden_neurons_xyz: 256
       n_layers_xyz: 8
@@ -580,631 +559,6 @@ model_factory_ImplicitronModelFactory_args:
           use_xavier_init: true
     view_metrics_ViewMetrics_args: {}
     regularization_metrics_RegularizationMetrics_args: {}
-  model_OverfitModel_args:
-    log_vars:
-    - loss_rgb_psnr_fg
-    - loss_rgb_psnr
-    - loss_rgb_mse
-    - loss_rgb_huber
-    - loss_depth_abs
-    - loss_depth_abs_fg
-    - loss_mask_neg_iou
-    - loss_mask_bce
-    - loss_mask_beta_prior
-    - loss_eikonal
-    - loss_density_tv
-    - loss_depth_neg_penalty
-    - loss_autodecoder_norm
-    - loss_prev_stage_rgb_mse
-    - loss_prev_stage_rgb_psnr_fg
-    - loss_prev_stage_rgb_psnr
-    - loss_prev_stage_mask_bce
-    - objective
-    - epoch
-    - sec/it
-    mask_images: true
-    mask_depths: true
-    render_image_width: 400
-    render_image_height: 400
-    mask_threshold: 0.5
-    output_rasterized_mc: false
-    bg_color:
-    - 0.0
-    - 0.0
-    - 0.0
-    chunk_size_grid: 4096
-    render_features_dimensions: 3
-    tqdm_trigger_threshold: 16
-    n_train_target_views: 1
-    sampling_mode_training: mask_sample
-    sampling_mode_evaluation: full_grid
-    global_encoder_class_type: null
-    raysampler_class_type: AdaptiveRaySampler
-    renderer_class_type: MultiPassEmissionAbsorptionRenderer
-    share_implicit_function_across_passes: false
-    implicit_function_class_type: NeuralRadianceFieldImplicitFunction
-    coarse_implicit_function_class_type: null
-    view_metrics_class_type: ViewMetrics
-    regularization_metrics_class_type: RegularizationMetrics
-    loss_weights:
-      loss_rgb_mse: 1.0
-      loss_prev_stage_rgb_mse: 1.0
-      loss_mask_bce: 0.0
-      loss_prev_stage_mask_bce: 0.0
-    global_encoder_HarmonicTimeEncoder_args:
-      n_harmonic_functions: 10
-      append_input: true
-      time_divisor: 1.0
-    global_encoder_SequenceAutodecoder_args:
-      autodecoder_args:
-        encoding_dim: 0
-        n_instances: 1
-        init_scale: 1.0
-        ignore_input: false
-    raysampler_AdaptiveRaySampler_args:
-      n_pts_per_ray_training: 64
-      n_pts_per_ray_evaluation: 64
-      n_rays_per_image_sampled_from_mask: 1024
-      n_rays_total_training: null
-      stratified_point_sampling_training: true
-      stratified_point_sampling_evaluation: false
-      cast_ray_bundle_as_cone: false
-      scene_extent: 8.0
-      scene_center:
-      - 0.0
-      - 0.0
-      - 0.0
-    raysampler_NearFarRaySampler_args:
-      n_pts_per_ray_training: 64
-      n_pts_per_ray_evaluation: 64
-      n_rays_per_image_sampled_from_mask: 1024
-      n_rays_total_training: null
-      stratified_point_sampling_training: true
-      stratified_point_sampling_evaluation: false
-      cast_ray_bundle_as_cone: false
-      min_depth: 0.1
-      max_depth: 8.0
-    renderer_LSTMRenderer_args:
-      num_raymarch_steps: 10
-      init_depth: 17.0
-      init_depth_noise_std: 0.0005
-      hidden_size: 16
-      n_feature_channels: 256
-      bg_color: null
-      verbose: false
-    renderer_MultiPassEmissionAbsorptionRenderer_args:
-      raymarcher_class_type: EmissionAbsorptionRaymarcher
-      n_pts_per_ray_fine_training: 64
-      n_pts_per_ray_fine_evaluation: 64
-      stratified_sampling_coarse_training: true
-      stratified_sampling_coarse_evaluation: false
-      append_coarse_samples_to_fine: true
-      density_noise_std_train: 0.0
-      return_weights: false
-      blurpool_weights: false
-      sample_pdf_eps: 1.0e-05
-      raymarcher_CumsumRaymarcher_args:
-        surface_thickness: 1
-        bg_color:
-        - 0.0
-        replicate_last_interval: false
-        background_opacity: 0.0
-        density_relu: true
-        blend_output: false
-      raymarcher_EmissionAbsorptionRaymarcher_args:
-        surface_thickness: 1
-        bg_color:
-        - 0.0
-        replicate_last_interval: false
-        background_opacity: 10000000000.0
-        density_relu: true
-        blend_output: false
-    renderer_SignedDistanceFunctionRenderer_args:
-      ray_normal_coloring_network_args:
-        feature_vector_size: 3
-        mode: idr
-        d_in: 9
-        d_out: 3
-        dims:
-        - 512
-        - 512
-        - 512
-        - 512
-        weight_norm: true
-        n_harmonic_functions_dir: 0
-        pooled_feature_dim: 0
-      bg_color:
-      - 0.0
-      soft_mask_alpha: 50.0
-      ray_tracer_args:
-        sdf_threshold: 5.0e-05
-        line_search_step: 0.5
-        line_step_iters: 1
-        sphere_tracing_iters: 10
-        n_steps: 100
-        n_secant_steps: 8
-    implicit_function_IdrFeatureField_args:
-      d_in: 3
-      d_out: 1
-      dims:
-      - 512
-      - 512
-      - 512
-      - 512
-      - 512
-      - 512
-      - 512
-      - 512
-      geometric_init: true
-      bias: 1.0
-      skip_in: []
-      weight_norm: true
-      n_harmonic_functions_xyz: 0
-      pooled_feature_dim: 0
-    implicit_function_NeRFormerImplicitFunction_args:
-      n_harmonic_functions_xyz: 10
-      n_harmonic_functions_dir: 4
-      n_hidden_neurons_dir: 128
-      input_xyz: true
-      xyz_ray_dir_in_camera_coords: false
-      use_integrated_positional_encoding: false
-      transformer_dim_down_factor: 2.0
-      n_hidden_neurons_xyz: 80
-      n_layers_xyz: 2
-      append_xyz:
-      - 1
-    implicit_function_NeuralRadianceFieldImplicitFunction_args:
-      n_harmonic_functions_xyz: 10
-      n_harmonic_functions_dir: 4
-      n_hidden_neurons_dir: 128
-      input_xyz: true
-      xyz_ray_dir_in_camera_coords: false
-      use_integrated_positional_encoding: false
-      transformer_dim_down_factor: 1.0
-      n_hidden_neurons_xyz: 256
-      n_layers_xyz: 8
-      append_xyz:
-      - 5
-    implicit_function_SRNHyperNetImplicitFunction_args:
-      latent_dim_hypernet: 0
-      hypernet_args:
-        n_harmonic_functions: 3
-        n_hidden_units: 256
-        n_layers: 2
-        n_hidden_units_hypernet: 256
-        n_layers_hypernet: 1
-        in_features: 3
-        out_features: 256
-        xyz_in_camera_coords: false
-      pixel_generator_args:
-        n_harmonic_functions: 4
-        n_hidden_units: 256
-        n_hidden_units_color: 128
-        n_layers: 2
-        in_features: 256
-        out_features: 3
-        ray_dir_in_camera_coords: false
-    implicit_function_SRNImplicitFunction_args:
-      raymarch_function_args:
-        n_harmonic_functions: 3
-        n_hidden_units: 256
-        n_layers: 2
-        in_features: 3
-        out_features: 256
-        xyz_in_camera_coords: false
-        raymarch_function: null
-      pixel_generator_args:
-        n_harmonic_functions: 4
-        n_hidden_units: 256
-        n_hidden_units_color: 128
-        n_layers: 2
-        in_features: 256
-        out_features: 3
-        ray_dir_in_camera_coords: false
-    implicit_function_VoxelGridImplicitFunction_args:
-      harmonic_embedder_xyz_density_args:
-        n_harmonic_functions: 6
-        omega_0: 1.0
-        logspace: true
-        append_input: true
-      harmonic_embedder_xyz_color_args:
-        n_harmonic_functions: 6
-        omega_0: 1.0
-        logspace: true
-        append_input: true
-      harmonic_embedder_dir_color_args:
-        n_harmonic_functions: 6
-        omega_0: 1.0
-        logspace: true
-        append_input: true
-      decoder_density_class_type: MLPDecoder
-      decoder_color_class_type: MLPDecoder
-      use_multiple_streams: true
-      xyz_ray_dir_in_camera_coords: false
-      scaffold_calculating_epochs: []
-      scaffold_resolution:
-      - 128
-      - 128
-      - 128
-      scaffold_empty_space_threshold: 0.001
-      scaffold_occupancy_chunk_size: -1
-      scaffold_max_pool_kernel_size: 3
-      scaffold_filter_points: true
-      volume_cropping_epochs: []
-      voxel_grid_density_args:
-        voxel_grid_class_type: FullResolutionVoxelGrid
-        extents:
-        - 2.0
-        - 2.0
-        - 2.0
-        translation:
-        - 0.0
-        - 0.0
-        - 0.0
-        init_std: 0.1
-        init_mean: 0.0
-        hold_voxel_grid_as_parameters: true
-        param_groups: {}
-        voxel_grid_CPFactorizedVoxelGrid_args:
-          align_corners: true
-          padding: zeros
-          mode: bilinear
-          n_features: 1
-          resolution_changes:
-            0:
-            - 128
-            - 128
-            - 128
-          n_components: 24
-          basis_matrix: true
-        voxel_grid_FullResolutionVoxelGrid_args:
-          align_corners: true
-          padding: zeros
-          mode: bilinear
-          n_features: 1
-          resolution_changes:
-            0:
-            - 128
-            - 128
-            - 128
-        voxel_grid_VMFactorizedVoxelGrid_args:
-          align_corners: true
-          padding: zeros
-          mode: bilinear
-          n_features: 1
-          resolution_changes:
-            0:
-            - 128
-            - 128
-            - 128
-          n_components: null
-          distribution_of_components: null
-          basis_matrix: true
-      voxel_grid_color_args:
-        voxel_grid_class_type: FullResolutionVoxelGrid
-        extents:
-        - 2.0
-        - 2.0
-        - 2.0
-        translation:
-        - 0.0
-        - 0.0
-        - 0.0
-        init_std: 0.1
-        init_mean: 0.0
-        hold_voxel_grid_as_parameters: true
-        param_groups: {}
-        voxel_grid_CPFactorizedVoxelGrid_args:
-          align_corners: true
-          padding: zeros
-          mode: bilinear
-          n_features: 1
-          resolution_changes:
-            0:
-            - 128
-            - 128
-            - 128
-          n_components: 24
-          basis_matrix: true
-        voxel_grid_FullResolutionVoxelGrid_args:
-          align_corners: true
-          padding: zeros
-          mode: bilinear
-          n_features: 1
-          resolution_changes:
-            0:
-            - 128
-            - 128
-            - 128
-        voxel_grid_VMFactorizedVoxelGrid_args:
-          align_corners: true
-          padding: zeros
-          mode: bilinear
-          n_features: 1
-          resolution_changes:
-            0:
-            - 128
-            - 128
-            - 128
-          n_components: null
-          distribution_of_components: null
-          basis_matrix: true
-      decoder_density_ElementwiseDecoder_args:
-        scale: 1.0
-        shift: 0.0
-        operation: IDENTITY
-      decoder_density_MLPDecoder_args:
-        param_groups: {}
-        network_args:
-          n_layers: 8
-          output_dim: 256
-          skip_dim: 39
-          hidden_dim: 256
-          input_skips:
-          - 5
-          skip_affine_trans: false
-          last_layer_bias_init: null
-          last_activation: RELU
-          use_xavier_init: true
-      decoder_color_ElementwiseDecoder_args:
-        scale: 1.0
-        shift: 0.0
-        operation: IDENTITY
-      decoder_color_MLPDecoder_args:
-        param_groups: {}
-        network_args:
-          n_layers: 8
-          output_dim: 256
-          skip_dim: 39
-          hidden_dim: 256
-          input_skips:
-          - 5
-          skip_affine_trans: false
-          last_layer_bias_init: null
-          last_activation: RELU
-          use_xavier_init: true
-    coarse_implicit_function_IdrFeatureField_args:
-      d_in: 3
-      d_out: 1
-      dims:
-      - 512
-      - 512
-      - 512
-      - 512
-      - 512
-      - 512
-      - 512
-      - 512
-      geometric_init: true
-      bias: 1.0
-      skip_in: []
-      weight_norm: true
-      n_harmonic_functions_xyz: 0
-      pooled_feature_dim: 0
-    coarse_implicit_function_NeRFormerImplicitFunction_args:
-      n_harmonic_functions_xyz: 10
-      n_harmonic_functions_dir: 4
-      n_hidden_neurons_dir: 128
-      input_xyz: true
-      xyz_ray_dir_in_camera_coords: false
-      use_integrated_positional_encoding: false
-      transformer_dim_down_factor: 2.0
-      n_hidden_neurons_xyz: 80
-      n_layers_xyz: 2
-      append_xyz:
-      - 1
-    coarse_implicit_function_NeuralRadianceFieldImplicitFunction_args:
-      n_harmonic_functions_xyz: 10
-      n_harmonic_functions_dir: 4
-      n_hidden_neurons_dir: 128
-      input_xyz: true
-      xyz_ray_dir_in_camera_coords: false
-      use_integrated_positional_encoding: false
-      transformer_dim_down_factor: 1.0
-      n_hidden_neurons_xyz: 256
-      n_layers_xyz: 8
-      append_xyz:
-      - 5
-    coarse_implicit_function_SRNHyperNetImplicitFunction_args:
-      latent_dim_hypernet: 0
-      hypernet_args:
-        n_harmonic_functions: 3
-        n_hidden_units: 256
-        n_layers: 2
-        n_hidden_units_hypernet: 256
-        n_layers_hypernet: 1
-        in_features: 3
-        out_features: 256
-        xyz_in_camera_coords: false
-      pixel_generator_args:
-        n_harmonic_functions: 4
-        n_hidden_units: 256
-        n_hidden_units_color: 128
-        n_layers: 2
-        in_features: 256
-        out_features: 3
-        ray_dir_in_camera_coords: false
-    coarse_implicit_function_SRNImplicitFunction_args:
-      raymarch_function_args:
-        n_harmonic_functions: 3
-        n_hidden_units: 256
-        n_layers: 2
-        in_features: 3
-        out_features: 256
-        xyz_in_camera_coords: false
-        raymarch_function: null
-      pixel_generator_args:
-        n_harmonic_functions: 4
-        n_hidden_units: 256
-        n_hidden_units_color: 128
-        n_layers: 2
-        in_features: 256
-        out_features: 3
-        ray_dir_in_camera_coords: false
-    coarse_implicit_function_VoxelGridImplicitFunction_args:
-      harmonic_embedder_xyz_density_args:
-        n_harmonic_functions: 6
-        omega_0: 1.0
-        logspace: true
-        append_input: true
-      harmonic_embedder_xyz_color_args:
-        n_harmonic_functions: 6
-        omega_0: 1.0
-        logspace: true
-        append_input: true
-      harmonic_embedder_dir_color_args:
-        n_harmonic_functions: 6
-        omega_0: 1.0
-        logspace: true
-        append_input: true
-      decoder_density_class_type: MLPDecoder
-      decoder_color_class_type: MLPDecoder
-      use_multiple_streams: true
-      xyz_ray_dir_in_camera_coords: false
-      scaffold_calculating_epochs: []
-      scaffold_resolution:
-      - 128
-      - 128
-      - 128
-      scaffold_empty_space_threshold: 0.001
-      scaffold_occupancy_chunk_size: -1
-      scaffold_max_pool_kernel_size: 3
-      scaffold_filter_points: true
-      volume_cropping_epochs: []
-      voxel_grid_density_args:
-        voxel_grid_class_type: FullResolutionVoxelGrid
-        extents:
-        - 2.0
-        - 2.0
-        - 2.0
-        translation:
-        - 0.0
-        - 0.0
-        - 0.0
-        init_std: 0.1
-        init_mean: 0.0
-        hold_voxel_grid_as_parameters: true
-        param_groups: {}
-        voxel_grid_CPFactorizedVoxelGrid_args:
-          align_corners: true
-          padding: zeros
-          mode: bilinear
-          n_features: 1
-          resolution_changes:
-            0:
-            - 128
-            - 128
-            - 128
-          n_components: 24
-          basis_matrix: true
-        voxel_grid_FullResolutionVoxelGrid_args:
-          align_corners: true
-          padding: zeros
-          mode: bilinear
-          n_features: 1
-          resolution_changes:
-            0:
-            - 128
-            - 128
-            - 128
-        voxel_grid_VMFactorizedVoxelGrid_args:
-          align_corners: true
-          padding: zeros
-          mode: bilinear
-          n_features: 1
-          resolution_changes:
-            0:
-            - 128
-            - 128
-            - 128
-          n_components: null
-          distribution_of_components: null
-          basis_matrix: true
-      voxel_grid_color_args:
-        voxel_grid_class_type: FullResolutionVoxelGrid
-        extents:
-        - 2.0
-        - 2.0
-        - 2.0
-        translation:
-        - 0.0
-        - 0.0
-        - 0.0
-        init_std: 0.1
-        init_mean: 0.0
-        hold_voxel_grid_as_parameters: true
-        param_groups: {}
-        voxel_grid_CPFactorizedVoxelGrid_args:
-          align_corners: true
-          padding: zeros
-          mode: bilinear
-          n_features: 1
-          resolution_changes:
-            0:
-            - 128
-            - 128
-            - 128
-          n_components: 24
-          basis_matrix: true
-        voxel_grid_FullResolutionVoxelGrid_args:
-          align_corners: true
-          padding: zeros
-          mode: bilinear
-          n_features: 1
-          resolution_changes:
-            0:
-            - 128
-            - 128
-            - 128
-        voxel_grid_VMFactorizedVoxelGrid_args:
-          align_corners: true
-          padding: zeros
-          mode: bilinear
-          n_features: 1
-          resolution_changes:
-            0:
-            - 128
-            - 128
-            - 128
-          n_components: null
-          distribution_of_components: null
-          basis_matrix: true
-      decoder_density_ElementwiseDecoder_args:
-        scale: 1.0
-        shift: 0.0
-        operation: IDENTITY
-      decoder_density_MLPDecoder_args:
-        param_groups: {}
-        network_args:
-          n_layers: 8
-          output_dim: 256
-          skip_dim: 39
-          hidden_dim: 256
-          input_skips:
-          - 5
-          skip_affine_trans: false
-          last_layer_bias_init: null
-          last_activation: RELU
-          use_xavier_init: true
-      decoder_color_ElementwiseDecoder_args:
-        scale: 1.0
-        shift: 0.0
-        operation: IDENTITY
-      decoder_color_MLPDecoder_args:
-        param_groups: {}
-        network_args:
-          n_layers: 8
-          output_dim: 256
-          skip_dim: 39
-          hidden_dim: 256
-          input_skips:
-          - 5
-          skip_affine_trans: false
-          last_layer_bias_init: null
-          last_activation: RELU
-          use_xavier_init: true
-    view_metrics_ViewMetrics_args: {}
-    regularization_metrics_RegularizationMetrics_args: {}
 optimizer_factory_ImplicitronOptimizerFactory_args:
   betas:
   - 0.9

projects/implicitron_trainer/tests/test_experiment.py

@@ -132,13 +132,6 @@ class TestExperiment(unittest.TestCase):
         # Check that the default config values, defined by Experiment and its
         # members, is what we expect it to be.
         cfg = OmegaConf.structured(experiment.Experiment)
-        # the following removes the possible effect of env variables
-        ds_arg = cfg.data_source_ImplicitronDataSource_args
-        ds_arg.dataset_map_provider_JsonIndexDatasetMapProvider_args.dataset_root = ""
-        ds_arg.dataset_map_provider_JsonIndexDatasetMapProviderV2_args.dataset_root = ""
-        if "dataset_map_provider_SqlIndexDatasetMapProvider_args" in ds_arg:
-            del ds_arg.dataset_map_provider_SqlIndexDatasetMapProvider_args
-        cfg.training_loop_ImplicitronTrainingLoop_args.visdom_port = 8097
         yaml = OmegaConf.to_yaml(cfg, sort_keys=False)
         if DEBUG:
             (DATA_DIR / "experiment.yaml").write_text(yaml)
@@ -148,11 +141,7 @@ class TestExperiment(unittest.TestCase):
         # Check that all the pre-prepared configs are valid.
         config_files = []
 
-        for pattern in (
-            "repro_singleseq*.yaml",
-            "repro_multiseq*.yaml",
-            "overfit_singleseq*.yaml",
-        ):
+        for pattern in ("repro_singleseq*.yaml", "repro_multiseq*.yaml"):
             config_files.extend(
                 [
                     f

projects/implicitron_trainer/visualize_reconstruction.py

@@ -39,7 +39,6 @@ def visualize_reconstruction(
     visdom_server: str = "http://127.0.0.1",
     visdom_port: int = 8097,
     visdom_env: Optional[str] = None,
-    **render_flyaround_kwargs,
 ) -> None:
     """
     Given an `exp_dir` containing a trained Implicitron model, generates videos consisting
@@ -61,8 +60,6 @@ def visualize_reconstruction(
         visdom_server: The address of the visdom server.
         visdom_port: The port of the visdom server.
         visdom_env: If set, defines a custom name for the visdom environment.
-        render_flyaround_kwargs: Keyword arguments passed to the invoked `render_flyaround`
-            function (see `pytorch3d.implicitron.models.visualization.render_flyaround`).
     """
 
     # In case an output directory is specified use it. If no output_directory
@@ -118,22 +115,20 @@ def visualize_reconstruction(
     # iterate over the sequences in the dataset
     for sequence_name in dataset.sequence_names():
         with torch.no_grad():
-            render_kwargs = {
-                "dataset": dataset,
-                "sequence_name": sequence_name,
-                "model": model,
-                "output_video_path": os.path.join(output_directory, "video"),
-                "n_source_views": n_source_views,
-                "visdom_show_preds": visdom_show_preds,
-                "n_flyaround_poses": n_eval_cameras,
-                "visdom_server": visdom_server,
-                "visdom_port": visdom_port,
-                "visdom_environment": visdom_env,
-                "video_resize": video_size,
-                "device": device,
-                **render_flyaround_kwargs,
-            }
-            render_flyaround(**render_kwargs)
+            render_flyaround(
+                dataset=dataset,
+                sequence_name=sequence_name,
+                model=model,
+                output_video_path=os.path.join(output_directory, "video"),
+                n_source_views=n_source_views,
+                visdom_show_preds=visdom_show_preds,
+                n_flyaround_poses=n_eval_cameras,
+                visdom_server=visdom_server,
+                visdom_port=visdom_port,
+                visdom_environment=visdom_env,
+                video_resize=video_size,
+                device=device,
+            )
 
 
 enable_get_default_args(visualize_reconstruction)

projects/nerf/nerf/raysampler.py

@@ -330,9 +330,9 @@ class NeRFRaysampler(torch.nn.Module):
 
             if self.training:
                 # During training we randomly subsample rays.
-                sel_rays = torch.randperm(
-                    n_pixels, device=full_ray_bundle.lengths.device
-                )[: self._mc_raysampler._n_rays_per_image]
+                sel_rays = torch.randperm(n_pixels, device=device)[
+                    : self._mc_raysampler._n_rays_per_image
+                ]
             else:
                 # In case we test, we take only the requested chunk.
                 if chunksize is None:

pytorch3d/__init__.py

@@ -4,4 +4,4 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-__version__ = "0.7.6"
+__version__ = "0.7.2"

pytorch3d/common/datatypes.py

@@ -4,6 +4,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import sys
 from typing import Optional, Union
 
 import torch
@@ -56,3 +57,19 @@ def get_device(x, device: Optional[Device] = None) -> torch.device:
 
     # Default device is cpu
     return torch.device("cpu")
+
+
+# Provide get_origin and get_args even in Python 3.7.
+
+if sys.version_info >= (3, 8, 0):
+    from typing import get_args, get_origin
+elif sys.version_info >= (3, 7, 0):
+
+    def get_origin(cls):  # pragma: no cover
+        return getattr(cls, "__origin__", None)
+
+    def get_args(cls):  # pragma: no cover
+        return getattr(cls, "__args__", None)
+
+else:
+    raise ImportError("This module requires Python 3.7+")

pytorch3d/csrc/face_areas_normals/face_areas_normals.cu

@@ -266,8 +266,6 @@ at::Tensor FaceAreasNormalsBackwardCuda(
       grad_normals_t{grad_normals, "grad_normals", 4};
   at::CheckedFrom c = "FaceAreasNormalsBackwardCuda";
   at::checkAllSameGPU(c, {verts_t, faces_t, grad_areas_t, grad_normals_t});
-  // This is nondeterministic because atomicAdd
-  at::globalContext().alertNotDeterministic("FaceAreasNormalsBackwardCuda");
 
   // Set the device for the kernel launch based on the device of verts
   at::cuda::CUDAGuard device_guard(verts.device());

pytorch3d/csrc/interp_face_attrs/interp_face_attrs.cu

@@ -130,9 +130,6 @@ std::tuple<at::Tensor, at::Tensor> InterpFaceAttrsBackwardCuda(
   at::checkAllSameType(
       c, {barycentric_coords_t, face_attrs_t, grad_pix_attrs_t});
 
-  // This is nondeterministic because atomicAdd
-  at::globalContext().alertNotDeterministic("InterpFaceAttrsBackwardCuda");
-
   // Set the device for the kernel launch based on the input
   at::cuda::CUDAGuard device_guard(pix_to_face.device());
   cudaStream_t stream = at::cuda::getCurrentCUDAStream();

pytorch3d/csrc/iou_box3d/iou_box3d.cu

@@ -12,6 +12,8 @@
 #include <math.h>
 #include <stdio.h>
 #include <stdlib.h>
+#include <thrust/device_vector.h>
+#include <thrust/tuple.h>
 #include "iou_box3d/iou_utils.cuh"
 
 // Parallelize over N*M computations which can each be done

pytorch3d/csrc/iou_box3d/iou_utils.cuh

@@ -8,6 +8,7 @@
 
 #include <float.h>
 #include <math.h>
+#include <thrust/device_vector.h>
 #include <cstdio>
 #include "utils/float_math.cuh"
 

pytorch3d/csrc/knn/knn.cu

@@ -495,7 +495,7 @@ __global__ void KNearestNeighborBackwardKernel(
     if ((p1_idx < num1) && (k < num2)) {
       const float grad_dist = grad_dists[n * P1 * K + p1_idx * K + k];
       // index of point in p2 corresponding to the k-th nearest neighbor
-      const int64_t p2_idx = idxs[n * P1 * K + p1_idx * K + k];
+      const size_t p2_idx = idxs[n * P1 * K + p1_idx * K + k];
       // If the index is the pad value of -1 then ignore it
       if (p2_idx == -1) {
         continue;
@@ -534,9 +534,6 @@ std::tuple<at::Tensor, at::Tensor> KNearestNeighborBackwardCuda(
       c, {p1_t, p2_t, lengths1_t, lengths2_t, idxs_t, grad_dists_t});
   at::checkAllSameType(c, {p1_t, p2_t, grad_dists_t});
 
-  // This is nondeterministic because atomicAdd
-  at::globalContext().alertNotDeterministic("KNearestNeighborBackwardCuda");
-
   // Set the device for the kernel launch based on the device of the input
   at::cuda::CUDAGuard device_guard(p1.device());
   cudaStream_t stream = at::cuda::getCurrentCUDAStream();

pytorch3d/csrc/marching_cubes/marching_cubes.cu

@@ -9,6 +9,8 @@
 #include <ATen/ATen.h>
 #include <ATen/cuda/CUDAContext.h>
 #include <c10/cuda/CUDAGuard.h>
+#include <thrust/device_vector.h>
+#include <thrust/scan.h>
 #include <cstdio>
 #include "marching_cubes/tables.h"
 
@@ -38,6 +40,20 @@ through" each cube in the grid.
 // EPS: Used to indicate if two float values are close
 __constant__ const float EPSILON = 1e-5;
 
+// Thrust wrapper for exclusive scan
+//
+// Args:
+//    output: pointer to on-device output array
+//    input: pointer to on-device input array, where scan is performed
+//    numElements: number of elements for the input array
+//
+void ThrustScanWrapper(int* output, int* input, int numElements) {
+  thrust::exclusive_scan(
+      thrust::device_ptr<int>(input),
+      thrust::device_ptr<int>(input + numElements),
+      thrust::device_ptr<int>(output));
+}
+
 // Linearly interpolate the position where an isosurface cuts an edge
 // between two vertices, based on their scalar values
 //
@@ -223,7 +239,7 @@ __global__ void CompactVoxelsKernel(
         compactedVoxelArray,
     const at::PackedTensorAccessor32<int, 1, at::RestrictPtrTraits>
         voxelOccupied,
-    const at::PackedTensorAccessor32<int64_t, 1, at::RestrictPtrTraits>
+    const at::PackedTensorAccessor32<int, 1, at::RestrictPtrTraits>
         voxelOccupiedScan,
     uint numVoxels) {
   uint id = blockIdx.x * blockDim.x + threadIdx.x;
@@ -255,8 +271,7 @@ __global__ void GenerateFacesKernel(
     at::PackedTensorAccessor<int64_t, 1, at::RestrictPtrTraits> ids,
     at::PackedTensorAccessor32<int, 1, at::RestrictPtrTraits>
         compactedVoxelArray,
-    at::PackedTensorAccessor32<int64_t, 1, at::RestrictPtrTraits>
-        numVertsScanned,
+    at::PackedTensorAccessor32<int, 1, at::RestrictPtrTraits> numVertsScanned,
     const uint activeVoxels,
     const at::PackedTensorAccessor32<float, 3, at::RestrictPtrTraits> vol,
     const at::PackedTensorAccessor32<int, 2, at::RestrictPtrTraits> faceTable,
@@ -440,24 +455,19 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> MarchingCubesCuda(
     grid.x = 65535;
   }
 
-  using at::indexing::None;
-  using at::indexing::Slice;
-
   auto d_voxelVerts =
-      at::zeros({numVoxels + 1}, at::TensorOptions().dtype(at::kInt))
+      at::zeros({numVoxels}, at::TensorOptions().dtype(at::kInt))
           .to(vol.device());
-  auto d_voxelVerts_ = d_voxelVerts.index({Slice(1, None)});
   auto d_voxelOccupied =
-      at::zeros({numVoxels + 1}, at::TensorOptions().dtype(at::kInt))
+      at::zeros({numVoxels}, at::TensorOptions().dtype(at::kInt))
           .to(vol.device());
-  auto d_voxelOccupied_ = d_voxelOccupied.index({Slice(1, None)});
 
   // Execute "ClassifyVoxelKernel" kernel to precompute
   // two arrays - d_voxelOccupied and d_voxelVertices to global memory,
   // which stores the occupancy state and number of voxel vertices per voxel.
   ClassifyVoxelKernel<<<grid, threads, 0, stream>>>(
-      d_voxelVerts_.packed_accessor32<int, 1, at::RestrictPtrTraits>(),
-      d_voxelOccupied_.packed_accessor32<int, 1, at::RestrictPtrTraits>(),
+      d_voxelVerts.packed_accessor32<int, 1, at::RestrictPtrTraits>(),
+      d_voxelOccupied.packed_accessor32<int, 1, at::RestrictPtrTraits>(),
       vol.packed_accessor32<float, 3, at::RestrictPtrTraits>(),
       isolevel);
   AT_CUDA_CHECK(cudaGetLastError());
@@ -467,12 +477,18 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> MarchingCubesCuda(
   // count for voxels in the grid and compute the number of active voxels.
   // If the number of active voxels is 0, return zero tensor for verts and
   // faces.
-
-  auto d_voxelOccupiedScan = at::cumsum(d_voxelOccupied, 0);
-  auto d_voxelOccupiedScan_ = d_voxelOccupiedScan.index({Slice(1, None)});
+  auto d_voxelOccupiedScan =
+      at::zeros({numVoxels}, at::TensorOptions().dtype(at::kInt))
+          .to(vol.device());
+  ThrustScanWrapper(
+      d_voxelOccupiedScan.data_ptr<int>(),
+      d_voxelOccupied.data_ptr<int>(),
+      numVoxels);
 
   // number of active voxels
-  int64_t activeVoxels = d_voxelOccupiedScan[numVoxels].cpu().item<int64_t>();
+  int lastElement = d_voxelVerts[numVoxels - 1].cpu().item<int>();
+  int lastScan = d_voxelOccupiedScan[numVoxels - 1].cpu().item<int>();
+  int activeVoxels = lastElement + lastScan;
 
   const int device_id = vol.device().index();
   auto opt = at::TensorOptions().dtype(at::kInt).device(at::kCUDA, device_id);
@@ -493,18 +509,22 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> MarchingCubesCuda(
   CompactVoxelsKernel<<<grid, threads, 0, stream>>>(
       d_compVoxelArray.packed_accessor32<int, 1, at::RestrictPtrTraits>(),
       d_voxelOccupied.packed_accessor32<int, 1, at::RestrictPtrTraits>(),
-      d_voxelOccupiedScan_
-          .packed_accessor32<int64_t, 1, at::RestrictPtrTraits>(),
+      d_voxelOccupiedScan.packed_accessor32<int, 1, at::RestrictPtrTraits>(),
       numVoxels);
   AT_CUDA_CHECK(cudaGetLastError());
   cudaDeviceSynchronize();
 
   // Scan d_voxelVerts array to generate offsets of vertices for each voxel
-  auto d_voxelVertsScan = at::cumsum(d_voxelVerts, 0);
-  auto d_voxelVertsScan_ = d_voxelVertsScan.index({Slice(1, None)});
+  auto d_voxelVertsScan = at::zeros({numVoxels}, opt);
+  ThrustScanWrapper(
+      d_voxelVertsScan.data_ptr<int>(),
+      d_voxelVerts.data_ptr<int>(),
+      numVoxels);
 
   // total number of vertices
-  int64_t totalVerts = d_voxelVertsScan[numVoxels].cpu().item<int64_t>();
+  lastElement = d_voxelVerts[numVoxels - 1].cpu().item<int>();
+  lastScan = d_voxelVertsScan[numVoxels - 1].cpu().item<int>();
+  int totalVerts = lastElement + lastScan;
 
   // Execute "GenerateFacesKernel" kernel
   // This runs only on the occupied voxels.
@@ -524,7 +544,7 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> MarchingCubesCuda(
       faces.packed_accessor<int64_t, 2, at::RestrictPtrTraits>(),
       ids.packed_accessor<int64_t, 1, at::RestrictPtrTraits>(),
       d_compVoxelArray.packed_accessor32<int, 1, at::RestrictPtrTraits>(),
-      d_voxelVertsScan_.packed_accessor32<int64_t, 1, at::RestrictPtrTraits>(),
+      d_voxelVertsScan.packed_accessor32<int, 1, at::RestrictPtrTraits>(),
       activeVoxels,
       vol.packed_accessor32<float, 3, at::RestrictPtrTraits>(),
       faceTable.packed_accessor32<int, 2, at::RestrictPtrTraits>(),

pytorch3d/csrc/marching_cubes/marching_cubes_cpu.cpp

@@ -71,8 +71,8 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> MarchingCubesCpu(
           if ((j + 1) % 3 == 0 && ps[0] != ps[1] && ps[1] != ps[2] &&
               ps[2] != ps[0]) {
             for (int k = 0; k < 3; k++) {
-              int64_t v = tri.at(k);
-              edge_id_to_v[v] = ps.at(k);
+              int v = tri[k];
+              edge_id_to_v[tri.at(k)] = ps.at(k);
               if (!uniq_edge_id.count(v)) {
                 uniq_edge_id[v] = verts.size();
                 verts.push_back(edge_id_to_v[v]);

pytorch3d/csrc/point_mesh/point_mesh_cuda.cu

@@ -305,8 +305,6 @@ std::tuple<at::Tensor, at::Tensor> DistanceBackwardCuda(
   at::CheckedFrom c = "DistanceBackwardCuda";
   at::checkAllSameGPU(c, {objects_t, targets_t, idx_objects_t, grad_dists_t});
   at::checkAllSameType(c, {objects_t, targets_t, grad_dists_t});
-  // This is nondeterministic because atomicAdd
-  at::globalContext().alertNotDeterministic("DistanceBackwardCuda");
 
   // Set the device for the kernel launch based on the device of the input
   at::cuda::CUDAGuard device_guard(objects.device());
@@ -626,9 +624,6 @@ std::tuple<at::Tensor, at::Tensor> PointFaceArrayDistanceBackwardCuda(
   at::CheckedFrom c = "PointFaceArrayDistanceBackwardCuda";
   at::checkAllSameGPU(c, {points_t, tris_t, grad_dists_t});
   at::checkAllSameType(c, {points_t, tris_t, grad_dists_t});
-  // This is nondeterministic because atomicAdd
-  at::globalContext().alertNotDeterministic(
-      "PointFaceArrayDistanceBackwardCuda");
 
   // Set the device for the kernel launch based on the device of the input
   at::cuda::CUDAGuard device_guard(points.device());
@@ -792,9 +787,6 @@ std::tuple<at::Tensor, at::Tensor> PointEdgeArrayDistanceBackwardCuda(
   at::CheckedFrom c = "PointEdgeArrayDistanceBackwardCuda";
   at::checkAllSameGPU(c, {points_t, segms_t, grad_dists_t});
   at::checkAllSameType(c, {points_t, segms_t, grad_dists_t});
-  // This is nondeterministic because atomicAdd
-  at::globalContext().alertNotDeterministic(
-      "PointEdgeArrayDistanceBackwardCuda");
 
   // Set the device for the kernel launch based on the device of the input
   at::cuda::CUDAGuard device_guard(points.device());

pytorch3d/csrc/points_to_volumes/points_to_volumes.cu

@@ -141,9 +141,6 @@ void PointsToVolumesForwardCuda(
        grid_sizes_t,
        mask_t});
 
-  // This is nondeterministic because atomicAdd
-  at::globalContext().alertNotDeterministic("PointsToVolumesForwardCuda");
-
   // Set the device for the kernel launch based on the device of the input
   at::cuda::CUDAGuard device_guard(points_3d.device());
   cudaStream_t stream = at::cuda::getCurrentCUDAStream();

pytorch3d/csrc/points_to_volumes/points_to_volumes.h

@@ -7,7 +7,6 @@
  */
 
 #pragma once
-#include <torch/csrc/autograd/VariableTypeUtils.h>
 #include <torch/extension.h>
 #include <cstdio>
 #include <tuple>
@@ -97,8 +96,6 @@ inline void PointsToVolumesForward(
         point_weight,
         align_corners,
         splat);
-    torch::autograd::increment_version(volume_features);
-    torch::autograd::increment_version(volume_densities);
     return;
 #else
     AT_ERROR("Not compiled with GPU support.");

pytorch3d/csrc/points_to_volumes/points_to_volumes_cpu.cpp

@@ -6,7 +6,6 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-#include <torch/csrc/autograd/VariableTypeUtils.h>
 #include <torch/extension.h>
 #include <algorithm>
 #include <cmath>
@@ -149,8 +148,6 @@ void PointsToVolumesForwardCpu(
       }
     }
   }
-  torch::autograd::increment_version(volume_features);
-  torch::autograd::increment_version(volume_densities);
 }
 
 // With nearest, the only smooth dependence is that volume features

pytorch3d/csrc/pulsar/global.h

@@ -30,18 +30,11 @@
 #define GLOBAL __global__
 #define RESTRICT __restrict__
 #define DEBUGBREAK()
-#ifdef __NVCC_DIAG_PRAGMA_SUPPORT__
-#pragma nv_diag_suppress 1866
-#pragma nv_diag_suppress 2941
-#pragma nv_diag_suppress 2951
-#pragma nv_diag_suppress 2967
-#else
 #pragma diag_suppress = attribute_not_allowed
 #pragma diag_suppress = 1866
 #pragma diag_suppress = 2941
 #pragma diag_suppress = 2951
 #pragma diag_suppress = 2967
-#endif
 #else // __CUDACC__
 #define INLINE inline
 #define HOST
@@ -56,7 +49,6 @@
 #pragma clang diagnostic pop
 #ifdef WITH_CUDA
 #include <ATen/cuda/CUDAContext.h>
-#include <vector_functions.h>
 #else
 #ifndef cudaStream_t
 typedef void* cudaStream_t;
@@ -73,6 +65,8 @@ struct float2 {
 struct float3 {
   float x, y, z;
 };
+#endif
+namespace py = pybind11;
 inline float3 make_float3(const float& x, const float& y, const float& z) {
   float3 res;
   res.x = x;
@@ -80,8 +74,6 @@ inline float3 make_float3(const float& x, const float& y, const float& z) {
   res.z = z;
   return res;
 }
-#endif
-namespace py = pybind11;
 
 inline bool operator==(const float3& a, const float3& b) {
   return a.x == b.x && a.y == b.y && a.z == b.z;

pytorch3d/csrc/pulsar/include/renderer.forward.device.h

@@ -102,7 +102,6 @@ void forward(
       self->workspace_d,
       self->workspace_size,
       stream);
-  CHECKLAUNCH();
   SORT_ASCENDING_WS(
       self->min_depth_d,
       self->min_depth_sorted_d,
@@ -112,7 +111,6 @@ void forward(
       self->workspace_d,
       self->workspace_size,
       stream);
-  CHECKLAUNCH();
   SORT_ASCENDING_WS(
       self->min_depth_d,
       self->min_depth_sorted_d,

pytorch3d/csrc/pulsar/pytorch/camera.h

@@ -37,7 +37,7 @@ inline void fill_cam_vecs(
   res->pixel_dir_y.x = pixel_dir_y.data_ptr<float>()[0];
   res->pixel_dir_y.y = pixel_dir_y.data_ptr<float>()[1];
   res->pixel_dir_y.z = pixel_dir_y.data_ptr<float>()[2];
-  auto sensor_dir_z = pixel_dir_y.cross(pixel_dir_x, -1);
+  auto sensor_dir_z = pixel_dir_y.cross(pixel_dir_x);
   sensor_dir_z /= sensor_dir_z.norm();
   if (right_handed) {
     sensor_dir_z *= -1.f;

pytorch3d/csrc/rasterize_meshes/rasterize_meshes.cu

@@ -583,9 +583,6 @@ at::Tensor RasterizeMeshesBackwardCuda(
   at::checkAllSameType(
       c, {face_verts_t, grad_zbuf_t, grad_bary_t, grad_dists_t});
 
-  // This is nondeterministic because atomicAdd
-  at::globalContext().alertNotDeterministic("RasterizeMeshesBackwardCuda");
-
   // Set the device for the kernel launch based on the device of the input
   at::cuda::CUDAGuard device_guard(face_verts.device());
   cudaStream_t stream = at::cuda::getCurrentCUDAStream();

pytorch3d/csrc/rasterize_points/rasterize_points.cu

@@ -423,8 +423,7 @@ at::Tensor RasterizePointsBackwardCuda(
   at::CheckedFrom c = "RasterizePointsBackwardCuda";
   at::checkAllSameGPU(c, {points_t, idxs_t, grad_zbuf_t, grad_dists_t});
   at::checkAllSameType(c, {points_t, grad_zbuf_t, grad_dists_t});
-  // This is nondeterministic because atomicAdd
-  at::globalContext().alertNotDeterministic("RasterizePointsBackwardCuda");
+
   // Set the device for the kernel launch based on the device of the input
   at::cuda::CUDAGuard device_guard(points.device());
   cudaStream_t stream = at::cuda::getCurrentCUDAStream();

pytorch3d/csrc/sample_farthest_points/sample_farthest_points.cu

@@ -155,7 +155,7 @@ at::Tensor FarthestPointSamplingCuda(
 
   // Max possible threads per block
   const int MAX_THREADS_PER_BLOCK = 1024;
-  const size_t threads = max(min(1 << points_pow_2, MAX_THREADS_PER_BLOCK), 2);
+  const size_t threads = max(min(1 << points_pow_2, MAX_THREADS_PER_BLOCK), 1);
 
   // Create the accessors
   auto points_a = points.packed_accessor64<float, 3, at::RestrictPtrTraits>();
@@ -215,6 +215,10 @@ at::Tensor FarthestPointSamplingCuda(
       FarthestPointSamplingKernel<2><<<threads, threads, shared_mem, stream>>>(
           points_a, lengths_a, K_a, idxs_a, min_point_dist_a, start_idxs_a);
       break;
+    case 1:
+      FarthestPointSamplingKernel<1><<<threads, threads, shared_mem, stream>>>(
+          points_a, lengths_a, K_a, idxs_a, min_point_dist_a, start_idxs_a);
+      break;
     default:
       FarthestPointSamplingKernel<1024>
           <<<blocks, threads, shared_mem, stream>>>(

pytorch3d/csrc/sample_pdf/sample_pdf.h

@@ -7,7 +7,6 @@
  */
 
 #pragma once
-#include <torch/csrc/autograd/VariableTypeUtils.h>
 #include <torch/extension.h>
 #include <cstdio>
 #include <tuple>
@@ -64,7 +63,6 @@ inline void SamplePdf(
 #ifdef WITH_CUDA
     CHECK_CUDA(weights);
     CHECK_CONTIGUOUS_CUDA(outputs);
-    torch::autograd::increment_version(outputs);
     SamplePdfCuda(bins, weights, outputs, eps);
     return;
 #else

pytorch3d/csrc/sample_pdf/sample_pdf_cpu.cpp

@@ -6,7 +6,6 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-#include <torch/csrc/autograd/VariableTypeUtils.h>
 #include <torch/extension.h>
 #include <algorithm>
 #include <thread>
@@ -138,5 +137,4 @@ void SamplePdfCpu(
   for (auto&& thread : threads) {
     thread.join();
   }
-  torch::autograd::increment_version(outputs);
 }

pytorch3d/implicitron/dataset/data_loader_map_provider.py

@@ -12,15 +12,14 @@ import torch
 from pytorch3d.implicitron.tools.config import registry, ReplaceableBase
 from torch.utils.data import (
     BatchSampler,
-    ConcatDataset,
+    ChainDataset,
     DataLoader,
     RandomSampler,
     Sampler,
 )
 
-from .dataset_base import DatasetBase
+from .dataset_base import DatasetBase, FrameData
 from .dataset_map_provider import DatasetMap
-from .frame_data import FrameData
 from .scene_batch_sampler import SceneBatchSampler
 from .utils import is_known_frame_scalar
 
@@ -483,7 +482,7 @@ class SequenceDataLoaderMapProvider(DataLoaderMapProviderBase):
             num_batches=num_batches,
         )
         return DataLoader(
-            ConcatDataset([dataset, train_dataset]),
+            ChainDataset([dataset, train_dataset]),
             batch_sampler=sampler,
             **data_loader_kwargs,
         )

pytorch3d/implicitron/dataset/data_source.py

@@ -13,8 +13,13 @@ from pytorch3d.implicitron.tools.config import (
 )
 from pytorch3d.renderer.cameras import CamerasBase
 
+from .blender_dataset_map_provider import BlenderDatasetMapProvider  # noqa
 from .data_loader_map_provider import DataLoaderMap, DataLoaderMapProviderBase
 from .dataset_map_provider import DatasetMap, DatasetMapProviderBase
+from .json_index_dataset_map_provider import JsonIndexDatasetMapProvider  # noqa
+from .json_index_dataset_map_provider_v2 import JsonIndexDatasetMapProviderV2  # noqa
+from .llff_dataset_map_provider import LlffDatasetMapProvider  # noqa
+from .rendered_mesh_dataset_map_provider import RenderedMeshDatasetMapProvider  # noqa
 
 
 class DataSourceBase(ReplaceableBase):
@@ -29,7 +34,6 @@ class DataSourceBase(ReplaceableBase):
     @property
     def all_train_cameras(self) -> Optional[CamerasBase]:
         """
-        DEPRECATED! The property will be removed in future versions.
         If the data is all for a single scene, a list
         of the known training cameras for that scene, which is
         used for evaluating the viewpoint difficulty of the
@@ -55,36 +59,6 @@ class ImplicitronDataSource(DataSourceBase):  # pyre-ignore[13]
     data_loader_map_provider: DataLoaderMapProviderBase
     data_loader_map_provider_class_type: str = "SequenceDataLoaderMapProvider"
 
-    @classmethod
-    def pre_expand(cls) -> None:
-        # use try/finally to bypass cinder's lazy imports
-        try:
-            from .blender_dataset_map_provider import (  # noqa: F401
-                BlenderDatasetMapProvider,
-            )
-            from .json_index_dataset_map_provider import (  # noqa: F401
-                JsonIndexDatasetMapProvider,
-            )
-            from .json_index_dataset_map_provider_v2 import (  # noqa: F401
-                JsonIndexDatasetMapProviderV2,
-            )
-            from .llff_dataset_map_provider import LlffDatasetMapProvider  # noqa: F401
-            from .rendered_mesh_dataset_map_provider import (  # noqa: F401
-                RenderedMeshDatasetMapProvider,
-            )
-            from .train_eval_data_loader_provider import (  # noqa: F401
-                TrainEvalDataLoaderMapProvider,
-            )
-
-            try:
-                from .sql_dataset_provider import (  # noqa: F401  # pyre-ignore
-                    SqlIndexDatasetMapProvider,
-                )
-            except ModuleNotFoundError:
-                pass  # environment without SQL dataset
-        finally:
-            pass
-
     def __post_init__(self):
         run_auto_creation(self)
         self._all_train_cameras_cache: Optional[Tuple[Optional[CamerasBase]]] = None
@@ -96,9 +70,6 @@ class ImplicitronDataSource(DataSourceBase):  # pyre-ignore[13]
 
     @property
     def all_train_cameras(self) -> Optional[CamerasBase]:
-        """
-        DEPRECATED! The property will be removed in future versions.
-        """
         if self._all_train_cameras_cache is None:  # pyre-ignore[16]
             all_train_cameras = self.dataset_map_provider.get_all_train_cameras()
             self._all_train_cameras_cache = (all_train_cameras,)

pytorch3d/implicitron/dataset/dataset_base.py

@@ -5,27 +5,217 @@
 # LICENSE file in the root directory of this source tree.
 
 from collections import defaultdict
-from dataclasses import dataclass
+from dataclasses import dataclass, field, fields
 from typing import (
+    Any,
     ClassVar,
     Dict,
     Iterable,
     Iterator,
     List,
+    Mapping,
     Optional,
     Sequence,
     Tuple,
     Type,
+    Union,
 )
 
+import numpy as np
 import torch
+from pytorch3d.renderer.camera_utils import join_cameras_as_batch
+from pytorch3d.renderer.cameras import CamerasBase, PerspectiveCameras
+from pytorch3d.structures.pointclouds import join_pointclouds_as_batch, Pointclouds
 
-from pytorch3d.implicitron.dataset.frame_data import FrameData
-from pytorch3d.implicitron.dataset.utils import GenericWorkaround
+
+@dataclass
+class FrameData(Mapping[str, Any]):
+    """
+    A type of the elements returned by indexing the dataset object.
+    It can represent both individual frames and batches of thereof;
+    in this documentation, the sizes of tensors refer to single frames;
+    add the first batch dimension for the collation result.
+
+    Args:
+        frame_number: The number of the frame within its sequence.
+            0-based continuous integers.
+        sequence_name: The unique name of the frame's sequence.
+        sequence_category: The object category of the sequence.
+        frame_timestamp: The time elapsed since the start of a sequence in sec.
+        image_size_hw: The size of the image in pixels; (height, width) tensor
+                        of shape (2,).
+        image_path: The qualified path to the loaded image (with dataset_root).
+        image_rgb: A Tensor of shape `(3, H, W)` holding the RGB image
+            of the frame; elements are floats in [0, 1].
+        mask_crop: A binary mask of shape `(1, H, W)` denoting the valid image
+            regions. Regions can be invalid (mask_crop[i,j]=0) in case they
+            are a result of zero-padding of the image after cropping around
+            the object bounding box; elements are floats in {0.0, 1.0}.
+        depth_path: The qualified path to the frame's depth map.
+        depth_map: A float Tensor of shape `(1, H, W)` holding the depth map
+            of the frame; values correspond to distances from the camera;
+            use `depth_mask` and `mask_crop` to filter for valid pixels.
+        depth_mask: A binary mask of shape `(1, H, W)` denoting pixels of the
+            depth map that are valid for evaluation, they have been checked for
+            consistency across views; elements are floats in {0.0, 1.0}.
+        mask_path: A qualified path to the foreground probability mask.
+        fg_probability: A Tensor of `(1, H, W)` denoting the probability of the
+            pixels belonging to the captured object; elements are floats
+            in [0, 1].
+        bbox_xywh: The bounding box tightly enclosing the foreground object in the
+            format (x0, y0, width, height). The convention assumes that
+            `x0+width` and `y0+height` includes the boundary of the box.
+            I.e., to slice out the corresponding crop from an image tensor `I`
+            we execute `crop = I[..., y0:y0+height, x0:x0+width]`
+        crop_bbox_xywh: The bounding box denoting the boundaries of `image_rgb`
+            in the original image coordinates in the format (x0, y0, width, height).
+            The convention is the same as for `bbox_xywh`. `crop_bbox_xywh` differs
+            from `bbox_xywh` due to padding (which can happen e.g. due to
+            setting `JsonIndexDataset.box_crop_context > 0`)
+        camera: A PyTorch3D camera object corresponding the frame's viewpoint,
+            corrected for cropping if it happened.
+        camera_quality_score: The score proportional to the confidence of the
+            frame's camera estimation (the higher the more accurate).
+        point_cloud_quality_score: The score proportional to the accuracy of the
+            frame's sequence point cloud (the higher the more accurate).
+        sequence_point_cloud_path: The path to the sequence's point cloud.
+        sequence_point_cloud: A PyTorch3D Pointclouds object holding the
+            point cloud corresponding to the frame's sequence. When the object
+            represents a batch of frames, point clouds may be deduplicated;
+            see `sequence_point_cloud_idx`.
+        sequence_point_cloud_idx: Integer indices mapping frame indices to the
+            corresponding point clouds in `sequence_point_cloud`; to get the
+            corresponding point cloud to `image_rgb[i]`, use
+            `sequence_point_cloud[sequence_point_cloud_idx[i]]`.
+        frame_type: The type of the loaded frame specified in
+            `subset_lists_file`, if provided.
+        meta: A dict for storing additional frame information.
+    """
+
+    frame_number: Optional[torch.LongTensor]
+    sequence_name: Union[str, List[str]]
+    sequence_category: Union[str, List[str]]
+    frame_timestamp: Optional[torch.Tensor] = None
+    image_size_hw: Optional[torch.Tensor] = None
+    image_path: Union[str, List[str], None] = None
+    image_rgb: Optional[torch.Tensor] = None
+    # masks out padding added due to cropping the square bit
+    mask_crop: Optional[torch.Tensor] = None
+    depth_path: Union[str, List[str], None] = None
+    depth_map: Optional[torch.Tensor] = None
+    depth_mask: Optional[torch.Tensor] = None
+    mask_path: Union[str, List[str], None] = None
+    fg_probability: Optional[torch.Tensor] = None
+    bbox_xywh: Optional[torch.Tensor] = None
+    crop_bbox_xywh: Optional[torch.Tensor] = None
+    camera: Optional[PerspectiveCameras] = None
+    camera_quality_score: Optional[torch.Tensor] = None
+    point_cloud_quality_score: Optional[torch.Tensor] = None
+    sequence_point_cloud_path: Union[str, List[str], None] = None
+    sequence_point_cloud: Optional[Pointclouds] = None
+    sequence_point_cloud_idx: Optional[torch.Tensor] = None
+    frame_type: Union[str, List[str], None] = None  # known | unseen
+    meta: dict = field(default_factory=lambda: {})
+
+    def to(self, *args, **kwargs):
+        new_params = {}
+        for f in fields(self):
+            value = getattr(self, f.name)
+            if isinstance(value, (torch.Tensor, Pointclouds, CamerasBase)):
+                new_params[f.name] = value.to(*args, **kwargs)
+            else:
+                new_params[f.name] = value
+        return type(self)(**new_params)
+
+    def cpu(self):
+        return self.to(device=torch.device("cpu"))
+
+    def cuda(self):
+        return self.to(device=torch.device("cuda"))
+
+    # the following functions make sure **frame_data can be passed to functions
+    def __iter__(self):
+        for f in fields(self):
+            yield f.name
+
+    def __getitem__(self, key):
+        return getattr(self, key)
+
+    def __len__(self):
+        return len(fields(self))
+
+    @classmethod
+    def collate(cls, batch):
+        """
+        Given a list objects `batch` of class `cls`, collates them into a batched
+        representation suitable for processing with deep networks.
+        """
+
+        elem = batch[0]
+
+        if isinstance(elem, cls):
+            pointcloud_ids = [id(el.sequence_point_cloud) for el in batch]
+            id_to_idx = defaultdict(list)
+            for i, pc_id in enumerate(pointcloud_ids):
+                id_to_idx[pc_id].append(i)
+
+            sequence_point_cloud = []
+            sequence_point_cloud_idx = -np.ones((len(batch),))
+            for i, ind in enumerate(id_to_idx.values()):
+                sequence_point_cloud_idx[ind] = i
+                sequence_point_cloud.append(batch[ind[0]].sequence_point_cloud)
+            assert (sequence_point_cloud_idx >= 0).all()
+
+            override_fields = {
+                "sequence_point_cloud": sequence_point_cloud,
+                "sequence_point_cloud_idx": sequence_point_cloud_idx.tolist(),
+            }
+            # note that the pre-collate value of sequence_point_cloud_idx is unused
+
+            collated = {}
+            for f in fields(elem):
+                list_values = override_fields.get(
+                    f.name, [getattr(d, f.name) for d in batch]
+                )
+                collated[f.name] = (
+                    cls.collate(list_values)
+                    if all(list_value is not None for list_value in list_values)
+                    else None
+                )
+            return cls(**collated)
+
+        elif isinstance(elem, Pointclouds):
+            return join_pointclouds_as_batch(batch)
+
+        elif isinstance(elem, CamerasBase):
+            # TODO: don't store K; enforce working in NDC space
+            return join_cameras_as_batch(batch)
+        else:
+            return torch.utils.data._utils.collate.default_collate(batch)
+
+
+class _GenericWorkaround:
+    """
+    OmegaConf.structured has a weirdness when you try to apply
+    it to a dataclass whose first base class is a Generic which is not
+    Dict. The issue is with a function called get_dict_key_value_types
+    in omegaconf/_utils.py.
+    For example this fails:
+
+        @dataclass(eq=False)
+        class D(torch.utils.data.Dataset[int]):
+            a: int = 3
+
+        OmegaConf.structured(D)
+
+    We avoid the problem by adding this class as an extra base class.
+    """
+
+    pass
 
 
 @dataclass(eq=False)
-class DatasetBase(GenericWorkaround, torch.utils.data.Dataset[FrameData]):
+class DatasetBase(_GenericWorkaround, torch.utils.data.Dataset[FrameData]):
     """
     Base class to describe a dataset to be used with Implicitron.
 
@@ -47,7 +237,7 @@ class DatasetBase(GenericWorkaround, torch.utils.data.Dataset[FrameData]):
         raise NotImplementedError()
 
     def get_frame_numbers_and_timestamps(
-        self, idxs: Sequence[int], subset_filter: Optional[Sequence[str]] = None
+        self, idxs: Sequence[int]
     ) -> List[Tuple[int, float]]:
         """
         If the sequences in the dataset are videos rather than
@@ -61,9 +251,7 @@ class DatasetBase(GenericWorkaround, torch.utils.data.Dataset[FrameData]):
         frames.
 
         Args:
-            idxs: frame index in self
-            subset_filter: If given, an index in idxs is ignored if the
-                corresponding frame is not in any of the named subsets.
+            idx: frame index in self
 
         Returns:
             tuple of
@@ -103,7 +291,7 @@ class DatasetBase(GenericWorkaround, torch.utils.data.Dataset[FrameData]):
         return dict(c2seq)
 
     def sequence_frames_in_order(
-        self, seq_name: str, subset_filter: Optional[Sequence[str]] = None
+        self, seq_name: str
     ) -> Iterator[Tuple[float, int, int]]:
         """Returns an iterator over the frame indices in a given sequence.
         We attempt to first sort by timestamp (if they are available),
@@ -120,9 +308,7 @@ class DatasetBase(GenericWorkaround, torch.utils.data.Dataset[FrameData]):
         """
         # pyre-ignore[16]
         seq_frame_indices = self._seq_to_idx[seq_name]
-        nos_timestamps = self.get_frame_numbers_and_timestamps(
-            seq_frame_indices, subset_filter
-        )
+        nos_timestamps = self.get_frame_numbers_and_timestamps(seq_frame_indices)
 
         yield from sorted(
             [
@@ -131,13 +317,11 @@ class DatasetBase(GenericWorkaround, torch.utils.data.Dataset[FrameData]):
             ]
         )
 
-    def sequence_indices_in_order(
-        self, seq_name: str, subset_filter: Optional[Sequence[str]] = None
-    ) -> Iterator[int]:
+    def sequence_indices_in_order(self, seq_name: str) -> Iterator[int]:
         """Same as `sequence_frames_in_order` but returns the iterator over
         only dataset indices.
         """
-        for _, _, idx in self.sequence_frames_in_order(seq_name, subset_filter):
+        for _, _, idx in self.sequence_frames_in_order(seq_name):
             yield idx
 
     # frame_data_type is the actual type of frames returned by the dataset.

pytorch3d/implicitron/dataset/dataset_map_provider.py

@@ -95,7 +95,6 @@ class DatasetMapProviderBase(ReplaceableBase):
 
     def get_all_train_cameras(self) -> Optional[CamerasBase]:
         """
-        DEPRECATED! The function will be removed in future versions.
         If the data is all for a single scene, returns a list
         of the known training cameras for that scene, which is
         used for evaluating the difficulty of the unknown

pytorch3d/implicitron/dataset/frame_data.py

@@ -1,777 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the BSD-style license found in the
-# LICENSE file in the root directory of this source tree.
-
-import os
-from abc import ABC, abstractmethod
-from collections import defaultdict
-from dataclasses import dataclass, field, fields
-from typing import (
-    Any,
-    ClassVar,
-    Generic,
-    List,
-    Mapping,
-    Optional,
-    Tuple,
-    Type,
-    TypeVar,
-    Union,
-)
-
-import numpy as np
-import torch
-
-from pytorch3d.implicitron.dataset import types
-from pytorch3d.implicitron.dataset.utils import (
-    adjust_camera_to_bbox_crop_,
-    adjust_camera_to_image_scale_,
-    bbox_xyxy_to_xywh,
-    clamp_box_to_image_bounds_and_round,
-    crop_around_box,
-    GenericWorkaround,
-    get_bbox_from_mask,
-    get_clamp_bbox,
-    load_depth,
-    load_depth_mask,
-    load_image,
-    load_mask,
-    load_pointcloud,
-    rescale_bbox,
-    resize_image,
-    safe_as_tensor,
-)
-from pytorch3d.implicitron.tools.config import registry, ReplaceableBase
-from pytorch3d.renderer.camera_utils import join_cameras_as_batch
-from pytorch3d.renderer.cameras import CamerasBase, PerspectiveCameras
-from pytorch3d.structures.pointclouds import join_pointclouds_as_batch, Pointclouds
-
-
-@dataclass
-class FrameData(Mapping[str, Any]):
-    """
-    A type of the elements returned by indexing the dataset object.
-    It can represent both individual frames and batches of thereof;
-    in this documentation, the sizes of tensors refer to single frames;
-    add the first batch dimension for the collation result.
-
-    Args:
-        frame_number: The number of the frame within its sequence.
-            0-based continuous integers.
-        sequence_name: The unique name of the frame's sequence.
-        sequence_category: The object category of the sequence.
-        frame_timestamp: The time elapsed since the start of a sequence in sec.
-        image_size_hw: The size of the original image in pixels; (height, width)
-            tensor of shape (2,). Note that it is optional, e.g. it can be `None`
-            if the frame annotation has no size ans image_rgb has not [yet] been
-            loaded. Image-less FrameData is valid but mutators like crop/resize
-            may fail if the original image size cannot be deduced.
-        effective_image_size_hw: The size of the image after mutations such as
-            crop/resize in pixels; (height, width). if the image has not been mutated,
-            it is equal to `image_size_hw`. Note that it is also optional, for the
-            same reason as `image_size_hw`.
-        image_path: The qualified path to the loaded image (with dataset_root).
-        image_rgb: A Tensor of shape `(3, H, W)` holding the RGB image
-            of the frame; elements are floats in [0, 1].
-        mask_crop: A binary mask of shape `(1, H, W)` denoting the valid image
-            regions. Regions can be invalid (mask_crop[i,j]=0) in case they
-            are a result of zero-padding of the image after cropping around
-            the object bounding box; elements are floats in {0.0, 1.0}.
-        depth_path: The qualified path to the frame's depth map.
-        depth_map: A float Tensor of shape `(1, H, W)` holding the depth map
-            of the frame; values correspond to distances from the camera;
-            use `depth_mask` and `mask_crop` to filter for valid pixels.
-        depth_mask: A binary mask of shape `(1, H, W)` denoting pixels of the
-            depth map that are valid for evaluation, they have been checked for
-            consistency across views; elements are floats in {0.0, 1.0}.
-        mask_path: A qualified path to the foreground probability mask.
-        fg_probability: A Tensor of `(1, H, W)` denoting the probability of the
-            pixels belonging to the captured object; elements are floats
-            in [0, 1].
-        bbox_xywh: The bounding box tightly enclosing the foreground object in the
-            format (x0, y0, width, height). The convention assumes that
-            `x0+width` and `y0+height` includes the boundary of the box.
-            I.e., to slice out the corresponding crop from an image tensor `I`
-            we execute `crop = I[..., y0:y0+height, x0:x0+width]`
-        crop_bbox_xywh: The bounding box denoting the boundaries of `image_rgb`
-            in the original image coordinates in the format (x0, y0, width, height).
-            The convention is the same as for `bbox_xywh`. `crop_bbox_xywh` differs
-            from `bbox_xywh` due to padding (which can happen e.g. due to
-            setting `JsonIndexDataset.box_crop_context > 0`)
-        camera: A PyTorch3D camera object corresponding the frame's viewpoint,
-            corrected for cropping if it happened.
-        camera_quality_score: The score proportional to the confidence of the
-            frame's camera estimation (the higher the more accurate).
-        point_cloud_quality_score: The score proportional to the accuracy of the
-            frame's sequence point cloud (the higher the more accurate).
-        sequence_point_cloud_path: The path to the sequence's point cloud.
-        sequence_point_cloud: A PyTorch3D Pointclouds object holding the
-            point cloud corresponding to the frame's sequence. When the object
-            represents a batch of frames, point clouds may be deduplicated;
-            see `sequence_point_cloud_idx`.
-        sequence_point_cloud_idx: Integer indices mapping frame indices to the
-            corresponding point clouds in `sequence_point_cloud`; to get the
-            corresponding point cloud to `image_rgb[i]`, use
-            `sequence_point_cloud[sequence_point_cloud_idx[i]]`.
-        frame_type: The type of the loaded frame specified in
-            `subset_lists_file`, if provided.
-        meta: A dict for storing additional frame information.
-    """
-
-    frame_number: Optional[torch.LongTensor]
-    sequence_name: Union[str, List[str]]
-    sequence_category: Union[str, List[str]]
-    frame_timestamp: Optional[torch.Tensor] = None
-    image_size_hw: Optional[torch.LongTensor] = None
-    effective_image_size_hw: Optional[torch.LongTensor] = None
-    image_path: Union[str, List[str], None] = None
-    image_rgb: Optional[torch.Tensor] = None
-    # masks out padding added due to cropping the square bit
-    mask_crop: Optional[torch.Tensor] = None
-    depth_path: Union[str, List[str], None] = None
-    depth_map: Optional[torch.Tensor] = None
-    depth_mask: Optional[torch.Tensor] = None
-    mask_path: Union[str, List[str], None] = None
-    fg_probability: Optional[torch.Tensor] = None
-    bbox_xywh: Optional[torch.Tensor] = None
-    crop_bbox_xywh: Optional[torch.Tensor] = None
-    camera: Optional[PerspectiveCameras] = None
-    camera_quality_score: Optional[torch.Tensor] = None
-    point_cloud_quality_score: Optional[torch.Tensor] = None
-    sequence_point_cloud_path: Union[str, List[str], None] = None
-    sequence_point_cloud: Optional[Pointclouds] = None
-    sequence_point_cloud_idx: Optional[torch.Tensor] = None
-    frame_type: Union[str, List[str], None] = None  # known | unseen
-    meta: dict = field(default_factory=lambda: {})
-
-    # NOTE that batching resets this attribute
-    _uncropped: bool = field(init=False, default=True)
-
-    def to(self, *args, **kwargs):
-        new_params = {}
-        for field_name in iter(self):
-            value = getattr(self, field_name)
-            if isinstance(value, (torch.Tensor, Pointclouds, CamerasBase)):
-                new_params[field_name] = value.to(*args, **kwargs)
-            else:
-                new_params[field_name] = value
-        frame_data = type(self)(**new_params)
-        frame_data._uncropped = self._uncropped
-        return frame_data
-
-    def cpu(self):
-        return self.to(device=torch.device("cpu"))
-
-    def cuda(self):
-        return self.to(device=torch.device("cuda"))
-
-    # the following functions make sure **frame_data can be passed to functions
-    def __iter__(self):
-        for f in fields(self):
-            if f.name.startswith("_"):
-                continue
-
-            yield f.name
-
-    def __getitem__(self, key):
-        return getattr(self, key)
-
-    def __len__(self):
-        return sum(1 for f in iter(self))
-
-    def crop_by_metadata_bbox_(
-        self,
-        box_crop_context: float,
-    ) -> None:
-        """Crops the frame data in-place by (possibly expanded) bounding box.
-        The bounding box is taken from the object state (usually taken from
-        the frame annotation or estimated from the foregroubnd mask).
-        If the expanded bounding box does not fit the image, it is clamped,
-        i.e. the image is *not* padded.
-
-        Args:
-            box_crop_context: rate of expansion for bbox; 0 means no expansion,
-
-        Raises:
-            ValueError: If the object does not contain a bounding box (usually when no
-                mask annotation is provided)
-            ValueError: If the frame data have been cropped or resized, thus the intrinsic
-                bounding box is not valid for the current image size.
-            ValueError: If the frame does not have an image size (usually a corner case
-                when no image has been loaded)
-        """
-        if self.bbox_xywh is None:
-            raise ValueError(
-                "Attempted cropping by metadata with empty bounding box. Consider either"
-                " to remove_empty_masks or turn off box_crop in the dataset config."
-            )
-
-        if not self._uncropped:
-            raise ValueError(
-                "Trying to apply the metadata bounding box to already cropped "
-                "or resized image; coordinates have changed."
-            )
-
-        self._crop_by_bbox_(
-            box_crop_context,
-            self.bbox_xywh,
-        )
-
-    def crop_by_given_bbox_(
-        self,
-        box_crop_context: float,
-        bbox_xywh: torch.Tensor,
-    ) -> None:
-        """Crops the frame data in-place by (possibly expanded) bounding box.
-        If the expanded bounding box does not fit the image, it is clamped,
-        i.e. the image is *not* padded.
-
-        Args:
-            box_crop_context: rate of expansion for bbox; 0 means no expansion,
-            bbox_xywh: bounding box in [x0, y0, width, height] format. If float
-                tensor, values are floored (after converting to [x0, y0, x1, y1]).
-
-        Raises:
-            ValueError: If the frame does not have an image size (usually a corner case
-                when no image has been loaded)
-        """
-        self._crop_by_bbox_(
-            box_crop_context,
-            bbox_xywh,
-        )
-
-    def _crop_by_bbox_(
-        self,
-        box_crop_context: float,
-        bbox_xywh: torch.Tensor,
-    ) -> None:
-        """Crops the frame data in-place by (possibly expanded) bounding box.
-        If the expanded bounding box does not fit the image, it is clamped,
-        i.e. the image is *not* padded.
-
-        Args:
-            box_crop_context: rate of expansion for bbox; 0 means no expansion,
-            bbox_xywh: bounding box in [x0, y0, width, height] format. If float
-                tensor, values are floored (after converting to [x0, y0, x1, y1]).
-
-        Raises:
-            ValueError: If the frame does not have an image size (usually a corner case
-                when no image has been loaded)
-        """
-        effective_image_size_hw = self.effective_image_size_hw
-        if effective_image_size_hw is None:
-            raise ValueError("Calling crop on image-less FrameData")
-
-        bbox_xyxy = get_clamp_bbox(
-            bbox_xywh,
-            image_path=self.image_path,  # pyre-ignore
-            box_crop_context=box_crop_context,
-        )
-        clamp_bbox_xyxy = clamp_box_to_image_bounds_and_round(
-            bbox_xyxy,
-            image_size_hw=tuple(self.effective_image_size_hw),  # pyre-ignore
-        )
-        crop_bbox_xywh = bbox_xyxy_to_xywh(clamp_bbox_xyxy)
-
-        if self.fg_probability is not None:
-            self.fg_probability = crop_around_box(
-                self.fg_probability,
-                clamp_bbox_xyxy,
-                self.mask_path,  # pyre-ignore
-            )
-        if self.image_rgb is not None:
-            self.image_rgb = crop_around_box(
-                self.image_rgb,
-                clamp_bbox_xyxy,
-                self.image_path,  # pyre-ignore
-            )
-
-        depth_map = self.depth_map
-        if depth_map is not None:
-            clamp_bbox_xyxy_depth = rescale_bbox(
-                clamp_bbox_xyxy, tuple(depth_map.shape[-2:]), effective_image_size_hw
-            ).long()
-            self.depth_map = crop_around_box(
-                depth_map,
-                clamp_bbox_xyxy_depth,
-                self.depth_path,  # pyre-ignore
-            )
-
-        depth_mask = self.depth_mask
-        if depth_mask is not None:
-            clamp_bbox_xyxy_depth = rescale_bbox(
-                clamp_bbox_xyxy, tuple(depth_mask.shape[-2:]), effective_image_size_hw
-            ).long()
-            self.depth_mask = crop_around_box(
-                depth_mask,
-                clamp_bbox_xyxy_depth,
-                self.mask_path,  # pyre-ignore
-            )
-
-        # changing principal_point according to bbox_crop
-        if self.camera is not None:
-            adjust_camera_to_bbox_crop_(
-                camera=self.camera,
-                image_size_wh=effective_image_size_hw.flip(dims=[-1]),
-                clamp_bbox_xywh=crop_bbox_xywh,
-            )
-
-        # pyre-ignore
-        self.effective_image_size_hw = crop_bbox_xywh[..., 2:].flip(dims=[-1])
-        self._uncropped = False
-
-    def resize_frame_(self, new_size_hw: torch.LongTensor) -> None:
-        """Resizes frame data in-place according to given dimensions.
-
-        Args:
-            new_size_hw: target image size [height, width], a LongTensor of shape (2,)
-
-        Raises:
-            ValueError: If the frame does not have an image size (usually a corner case
-                when no image has been loaded)
-        """
-
-        effective_image_size_hw = self.effective_image_size_hw
-        if effective_image_size_hw is None:
-            raise ValueError("Calling resize on image-less FrameData")
-
-        image_height, image_width = new_size_hw.tolist()
-
-        if self.fg_probability is not None:
-            self.fg_probability, _, _ = resize_image(
-                self.fg_probability,
-                image_height=image_height,
-                image_width=image_width,
-                mode="nearest",
-            )
-
-        if self.image_rgb is not None:
-            self.image_rgb, _, self.mask_crop = resize_image(
-                self.image_rgb, image_height=image_height, image_width=image_width
-            )
-
-        if self.depth_map is not None:
-            self.depth_map, _, _ = resize_image(
-                self.depth_map,
-                image_height=image_height,
-                image_width=image_width,
-                mode="nearest",
-            )
-
-        if self.depth_mask is not None:
-            self.depth_mask, _, _ = resize_image(
-                self.depth_mask,
-                image_height=image_height,
-                image_width=image_width,
-                mode="nearest",
-            )
-
-        if self.camera is not None:
-            if self.image_size_hw is None:
-                raise ValueError(
-                    "image_size_hw has to be defined for resizing FrameData with cameras."
-                )
-            adjust_camera_to_image_scale_(
-                camera=self.camera,
-                original_size_wh=effective_image_size_hw.flip(dims=[-1]),
-                new_size_wh=new_size_hw.flip(dims=[-1]),  # pyre-ignore
-            )
-
-        self.effective_image_size_hw = new_size_hw
-        self._uncropped = False
-
-    @classmethod
-    def collate(cls, batch):
-        """
-        Given a list objects `batch` of class `cls`, collates them into a batched
-        representation suitable for processing with deep networks.
-        """
-
-        elem = batch[0]
-
-        if isinstance(elem, cls):
-            pointcloud_ids = [id(el.sequence_point_cloud) for el in batch]
-            id_to_idx = defaultdict(list)
-            for i, pc_id in enumerate(pointcloud_ids):
-                id_to_idx[pc_id].append(i)
-
-            sequence_point_cloud = []
-            sequence_point_cloud_idx = -np.ones((len(batch),))
-            for i, ind in enumerate(id_to_idx.values()):
-                sequence_point_cloud_idx[ind] = i
-                sequence_point_cloud.append(batch[ind[0]].sequence_point_cloud)
-            assert (sequence_point_cloud_idx >= 0).all()
-
-            override_fields = {
-                "sequence_point_cloud": sequence_point_cloud,
-                "sequence_point_cloud_idx": sequence_point_cloud_idx.tolist(),
-            }
-            # note that the pre-collate value of sequence_point_cloud_idx is unused
-
-            collated = {}
-            for f in fields(elem):
-                if not f.init:
-                    continue
-
-                list_values = override_fields.get(
-                    f.name, [getattr(d, f.name) for d in batch]
-                )
-                collated[f.name] = (
-                    cls.collate(list_values)
-                    if all(list_value is not None for list_value in list_values)
-                    else None
-                )
-            return cls(**collated)
-
-        elif isinstance(elem, Pointclouds):
-            return join_pointclouds_as_batch(batch)
-
-        elif isinstance(elem, CamerasBase):
-            # TODO: don't store K; enforce working in NDC space
-            return join_cameras_as_batch(batch)
-        else:
-            return torch.utils.data._utils.collate.default_collate(batch)
-
-
-FrameDataSubtype = TypeVar("FrameDataSubtype", bound=FrameData)
-
-
-class FrameDataBuilderBase(ReplaceableBase, Generic[FrameDataSubtype], ABC):
-    """A base class for FrameDataBuilders that build a FrameData object, load and
-    process the binary data (crop and resize). Implementations should parametrize
-    the class with a subtype of FrameData and set frame_data_type class variable to
-    that type. They have to also implement `build` method.
-    """
-
-    # To be initialised to FrameDataSubtype
-    frame_data_type: ClassVar[Type[FrameDataSubtype]]
-
-    @abstractmethod
-    def build(
-        self,
-        frame_annotation: types.FrameAnnotation,
-        sequence_annotation: types.SequenceAnnotation,
-        *,
-        load_blobs: bool = True,
-        **kwargs,
-    ) -> FrameDataSubtype:
-        """An abstract method to build the frame data based on raw frame/sequence
-        annotations, load the binary data and adjust them according to the metadata.
-        """
-        raise NotImplementedError()
-
-
-class GenericFrameDataBuilder(FrameDataBuilderBase[FrameDataSubtype], ABC):
-    """
-    A class to build a FrameData object, load and process the binary data (crop and
-    resize). This is an abstract class for extending to build FrameData subtypes. Most
-    users need to use concrete `FrameDataBuilder` class instead.
-    Beware that modifications of frame data are done in-place.
-
-    Args:
-        dataset_root: The root folder of the dataset; all paths in frame / sequence
-            annotations are defined w.r.t. this root. Has to be set if any of the
-            load_* flabs below is true.
-        load_images: Enable loading the frame RGB data.
-        load_depths: Enable loading the frame depth maps.
-        load_depth_masks: Enable loading the frame depth map masks denoting the
-            depth values used for evaluation (the points consistent across views).
-        load_masks: Enable loading frame foreground masks.
-        load_point_clouds: Enable loading sequence-level point clouds.
-        max_points: Cap on the number of loaded points in the point cloud;
-            if reached, they are randomly sampled without replacement.
-        mask_images: Whether to mask the images with the loaded foreground masks;
-            0 value is used for background.
-        mask_depths: Whether to mask the depth maps with the loaded foreground
-            masks; 0 value is used for background.
-        image_height: The height of the returned images, masks, and depth maps;
-            aspect ratio is preserved during cropping/resizing.
-        image_width: The width of the returned images, masks, and depth maps;
-            aspect ratio is preserved during cropping/resizing.
-        box_crop: Enable cropping of the image around the bounding box inferred
-            from the foreground region of the loaded segmentation mask; masks
-            and depth maps are cropped accordingly; cameras are corrected.
-        box_crop_mask_thr: The threshold used to separate pixels into foreground
-            and background based on the foreground_probability mask; if no value
-            is greater than this threshold, the loader lowers it and repeats.
-        box_crop_context: The amount of additional padding added to each
-            dimension of the cropping bounding box, relative to box size.
-        path_manager: Optionally a PathManager for interpreting paths in a special way.
-    """
-
-    dataset_root: Optional[str] = None
-    load_images: bool = True
-    load_depths: bool = True
-    load_depth_masks: bool = True
-    load_masks: bool = True
-    load_point_clouds: bool = False
-    max_points: int = 0
-    mask_images: bool = False
-    mask_depths: bool = False
-    image_height: Optional[int] = 800
-    image_width: Optional[int] = 800
-    box_crop: bool = True
-    box_crop_mask_thr: float = 0.4
-    box_crop_context: float = 0.3
-    path_manager: Any = None
-
-    def __post_init__(self) -> None:
-        load_any_blob = (
-            self.load_images
-            or self.load_depths
-            or self.load_depth_masks
-            or self.load_masks
-            or self.load_point_clouds
-        )
-        if load_any_blob and self.dataset_root is None:
-            raise ValueError(
-                "dataset_root must be set to load any blob data. "
-                "Make sure it is set in either FrameDataBuilder or Dataset params."
-            )
-
-        if load_any_blob and not self._exists_in_dataset_root(""):
-            raise ValueError(
-                f"dataset_root is passed but {self.dataset_root} does not exist."
-            )
-
-    def build(
-        self,
-        frame_annotation: types.FrameAnnotation,
-        sequence_annotation: types.SequenceAnnotation,
-        *,
-        load_blobs: bool = True,
-        **kwargs,
-    ) -> FrameDataSubtype:
-        """Builds the frame data based on raw frame/sequence annotations, loads the
-        binary data and adjust them according to the metadata. The processing includes:
-            * if box_crop is set, the image/mask/depth are cropped with the bounding
-                box provided or estimated from MaskAnnotation,
-            * if image_height/image_width are set, the image/mask/depth are resized to
-                fit that resolution. Note that the aspect ratio is preserved, and the
-                (possibly cropped) image is pasted into the top-left corner. In the
-                resulting frame_data, mask_crop field corresponds to the mask of the
-                pasted image.
-
-        Args:
-            frame_annotation: frame annotation
-            sequence_annotation: sequence annotation
-            load_blobs: if the function should attempt loading the image, depth map
-                and mask, and foreground mask
-
-        Returns:
-            The constructed FrameData object.
-        """
-
-        point_cloud = sequence_annotation.point_cloud
-
-        frame_data = self.frame_data_type(
-            frame_number=safe_as_tensor(frame_annotation.frame_number, torch.long),
-            frame_timestamp=safe_as_tensor(
-                frame_annotation.frame_timestamp, torch.float
-            ),
-            sequence_name=frame_annotation.sequence_name,
-            sequence_category=sequence_annotation.category,
-            camera_quality_score=safe_as_tensor(
-                sequence_annotation.viewpoint_quality_score, torch.float
-            ),
-            point_cloud_quality_score=safe_as_tensor(
-                point_cloud.quality_score, torch.float
-            )
-            if point_cloud is not None
-            else None,
-        )
-
-        fg_mask_np: Optional[np.ndarray] = None
-        mask_annotation = frame_annotation.mask
-        if mask_annotation is not None:
-            if load_blobs and self.load_masks:
-                fg_mask_np, mask_path = self._load_fg_probability(frame_annotation)
-                frame_data.mask_path = mask_path
-                frame_data.fg_probability = safe_as_tensor(fg_mask_np, torch.float)
-
-            bbox_xywh = mask_annotation.bounding_box_xywh
-            if bbox_xywh is None and fg_mask_np is not None:
-                bbox_xywh = get_bbox_from_mask(fg_mask_np, self.box_crop_mask_thr)
-
-            frame_data.bbox_xywh = safe_as_tensor(bbox_xywh, torch.float)
-
-        if frame_annotation.image is not None:
-            image_size_hw = safe_as_tensor(frame_annotation.image.size, torch.long)
-            frame_data.image_size_hw = image_size_hw  # original image size
-            # image size after crop/resize
-            frame_data.effective_image_size_hw = image_size_hw
-            image_path = None
-            dataset_root = self.dataset_root
-            if frame_annotation.image.path is not None and dataset_root is not None:
-                image_path = os.path.join(dataset_root, frame_annotation.image.path)
-                frame_data.image_path = image_path
-
-            if load_blobs and self.load_images:
-                if image_path is None:
-                    raise ValueError("Image path is required to load images.")
-
-                image_np = load_image(self._local_path(image_path))
-                frame_data.image_rgb = self._postprocess_image(
-                    image_np, frame_annotation.image.size, frame_data.fg_probability
-                )
-
-        if (
-            load_blobs
-            and self.load_depths
-            and frame_annotation.depth is not None
-            and frame_annotation.depth.path is not None
-        ):
-            (
-                frame_data.depth_map,
-                frame_data.depth_path,
-                frame_data.depth_mask,
-            ) = self._load_mask_depth(frame_annotation, fg_mask_np)
-
-        if load_blobs and self.load_point_clouds and point_cloud is not None:
-            pcl_path = self._fix_point_cloud_path(point_cloud.path)
-            frame_data.sequence_point_cloud = load_pointcloud(
-                self._local_path(pcl_path), max_points=self.max_points
-            )
-            frame_data.sequence_point_cloud_path = pcl_path
-
-        if frame_annotation.viewpoint is not None:
-            frame_data.camera = self._get_pytorch3d_camera(frame_annotation)
-
-        if self.box_crop:
-            frame_data.crop_by_metadata_bbox_(self.box_crop_context)
-
-        if self.image_height is not None and self.image_width is not None:
-            new_size = (self.image_height, self.image_width)
-            frame_data.resize_frame_(
-                new_size_hw=torch.tensor(new_size, dtype=torch.long),  # pyre-ignore
-            )
-
-        return frame_data
-
-    def _load_fg_probability(
-        self, entry: types.FrameAnnotation
-    ) -> Tuple[np.ndarray, str]:
-        assert self.dataset_root is not None and entry.mask is not None
-        full_path = os.path.join(self.dataset_root, entry.mask.path)
-        fg_probability = load_mask(self._local_path(full_path))
-        if fg_probability.shape[-2:] != entry.image.size:
-            raise ValueError(
-                f"bad mask size: {fg_probability.shape[-2:]} vs {entry.image.size}!"
-            )
-
-        return fg_probability, full_path
-
-    def _postprocess_image(
-        self,
-        image_np: np.ndarray,
-        image_size: Tuple[int, int],
-        fg_probability: Optional[torch.Tensor],
-    ) -> torch.Tensor:
-        image_rgb = safe_as_tensor(image_np, torch.float)
-
-        if image_rgb.shape[-2:] != image_size:
-            raise ValueError(f"bad image size: {image_rgb.shape[-2:]} vs {image_size}!")
-
-        if self.mask_images:
-            assert fg_probability is not None
-            image_rgb *= fg_probability
-
-        return image_rgb
-
-    def _load_mask_depth(
-        self,
-        entry: types.FrameAnnotation,
-        fg_mask: Optional[np.ndarray],
-    ) -> Tuple[torch.Tensor, str, torch.Tensor]:
-        entry_depth = entry.depth
-        dataset_root = self.dataset_root
-        assert dataset_root is not None
-        assert entry_depth is not None and entry_depth.path is not None
-        path = os.path.join(dataset_root, entry_depth.path)
-        depth_map = load_depth(self._local_path(path), entry_depth.scale_adjustment)
-
-        if self.mask_depths:
-            assert fg_mask is not None
-            depth_map *= fg_mask
-
-        mask_path = entry_depth.mask_path
-        if self.load_depth_masks and mask_path is not None:
-            mask_path = os.path.join(dataset_root, mask_path)
-            depth_mask = load_depth_mask(self._local_path(mask_path))
-        else:
-            depth_mask = (depth_map > 0.0).astype(np.float32)
-
-        return torch.tensor(depth_map), path, torch.tensor(depth_mask)
-
-    def _get_pytorch3d_camera(
-        self,
-        entry: types.FrameAnnotation,
-    ) -> PerspectiveCameras:
-        entry_viewpoint = entry.viewpoint
-        assert entry_viewpoint is not None
-        # principal point and focal length
-        principal_point = torch.tensor(
-            entry_viewpoint.principal_point, dtype=torch.float
-        )
-        focal_length = torch.tensor(entry_viewpoint.focal_length, dtype=torch.float)
-
-        format = entry_viewpoint.intrinsics_format
-        if entry_viewpoint.intrinsics_format == "ndc_norm_image_bounds":
-            # legacy PyTorch3D NDC format
-            # convert to pixels unequally and convert to ndc equally
-            image_size_as_list = list(reversed(entry.image.size))
-            image_size_wh = torch.tensor(image_size_as_list, dtype=torch.float)
-            per_axis_scale = image_size_wh / image_size_wh.min()
-            focal_length = focal_length * per_axis_scale
-            principal_point = principal_point * per_axis_scale
-        elif entry_viewpoint.intrinsics_format != "ndc_isotropic":
-            raise ValueError(f"Unknown intrinsics format: {format}")
-
-        return PerspectiveCameras(
-            focal_length=focal_length[None],
-            principal_point=principal_point[None],
-            R=torch.tensor(entry_viewpoint.R, dtype=torch.float)[None],
-            T=torch.tensor(entry_viewpoint.T, dtype=torch.float)[None],
-        )
-
-    def _fix_point_cloud_path(self, path: str) -> str:
-        """
-        Fix up a point cloud path from the dataset.
-        Some files in Co3Dv2 have an accidental absolute path stored.
-        """
-        unwanted_prefix = (
-            "/large_experiments/p3/replay/datasets/co3d/co3d45k_220512/export_v23/"
-        )
-        if path.startswith(unwanted_prefix):
-            path = path[len(unwanted_prefix) :]
-        assert self.dataset_root is not None
-        return os.path.join(self.dataset_root, path)
-
-    def _local_path(self, path: str) -> str:
-        if self.path_manager is None:
-            return path
-        return self.path_manager.get_local_path(path)
-
-    def _exists_in_dataset_root(self, relpath) -> bool:
-        if not self.dataset_root:
-            return False
-
-        full_path = os.path.join(self.dataset_root, relpath)
-        if self.path_manager is None:
-            return os.path.exists(full_path)
-        else:
-            return self.path_manager.exists(full_path)
-
-
-@registry.register
-class FrameDataBuilder(GenericWorkaround, GenericFrameDataBuilder[FrameData]):
-    """
-    A concrete class to build a FrameData object, load and process the binary data (crop
-    and resize). Beware that modifications of frame data are done in-place. Please see
-    the documentation for `GenericFrameDataBuilder` for the description of parameters
-    and methods.
-    """
-
-    frame_data_type: ClassVar[Type[FrameData]] = FrameData

pytorch3d/implicitron/dataset/json_index_dataset.py

@@ -15,6 +15,7 @@ import random
 import warnings
 from collections import defaultdict
 from itertools import islice
+from pathlib import Path
 from typing import (
     Any,
     ClassVar,
@@ -29,16 +30,20 @@ from typing import (
     Union,
 )
 
-from pytorch3d.implicitron.dataset import types
-from pytorch3d.implicitron.dataset.dataset_base import DatasetBase
-from pytorch3d.implicitron.dataset.frame_data import FrameData, FrameDataBuilder
-from pytorch3d.implicitron.dataset.utils import is_known_frame_scalar
+import numpy as np
+import torch
+from PIL import Image
 from pytorch3d.implicitron.tools.config import registry, ReplaceableBase
+from pytorch3d.io import IO
 from pytorch3d.renderer.camera_utils import join_cameras_as_batch
-from pytorch3d.renderer.cameras import CamerasBase
-
+from pytorch3d.renderer.cameras import CamerasBase, PerspectiveCameras
+from pytorch3d.structures.pointclouds import Pointclouds
 from tqdm import tqdm
 
+from . import types
+from .dataset_base import DatasetBase, FrameData
+from .utils import is_known_frame_scalar
+
 
 logger = logging.getLogger(__name__)
 
@@ -60,7 +65,7 @@ class JsonIndexDataset(DatasetBase, ReplaceableBase):
     A dataset with annotations in json files like the Common Objects in 3D
     (CO3D) dataset.
 
-    Metadata-related args::
+    Args:
         frame_annotations_file: A zipped json file containing metadata of the
             frames in the dataset, serialized List[types.FrameAnnotation].
         sequence_annotations_file: A zipped json file containing metadata of the
@@ -78,24 +83,6 @@ class JsonIndexDataset(DatasetBase, ReplaceableBase):
         pick_sequence: A list of sequence names to restrict the dataset to.
         exclude_sequence: A list of the names of the sequences to exclude.
         limit_category_to: Restrict the dataset to the given list of categories.
-        remove_empty_masks: Removes the frames with no active foreground pixels
-            in the segmentation mask after thresholding (see box_crop_mask_thr).
-        n_frames_per_sequence: If > 0, randomly samples #n_frames_per_sequence
-            frames in each sequences uniformly without replacement if it has
-            more frames than that; applied before other frame-level filters.
-        seed: The seed of the random generator sampling #n_frames_per_sequence
-            random frames per sequence.
-        sort_frames: Enable frame annotations sorting to group frames from the
-            same sequences together and order them by timestamps
-        eval_batches: A list of batches that form the evaluation set;
-            list of batch-sized lists of indices corresponding to __getitem__
-            of this class, thus it can be used directly as a batch sampler.
-        eval_batch_index:
-            ( Optional[List[List[Union[Tuple[str, int, str], Tuple[str, int]]]] )
-            A list of batches of frames described as (sequence_name, frame_idx)
-            that can form the evaluation set, `eval_batches` will be set from this.
-
-    Blob-loading parameters:
         dataset_root: The root folder of the dataset; all the paths in jsons are
             specified relative to this root (but not json paths themselves).
         load_images: Enable loading the frame RGB data.
@@ -122,6 +109,23 @@ class JsonIndexDataset(DatasetBase, ReplaceableBase):
             is greater than this threshold, the loader lowers it and repeats.
         box_crop_context: The amount of additional padding added to each
             dimension of the cropping bounding box, relative to box size.
+        remove_empty_masks: Removes the frames with no active foreground pixels
+            in the segmentation mask after thresholding (see box_crop_mask_thr).
+        n_frames_per_sequence: If > 0, randomly samples #n_frames_per_sequence
+            frames in each sequences uniformly without replacement if it has
+            more frames than that; applied before other frame-level filters.
+        seed: The seed of the random generator sampling #n_frames_per_sequence
+            random frames per sequence.
+        sort_frames: Enable frame annotations sorting to group frames from the
+            same sequences together and order them by timestamps
+        eval_batches: A list of batches that form the evaluation set;
+            list of batch-sized lists of indices corresponding to __getitem__
+            of this class, thus it can be used directly as a batch sampler.
+        eval_batch_index:
+            ( Optional[List[List[Union[Tuple[str, int, str], Tuple[str, int]]]] )
+            A list of batches of frames described as (sequence_name, frame_idx)
+            that can form the evaluation set, `eval_batches` will be set from this.
+
     """
 
     frame_annotations_type: ClassVar[
@@ -158,14 +162,12 @@ class JsonIndexDataset(DatasetBase, ReplaceableBase):
     sort_frames: bool = False
     eval_batches: Any = None
     eval_batch_index: Any = None
-    # initialised in __post_init__
-    # commented because of OmegaConf (for tests to pass)
-    # _frame_data_builder: FrameDataBuilder = field(init=False)
     # frame_annots: List[FrameAnnotsEntry] = field(init=False)
     # seq_annots: Dict[str, types.SequenceAnnotation] = field(init=False)
-    # _seq_to_idx: Dict[str, List[int]] = field(init=False)
 
     def __post_init__(self) -> None:
+        # pyre-fixme[16]: `JsonIndexDataset` has no attribute `subset_to_image_path`.
+        self.subset_to_image_path = None
         self._load_frames()
         self._load_sequences()
         if self.sort_frames:
@@ -173,28 +175,9 @@ class JsonIndexDataset(DatasetBase, ReplaceableBase):
         self._load_subset_lists()
         self._filter_db()  # also computes sequence indices
         self._extract_and_set_eval_batches()
-
-        # pyre-ignore
-        self._frame_data_builder = FrameDataBuilder(
-            dataset_root=self.dataset_root,
-            load_images=self.load_images,
-            load_depths=self.load_depths,
-            load_depth_masks=self.load_depth_masks,
-            load_masks=self.load_masks,
-            load_point_clouds=self.load_point_clouds,
-            max_points=self.max_points,
-            mask_images=self.mask_images,
-            mask_depths=self.mask_depths,
-            image_height=self.image_height,
-            image_width=self.image_width,
-            box_crop=self.box_crop,
-            box_crop_mask_thr=self.box_crop_mask_thr,
-            box_crop_context=self.box_crop_context,
-            path_manager=self.path_manager,
-        )
         logger.info(str(self))
 
-    def _extract_and_set_eval_batches(self) -> None:
+    def _extract_and_set_eval_batches(self):
         """
         Sets eval_batches based on input eval_batch_index.
         """
@@ -224,13 +207,13 @@ class JsonIndexDataset(DatasetBase, ReplaceableBase):
             # https://gist.github.com/treyhunner/f35292e676efa0be1728
             functools.reduce(
                 lambda a, b: {**a, **b},
-                # pyre-ignore[16]
-                [d.seq_annots for d in other_datasets],
+                [d.seq_annots for d in other_datasets],  # pyre-ignore[16]
             )
         )
         all_eval_batches = [
             self.eval_batches,
-            *[d.eval_batches for d in other_datasets],  # pyre-ignore[16]
+            # pyre-ignore
+            *[d.eval_batches for d in other_datasets],
         ]
         if not (
             all(ba is None for ba in all_eval_batches)
@@ -268,7 +251,7 @@ class JsonIndexDataset(DatasetBase, ReplaceableBase):
         allow_missing_indices: bool = False,
         remove_missing_indices: bool = False,
         suppress_missing_index_warning: bool = True,
-    ) -> Union[List[List[Optional[int]]], List[List[int]]]:
+    ) -> List[List[Union[Optional[int], int]]]:
         """
         Obtain indices into the dataset object given a list of frame ids.
 
@@ -340,7 +323,9 @@ class JsonIndexDataset(DatasetBase, ReplaceableBase):
             valid_dataset_idx = [
                 [b for b in batch if b is not None] for batch in dataset_idx
             ]
-            return [batch for batch in valid_dataset_idx if len(batch) > 0]
+            return [  # pyre-ignore[7]
+                batch for batch in valid_dataset_idx if len(batch) > 0
+            ]
 
         return dataset_idx
 
@@ -432,18 +417,255 @@ class JsonIndexDataset(DatasetBase, ReplaceableBase):
             raise IndexError(f"index {index} out of range {len(self.frame_annots)}")
 
         entry = self.frame_annots[index]["frame_annotation"]
-
-        # pyre-ignore
-        frame_data = self._frame_data_builder.build(
-            entry,
-            # pyre-ignore
-            self.seq_annots[entry.sequence_name],
+        # pyre-ignore[16]
+        point_cloud = self.seq_annots[entry.sequence_name].point_cloud
+        frame_data = FrameData(
+            frame_number=_safe_as_tensor(entry.frame_number, torch.long),
+            frame_timestamp=_safe_as_tensor(entry.frame_timestamp, torch.float),
+            sequence_name=entry.sequence_name,
+            sequence_category=self.seq_annots[entry.sequence_name].category,
+            camera_quality_score=_safe_as_tensor(
+                self.seq_annots[entry.sequence_name].viewpoint_quality_score,
+                torch.float,
+            ),
+            point_cloud_quality_score=_safe_as_tensor(
+                point_cloud.quality_score, torch.float
+            )
+            if point_cloud is not None
+            else None,
         )
-        # Optional field
+
+        # The rest of the fields are optional
         frame_data.frame_type = self._get_frame_type(self.frame_annots[index])
 
+        (
+            frame_data.fg_probability,
+            frame_data.mask_path,
+            frame_data.bbox_xywh,
+            clamp_bbox_xyxy,
+            frame_data.crop_bbox_xywh,
+        ) = self._load_crop_fg_probability(entry)
+
+        scale = 1.0
+        if self.load_images and entry.image is not None:
+            # original image size
+            frame_data.image_size_hw = _safe_as_tensor(entry.image.size, torch.long)
+
+            (
+                frame_data.image_rgb,
+                frame_data.image_path,
+                frame_data.mask_crop,
+                scale,
+            ) = self._load_crop_images(
+                entry, frame_data.fg_probability, clamp_bbox_xyxy
+            )
+
+        if self.load_depths and entry.depth is not None:
+            (
+                frame_data.depth_map,
+                frame_data.depth_path,
+                frame_data.depth_mask,
+            ) = self._load_mask_depth(entry, clamp_bbox_xyxy, frame_data.fg_probability)
+
+        if entry.viewpoint is not None:
+            frame_data.camera = self._get_pytorch3d_camera(
+                entry,
+                scale,
+                clamp_bbox_xyxy,
+            )
+
+        if self.load_point_clouds and point_cloud is not None:
+            pcl_path = self._fix_point_cloud_path(point_cloud.path)
+            frame_data.sequence_point_cloud = _load_pointcloud(
+                self._local_path(pcl_path), max_points=self.max_points
+            )
+            frame_data.sequence_point_cloud_path = pcl_path
+
         return frame_data
 
+    def _fix_point_cloud_path(self, path: str) -> str:
+        """
+        Fix up a point cloud path from the dataset.
+        Some files in Co3Dv2 have an accidental absolute path stored.
+        """
+        unwanted_prefix = (
+            "/large_experiments/p3/replay/datasets/co3d/co3d45k_220512/export_v23/"
+        )
+        if path.startswith(unwanted_prefix):
+            path = path[len(unwanted_prefix) :]
+        return os.path.join(self.dataset_root, path)
+
+    def _load_crop_fg_probability(
+        self, entry: types.FrameAnnotation
+    ) -> Tuple[
+        Optional[torch.Tensor],
+        Optional[str],
+        Optional[torch.Tensor],
+        Optional[torch.Tensor],
+        Optional[torch.Tensor],
+    ]:
+        fg_probability = None
+        full_path = None
+        bbox_xywh = None
+        clamp_bbox_xyxy = None
+        crop_box_xywh = None
+
+        if (self.load_masks or self.box_crop) and entry.mask is not None:
+            full_path = os.path.join(self.dataset_root, entry.mask.path)
+            mask = _load_mask(self._local_path(full_path))
+
+            if mask.shape[-2:] != entry.image.size:
+                raise ValueError(
+                    f"bad mask size: {mask.shape[-2:]} vs {entry.image.size}!"
+                )
+
+            bbox_xywh = torch.tensor(_get_bbox_from_mask(mask, self.box_crop_mask_thr))
+
+            if self.box_crop:
+                clamp_bbox_xyxy = _clamp_box_to_image_bounds_and_round(
+                    _get_clamp_bbox(
+                        bbox_xywh,
+                        image_path=entry.image.path,
+                        box_crop_context=self.box_crop_context,
+                    ),
+                    image_size_hw=tuple(mask.shape[-2:]),
+                )
+                crop_box_xywh = _bbox_xyxy_to_xywh(clamp_bbox_xyxy)
+
+                mask = _crop_around_box(mask, clamp_bbox_xyxy, full_path)
+
+            fg_probability, _, _ = self._resize_image(mask, mode="nearest")
+
+        return fg_probability, full_path, bbox_xywh, clamp_bbox_xyxy, crop_box_xywh
+
+    def _load_crop_images(
+        self,
+        entry: types.FrameAnnotation,
+        fg_probability: Optional[torch.Tensor],
+        clamp_bbox_xyxy: Optional[torch.Tensor],
+    ) -> Tuple[torch.Tensor, str, torch.Tensor, float]:
+        assert self.dataset_root is not None and entry.image is not None
+        path = os.path.join(self.dataset_root, entry.image.path)
+        image_rgb = _load_image(self._local_path(path))
+
+        if image_rgb.shape[-2:] != entry.image.size:
+            raise ValueError(
+                f"bad image size: {image_rgb.shape[-2:]} vs {entry.image.size}!"
+            )
+
+        if self.box_crop:
+            assert clamp_bbox_xyxy is not None
+            image_rgb = _crop_around_box(image_rgb, clamp_bbox_xyxy, path)
+
+        image_rgb, scale, mask_crop = self._resize_image(image_rgb)
+
+        if self.mask_images:
+            assert fg_probability is not None
+            image_rgb *= fg_probability
+
+        return image_rgb, path, mask_crop, scale
+
+    def _load_mask_depth(
+        self,
+        entry: types.FrameAnnotation,
+        clamp_bbox_xyxy: Optional[torch.Tensor],
+        fg_probability: Optional[torch.Tensor],
+    ) -> Tuple[torch.Tensor, str, torch.Tensor]:
+        entry_depth = entry.depth
+        assert entry_depth is not None
+        path = os.path.join(self.dataset_root, entry_depth.path)
+        depth_map = _load_depth(self._local_path(path), entry_depth.scale_adjustment)
+
+        if self.box_crop:
+            assert clamp_bbox_xyxy is not None
+            depth_bbox_xyxy = _rescale_bbox(
+                clamp_bbox_xyxy, entry.image.size, depth_map.shape[-2:]
+            )
+            depth_map = _crop_around_box(depth_map, depth_bbox_xyxy, path)
+
+        depth_map, _, _ = self._resize_image(depth_map, mode="nearest")
+
+        if self.mask_depths:
+            assert fg_probability is not None
+            depth_map *= fg_probability
+
+        if self.load_depth_masks:
+            assert entry_depth.mask_path is not None
+            mask_path = os.path.join(self.dataset_root, entry_depth.mask_path)
+            depth_mask = _load_depth_mask(self._local_path(mask_path))
+
+            if self.box_crop:
+                assert clamp_bbox_xyxy is not None
+                depth_mask_bbox_xyxy = _rescale_bbox(
+                    clamp_bbox_xyxy, entry.image.size, depth_mask.shape[-2:]
+                )
+                depth_mask = _crop_around_box(
+                    depth_mask, depth_mask_bbox_xyxy, mask_path
+                )
+
+            depth_mask, _, _ = self._resize_image(depth_mask, mode="nearest")
+        else:
+            depth_mask = torch.ones_like(depth_map)
+
+        return depth_map, path, depth_mask
+
+    def _get_pytorch3d_camera(
+        self,
+        entry: types.FrameAnnotation,
+        scale: float,
+        clamp_bbox_xyxy: Optional[torch.Tensor],
+    ) -> PerspectiveCameras:
+        entry_viewpoint = entry.viewpoint
+        assert entry_viewpoint is not None
+        # principal point and focal length
+        principal_point = torch.tensor(
+            entry_viewpoint.principal_point, dtype=torch.float
+        )
+        focal_length = torch.tensor(entry_viewpoint.focal_length, dtype=torch.float)
+
+        half_image_size_wh_orig = (
+            torch.tensor(list(reversed(entry.image.size)), dtype=torch.float) / 2.0
+        )
+
+        # first, we convert from the dataset's NDC convention to pixels
+        format = entry_viewpoint.intrinsics_format
+        if format.lower() == "ndc_norm_image_bounds":
+            # this is e.g. currently used in CO3D for storing intrinsics
+            rescale = half_image_size_wh_orig
+        elif format.lower() == "ndc_isotropic":
+            rescale = half_image_size_wh_orig.min()
+        else:
+            raise ValueError(f"Unknown intrinsics format: {format}")
+
+        # principal point and focal length in pixels
+        principal_point_px = half_image_size_wh_orig - principal_point * rescale
+        focal_length_px = focal_length * rescale
+        if self.box_crop:
+            assert clamp_bbox_xyxy is not None
+            principal_point_px -= clamp_bbox_xyxy[:2]
+
+        # now, convert from pixels to PyTorch3D v0.5+ NDC convention
+        if self.image_height is None or self.image_width is None:
+            out_size = list(reversed(entry.image.size))
+        else:
+            out_size = [self.image_width, self.image_height]
+
+        half_image_size_output = torch.tensor(out_size, dtype=torch.float) / 2.0
+        half_min_image_size_output = half_image_size_output.min()
+
+        # rescaled principal point and focal length in ndc
+        principal_point = (
+            half_image_size_output - principal_point_px * scale
+        ) / half_min_image_size_output
+        focal_length = focal_length_px * scale / half_min_image_size_output
+
+        return PerspectiveCameras(
+            focal_length=focal_length[None],
+            principal_point=principal_point[None],
+            R=torch.tensor(entry_viewpoint.R, dtype=torch.float)[None],
+            T=torch.tensor(entry_viewpoint.T, dtype=torch.float)[None],
+        )
+
     def _load_frames(self) -> None:
         logger.info(f"Loading Co3D frames from {self.frame_annotations_file}.")
         local_file = self._local_path(self.frame_annotations_file)
@@ -631,23 +853,46 @@ class JsonIndexDataset(DatasetBase, ReplaceableBase):
         # pyre-ignore[16]
         self._seq_to_idx = seq_to_idx
 
+    def _resize_image(
+        self, image, mode="bilinear"
+    ) -> Tuple[torch.Tensor, float, torch.Tensor]:
+        image_height, image_width = self.image_height, self.image_width
+        if image_height is None or image_width is None:
+            # skip the resizing
+            imre_ = torch.from_numpy(image)
+            return imre_, 1.0, torch.ones_like(imre_[:1])
+        # takes numpy array, returns pytorch tensor
+        minscale = min(
+            image_height / image.shape[-2],
+            image_width / image.shape[-1],
+        )
+        imre = torch.nn.functional.interpolate(
+            torch.from_numpy(image)[None],
+            scale_factor=minscale,
+            mode=mode,
+            align_corners=False if mode == "bilinear" else None,
+            recompute_scale_factor=True,
+        )[0]
+        # pyre-fixme[19]: Expected 1 positional argument.
+        imre_ = torch.zeros(image.shape[0], self.image_height, self.image_width)
+        imre_[:, 0 : imre.shape[1], 0 : imre.shape[2]] = imre
+        # pyre-fixme[6]: For 2nd param expected `int` but got `Optional[int]`.
+        # pyre-fixme[6]: For 3rd param expected `int` but got `Optional[int]`.
+        mask = torch.zeros(1, self.image_height, self.image_width)
+        mask[:, 0 : imre.shape[1], 0 : imre.shape[2]] = 1.0
+        return imre_, minscale, mask
+
     def _local_path(self, path: str) -> str:
         if self.path_manager is None:
             return path
         return self.path_manager.get_local_path(path)
 
     def get_frame_numbers_and_timestamps(
-        self, idxs: Sequence[int], subset_filter: Optional[Sequence[str]] = None
+        self, idxs: Sequence[int]
     ) -> List[Tuple[int, float]]:
         out: List[Tuple[int, float]] = []
         for idx in idxs:
-            if (
-                subset_filter is not None
-                # pyre-fixme[16]: `JsonIndexDataset` has no attribute `frame_annots`.
-                and self.frame_annots[idx]["subset"] not in subset_filter
-            ):
-                continue
-
+            # pyre-ignore[16]
             frame_annotation = self.frame_annots[idx]["frame_annotation"]
             out.append(
                 (frame_annotation.frame_number, frame_annotation.frame_timestamp)
@@ -667,3 +912,169 @@ class JsonIndexDataset(DatasetBase, ReplaceableBase):
 
 def _seq_name_to_seed(seq_name) -> int:
     return int(hashlib.sha1(seq_name.encode("utf-8")).hexdigest(), 16)
+
+
+def _load_image(path) -> np.ndarray:
+    with Image.open(path) as pil_im:
+        im = np.array(pil_im.convert("RGB"))
+    im = im.transpose((2, 0, 1))
+    im = im.astype(np.float32) / 255.0
+    return im
+
+
+def _load_16big_png_depth(depth_png) -> np.ndarray:
+    with Image.open(depth_png) as depth_pil:
+        # the image is stored with 16-bit depth but PIL reads it as I (32 bit).
+        # we cast it to uint16, then reinterpret as float16, then cast to float32
+        depth = (
+            np.frombuffer(np.array(depth_pil, dtype=np.uint16), dtype=np.float16)
+            .astype(np.float32)
+            .reshape((depth_pil.size[1], depth_pil.size[0]))
+        )
+    return depth
+
+
+def _load_1bit_png_mask(file: str) -> np.ndarray:
+    with Image.open(file) as pil_im:
+        mask = (np.array(pil_im.convert("L")) > 0.0).astype(np.float32)
+    return mask
+
+
+def _load_depth_mask(path: str) -> np.ndarray:
+    if not path.lower().endswith(".png"):
+        raise ValueError('unsupported depth mask file name "%s"' % path)
+    m = _load_1bit_png_mask(path)
+    return m[None]  # fake feature channel
+
+
+def _load_depth(path, scale_adjustment) -> np.ndarray:
+    if not path.lower().endswith(".png"):
+        raise ValueError('unsupported depth file name "%s"' % path)
+
+    d = _load_16big_png_depth(path) * scale_adjustment
+    d[~np.isfinite(d)] = 0.0
+    return d[None]  # fake feature channel
+
+
+def _load_mask(path) -> np.ndarray:
+    with Image.open(path) as pil_im:
+        mask = np.array(pil_im)
+    mask = mask.astype(np.float32) / 255.0
+    return mask[None]  # fake feature channel
+
+
+def _get_1d_bounds(arr) -> Tuple[int, int]:
+    nz = np.flatnonzero(arr)
+    return nz[0], nz[-1] + 1
+
+
+def _get_bbox_from_mask(
+    mask, thr, decrease_quant: float = 0.05
+) -> Tuple[int, int, int, int]:
+    # bbox in xywh
+    masks_for_box = np.zeros_like(mask)
+    while masks_for_box.sum() <= 1.0:
+        masks_for_box = (mask > thr).astype(np.float32)
+        thr -= decrease_quant
+    if thr <= 0.0:
+        warnings.warn(f"Empty masks_for_bbox (thr={thr}) => using full image.")
+
+    x0, x1 = _get_1d_bounds(masks_for_box.sum(axis=-2))
+    y0, y1 = _get_1d_bounds(masks_for_box.sum(axis=-1))
+
+    return x0, y0, x1 - x0, y1 - y0
+
+
+def _get_clamp_bbox(
+    bbox: torch.Tensor,
+    box_crop_context: float = 0.0,
+    image_path: str = "",
+) -> torch.Tensor:
+    # box_crop_context: rate of expansion for bbox
+    # returns possibly expanded bbox xyxy as float
+
+    bbox = bbox.clone()  # do not edit bbox in place
+
+    # increase box size
+    if box_crop_context > 0.0:
+        c = box_crop_context
+        bbox = bbox.float()
+        bbox[0] -= bbox[2] * c / 2
+        bbox[1] -= bbox[3] * c / 2
+        bbox[2] += bbox[2] * c
+        bbox[3] += bbox[3] * c
+
+    if (bbox[2:] <= 1.0).any():
+        raise ValueError(
+            f"squashed image {image_path}!! The bounding box contains no pixels."
+        )
+
+    bbox[2:] = torch.clamp(bbox[2:], 2)  # set min height, width to 2 along both axes
+    bbox_xyxy = _bbox_xywh_to_xyxy(bbox, clamp_size=2)
+
+    return bbox_xyxy
+
+
+def _crop_around_box(tensor, bbox, impath: str = ""):
+    # bbox is xyxy, where the upper bound is corrected with +1
+    bbox = _clamp_box_to_image_bounds_and_round(
+        bbox,
+        image_size_hw=tensor.shape[-2:],
+    )
+    tensor = tensor[..., bbox[1] : bbox[3], bbox[0] : bbox[2]]
+    assert all(c > 0 for c in tensor.shape), f"squashed image {impath}"
+    return tensor
+
+
+def _clamp_box_to_image_bounds_and_round(
+    bbox_xyxy: torch.Tensor,
+    image_size_hw: Tuple[int, int],
+) -> torch.LongTensor:
+    bbox_xyxy = bbox_xyxy.clone()
+    bbox_xyxy[[0, 2]] = torch.clamp(bbox_xyxy[[0, 2]], 0, image_size_hw[-1])
+    bbox_xyxy[[1, 3]] = torch.clamp(bbox_xyxy[[1, 3]], 0, image_size_hw[-2])
+    if not isinstance(bbox_xyxy, torch.LongTensor):
+        bbox_xyxy = bbox_xyxy.round().long()
+    return bbox_xyxy  # pyre-ignore [7]
+
+
+def _rescale_bbox(bbox: torch.Tensor, orig_res, new_res) -> torch.Tensor:
+    assert bbox is not None
+    assert np.prod(orig_res) > 1e-8
+    # average ratio of dimensions
+    rel_size = (new_res[0] / orig_res[0] + new_res[1] / orig_res[1]) / 2.0
+    return bbox * rel_size
+
+
+def _bbox_xyxy_to_xywh(xyxy: torch.Tensor) -> torch.Tensor:
+    wh = xyxy[2:] - xyxy[:2]
+    xywh = torch.cat([xyxy[:2], wh])
+    return xywh
+
+
+def _bbox_xywh_to_xyxy(
+    xywh: torch.Tensor, clamp_size: Optional[int] = None
+) -> torch.Tensor:
+    xyxy = xywh.clone()
+    if clamp_size is not None:
+        xyxy[2:] = torch.clamp(xyxy[2:], clamp_size)
+    xyxy[2:] += xyxy[:2]
+    return xyxy
+
+
+def _safe_as_tensor(data, dtype):
+    if data is None:
+        return None
+    return torch.tensor(data, dtype=dtype)
+
+
+# NOTE this cache is per-worker; they are implemented as processes.
+# each batch is loaded and collated by a single worker;
+# since sequences tend to co-occur within batches, this is useful.
+@functools.lru_cache(maxsize=256)
+def _load_pointcloud(pcl_path: Union[str, Path], max_points: int = 0) -> Pointclouds:
+    pcl = IO().load_pointcloud(pcl_path)
+    if max_points > 0:
+        pcl = pcl.subsample(max_points)
+
+    return pcl

pytorch3d/implicitron/dataset/load_llff.py

@@ -34,7 +34,11 @@ def _minify(basedir, path_manager, factors=(), resolutions=()):
 
     imgdir = os.path.join(basedir, "images")
     imgs = [os.path.join(imgdir, f) for f in sorted(_ls(path_manager, imgdir))]
-    imgs = [f for f in imgs if f.endswith("JPG", "jpg", "png", "jpeg", "PNG")]
+    imgs = [
+        f
+        for f in imgs
+        if any([f.endswith(ex) for ex in ["JPG", "jpg", "png", "jpeg", "PNG"]])
+    ]
     imgdir_orig = imgdir
 
     wd = os.getcwd()

pytorch3d/implicitron/dataset/orm_types.py

@@ -1,189 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the BSD-style license found in the
-# LICENSE file in the root directory of this source tree.
-
-# This functionality requires SQLAlchemy 2.0 or later.
-
-import math
-import struct
-from typing import Optional, Tuple
-
-import numpy as np
-
-from pytorch3d.implicitron.dataset.types import (
-    DepthAnnotation,
-    ImageAnnotation,
-    MaskAnnotation,
-    PointCloudAnnotation,
-    VideoAnnotation,
-    ViewpointAnnotation,
-)
-
-from sqlalchemy import LargeBinary
-from sqlalchemy.orm import (
-    composite,
-    DeclarativeBase,
-    Mapped,
-    mapped_column,
-    MappedAsDataclass,
-)
-from sqlalchemy.types import TypeDecorator
-
-
-# these produce policies to serialize structured types to blobs
-def ArrayTypeFactory(shape=None):
-    if shape is None:
-
-        class VariableShapeNumpyArrayType(TypeDecorator):
-            impl = LargeBinary
-
-            def process_bind_param(self, value, dialect):
-                if value is None:
-                    return None
-
-                ndim_bytes = np.int32(value.ndim).tobytes()
-                shape_bytes = np.array(value.shape, dtype=np.int64).tobytes()
-                value_bytes = value.astype(np.float32).tobytes()
-                return ndim_bytes + shape_bytes + value_bytes
-
-            def process_result_value(self, value, dialect):
-                if value is None:
-                    return None
-
-                ndim = np.frombuffer(value[:4], dtype=np.int32)[0]
-                value_start = 4 + 8 * ndim
-                shape = np.frombuffer(value[4:value_start], dtype=np.int64)
-                assert shape.shape == (ndim,)
-                return np.frombuffer(value[value_start:], dtype=np.float32).reshape(
-                    shape
-                )
-
-        return VariableShapeNumpyArrayType
-
-    class NumpyArrayType(TypeDecorator):
-        impl = LargeBinary
-
-        def process_bind_param(self, value, dialect):
-            if value is not None:
-                if value.shape != shape:
-                    raise ValueError(f"Passed an array of wrong shape: {value.shape}")
-                return value.astype(np.float32).tobytes()
-            return None
-
-        def process_result_value(self, value, dialect):
-            if value is not None:
-                return np.frombuffer(value, dtype=np.float32).reshape(shape)
-            return None
-
-    return NumpyArrayType
-
-
-def TupleTypeFactory(dtype=float, shape: Tuple[int, ...] = (2,)):
-    format_symbol = {
-        float: "f",  # float32
-        int: "i",  # int32
-    }[dtype]
-
-    class TupleType(TypeDecorator):
-        impl = LargeBinary
-        _format = format_symbol * math.prod(shape)
-
-        def process_bind_param(self, value, _):
-            if value is None:
-                return None
-
-            if len(shape) > 1:
-                value = np.array(value, dtype=dtype).reshape(-1)
-
-            return struct.pack(TupleType._format, *value)
-
-        def process_result_value(self, value, _):
-            if value is None:
-                return None
-
-            loaded = struct.unpack(TupleType._format, value)
-            if len(shape) > 1:
-                loaded = _rec_totuple(
-                    np.array(loaded, dtype=dtype).reshape(shape).tolist()
-                )
-
-            return loaded
-
-    return TupleType
-
-
-def _rec_totuple(t):
-    if isinstance(t, list):
-        return tuple(_rec_totuple(x) for x in t)
-
-    return t
-
-
-class Base(MappedAsDataclass, DeclarativeBase):
-    """subclasses will be converted to dataclasses"""
-
-
-class SqlFrameAnnotation(Base):
-    __tablename__ = "frame_annots"
-
-    sequence_name: Mapped[str] = mapped_column(primary_key=True)
-    frame_number: Mapped[int] = mapped_column(primary_key=True)
-    frame_timestamp: Mapped[float] = mapped_column(index=True)
-
-    image: Mapped[ImageAnnotation] = composite(
-        mapped_column("_image_path"),
-        mapped_column("_image_size", TupleTypeFactory(int)),
-    )
-
-    depth: Mapped[DepthAnnotation] = composite(
-        mapped_column("_depth_path", nullable=True),
-        mapped_column("_depth_scale_adjustment", nullable=True),
-        mapped_column("_depth_mask_path", nullable=True),
-    )
-
-    mask: Mapped[MaskAnnotation] = composite(
-        mapped_column("_mask_path", nullable=True),
-        mapped_column("_mask_mass", index=True, nullable=True),
-        mapped_column(
-            "_mask_bounding_box_xywh",
-            TupleTypeFactory(float, shape=(4,)),
-            nullable=True,
-        ),
-    )
-
-    viewpoint: Mapped[ViewpointAnnotation] = composite(
-        mapped_column(
-            "_viewpoint_R", TupleTypeFactory(float, shape=(3, 3)), nullable=True
-        ),
-        mapped_column(
-            "_viewpoint_T", TupleTypeFactory(float, shape=(3,)), nullable=True
-        ),
-        mapped_column(
-            "_viewpoint_focal_length", TupleTypeFactory(float), nullable=True
-        ),
-        mapped_column(
-            "_viewpoint_principal_point", TupleTypeFactory(float), nullable=True
-        ),
-        mapped_column("_viewpoint_intrinsics_format", nullable=True),
-    )
-
-
-class SqlSequenceAnnotation(Base):
-    __tablename__ = "sequence_annots"
-
-    sequence_name: Mapped[str] = mapped_column(primary_key=True)
-    category: Mapped[str] = mapped_column(index=True)
-
-    video: Mapped[VideoAnnotation] = composite(
-        mapped_column("_video_path", nullable=True),
-        mapped_column("_video_length", nullable=True),
-    )
-    point_cloud: Mapped[PointCloudAnnotation] = composite(
-        mapped_column("_point_cloud_path", nullable=True),
-        mapped_column("_point_cloud_quality_score", nullable=True),
-        mapped_column("_point_cloud_n_points", nullable=True),
-    )
-    # the bigger the better
-    viewpoint_quality_score: Mapped[Optional[float]] = mapped_column()

pytorch3d/implicitron/dataset/rendered_mesh_dataset_map_provider.py

@@ -8,7 +8,11 @@ from os.path import dirname, join, realpath
 from typing import Optional, Tuple
 
 import torch
-from pytorch3d.implicitron.tools.config import registry, run_auto_creation
+from pytorch3d.implicitron.tools.config import (
+    expand_args_fields,
+    registry,
+    run_auto_creation,
+)
 from pytorch3d.io import IO
 from pytorch3d.renderer import (
     AmbientLights,
@@ -45,7 +49,7 @@ class RenderedMeshDatasetMapProvider(DatasetMapProviderBase):  # pyre-ignore [13
     if one is available, the data it produces is on the CPU just like
     the data returned by implicitron's other dataset map providers.
     This is because both datasets and models can be large, so implicitron's
-    training loop expects data on the CPU and only moves
+    GenericModel.forward (etc) expects data on the CPU and only moves
     what it needs to the device.
 
     For a more detailed explanation of this code, please refer to the
@@ -57,23 +61,16 @@ class RenderedMeshDatasetMapProvider(DatasetMapProviderBase):  # pyre-ignore [13
             the cow mesh in the same repo as this code.
         azimuth_range: number of degrees on each side of the start position to
             take samples
-        distance: distance from camera centres to the origin.
         resolution: the common height and width of the output images.
         use_point_light: whether to use a particular point light as opposed
             to ambient white.
-        gpu_idx: which gpu to use for rendering the mesh.
-        path_manager_factory: (Optional) An object that generates an instance of
-            PathManager that can translate provided file paths.
-        path_manager_factory_class_type: The class type of `path_manager_factory`.
     """
 
     num_views: int = 40
     data_file: Optional[str] = None
     azimuth_range: float = 180
-    distance: float = 2.7
     resolution: int = 128
     use_point_light: bool = True
-    gpu_idx: Optional[int] = 0
     path_manager_factory: PathManagerFactory
     path_manager_factory_class_type: str = "PathManagerFactory"
 
@@ -88,8 +85,8 @@ class RenderedMeshDatasetMapProvider(DatasetMapProviderBase):  # pyre-ignore [13
     def __post_init__(self) -> None:
         super().__init__()
         run_auto_creation(self)
-        if torch.cuda.is_available() and self.gpu_idx is not None:
-            device = torch.device(f"cuda:{self.gpu_idx}")
+        if torch.cuda.is_available():
+            device = torch.device("cuda:0")
         else:
             device = torch.device("cpu")
         if self.data_file is None:
@@ -109,13 +106,13 @@ class RenderedMeshDatasetMapProvider(DatasetMapProviderBase):  # pyre-ignore [13
             num_views=self.num_views,
             mesh=mesh,
             azimuth_range=self.azimuth_range,
-            distance=self.distance,
             resolution=self.resolution,
             device=device,
             use_point_light=self.use_point_light,
         )
         # pyre-ignore[16]
         self.poses = poses.cpu()
+        expand_args_fields(SingleSceneDataset)
         # pyre-ignore[16]
         self.train_dataset = SingleSceneDataset(  # pyre-ignore[28]
             object_name="cow",
@@ -133,7 +130,6 @@ def _generate_cow_renders(
     num_views: int,
     mesh: Meshes,
     azimuth_range: float,
-    distance: float,
     resolution: int,
     device: torch.device,
     use_point_light: bool,
@@ -172,11 +168,11 @@ def _generate_cow_renders(
     else:
         lights = AmbientLights(device=device)
 
-    # Initialize a perspective camera that represents a batch of different
+    # Initialize an OpenGL perspective camera that represents a batch of different
     # viewing angles. All the cameras helper methods support mixed type inputs and
-    # broadcasting. So we can view the camera from a fixed distance, and
+    # broadcasting. So we can view the camera from the a distance of dist=2.7, and
     # then specify elevation and azimuth angles for each viewpoint as tensors.
-    R, T = look_at_view_transform(dist=distance, elev=elev, azim=azim)
+    R, T = look_at_view_transform(dist=2.7, elev=elev, azim=azim)
     cameras = FoVPerspectiveCameras(device=device, R=R, T=T)
 
     # Define the settings for rasterization and shading.

pytorch3d/implicitron/dataset/single_sequence_dataset.py

@@ -9,7 +9,7 @@
 # provide data for a single scene.
 
 from dataclasses import field
-from typing import Iterable, Iterator, List, Optional, Sequence, Tuple
+from typing import Iterable, Iterator, List, Optional, Tuple
 
 import numpy as np
 import torch
@@ -20,9 +20,8 @@ from pytorch3d.implicitron.tools.config import (
 )
 from pytorch3d.renderer import CamerasBase, join_cameras_as_batch, PerspectiveCameras
 
-from .dataset_base import DatasetBase
+from .dataset_base import DatasetBase, FrameData
 from .dataset_map_provider import DatasetMap, DatasetMapProviderBase, PathManagerFactory
-from .frame_data import FrameData
 from .utils import DATASET_TYPE_KNOWN, DATASET_TYPE_UNKNOWN
 
 _SINGLE_SEQUENCE_NAME: str = "one_sequence"
@@ -48,11 +47,10 @@ class SingleSceneDataset(DatasetBase, Configurable):
         return len(self.poses)
 
     def sequence_frames_in_order(
-        self, seq_name: str, subset_filter: Optional[Sequence[str]] = None
+        self, seq_name: str
     ) -> Iterator[Tuple[float, int, int]]:
         for i in range(len(self)):
-            if subset_filter is None or self.frame_types[i] in subset_filter:
-                yield 0.0, i, i
+            yield (0.0, i, i)
 
     def __getitem__(self, index) -> FrameData:
         if index >= len(self):
@@ -69,8 +67,7 @@ class SingleSceneDataset(DatasetBase, Configurable):
             sequence_name=_SINGLE_SEQUENCE_NAME,
             sequence_category=self.object_name,
             camera=pose,
-            # pyre-ignore
-            image_size_hw=torch.tensor(image.shape[1:], dtype=torch.long),
+            image_size_hw=torch.tensor(image.shape[1:]),
             image_rgb=image,
             fg_probability=fg_probability,
             frame_type=frame_type,

pytorch3d/implicitron/dataset/sql_dataset.py

@@ -1,768 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the BSD-style license found in the
-# LICENSE file in the root directory of this source tree.
-
-import hashlib
-import json
-import logging
-import os
-from dataclasses import dataclass
-from typing import (
-    Any,
-    ClassVar,
-    Dict,
-    Iterable,
-    Iterator,
-    List,
-    Optional,
-    Sequence,
-    Tuple,
-    Type,
-    Union,
-)
-
-import numpy as np
-import pandas as pd
-import sqlalchemy as sa
-import torch
-from pytorch3d.implicitron.dataset.dataset_base import DatasetBase
-
-from pytorch3d.implicitron.dataset.frame_data import (  # noqa
-    FrameData,
-    FrameDataBuilder,
-    FrameDataBuilderBase,
-)
-from pytorch3d.implicitron.tools.config import (
-    registry,
-    ReplaceableBase,
-    run_auto_creation,
-)
-from sqlalchemy.orm import Session
-
-from .orm_types import SqlFrameAnnotation, SqlSequenceAnnotation
-
-
-logger = logging.getLogger(__name__)
-
-
-_SET_LISTS_TABLE: str = "set_lists"
-
-
-@registry.register
-class SqlIndexDataset(DatasetBase, ReplaceableBase):  # pyre-ignore
-    """
-    A dataset with annotations stored as SQLite tables. This is an index-based dataset.
-    The length is returned after all sequence and frame filters are applied (see param
-    definitions below). Indices can either be ordinal in [0, len), or pairs of
-    (sequence_name, frame_number); with the performance of `dataset[i]` and
-    `dataset[sequence_name, frame_number]` being same. A faster way to get metadata only
-    (without blobs) is `dataset.meta[idx]` indexing; it requires box_crop==False.
-    With ordinal indexing, the sequences are NOT guaranteed to span contiguous index
-    ranges, and frame numbers are NOT guaranteed to be increasing within a sequence.
-    Sequence-aware batch samplers have to use `sequence_[frames|indices]_in_order`
-    iterators, which are efficient.
-
-    This functionality requires SQLAlchemy 2.0 or later.
-
-    Metadata-related args:
-        sqlite_metadata_file: A SQLite file containing frame and sequence annotation
-            tables (mapping to SqlFrameAnnotation and SqlSequenceAnnotation,
-            respectively).
-        dataset_root: A root directory to look for images, masks, etc. It can be
-            alternatively set in `frame_data_builder` args, but this takes precedence.
-        subset_lists_file: A JSON/sqlite file containing the lists of frames
-            corresponding to different subsets (e.g. train/val/test) of the dataset;
-            format: {subset: [(sequence_name, frame_id, file_path)]}. All entries
-            must be present in frame_annotation metadata table.
-        path_manager: a facade for non-POSIX filesystems.
-        subsets: Restrict frames/sequences only to the given list of subsets
-            as defined in subset_lists_file (see above). Applied before all other
-            filters.
-        remove_empty_masks: Removes the frames with no active foreground pixels
-            in the segmentation mask (needs frame_annotation.mask.mass to be set;
-            null values are retained).
-        pick_frames_sql_clause: SQL WHERE clause to constrain frame annotations
-            NOTE: This is a potential security risk! The string is passed to the SQL
-            engine verbatim. Don’t expose it to end users of your application!
-        pick_categories: Restrict the dataset to the given list of categories.
-        pick_sequences: A Sequence of sequence names to restrict the dataset to.
-        exclude_sequences: A Sequence of the names of the sequences to exclude.
-        limit_sequences_per_category_to: Limit the dataset to the first up to N
-            sequences within each category (applies after all other sequence filters
-            but before `limit_sequences_to`).
-        limit_sequences_to: Limit the dataset to the first `limit_sequences_to`
-            sequences (after other sequence filters have been applied but before
-            frame-based filters).
-        limit_to: Limit the dataset to the first #limit_to frames (after other
-            filters have been applied, except n_frames_per_sequence).
-        n_frames_per_sequence: If > 0, randomly samples `n_frames_per_sequence`
-            frames in each sequences uniformly without replacement if it has
-            more frames than that; applied after other frame-level filters.
-        seed: The seed of the random generator sampling `n_frames_per_sequence`
-            random frames per sequence.
-    """
-
-    frame_annotations_type: ClassVar[Type[SqlFrameAnnotation]] = SqlFrameAnnotation
-
-    sqlite_metadata_file: str = ""
-    dataset_root: Optional[str] = None
-    subset_lists_file: str = ""
-    eval_batches_file: Optional[str] = None
-    path_manager: Any = None
-    subsets: Optional[List[str]] = None
-    remove_empty_masks: bool = True
-    pick_frames_sql_clause: Optional[str] = None
-    pick_categories: Tuple[str, ...] = ()
-
-    pick_sequences: Tuple[str, ...] = ()
-    exclude_sequences: Tuple[str, ...] = ()
-    limit_sequences_per_category_to: int = 0
-    limit_sequences_to: int = 0
-    limit_to: int = 0
-    n_frames_per_sequence: int = -1
-    seed: int = 0
-    remove_empty_masks_poll_whole_table_threshold: int = 300_000
-    # we set it manually in the constructor
-    # _index: pd.DataFrame = field(init=False)
-
-    frame_data_builder: FrameDataBuilderBase
-    frame_data_builder_class_type: str = "FrameDataBuilder"
-
-    def __post_init__(self) -> None:
-        if sa.__version__ < "2.0":
-            raise ImportError("This class requires SQL Alchemy 2.0 or later")
-
-        if not self.sqlite_metadata_file:
-            raise ValueError("sqlite_metadata_file must be set")
-
-        if self.dataset_root:
-            frame_builder_type = self.frame_data_builder_class_type
-            getattr(self, f"frame_data_builder_{frame_builder_type}_args")[
-                "dataset_root"
-            ] = self.dataset_root
-
-        run_auto_creation(self)
-        self.frame_data_builder.path_manager = self.path_manager
-
-        # pyre-ignore  # NOTE: sqlite-specific args (read-only mode).
-        self._sql_engine = sa.create_engine(
-            f"sqlite:///file:{self.sqlite_metadata_file}?mode=ro&uri=true"
-        )
-
-        sequences = self._get_filtered_sequences_if_any()
-
-        if self.subsets:
-            index = self._build_index_from_subset_lists(sequences)
-        else:
-            # TODO: if self.subset_lists_file and not self.subsets, it might be faster to
-            # still use the concatenated lists, assuming they cover the whole dataset
-            index = self._build_index_from_db(sequences)
-
-        if self.n_frames_per_sequence >= 0:
-            index = self._stratified_sample_index(index)
-
-        if len(index) == 0:
-            raise ValueError(f"There are no frames in the subsets: {self.subsets}!")
-
-        self._index = index.set_index(["sequence_name", "frame_number"])  # pyre-ignore
-
-        self.eval_batches = None  # pyre-ignore
-        if self.eval_batches_file:
-            self.eval_batches = self._load_filter_eval_batches()
-
-        logger.info(str(self))
-
-    def __len__(self) -> int:
-        # pyre-ignore[16]
-        return len(self._index)
-
-    def __getitem__(self, frame_idx: Union[int, Tuple[str, int]]) -> FrameData:
-        """
-        Fetches FrameData by either iloc in the index or by (sequence, frame_no) pair
-        """
-        return self._get_item(frame_idx, True)
-
-    @property
-    def meta(self):
-        """
-        Allows accessing metadata only without loading blobs using `dataset.meta[idx]`.
-        Requires box_crop==False, since in that case, cameras cannot be adjusted
-        without loading masks.
-
-        Returns:
-            FrameData objects with blob fields like `image_rgb` set to None.
-
-        Raises:
-            ValueError if dataset.box_crop is set.
-        """
-        return SqlIndexDataset._MetadataAccessor(self)
-
-    @dataclass
-    class _MetadataAccessor:
-        dataset: "SqlIndexDataset"
-
-        def __getitem__(self, frame_idx: Union[int, Tuple[str, int]]) -> FrameData:
-            return self.dataset._get_item(frame_idx, False)
-
-    def _get_item(
-        self, frame_idx: Union[int, Tuple[str, int]], load_blobs: bool = True
-    ) -> FrameData:
-        if isinstance(frame_idx, int):
-            if frame_idx >= len(self._index):
-                raise IndexError(f"index {frame_idx} out of range {len(self._index)}")
-
-            seq, frame = self._index.index[frame_idx]
-        else:
-            seq, frame, *rest = frame_idx
-            if isinstance(frame, torch.LongTensor):
-                frame = frame.item()
-
-            if (seq, frame) not in self._index.index:
-                raise IndexError(
-                    f"Sequence-frame index {frame_idx} not found; was it filtered out?"
-                )
-
-            if rest and rest[0] != self._index.loc[(seq, frame), "_image_path"]:
-                raise IndexError(f"Non-matching image path in {frame_idx}.")
-
-        stmt = sa.select(self.frame_annotations_type).where(
-            self.frame_annotations_type.sequence_name == seq,
-            self.frame_annotations_type.frame_number
-            == int(frame),  # cast from np.int64
-        )
-        seq_stmt = sa.select(SqlSequenceAnnotation).where(
-            SqlSequenceAnnotation.sequence_name == seq
-        )
-        with Session(self._sql_engine) as session:
-            entry = session.scalars(stmt).one()
-            seq_metadata = session.scalars(seq_stmt).one()
-
-        assert entry.image.path == self._index.loc[(seq, frame), "_image_path"]
-
-        frame_data = self.frame_data_builder.build(
-            entry, seq_metadata, load_blobs=load_blobs
-        )
-
-        # The rest of the fields are optional
-        frame_data.frame_type = self._get_frame_type(entry)
-        return frame_data
-
-    def __str__(self) -> str:
-        # pyre-ignore[16]
-        return f"SqlIndexDataset #frames={len(self._index)}"
-
-    def sequence_names(self) -> Iterable[str]:
-        """Returns an iterator over sequence names in the dataset."""
-        return self._index.index.unique("sequence_name")
-
-    # override
-    def category_to_sequence_names(self) -> Dict[str, List[str]]:
-        stmt = sa.select(
-            SqlSequenceAnnotation.category, SqlSequenceAnnotation.sequence_name
-        ).where(  # we limit results to sequences that have frames after all filters
-            SqlSequenceAnnotation.sequence_name.in_(self.sequence_names())
-        )
-        with self._sql_engine.connect() as connection:
-            cat_to_seqs = pd.read_sql(stmt, connection)
-
-        return cat_to_seqs.groupby("category")["sequence_name"].apply(list).to_dict()
-
-    # override
-    def get_frame_numbers_and_timestamps(
-        self, idxs: Sequence[int], subset_filter: Optional[Sequence[str]] = None
-    ) -> List[Tuple[int, float]]:
-        """
-        Implements the DatasetBase method.
-
-        NOTE: Avoid this function as there are more efficient alternatives such as
-        querying `dataset[idx]` directly or getting all sequence frames with
-        `sequence_[frames|indices]_in_order`.
-
-        Return the index and timestamp in their videos of the frames whose
-        indices are given in `idxs`. They need to belong to the same sequence!
-        If timestamps are absent, they are replaced with zeros.
-        This is used for letting SceneBatchSampler identify consecutive
-        frames.
-
-        Args:
-            idxs: a sequence int frame index in the dataset (it can be a slice)
-            subset_filter: must remain None
-
-        Returns:
-            list of tuples of
-                - frame index in video
-                - timestamp of frame in video, coalesced with 0s
-
-        Raises:
-            ValueError if idxs belong to more than one sequence.
-        """
-
-        if subset_filter is not None:
-            raise NotImplementedError(
-                "Subset filters are not supported in SQL Dataset. "
-                "We encourage creating a dataset per subset."
-            )
-
-        index_slice, _ = self._get_frame_no_coalesced_ts_by_row_indices(idxs)
-        # alternatively, we can use `.values.tolist()`, which may be faster
-        # but returns a list of lists
-        return list(index_slice.itertuples())
-
-    # override
-    def sequence_frames_in_order(
-        self, seq_name: str, subset_filter: Optional[Sequence[str]] = None
-    ) -> Iterator[Tuple[float, int, int]]:
-        """
-        Overrides the default DatasetBase implementation (we don’t use `_seq_to_idx`).
-        Returns an iterator over the frame indices in a given sequence.
-        We attempt to first sort by timestamp (if they are available),
-        then by frame number.
-
-        Args:
-            seq_name: the name of the sequence.
-            subset_filter: subset names to filter to
-
-        Returns:
-            an iterator over triplets `(timestamp, frame_no, dataset_idx)`,
-                where `frame_no` is the index within the sequence, and
-                `dataset_idx` is the index within the dataset.
-                `None` timestamps are replaced with 0s.
-        """
-        # TODO: implement sort_timestamp_first? (which would matter if the orders
-        # of frame numbers and timestamps are different)
-        rows = self._index.index.get_loc(seq_name)
-        if isinstance(rows, slice):
-            assert rows.stop is not None, "Unexpected result from pandas"
-            rows = range(rows.start or 0, rows.stop, rows.step or 1)
-        else:
-            rows = np.where(rows)[0]
-
-        index_slice, idx = self._get_frame_no_coalesced_ts_by_row_indices(
-            rows, seq_name, subset_filter
-        )
-        index_slice["idx"] = idx
-
-        yield from index_slice.itertuples(index=False)
-
-    # override
-    def get_eval_batches(self) -> Optional[List[Any]]:
-        """
-        This class does not support eval batches with ordinal indices. You can pass
-        eval_batches as a batch_sampler to a data_loader since the dataset supports
-        `dataset[seq_name, frame_no]` indexing.
-        """
-        return self.eval_batches
-
-    # override
-    def join(self, other_datasets: Iterable[DatasetBase]) -> None:
-        raise ValueError("Not supported! Preprocess the data by merging them instead.")
-
-    # override
-    @property
-    def frame_data_type(self) -> Type[FrameData]:
-        return self.frame_data_builder.frame_data_type
-
-    def is_filtered(self) -> bool:
-        """
-        Returns `True` in case the dataset has been filtered and thus some frame
-        annotations stored on the disk might be missing in the dataset object.
-        Does not account for subsets.
-
-        Returns:
-            is_filtered: `True` if the dataset has been filtered, else `False`.
-        """
-        return (
-            self.remove_empty_masks
-            or self.limit_to > 0
-            or self.limit_sequences_to > 0
-            or self.limit_sequences_per_category_to > 0
-            or len(self.pick_sequences) > 0
-            or len(self.exclude_sequences) > 0
-            or len(self.pick_categories) > 0
-            or self.n_frames_per_sequence > 0
-        )
-
-    def _get_filtered_sequences_if_any(self) -> Optional[pd.Series]:
-        # maximum possible filter (if limit_sequences_per_category_to == 0):
-        # WHERE category IN 'self.pick_categories'
-        # AND sequence_name IN 'self.pick_sequences'
-        # AND sequence_name NOT IN 'self.exclude_sequences'
-        # LIMIT 'self.limit_sequence_to'
-
-        where_conditions = [
-            *self._get_category_filters(),
-            *self._get_pick_filters(),
-            *self._get_exclude_filters(),
-        ]
-
-        def add_where(stmt):
-            return stmt.where(*where_conditions) if where_conditions else stmt
-
-        if self.limit_sequences_per_category_to <= 0:
-            stmt = add_where(sa.select(SqlSequenceAnnotation.sequence_name))
-        else:
-            subquery = sa.select(
-                SqlSequenceAnnotation.sequence_name,
-                sa.func.row_number()
-                .over(
-                    order_by=sa.text("ROWID"),  # NOTE: ROWID is SQLite-specific
-                    partition_by=SqlSequenceAnnotation.category,
-                )
-                .label("row_number"),
-            )
-
-            subquery = add_where(subquery).subquery()
-            stmt = sa.select(subquery.c.sequence_name).where(
-                subquery.c.row_number <= self.limit_sequences_per_category_to
-            )
-
-        if self.limit_sequences_to > 0:
-            logger.info(
-                f"Limiting dataset to first {self.limit_sequences_to} sequences"
-            )
-            # NOTE: ROWID is SQLite-specific
-            stmt = stmt.order_by(sa.text("ROWID")).limit(self.limit_sequences_to)
-
-        if (
-            not where_conditions
-            and self.limit_sequences_to <= 0
-            and self.limit_sequences_per_category_to <= 0
-        ):
-            # we will not need to filter by sequences
-            return None
-
-        with self._sql_engine.connect() as connection:
-            sequences = pd.read_sql_query(stmt, connection)["sequence_name"]
-        logger.info("... retained %d sequences" % len(sequences))
-
-        return sequences
-
-    def _get_category_filters(self) -> List[sa.ColumnElement]:
-        if not self.pick_categories:
-            return []
-
-        logger.info(f"Limiting dataset to categories: {self.pick_categories}")
-        return [SqlSequenceAnnotation.category.in_(self.pick_categories)]
-
-    def _get_pick_filters(self) -> List[sa.ColumnElement]:
-        if not self.pick_sequences:
-            return []
-
-        logger.info(f"Limiting dataset to sequences: {self.pick_sequences}")
-        return [SqlSequenceAnnotation.sequence_name.in_(self.pick_sequences)]
-
-    def _get_exclude_filters(self) -> List[sa.ColumnOperators]:
-        if not self.exclude_sequences:
-            return []
-
-        logger.info(f"Removing sequences from the dataset: {self.exclude_sequences}")
-        return [SqlSequenceAnnotation.sequence_name.notin_(self.exclude_sequences)]
-
-    def _load_subsets_from_json(self, subset_lists_path: str) -> pd.DataFrame:
-        assert self.subsets is not None
-        with open(subset_lists_path, "r") as f:
-            subset_to_seq_frame = json.load(f)
-
-        seq_frame_list = sum(
-            (
-                [(*row, subset) for row in subset_to_seq_frame[subset]]
-                for subset in self.subsets
-            ),
-            [],
-        )
-        index = pd.DataFrame(
-            seq_frame_list,
-            columns=["sequence_name", "frame_number", "_image_path", "subset"],
-        )
-        return index
-
-    def _load_subsets_from_sql(self, subset_lists_path: str) -> pd.DataFrame:
-        subsets = self.subsets
-        assert subsets is not None
-        # we need a new engine since we store the subsets in a separate DB
-        engine = sa.create_engine(f"sqlite:///{subset_lists_path}")
-        table = sa.Table(_SET_LISTS_TABLE, sa.MetaData(), autoload_with=engine)
-        stmt = sa.select(table).where(table.c.subset.in_(subsets))
-        with engine.connect() as connection:
-            index = pd.read_sql(stmt, connection)
-
-        return index
-
-    def _build_index_from_subset_lists(
-        self, sequences: Optional[pd.Series]
-    ) -> pd.DataFrame:
-        if not self.subset_lists_file:
-            raise ValueError("Requested subsets but subset_lists_file not given")
-
-        logger.info(f"Loading subset lists from {self.subset_lists_file}.")
-
-        subset_lists_path = self._local_path(self.subset_lists_file)
-        if subset_lists_path.lower().endswith(".json"):
-            index = self._load_subsets_from_json(subset_lists_path)
-        else:
-            index = self._load_subsets_from_sql(subset_lists_path)
-        index = index.set_index(["sequence_name", "frame_number"])
-        logger.info(f"  -> loaded {len(index)} samples of {self.subsets}.")
-
-        if sequences is not None:
-            logger.info("Applying filtered sequences.")
-            sequence_values = index.index.get_level_values("sequence_name")
-            index = index.loc[sequence_values.isin(sequences)]
-            logger.info(f"  -> retained {len(index)} samples.")
-
-        pick_frames_criteria = []
-        if self.remove_empty_masks:
-            logger.info("Culling samples with empty masks.")
-
-            if len(index) > self.remove_empty_masks_poll_whole_table_threshold:
-                # APPROACH 1: find empty masks and drop indices.
-                # dev load: 17s / 15 s (3.1M / 500K)
-                stmt = sa.select(
-                    self.frame_annotations_type.sequence_name,
-                    self.frame_annotations_type.frame_number,
-                ).where(self.frame_annotations_type._mask_mass == 0)
-                with Session(self._sql_engine) as session:
-                    to_remove = session.execute(stmt).all()
-
-                # Pandas uses np.int64 for integer types, so we have to case
-                # we might want to read it to pandas DataFrame directly to avoid the loop
-                to_remove = [(seq, np.int64(fr)) for seq, fr in to_remove]
-                index.drop(to_remove, errors="ignore", inplace=True)
-            else:
-                # APPROACH 3: load index into a temp table and join with annotations
-                # dev load: 94 s / 23 s (3.1M / 500K)
-                pick_frames_criteria.append(
-                    sa.or_(
-                        self.frame_annotations_type._mask_mass.is_(None),
-                        self.frame_annotations_type._mask_mass != 0,
-                    )
-                )
-
-        if self.pick_frames_sql_clause:
-            logger.info("Applying the custom SQL clause.")
-            pick_frames_criteria.append(sa.text(self.pick_frames_sql_clause))
-
-        if pick_frames_criteria:
-            index = self._pick_frames_by_criteria(index, pick_frames_criteria)
-
-        logger.info(f"  -> retained {len(index)} samples.")
-
-        if self.limit_to > 0:
-            logger.info(f"Limiting dataset to first {self.limit_to} frames")
-            index = index.sort_index().iloc[: self.limit_to]
-
-        return index.reset_index()
-
-    def _pick_frames_by_criteria(self, index: pd.DataFrame, criteria) -> pd.DataFrame:
-        IndexTable = self._get_temp_index_table_instance()
-        with self._sql_engine.connect() as connection:
-            IndexTable.create(connection)
-            # we don’t let pandas’s `to_sql` create the table automatically as
-            # the table would be permanent, so we create it and append with pandas
-            n_rows = index.to_sql(IndexTable.name, connection, if_exists="append")
-            assert n_rows == len(index)
-            sa_type = self.frame_annotations_type
-            stmt = (
-                sa.select(IndexTable)
-                .select_from(
-                    IndexTable.join(
-                        self.frame_annotations_type,
-                        sa.and_(
-                            sa_type.sequence_name == IndexTable.c.sequence_name,
-                            sa_type.frame_number == IndexTable.c.frame_number,
-                        ),
-                    )
-                )
-                .where(*criteria)
-            )
-            return pd.read_sql_query(stmt, connection).set_index(
-                ["sequence_name", "frame_number"]
-            )
-
-    def _build_index_from_db(self, sequences: Optional[pd.Series]):
-        logger.info("Loading sequcence-frame index from the database")
-        stmt = sa.select(
-            self.frame_annotations_type.sequence_name,
-            self.frame_annotations_type.frame_number,
-            self.frame_annotations_type._image_path,
-            sa.null().label("subset"),
-        )
-        where_conditions = []
-        if sequences is not None:
-            logger.info("  applying filtered sequences")
-            where_conditions.append(
-                self.frame_annotations_type.sequence_name.in_(sequences.tolist())
-            )
-
-        if self.remove_empty_masks:
-            logger.info("  excluding samples with empty masks")
-            where_conditions.append(
-                sa.or_(
-                    self.frame_annotations_type._mask_mass.is_(None),
-                    self.frame_annotations_type._mask_mass != 0,
-                )
-            )
-
-        if self.pick_frames_sql_clause:
-            logger.info("  applying custom SQL clause")
-            where_conditions.append(sa.text(self.pick_frames_sql_clause))
-
-        if where_conditions:
-            stmt = stmt.where(*where_conditions)
-
-        if self.limit_to > 0:
-            logger.info(f"Limiting dataset to first {self.limit_to} frames")
-            stmt = stmt.order_by(
-                self.frame_annotations_type.sequence_name,
-                self.frame_annotations_type.frame_number,
-            ).limit(self.limit_to)
-
-        with self._sql_engine.connect() as connection:
-            index = pd.read_sql_query(stmt, connection)
-
-        logger.info(f"  -> loaded {len(index)} samples.")
-        return index
-
-    def _sort_index_(self, index):
-        logger.info("Sorting the index by sequence and frame number.")
-        index.sort_values(["sequence_name", "frame_number"], inplace=True)
-        logger.info("  -> Done.")
-
-    def _load_filter_eval_batches(self):
-        assert self.eval_batches_file
-        logger.info(f"Loading eval batches from {self.eval_batches_file}")
-
-        if not os.path.isfile(self.eval_batches_file):
-            # The batch indices file does not exist.
-            # Most probably the user has not specified the root folder.
-            raise ValueError(
-                f"Looking for dataset json file in {self.eval_batches_file}. "
-                + "Please specify a correct dataset_root folder."
-            )
-
-        with open(self.eval_batches_file, "r") as f:
-            eval_batches = json.load(f)
-
-        # limit the dataset to sequences to allow multiple evaluations in one file
-        pick_sequences = set(self.pick_sequences)
-        if self.pick_categories:
-            cat_to_seq = self.category_to_sequence_names()
-            pick_sequences.update(
-                seq for cat in self.pick_categories for seq in cat_to_seq[cat]
-            )
-
-        if pick_sequences:
-            old_len = len(eval_batches)
-            eval_batches = [b for b in eval_batches if b[0][0] in pick_sequences]
-            logger.warn(
-                f"Picked eval batches by sequence/cat: {old_len} -> {len(eval_batches)}"
-            )
-
-        if self.exclude_sequences:
-            old_len = len(eval_batches)
-            exclude_sequences = set(self.exclude_sequences)
-            eval_batches = [b for b in eval_batches if b[0][0] not in exclude_sequences]
-            logger.warn(
-                f"Excluded eval batches by sequence: {old_len} -> {len(eval_batches)}"
-            )
-
-        return eval_batches
-
-    def _stratified_sample_index(self, index):
-        # NOTE this stratified sampling can be done more efficiently in
-        # the no-subset case above if it is added to the SQL query.
-        # We keep this generic implementation since no-subset case is uncommon
-        index = index.groupby("sequence_name", group_keys=False).apply(
-            lambda seq_frames: seq_frames.sample(
-                min(len(seq_frames), self.n_frames_per_sequence),
-                random_state=(
-                    _seq_name_to_seed(seq_frames.iloc[0]["sequence_name"]) + self.seed
-                ),
-            )
-        )
-        logger.info(f"  -> retained {len(index)} samples aster stratified sampling.")
-        return index
-
-    def _get_frame_type(self, entry: SqlFrameAnnotation) -> Optional[str]:
-        return self._index.loc[(entry.sequence_name, entry.frame_number), "subset"]
-
-    def _get_frame_no_coalesced_ts_by_row_indices(
-        self,
-        idxs: Sequence[int],
-        seq_name: Optional[str] = None,
-        subset_filter: Union[Sequence[str], str, None] = None,
-    ) -> Tuple[pd.DataFrame, Sequence[int]]:
-        """
-        Loads timestamps for given index rows belonging to the same sequence.
-        If seq_name is known, it speeds up the computation.
-        Raises ValueError if `idxs` do not all belong to a single sequences .
-        """
-        index_slice = self._index.iloc[idxs]
-        if subset_filter is not None:
-            if isinstance(subset_filter, str):
-                subset_filter = [subset_filter]
-            indicator = index_slice["subset"].isin(subset_filter)
-            index_slice = index_slice.loc[indicator]
-            idxs = [i for i, isin in zip(idxs, indicator) if isin]
-
-        frames = index_slice.index.get_level_values("frame_number").tolist()
-        if seq_name is None:
-            seq_name_list = index_slice.index.get_level_values("sequence_name").tolist()
-            seq_name_set = set(seq_name_list)
-            if len(seq_name_set) > 1:
-                raise ValueError("Given indices belong to more than one sequence.")
-            elif len(seq_name_set) == 1:
-                seq_name = seq_name_list[0]
-
-        coalesced_ts = sa.sql.functions.coalesce(
-            self.frame_annotations_type.frame_timestamp, 0
-        )
-        stmt = sa.select(
-            coalesced_ts.label("frame_timestamp"),
-            self.frame_annotations_type.frame_number,
-        ).where(
-            self.frame_annotations_type.sequence_name == seq_name,
-            self.frame_annotations_type.frame_number.in_(frames),
-        )
-
-        with self._sql_engine.connect() as connection:
-            frame_no_ts = pd.read_sql_query(stmt, connection)
-
-        if len(frame_no_ts) != len(index_slice):
-            raise ValueError(
-                "Not all indices are found in the database; "
-                "do they belong to more than one sequence?"
-            )
-
-        return frame_no_ts, idxs
-
-    def _local_path(self, path: str) -> str:
-        if self.path_manager is None:
-            return path
-        return self.path_manager.get_local_path(path)
-
-    def _get_temp_index_table_instance(self, table_name: str = "__index"):
-        CachedTable = self.frame_annotations_type.metadata.tables.get(table_name)
-        if CachedTable is not None:  # table definition is not idempotent
-            return CachedTable
-
-        return sa.Table(
-            table_name,
-            self.frame_annotations_type.metadata,
-            sa.Column("sequence_name", sa.String, primary_key=True),
-            sa.Column("frame_number", sa.Integer, primary_key=True),
-            sa.Column("_image_path", sa.String),
-            sa.Column("subset", sa.String),
-            prefixes=["TEMP"],  # NOTE SQLite specific!
-        )
-
-
-def _seq_name_to_seed(seq_name) -> int:
-    """Generates numbers in [0, 2 ** 28)"""
-    return int(hashlib.sha1(seq_name.encode("utf-8")).hexdigest()[:7], 16)
-
-
-def _safe_as_tensor(data, dtype):
-    return torch.tensor(data, dtype=dtype) if data is not None else None

pytorch3d/implicitron/dataset/sql_dataset_provider.py

@@ -1,424 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the BSD-style license found in the
-# LICENSE file in the root directory of this source tree.
-
-
-import logging
-import os
-from typing import List, Optional, Tuple, Type
-
-import numpy as np
-
-from omegaconf import DictConfig, OmegaConf
-
-from pytorch3d.implicitron.dataset.dataset_map_provider import (
-    DatasetMap,
-    DatasetMapProviderBase,
-    PathManagerFactory,
-)
-from pytorch3d.implicitron.tools.config import (
-    expand_args_fields,
-    registry,
-    run_auto_creation,
-)
-
-from .sql_dataset import SqlIndexDataset
-
-
-_CO3D_SQL_DATASET_ROOT: str = os.getenv("CO3D_SQL_DATASET_ROOT", "")
-
-# _NEED_CONTROL is a list of those elements of SqlIndexDataset which
-# are not directly specified for it in the config but come from the
-# DatasetMapProvider.
-_NEED_CONTROL: Tuple[str, ...] = (
-    "path_manager",
-    "subsets",
-    "sqlite_metadata_file",
-    "subset_lists_file",
-)
-
-logger = logging.getLogger(__name__)
-
-
-@registry.register
-class SqlIndexDatasetMapProvider(DatasetMapProviderBase):  # pyre-ignore [13]
-    """
-    Generates the training, validation, and testing dataset objects for
-    a dataset laid out on disk like SQL-CO3D, with annotations in an SQLite data base.
-
-    The dataset is organized in the filesystem as follows::
-
-        self.dataset_root
-            ├── <possible/partition/0>
-            │   ├── <sequence_name_0>
-            │   │   ├── depth_masks
-            │   │   ├── depths
-            │   │   ├── images
-            │   │   ├── masks
-            │   │   └── pointcloud.ply
-            │   ├── <sequence_name_1>
-            │   │   ├── depth_masks
-            │   │   ├── depths
-            │   │   ├── images
-            │   │   ├── masks
-            │   │   └── pointcloud.ply
-            │   ├── ...
-            │   ├── <sequence_name_N>
-            │   ├── set_lists
-            │       ├── <subset_base_name_0>.json
-            │       ├── <subset_base_name_1>.json
-            │       ├── ...
-            │       ├── <subset_base_name_2>.json
-            │   ├── eval_batches
-            │   │   ├── <eval_batches_base_name_0>.json
-            │   │   ├── <eval_batches_base_name_1>.json
-            │   │   ├── ...
-            │   │   ├── <eval_batches_base_name_M>.json
-            │   ├── frame_annotations.jgz
-            │   ├── sequence_annotations.jgz
-            ├── <possible/partition/1>
-            ├── ...
-            ├── <possible/partition/K>
-            ├── set_lists
-                ├── <subset_base_name_0>.sqlite
-                ├── <subset_base_name_1>.sqlite
-                ├── ...
-                ├── <subset_base_name_2>.sqlite
-            ├── eval_batches
-            │   ├── <eval_batches_base_name_0>.json
-            │   ├── <eval_batches_base_name_1>.json
-            │   ├── ...
-            │   ├── <eval_batches_base_name_M>.json
-
-    The dataset contains sequences named `<sequence_name_i>` that may be partitioned by
-    directories such as `<possible/partition/0>` e.g. representing categories but they
-    can also be stored in a flat structure. Each sequence folder contains the list of
-    sequence images, depth maps, foreground masks, and valid-depth masks
-    `images`, `depths`, `masks`, and `depth_masks` respectively. Furthermore,
-    `set_lists/` dirtectories (with partitions or global) store json or sqlite files
-    `<subset_base_name_l>.<ext>`, each describing a certain sequence subset.
-    These subset path conventions are not hard-coded and arbitrary relative path can be
-    specified by setting `self.subset_lists_path` to the relative path w.r.t.
-    dataset root.
-
-    Each `<subset_base_name_l>.json` file contains the following dictionary::
-
-        {
-            "train": [
-                (sequence_name: str, frame_number: int, image_path: str),
-                ...
-            ],
-            "val": [
-                (sequence_name: str, frame_number: int, image_path: str),
-                ...
-            ],
-            "test": [
-                (sequence_name: str, frame_number: int, image_path: str),
-                ...
-            ],
-        ]
-
-    defining the list of frames (identified with their `sequence_name` and
-    `frame_number`) in the "train", "val", and "test" subsets of the dataset. In case of
-    SQLite format, `<subset_base_name_l>.sqlite` contains a table with the header::
-
-        | sequence_name | frame_number | image_path | subset |
-
-    Note that `frame_number` can be obtained only from the metadata and
-    does not necesarrily correspond to the numeric suffix of the corresponding image
-    file name (e.g. a file `<partition_0>/<sequence_name_0>/images/frame00005.jpg` can
-    have its frame number set to `20`, not 5).
-
-    Each `<eval_batches_base_name_M>.json` file contains a list of evaluation examples
-    in the following form::
-
-        [
-            [  # batch 1
-                (sequence_name: str, frame_number: int, image_path: str),
-                ...
-            ],
-            [  # batch 2
-                (sequence_name: str, frame_number: int, image_path: str),
-                ...
-            ],
-        ]
-
-    Note that the evaluation examples always come from the `"test"` subset of the dataset.
-    (test frames can repeat across batches). The batches can contain single element,
-    which is typical in case of regular radiance field fitting.
-
-    Args:
-        subset_lists_path: The relative path to the dataset subset definition.
-            For CO3D, these include e.g. "skateboard/set_lists/set_lists_manyview_dev_0.json".
-            By default (None), dataset is not partitioned to subsets (in that case, setting
-            `ignore_subsets` will speed up construction)
-        dataset_root: The root folder of the dataset.
-        metadata_basename: name of the SQL metadata file in dataset_root;
-            not expected to be changed by users
-        test_on_train: Construct validation and test datasets from
-            the training subset; note that in practice, in this
-            case all subset dataset objects will be same
-        only_test_set: Load only the test set. Incompatible with `test_on_train`.
-        ignore_subsets: Don’t filter by subsets in the dataset; note that in this
-            case all subset datasets will be same
-        eval_batch_num_training_frames: Add a certain number of training frames to each
-            eval batch. Useful for evaluating models that require
-            source views as input (e.g. NeRF-WCE / PixelNeRF).
-        dataset_args: Specifies additional arguments to the
-            JsonIndexDataset constructor call.
-        path_manager_factory: (Optional) An object that generates an instance of
-            PathManager that can translate provided file paths.
-        path_manager_factory_class_type: The class type of `path_manager_factory`.
-    """
-
-    category: Optional[str] = None
-    subset_list_name: Optional[str] = None  # TODO: docs
-    # OR
-    subset_lists_path: Optional[str] = None
-    eval_batches_path: Optional[str] = None
-
-    dataset_root: str = _CO3D_SQL_DATASET_ROOT
-    metadata_basename: str = "metadata.sqlite"
-
-    test_on_train: bool = False
-    only_test_set: bool = False
-    ignore_subsets: bool = False
-    train_subsets: Tuple[str, ...] = ("train",)
-    val_subsets: Tuple[str, ...] = ("val",)
-    test_subsets: Tuple[str, ...] = ("test",)
-
-    eval_batch_num_training_frames: int = 0
-
-    # this is a mould that is never constructed, used to build self._dataset_map values
-    dataset_class_type: str = "SqlIndexDataset"
-    dataset: SqlIndexDataset
-
-    path_manager_factory: PathManagerFactory
-    path_manager_factory_class_type: str = "PathManagerFactory"
-
-    def __post_init__(self):
-        super().__init__()
-        run_auto_creation(self)
-
-        if self.only_test_set and self.test_on_train:
-            raise ValueError("Cannot have only_test_set and test_on_train")
-
-        if self.ignore_subsets and not self.only_test_set:
-            self.test_on_train = True  # no point in loading same data 3 times
-
-        path_manager = self.path_manager_factory.get()
-
-        sqlite_metadata_file = os.path.join(self.dataset_root, self.metadata_basename)
-        sqlite_metadata_file = _local_path(path_manager, sqlite_metadata_file)
-
-        if not os.path.isfile(sqlite_metadata_file):
-            # The sqlite_metadata_file does not exist.
-            # Most probably the user has not specified the root folder.
-            raise ValueError(
-                f"Looking for frame annotations in {sqlite_metadata_file}."
-                + " Please specify a correct dataset_root folder."
-                + " Note: By default the root folder is taken from the"
-                + " CO3D_SQL_DATASET_ROOT environment variable."
-            )
-
-        if self.subset_lists_path and self.subset_list_name:
-            raise ValueError(
-                "subset_lists_path and subset_list_name cannot be both set"
-            )
-
-        subset_lists_file = self._get_lists_file("set_lists")
-
-        # setup the common dataset arguments
-        common_dataset_kwargs = {
-            **getattr(self, f"dataset_{self.dataset_class_type}_args"),
-            "sqlite_metadata_file": sqlite_metadata_file,
-            "dataset_root": self.dataset_root,
-            "subset_lists_file": subset_lists_file,
-            "path_manager": path_manager,
-        }
-
-        if self.category:
-            logger.info(f"Forcing category filter in the datasets to {self.category}")
-            common_dataset_kwargs["pick_categories"] = self.category.split(",")
-
-        # get the used dataset type
-        dataset_type: Type[SqlIndexDataset] = registry.get(
-            SqlIndexDataset, self.dataset_class_type
-        )
-        expand_args_fields(dataset_type)
-
-        if subset_lists_file is not None and not os.path.isfile(subset_lists_file):
-            available_subsets = self._get_available_subsets(
-                OmegaConf.to_object(common_dataset_kwargs["pick_categories"])
-            )
-            msg = f"Cannot find subset list file {self.subset_lists_path}."
-            if available_subsets:
-                msg += f" Some of the available subsets: {str(available_subsets)}."
-            raise ValueError(msg)
-
-        train_dataset = None
-        val_dataset = None
-        if not self.only_test_set:
-            # load the training set
-            logger.debug("Constructing train dataset.")
-            train_dataset = dataset_type(
-                **common_dataset_kwargs, subsets=self._get_subsets(self.train_subsets)
-            )
-            logger.info(f"Train dataset: {str(train_dataset)}")
-
-        if self.test_on_train:
-            assert train_dataset is not None
-            val_dataset = test_dataset = train_dataset
-        else:
-            # load the val and test sets
-            if not self.only_test_set:
-                # NOTE: this is always loaded in JsonProviderV2
-                logger.debug("Extracting val dataset.")
-                val_dataset = dataset_type(
-                    **common_dataset_kwargs, subsets=self._get_subsets(self.val_subsets)
-                )
-                logger.info(f"Val dataset: {str(val_dataset)}")
-
-            logger.debug("Extracting test dataset.")
-            eval_batches_file = self._get_lists_file("eval_batches")
-            del common_dataset_kwargs["eval_batches_file"]
-            test_dataset = dataset_type(
-                **common_dataset_kwargs,
-                subsets=self._get_subsets(self.test_subsets, True),
-                eval_batches_file=eval_batches_file,
-            )
-            logger.info(f"Test dataset: {str(test_dataset)}")
-
-            if (
-                eval_batches_file is not None
-                and self.eval_batch_num_training_frames > 0
-            ):
-                self._extend_eval_batches(test_dataset)
-
-        self._dataset_map = DatasetMap(
-            train=train_dataset, val=val_dataset, test=test_dataset
-        )
-
-    def _get_subsets(self, subsets, is_eval: bool = False):
-        if self.ignore_subsets:
-            return None
-
-        if is_eval and self.eval_batch_num_training_frames > 0:
-            # we will need to have training frames for extended batches
-            return list(subsets) + list(self.train_subsets)
-
-        return subsets
-
-    def _extend_eval_batches(self, test_dataset: SqlIndexDataset) -> None:
-        rng = np.random.default_rng(seed=0)
-        eval_batches = test_dataset.get_eval_batches()
-        if eval_batches is None:
-            raise ValueError("Eval batches were not loaded!")
-
-        for batch in eval_batches:
-            sequence = batch[0][0]
-            seq_frames = list(
-                test_dataset.sequence_frames_in_order(sequence, self.train_subsets)
-            )
-            idx_to_add = rng.permutation(len(seq_frames))[
-                : self.eval_batch_num_training_frames
-            ]
-            batch.extend((sequence, seq_frames[a][1]) for a in idx_to_add)
-
-    @classmethod
-    def dataset_tweak_args(cls, type, args: DictConfig) -> None:
-        """
-        Called by get_default_args.
-        Certain fields are not exposed on each dataset class
-        but rather are controlled by this provider class.
-        """
-        for key in _NEED_CONTROL:
-            del args[key]
-
-    def create_dataset(self):
-        # No `dataset` member of this class is created.
-        # The dataset(s) live in `self.get_dataset_map`.
-        pass
-
-    def get_dataset_map(self) -> DatasetMap:
-        return self._dataset_map  # pyre-ignore [16]
-
-    def _get_available_subsets(self, categories: List[str]):
-        """
-        Get the available subset names for a given category folder (if given) inside
-        a root dataset folder `dataset_root`.
-        """
-        path_manager = self.path_manager_factory.get()
-
-        subsets: List[str] = []
-        for prefix in [""] + categories:
-            set_list_dir = os.path.join(self.dataset_root, prefix, "set_lists")
-            if not (
-                (path_manager is not None) and path_manager.isdir(set_list_dir)
-            ) and not os.path.isdir(set_list_dir):
-                continue
-
-            set_list_files = (os.listdir if path_manager is None else path_manager.ls)(
-                set_list_dir
-            )
-            subsets.extend(os.path.join(prefix, "set_lists", f) for f in set_list_files)
-
-        return subsets
-
-    def _get_lists_file(self, flavor: str) -> Optional[str]:
-        if flavor == "eval_batches":
-            subset_lists_path = self.eval_batches_path
-        else:
-            subset_lists_path = self.subset_lists_path
-
-        if not subset_lists_path and not self.subset_list_name:
-            return None
-
-        category_elem = ""
-        if self.category and "," not in self.category:
-            # if multiple categories are given, looking for global set lists
-            category_elem = self.category
-
-        subset_lists_path = subset_lists_path or (
-            os.path.join(
-                category_elem, f"{flavor}", f"{flavor}_{self.subset_list_name}"
-            )
-        )
-
-        assert subset_lists_path
-        path_manager = self.path_manager_factory.get()
-        # try absolute path first
-        subset_lists_file = _get_local_path_check_extensions(
-            subset_lists_path, path_manager
-        )
-        if subset_lists_file:
-            return subset_lists_file
-
-        full_path = os.path.join(self.dataset_root, subset_lists_path)
-        subset_lists_file = _get_local_path_check_extensions(full_path, path_manager)
-
-        if not subset_lists_file:
-            raise FileNotFoundError(
-                f"Subset lists path given but not found: {full_path}"
-            )
-
-        return subset_lists_file
-
-
-def _get_local_path_check_extensions(
-    path, path_manager, extensions=("", ".sqlite", ".json")
-) -> Optional[str]:
-    for ext in extensions:
-        local = _local_path(path_manager, path + ext)
-        if os.path.isfile(local):
-            return local
-
-    return None
-
-
-def _local_path(path_manager, path: str) -> str:
-    if path_manager is None:
-        return path
-    return path_manager.get_local_path(path)

pytorch3d/implicitron/dataset/train_eval_data_loader_provider.py

@@ -1,189 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the BSD-style license found in the
-# LICENSE file in the root directory of this source tree.
-
-import logging
-from typing import Any, Dict, Optional, Tuple
-
-from pytorch3d.implicitron.dataset.data_loader_map_provider import (
-    DataLoaderMap,
-    SceneBatchSampler,
-    SequenceDataLoaderMapProvider,
-)
-from pytorch3d.implicitron.dataset.dataset_base import DatasetBase
-from pytorch3d.implicitron.dataset.dataset_map_provider import DatasetMap
-from pytorch3d.implicitron.dataset.frame_data import FrameData
-from pytorch3d.implicitron.tools.config import registry, run_auto_creation
-
-from torch.utils.data import DataLoader
-
-logger = logging.getLogger(__name__)
-
-
-# TODO: we can merge it with SequenceDataLoaderMapProvider in PyTorch3D
-# and support both eval_batches protocols
-@registry.register
-class TrainEvalDataLoaderMapProvider(SequenceDataLoaderMapProvider):
-    """
-    Implementation of DataLoaderMapProviderBase that may use internal eval batches for
-    the test dataset. In particular, if `eval_batches_relpath` is set, it loads
-    eval batches from that json file, otherwise test set is treated in the same way as
-    train and val, i.e. the parameters `dataset_length_test` and `test_conditioning_type`
-    are respected.
-
-    If conditioning is not required, then the batch size should
-    be set as 1, and most of the fields do not matter.
-
-    If conditioning is required, each batch will contain one main
-    frame first to predict and the, rest of the elements are for
-    conditioning.
-
-    If images_per_seq_options is left empty, the conditioning
-    frames are picked according to the conditioning type given.
-    This does not have regard to the order of frames in a
-    scene, or which frames belong to what scene.
-
-    If images_per_seq_options is given, then the conditioning types
-    must be SAME and the remaining fields are used.
-
-    Members:
-        batch_size: The size of the batch of the data loader.
-        num_workers: Number of data-loading threads in each data loader.
-        dataset_length_train: The number of batches in a training epoch. Or 0 to mean
-            an epoch is the length of the training set.
-        dataset_length_val: The number of batches in a validation epoch. Or 0 to mean
-            an epoch is the length of the validation set.
-        dataset_length_test: used if test_dataset.eval_batches is NOT set. The number of
-            batches in a testing epoch. Or 0 to mean an epoch is the length of the test
-            set.
-        images_per_seq_options: Possible numbers of frames sampled per sequence in a batch.
-            If a conditioning_type is KNOWN or TRAIN, then this must be left at its initial
-            value. Empty (the default) means that we are not careful about which frames
-            come from which scene.
-        sample_consecutive_frames: if True, will sample a contiguous interval of frames
-            in the sequence. It first sorts the frames by timestimps when available,
-            otherwise by frame numbers, finds the connected segments within the sequence
-            of sufficient length, then samples a random pivot element among them and
-            ideally uses it as a middle of the temporal window, shifting the borders
-            where necessary. This strategy mitigates the bias against shorter segments
-            and their boundaries.
-        consecutive_frames_max_gap: if a number > 0, then used to define the maximum
-            difference in frame_number of neighbouring frames when forming connected
-            segments; if both this and consecutive_frames_max_gap_seconds are 0s,
-            the whole sequence is considered a segment regardless of frame numbers.
-        consecutive_frames_max_gap_seconds: if a number > 0.0, then used to define the
-            maximum difference in frame_timestamp of neighbouring frames when forming
-            connected segments; if both this and consecutive_frames_max_gap are 0s,
-            the whole sequence is considered a segment regardless of frame timestamps.
-    """
-
-    batch_size: int = 1
-    num_workers: int = 0
-
-    dataset_length_train: int = 0
-    dataset_length_val: int = 0
-    dataset_length_test: int = 0
-
-    images_per_seq_options: Tuple[int, ...] = ()
-    sample_consecutive_frames: bool = False
-    consecutive_frames_max_gap: int = 0
-    consecutive_frames_max_gap_seconds: float = 0.1
-
-    def __post_init__(self):
-        run_auto_creation(self)
-
-    def get_data_loader_map(self, datasets: DatasetMap) -> DataLoaderMap:
-        """
-        Returns a collection of data loaders for a given collection of datasets.
-        """
-        train = self._make_generic_data_loader(
-            datasets.train,
-            self.dataset_length_train,
-            datasets.train,
-        )
-
-        val = self._make_generic_data_loader(
-            datasets.val,
-            self.dataset_length_val,
-            datasets.train,
-        )
-
-        if datasets.test is not None and datasets.test.get_eval_batches() is not None:
-            test = self._make_eval_data_loader(datasets.test)
-        else:
-            test = self._make_generic_data_loader(
-                datasets.test,
-                self.dataset_length_test,
-                datasets.train,
-            )
-
-        return DataLoaderMap(train=train, val=val, test=test)
-
-    def _make_eval_data_loader(
-        self,
-        dataset: Optional[DatasetBase],
-    ) -> Optional[DataLoader[FrameData]]:
-        if dataset is None:
-            return None
-
-        return DataLoader(
-            dataset,
-            batch_sampler=dataset.get_eval_batches(),
-            **self._get_data_loader_common_kwargs(dataset),
-        )
-
-    def _make_generic_data_loader(
-        self,
-        dataset: Optional[DatasetBase],
-        num_batches: int,
-        train_dataset: Optional[DatasetBase],
-    ) -> Optional[DataLoader[FrameData]]:
-        """
-        Returns the dataloader for a dataset.
-
-        Args:
-            dataset: the dataset
-            num_batches: possible ceiling on number of batches per epoch
-            train_dataset: the training dataset, used if conditioning_type==TRAIN
-            conditioning_type: source for padding of batches
-        """
-        if dataset is None:
-            return None
-
-        data_loader_kwargs = self._get_data_loader_common_kwargs(dataset)
-
-        if len(self.images_per_seq_options) > 0:
-            # this is a typical few-view setup
-            # conditioning comes from the same subset since subsets are split by seqs
-            batch_sampler = SceneBatchSampler(
-                dataset,
-                self.batch_size,
-                num_batches=len(dataset) if num_batches <= 0 else num_batches,
-                images_per_seq_options=self.images_per_seq_options,
-                sample_consecutive_frames=self.sample_consecutive_frames,
-                consecutive_frames_max_gap=self.consecutive_frames_max_gap,
-                consecutive_frames_max_gap_seconds=self.consecutive_frames_max_gap_seconds,
-            )
-            return DataLoader(
-                dataset,
-                batch_sampler=batch_sampler,
-                **data_loader_kwargs,
-            )
-
-        if self.batch_size == 1:
-            # this is a typical many-view setup (without conditioning)
-            return self._simple_loader(dataset, num_batches, data_loader_kwargs)
-
-        # edge case: conditioning on train subset, typical for Nerformer-like many-view
-        # there is only one sequence in all datasets, so we condition on another subset
-        return self._train_loader(
-            dataset, train_dataset, num_batches, data_loader_kwargs
-        )
-
-    def _get_data_loader_common_kwargs(self, dataset: DatasetBase) -> Dict[str, Any]:
-        return {
-            "num_workers": self.num_workers,
-            "collate_fn": dataset.frame_data_type.collate,
-        }

pytorch3d/implicitron/dataset/types.py

@@ -9,21 +9,10 @@ import dataclasses
 import gzip
 import json
 from dataclasses import dataclass, Field, MISSING
-from typing import (
-    Any,
-    cast,
-    Dict,
-    get_args,
-    get_origin,
-    IO,
-    Optional,
-    Tuple,
-    Type,
-    TypeVar,
-    Union,
-)
+from typing import Any, cast, Dict, IO, Optional, Tuple, Type, TypeVar, Union
 
 import numpy as np
+from pytorch3d.common.datatypes import get_args, get_origin
 
 
 _X = TypeVar("_X")
@@ -55,8 +44,6 @@ class MaskAnnotation:
     path: str
     # (soft) number of pixels in the mask; sum(Prob(fg | pixel))
     mass: Optional[float] = None
-    # tight bounding box around the foreground mask
-    bounding_box_xywh: Optional[Tuple[float, float, float, float]] = None
 
 
 @dataclass
@@ -204,7 +191,7 @@ def _dataclass_list_from_dict_list(dlist, typeannot):
     # otherwise, we dispatch by the type of the provided annotation to convert to
     if issubclass(cls, tuple) and hasattr(cls, "_fields"):  # namedtuple
         # For namedtuple, call the function recursively on the lists of corresponding keys
-        types = cls.__annotations__.values()
+        types = cls._field_types.values()
         dlist_T = zip(*dlist)
         res_T = [
             _dataclass_list_from_dict_list(key_list, tp)
@@ -270,7 +257,7 @@ def _dataclass_from_dict(d, typeannot):
 
     cls = get_origin(typeannot) or typeannot
     if issubclass(cls, tuple) and hasattr(cls, "_fields"):  # namedtuple
-        types = cls.__annotations__.values()
+        types = cls._field_types.values()
         return cls(*[_dataclass_from_dict(v, tp) for v, tp in zip(d, types)])
     elif issubclass(cls, (list, tuple)):
         types = get_args(typeannot)

pytorch3d/implicitron/dataset/utils.py

@@ -5,18 +5,10 @@
 # LICENSE file in the root directory of this source tree.
 
 
-import functools
-import warnings
-from pathlib import Path
-from typing import List, Optional, Tuple, TypeVar, Union
+from typing import List, Optional
 
-import numpy as np
 import torch
-from PIL import Image
 
-from pytorch3d.io import IO
-from pytorch3d.renderer.cameras import PerspectiveCameras
-from pytorch3d.structures.pointclouds import Pointclouds
 
 DATASET_TYPE_TRAIN = "train"
 DATASET_TYPE_TEST = "test"
@@ -24,26 +16,6 @@ DATASET_TYPE_KNOWN = "known"
 DATASET_TYPE_UNKNOWN = "unseen"
 
 
-class GenericWorkaround:
-    """
-    OmegaConf.structured has a weirdness when you try to apply
-    it to a dataclass whose first base class is a Generic which is not
-    Dict. The issue is with a function called get_dict_key_value_types
-    in omegaconf/_utils.py.
-    For example this fails:
-
-        @dataclass(eq=False)
-        class D(torch.utils.data.Dataset[int]):
-            a: int = 3
-
-        OmegaConf.structured(D)
-
-    We avoid the problem by adding this class as an extra base class.
-    """
-
-    pass
-
-
 def is_known_frame_scalar(frame_type: str) -> bool:
     """
     Given a single frame type corresponding to a single frame, return whether
@@ -80,301 +52,3 @@ def is_train_frame(
         dtype=torch.bool,
         device=device,
     )
-
-
-def get_bbox_from_mask(
-    mask: np.ndarray, thr: float, decrease_quant: float = 0.05
-) -> Tuple[int, int, int, int]:
-    # bbox in xywh
-    masks_for_box = np.zeros_like(mask)
-    while masks_for_box.sum() <= 1.0:
-        masks_for_box = (mask > thr).astype(np.float32)
-        thr -= decrease_quant
-    if thr <= 0.0:
-        warnings.warn(
-            f"Empty masks_for_bbox (thr={thr}) => using full image.", stacklevel=1
-        )
-
-    x0, x1 = get_1d_bounds(masks_for_box.sum(axis=-2))
-    y0, y1 = get_1d_bounds(masks_for_box.sum(axis=-1))
-
-    return x0, y0, x1 - x0, y1 - y0
-
-
-def crop_around_box(
-    tensor: torch.Tensor, bbox: torch.Tensor, impath: str = ""
-) -> torch.Tensor:
-    # bbox is xyxy, where the upper bound is corrected with +1
-    bbox = clamp_box_to_image_bounds_and_round(
-        bbox,
-        image_size_hw=tuple(tensor.shape[-2:]),
-    )
-    tensor = tensor[..., bbox[1] : bbox[3], bbox[0] : bbox[2]]
-    assert all(c > 0 for c in tensor.shape), f"squashed image {impath}"
-    return tensor
-
-
-def clamp_box_to_image_bounds_and_round(
-    bbox_xyxy: torch.Tensor,
-    image_size_hw: Tuple[int, int],
-) -> torch.LongTensor:
-    bbox_xyxy = bbox_xyxy.clone()
-    bbox_xyxy[[0, 2]] = torch.clamp(bbox_xyxy[[0, 2]], 0, image_size_hw[-1])
-    bbox_xyxy[[1, 3]] = torch.clamp(bbox_xyxy[[1, 3]], 0, image_size_hw[-2])
-    if not isinstance(bbox_xyxy, torch.LongTensor):
-        bbox_xyxy = bbox_xyxy.round().long()
-    return bbox_xyxy  # pyre-ignore [7]
-
-
-T = TypeVar("T", bound=torch.Tensor)
-
-
-def bbox_xyxy_to_xywh(xyxy: T) -> T:
-    wh = xyxy[2:] - xyxy[:2]
-    xywh = torch.cat([xyxy[:2], wh])
-    return xywh  # pyre-ignore
-
-
-def get_clamp_bbox(
-    bbox: torch.Tensor,
-    box_crop_context: float = 0.0,
-    image_path: str = "",
-) -> torch.Tensor:
-    # box_crop_context: rate of expansion for bbox
-    # returns possibly expanded bbox xyxy as float
-
-    bbox = bbox.clone()  # do not edit bbox in place
-
-    # increase box size
-    if box_crop_context > 0.0:
-        c = box_crop_context
-        bbox = bbox.float()
-        bbox[0] -= bbox[2] * c / 2
-        bbox[1] -= bbox[3] * c / 2
-        bbox[2] += bbox[2] * c
-        bbox[3] += bbox[3] * c
-
-    if (bbox[2:] <= 1.0).any():
-        raise ValueError(
-            f"squashed image {image_path}!! The bounding box contains no pixels."
-        )
-
-    bbox[2:] = torch.clamp(bbox[2:], 2)  # set min height, width to 2 along both axes
-    bbox_xyxy = bbox_xywh_to_xyxy(bbox, clamp_size=2)
-
-    return bbox_xyxy
-
-
-def rescale_bbox(
-    bbox: torch.Tensor,
-    orig_res: Union[Tuple[int, int], torch.LongTensor],
-    new_res: Union[Tuple[int, int], torch.LongTensor],
-) -> torch.Tensor:
-    assert bbox is not None
-    assert np.prod(orig_res) > 1e-8
-    # average ratio of dimensions
-    # pyre-ignore
-    rel_size = (new_res[0] / orig_res[0] + new_res[1] / orig_res[1]) / 2.0
-    return bbox * rel_size
-
-
-def bbox_xywh_to_xyxy(
-    xywh: torch.Tensor, clamp_size: Optional[int] = None
-) -> torch.Tensor:
-    xyxy = xywh.clone()
-    if clamp_size is not None:
-        xyxy[2:] = torch.clamp(xyxy[2:], clamp_size)
-    xyxy[2:] += xyxy[:2]
-    return xyxy
-
-
-def get_1d_bounds(arr: np.ndarray) -> Tuple[int, int]:
-    nz = np.flatnonzero(arr)
-    return nz[0], nz[-1] + 1
-
-
-def resize_image(
-    image: Union[np.ndarray, torch.Tensor],
-    image_height: Optional[int],
-    image_width: Optional[int],
-    mode: str = "bilinear",
-) -> Tuple[torch.Tensor, float, torch.Tensor]:
-
-    if isinstance(image, np.ndarray):
-        image = torch.from_numpy(image)
-
-    if image_height is None or image_width is None:
-        # skip the resizing
-        return image, 1.0, torch.ones_like(image[:1])
-    # takes numpy array or tensor, returns pytorch tensor
-    minscale = min(
-        image_height / image.shape[-2],
-        image_width / image.shape[-1],
-    )
-    imre = torch.nn.functional.interpolate(
-        image[None],
-        scale_factor=minscale,
-        mode=mode,
-        align_corners=False if mode == "bilinear" else None,
-        recompute_scale_factor=True,
-    )[0]
-    imre_ = torch.zeros(image.shape[0], image_height, image_width)
-    imre_[:, 0 : imre.shape[1], 0 : imre.shape[2]] = imre
-    mask = torch.zeros(1, image_height, image_width)
-    mask[:, 0 : imre.shape[1], 0 : imre.shape[2]] = 1.0
-    return imre_, minscale, mask
-
-
-def transpose_normalize_image(image: np.ndarray) -> np.ndarray:
-    im = np.atleast_3d(image).transpose((2, 0, 1))
-    return im.astype(np.float32) / 255.0
-
-
-def load_image(path: str) -> np.ndarray:
-    with Image.open(path) as pil_im:
-        im = np.array(pil_im.convert("RGB"))
-
-    return transpose_normalize_image(im)
-
-
-def load_mask(path: str) -> np.ndarray:
-    with Image.open(path) as pil_im:
-        mask = np.array(pil_im)
-
-    return transpose_normalize_image(mask)
-
-
-def load_depth(path: str, scale_adjustment: float) -> np.ndarray:
-    if path.lower().endswith(".exr"):
-        # NOTE: environment variable OPENCV_IO_ENABLE_OPENEXR must be set to 1
-        # You will have to accept these vulnerabilities by using OpenEXR:
-        # https://github.com/opencv/opencv/issues/21326
-        import cv2
-
-        d = cv2.imread(path, cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH)[..., 0]
-        d[d > 1e9] = 0.0
-    elif path.lower().endswith(".png"):
-        d = load_16big_png_depth(path)
-    else:
-        raise ValueError('unsupported depth file name "%s"' % path)
-
-    d = d * scale_adjustment
-
-    d[~np.isfinite(d)] = 0.0
-    return d[None]  # fake feature channel
-
-
-def load_16big_png_depth(depth_png: str) -> np.ndarray:
-    with Image.open(depth_png) as depth_pil:
-        # the image is stored with 16-bit depth but PIL reads it as I (32 bit).
-        # we cast it to uint16, then reinterpret as float16, then cast to float32
-        depth = (
-            np.frombuffer(np.array(depth_pil, dtype=np.uint16), dtype=np.float16)
-            .astype(np.float32)
-            .reshape((depth_pil.size[1], depth_pil.size[0]))
-        )
-    return depth
-
-
-def load_1bit_png_mask(file: str) -> np.ndarray:
-    with Image.open(file) as pil_im:
-        mask = (np.array(pil_im.convert("L")) > 0.0).astype(np.float32)
-    return mask
-
-
-def load_depth_mask(path: str) -> np.ndarray:
-    if not path.lower().endswith(".png"):
-        raise ValueError('unsupported depth mask file name "%s"' % path)
-    m = load_1bit_png_mask(path)
-    return m[None]  # fake feature channel
-
-
-def safe_as_tensor(data, dtype):
-    return torch.tensor(data, dtype=dtype) if data is not None else None
-
-
-def _convert_ndc_to_pixels(
-    focal_length: torch.Tensor,
-    principal_point: torch.Tensor,
-    image_size_wh: torch.Tensor,
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    half_image_size = image_size_wh / 2
-    rescale = half_image_size.min()
-    principal_point_px = half_image_size - principal_point * rescale
-    focal_length_px = focal_length * rescale
-    return focal_length_px, principal_point_px
-
-
-def _convert_pixels_to_ndc(
-    focal_length_px: torch.Tensor,
-    principal_point_px: torch.Tensor,
-    image_size_wh: torch.Tensor,
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    half_image_size = image_size_wh / 2
-    rescale = half_image_size.min()
-    principal_point = (half_image_size - principal_point_px) / rescale
-    focal_length = focal_length_px / rescale
-    return focal_length, principal_point
-
-
-def adjust_camera_to_bbox_crop_(
-    camera: PerspectiveCameras,
-    image_size_wh: torch.Tensor,
-    clamp_bbox_xywh: torch.Tensor,
-) -> None:
-    if len(camera) != 1:
-        raise ValueError("Adjusting currently works with singleton cameras camera only")
-
-    focal_length_px, principal_point_px = _convert_ndc_to_pixels(
-        camera.focal_length[0],
-        camera.principal_point[0],
-        image_size_wh,
-    )
-    principal_point_px_cropped = principal_point_px - clamp_bbox_xywh[:2]
-
-    focal_length, principal_point_cropped = _convert_pixels_to_ndc(
-        focal_length_px,
-        principal_point_px_cropped,
-        clamp_bbox_xywh[2:],
-    )
-
-    camera.focal_length = focal_length[None]
-    camera.principal_point = principal_point_cropped[None]
-
-
-def adjust_camera_to_image_scale_(
-    camera: PerspectiveCameras,
-    original_size_wh: torch.Tensor,
-    new_size_wh: torch.LongTensor,
-) -> PerspectiveCameras:
-    focal_length_px, principal_point_px = _convert_ndc_to_pixels(
-        camera.focal_length[0],
-        camera.principal_point[0],
-        original_size_wh,
-    )
-
-    # now scale and convert from pixels to NDC
-    image_size_wh_output = new_size_wh.float()
-    scale = (image_size_wh_output / original_size_wh).min(dim=-1, keepdim=True).values
-    focal_length_px_scaled = focal_length_px * scale
-    principal_point_px_scaled = principal_point_px * scale
-
-    focal_length_scaled, principal_point_scaled = _convert_pixels_to_ndc(
-        focal_length_px_scaled,
-        principal_point_px_scaled,
-        image_size_wh_output,
-    )
-    camera.focal_length = focal_length_scaled[None]
-    camera.principal_point = principal_point_scaled[None]  # pyre-ignore
-
-
-# NOTE this cache is per-worker; they are implemented as processes.
-# each batch is loaded and collated by a single worker;
-# since sequences tend to co-occur within batches, this is useful.
-@functools.lru_cache(maxsize=256)
-def load_pointcloud(pcl_path: Union[str, Path], max_points: int = 0) -> Pointclouds:
-    pcl = IO().load_pointcloud(pcl_path)
-    if max_points > 0:
-        pcl = pcl.subsample(max_points)
-
-    return pcl

pytorch3d/implicitron/dataset/visualize.py

@@ -10,7 +10,7 @@ import torch
 from pytorch3d.implicitron.tools.point_cloud_utils import get_rgbd_point_cloud
 from pytorch3d.structures import Pointclouds
 
-from .frame_data import FrameData
+from .dataset_base import FrameData
 from .json_index_dataset import JsonIndexDataset
 
 
@@ -89,8 +89,9 @@ def get_implicitron_sequence_pointcloud(
             frame_data.image_rgb,
             frame_data.depth_map,
             (cast(torch.Tensor, frame_data.fg_probability) > 0.5).float()
-            if mask_points and frame_data.fg_probability is not None
+            if frame_data.fg_probability is not None
             else None,
+            mask_points=mask_points,
         )
 
     return point_cloud, frame_data

pytorch3d/implicitron/eval_demo.py

@@ -130,7 +130,7 @@ def evaluate_dbir_for_category(
         raise ValueError("Image size should be set in the dataset")
 
     # init the simple DBIR model
-    model = ModelDBIR(
+    model = ModelDBIR(  # pyre-ignore[28]: c’tor implicitly overridden
         render_image_width=image_size,
         render_image_height=image_size,
         bg_color=bg_color,
@@ -153,12 +153,21 @@ def evaluate_dbir_for_category(
                 preds["implicitron_render"],
                 bg_color=bg_color,
                 lpips_model=lpips_model,
+                source_cameras=data_source.all_train_cameras,
             )
         )
 
+    if task == Task.SINGLE_SEQUENCE:
+        camera_difficulty_bin_breaks = 0.97, 0.98
+        multisequence_evaluation = False
+    else:
+        camera_difficulty_bin_breaks = 2.0 / 3, 5.0 / 6
+        multisequence_evaluation = True
+
     category_result_flat, category_result = summarize_nvs_eval_results(
         per_batch_eval_results,
-        is_multisequence=task != Task.SINGLE_SEQUENCE,
+        camera_difficulty_bin_breaks=camera_difficulty_bin_breaks,
+        is_multisequence=multisequence_evaluation,
     )
 
     return category_result["results"]

pytorch3d/implicitron/evaluation/evaluate_new_view_synthesis.py

@@ -14,15 +14,17 @@ from typing import Any, Dict, List, Optional, Sequence, Tuple, TYPE_CHECKING, Un
 import numpy as np
 import torch
 import torch.nn.functional as F
-from pytorch3d.implicitron.dataset.frame_data import FrameData
-from pytorch3d.implicitron.dataset.utils import is_train_frame
+from pytorch3d.implicitron.dataset.dataset_base import FrameData
+from pytorch3d.implicitron.dataset.utils import is_known_frame, is_train_frame
 from pytorch3d.implicitron.models.base_model import ImplicitronRender
 from pytorch3d.implicitron.tools import vis_utils
+from pytorch3d.implicitron.tools.camera_utils import volumetric_camera_overlaps
 from pytorch3d.implicitron.tools.image_utils import mask_background
 from pytorch3d.implicitron.tools.metric_utils import calc_psnr, eval_depth, iou, rgb_l1
 from pytorch3d.implicitron.tools.point_cloud_utils import get_rgbd_point_cloud
 from pytorch3d.implicitron.tools.vis_utils import make_depth_image
-from pytorch3d.renderer.cameras import PerspectiveCameras
+from pytorch3d.renderer.camera_utils import join_cameras_as_batch
+from pytorch3d.renderer.cameras import CamerasBase, PerspectiveCameras
 from pytorch3d.vis.plotly_vis import plot_scene
 from tabulate import tabulate
 
@@ -38,8 +40,8 @@ class _Visualizer:
     image_render: torch.Tensor
     image_rgb_masked: torch.Tensor
     depth_render: torch.Tensor
-    depth_map: Optional[torch.Tensor]
-    depth_mask: Optional[torch.Tensor]
+    depth_map: torch.Tensor
+    depth_mask: torch.Tensor
 
     visdom_env: str = "eval_debug"
 
@@ -75,11 +77,9 @@ class _Visualizer:
         viz = self._viz
         viz.images(
             torch.cat(
-                (make_depth_image(self.depth_render, loss_mask_now),)
-                + (
-                    (make_depth_image(self.depth_map, loss_mask_now),)
-                    if self.depth_map is not None
-                    else ()
+                (
+                    make_depth_image(self.depth_render, loss_mask_now),
+                    make_depth_image(self.depth_map, loss_mask_now),
                 ),
                 dim=3,
             ),
@@ -93,13 +93,12 @@ class _Visualizer:
             win="depth_abs" + name_postfix + "_mask",
             opts={"title": f"depth_abs_{name_postfix}_{depth_loss:1.2f}_mask"},
         )
-        if self.depth_mask is not None:
-            viz.images(
-                self.depth_mask,
-                env=self.visdom_env,
-                win="depth_abs" + name_postfix + "_maskd",
-                opts={"title": f"depth_abs_{name_postfix}_{depth_loss:1.2f}_maskd"},
-            )
+        viz.images(
+            self.depth_mask,
+            env=self.visdom_env,
+            win="depth_abs" + name_postfix + "_maskd",
+            opts={"title": f"depth_abs_{name_postfix}_{depth_loss:1.2f}_maskd"},
+        )
 
         # show the 3D plot
         # pyre-fixme[9]: viewpoint_trivial has type `PerspectiveCameras`; used as
@@ -107,30 +106,29 @@ class _Visualizer:
         viewpoint_trivial: PerspectiveCameras = PerspectiveCameras().to(
             loss_mask_now.device
         )
+        pcl_pred = get_rgbd_point_cloud(
+            viewpoint_trivial,
+            self.image_render,
+            self.depth_render,
+            # mask_crop,
+            torch.ones_like(self.depth_render),
+            # loss_mask_now,
+        )
+        pcl_gt = get_rgbd_point_cloud(
+            viewpoint_trivial,
+            self.image_rgb_masked,
+            self.depth_map,
+            # mask_crop,
+            torch.ones_like(self.depth_map),
+            # loss_mask_now,
+        )
         _pcls = {
-            "pred_depth": get_rgbd_point_cloud(
-                viewpoint_trivial,
-                self.image_render,
-                self.depth_render,
-                # mask_crop,
-                torch.ones_like(self.depth_render),
-                # loss_mask_now,
-            )
+            pn: p
+            for pn, p in zip(("pred_depth", "gt_depth"), (pcl_pred, pcl_gt))
+            if int(p.num_points_per_cloud()) > 0
         }
-        if self.depth_map is not None:
-            _pcls["gt_depth"] = get_rgbd_point_cloud(
-                viewpoint_trivial,
-                self.image_rgb_masked,
-                self.depth_map,
-                # mask_crop,
-                torch.ones_like(self.depth_map),
-                # loss_mask_now,
-            )
-
-        _pcls = {pn: p for pn, p in _pcls.items() if int(p.num_points_per_cloud()) > 0}
-
         plotlyplot = plot_scene(
-            {f"pcl{name_postfix}": _pcls},  # pyre-ignore
+            {f"pcl{name_postfix}": _pcls},
             camera_scale=1.0,
             pointcloud_max_points=10000,
             pointcloud_marker_size=1,
@@ -151,6 +149,7 @@ def eval_batch(
     visualize: bool = False,
     visualize_visdom_env: str = "eval_debug",
     break_after_visualising: bool = True,
+    source_cameras: Optional[CamerasBase] = None,
 ) -> Dict[str, Any]:
     """
     Produce performance metrics for a single batch of new-view synthesis
@@ -172,6 +171,8 @@ def eval_batch(
             ground truth.
         lpips_model: A pre-trained model for evaluating the LPIPS metric.
         visualize: If True, visualizes the results to Visdom.
+        source_cameras: A list of all training cameras for evaluating the
+            difficulty of the target views.
 
     Returns:
         results: A dictionary holding evaluation metrics.
@@ -223,10 +224,17 @@ def eval_batch(
         frame_type = [frame_type]
 
     is_train = is_train_frame(frame_type)
-    if len(is_train) > 1 and (is_train[1] != is_train[1:]).any():
+    if not (is_train[0] == is_train).all():
+        raise ValueError("All frames in the eval batch have to be either train/test.")
+
+    # pyre-fixme[16]: `Optional` has no attribute `device`.
+    is_known = is_known_frame(frame_type, device=frame_data.image_rgb.device)
+
+    if not ((is_known[1:] == 1).all() and (is_known[0] == 0).all()):
         raise ValueError(
-            "All (conditioning) frames in the eval batch have to be either train/test."
-        )
+            "For evaluation the first element of the batch has to be"
+            + " a target view while the rest should be source views."
+        )  # TODO: do we need to enforce this?
 
     for k in [
         "depth_map",
@@ -281,10 +289,10 @@ def eval_batch(
             image_render=image_render,
             image_rgb_masked=image_rgb_masked,
             depth_render=cloned_render["depth_render"],
+            # pyre-fixme[6]: Expected `Tensor` for 4th param but got
+            #  `Optional[torch.Tensor]`.
             depth_map=frame_data.depth_map,
-            depth_mask=frame_data.depth_mask[:1]
-            if frame_data.depth_mask is not None
-            else None,
+            depth_mask=frame_data.depth_mask[:1],
             visdom_env=visualize_visdom_env,
         )
 
@@ -339,8 +347,6 @@ def eval_batch(
     ):
         results[rgb_metric_name] = rgb_metric_fun(
             image_render,
-            # pyre-fixme[6]: For 2nd argument expected `Tensor` but got
-            #  `Optional[Tensor]`.
             image_rgb,
             mask=mask_crop,
         )
@@ -359,9 +365,18 @@ def eval_batch(
     # convert all metrics to floats
     results = {k: float(v) for k, v in results.items()}
 
+    if source_cameras is None:
+        # pyre-fixme[16]: Optional has no attribute __getitem__
+        source_cameras = frame_data.camera[torch.where(is_known)[0]]
+
     results["meta"] = {
+        # calculate the camera difficulties and add to results
+        "camera_difficulty": calculate_camera_difficulties(
+            frame_data.camera[0],
+            source_cameras,
+        )[0].item(),
         # store the size of the batch (corresponds to n_src_views+1)
-        "batch_size": len(frame_type),
+        "batch_size": int(is_known.numel()),
         # store the type of the target frame
         # pyre-fixme[16]: `None` has no attribute `__getitem__`.
         "frame_type": str(frame_data.frame_type[0]),
@@ -391,6 +406,33 @@ def average_per_batch_results(
     }
 
 
+def calculate_camera_difficulties(
+    cameras_target: CamerasBase,
+    cameras_source: CamerasBase,
+) -> torch.Tensor:
+    """
+    Calculate the difficulties of the target cameras, given a set of known
+    cameras `cameras_source`.
+
+    Returns:
+        a tensor of shape (len(cameras_target),)
+    """
+    ious = [
+        volumetric_camera_overlaps(
+            join_cameras_as_batch(
+                # pyre-fixme[6]: Expected `CamerasBase` for 1st param but got
+                #  `Optional[pytorch3d.renderer.utils.TensorProperties]`.
+                [cameras_target[cami], cameras_source.to(cameras_target.device)]
+            )
+        )[0, :]
+        for cami in range(cameras_target.R.shape[0])
+    ]
+    camera_difficulties = torch.stack(
+        [_reduce_camera_iou_overlap(iou[1:]) for iou in ious]
+    )
+    return camera_difficulties
+
+
 def _reduce_camera_iou_overlap(ious: torch.Tensor, topk: int = 2) -> torch.Tensor:
     """
     Calculate the final camera difficulty by computing the average of the
@@ -416,7 +458,8 @@ def _get_camera_difficulty_bin_edges(camera_difficulty_bin_breaks: Tuple[float,
 def summarize_nvs_eval_results(
     per_batch_eval_results: List[Dict[str, Any]],
     is_multisequence: bool,
-) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+    camera_difficulty_bin_breaks: Tuple[float, float],
+):
     """
     Compile the per-batch evaluation results `per_batch_eval_results` into
     a set of aggregate metrics. The produced metrics depend on is_multisequence.
@@ -439,12 +482,19 @@ def summarize_nvs_eval_results(
     batch_sizes = torch.tensor(
         [r["meta"]["batch_size"] for r in per_batch_eval_results]
     ).long()
-
+    camera_difficulty = torch.tensor(
+        [r["meta"]["camera_difficulty"] for r in per_batch_eval_results]
+    ).float()
     is_train = is_train_frame([r["meta"]["frame_type"] for r in per_batch_eval_results])
 
     # init the result database dict
     results = []
 
+    diff_bin_edges, diff_bin_names = _get_camera_difficulty_bin_edges(
+        camera_difficulty_bin_breaks
+    )
+    n_diff_edges = diff_bin_edges.numel()
+
     # add per set averages
     for SET in eval_sets:
         if SET is None:
@@ -454,17 +504,26 @@ def summarize_nvs_eval_results(
             ok_set = is_train == int(SET == "train")
             set_name = SET
 
-        # average over all results
-        bin_results = average_per_batch_results(
-            per_batch_eval_results, idx=torch.where(ok_set)[0]
-        )
-        results.append(
-            {
-                "subset": set_name,
-                "subsubset": "diff=all",
-                "metrics": bin_results,
-            }
-        )
+        # eval each difficulty bin, including a full average result (diff_bin=None)
+        for diff_bin in [None, *list(range(n_diff_edges - 1))]:
+            if diff_bin is None:
+                # average over all results
+                in_bin = ok_set
+                diff_bin_name = "all"
+            else:
+                b1, b2 = diff_bin_edges[diff_bin : (diff_bin + 2)]
+                in_bin = ok_set & (camera_difficulty > b1) & (camera_difficulty <= b2)
+                diff_bin_name = diff_bin_names[diff_bin]
+            bin_results = average_per_batch_results(
+                per_batch_eval_results, idx=torch.where(in_bin)[0]
+            )
+            results.append(
+                {
+                    "subset": set_name,
+                    "subsubset": f"diff={diff_bin_name}",
+                    "metrics": bin_results,
+                }
+            )
 
         if is_multisequence:
             # split based on n_src_views
@@ -493,7 +552,7 @@ def _get_flat_nvs_metric_key(result, metric_name) -> str:
     return metric_key
 
 
-def flatten_nvs_results(results) -> Dict[str, Any]:
+def flatten_nvs_results(results):
     """
     Takes input `results` list of dicts of the form::
 
@@ -512,6 +571,7 @@ def flatten_nvs_results(results) -> Dict[str, Any]:
             'subset=train/test/...|subsubset=src=1/src=2/...': nvs_eval_metrics,
             ...
         }
+
     """
     results_flat = {}
     for result in results:

pytorch3d/implicitron/evaluation/evaluator.py

@@ -14,6 +14,8 @@ from typing import Any, Dict, List, Optional, Tuple
 import torch
 
 import tqdm
+from pytorch3d.implicitron.dataset import utils as ds_utils
+
 from pytorch3d.implicitron.evaluation import evaluate_new_view_synthesis as evaluate
 from pytorch3d.implicitron.models.base_model import EvaluationMode, ImplicitronModelBase
 from pytorch3d.implicitron.tools.config import (
@@ -21,6 +23,7 @@ from pytorch3d.implicitron.tools.config import (
     ReplaceableBase,
     run_auto_creation,
 )
+from pytorch3d.renderer.cameras import CamerasBase
 from torch.utils.data import DataLoader
 
 logger = logging.getLogger(__name__)
@@ -47,19 +50,22 @@ class EvaluatorBase(ReplaceableBase):
 class ImplicitronEvaluator(EvaluatorBase):
     """
     Evaluate the results of Implicitron training.
+
+    Members:
+        camera_difficulty_bin_breaks: low/medium vals to divide camera difficulties into
+            [0-eps, low, medium, 1+eps].
     """
 
-    # UNUSED; preserved for compatibility purposes
     camera_difficulty_bin_breaks: Tuple[float, ...] = 0.97, 0.98
 
     def __post_init__(self):
         run_auto_creation(self)
 
-    # pyre-fixme[14]: `run` overrides method defined in `EvaluatorBase` inconsistently.
     def run(
         self,
         model: ImplicitronModelBase,
         dataloader: DataLoader,
+        all_train_cameras: Optional[CamerasBase],
         device: torch.device,
         dump_to_json: bool = False,
         exp_dir: Optional[str] = None,
@@ -73,6 +79,7 @@ class ImplicitronEvaluator(EvaluatorBase):
         Args:
             model: A (trained) model to evaluate.
             dataloader: A test dataloader.
+            all_train_cameras: Camera instances we used for training.
             device: A torch device.
             dump_to_json: If True, will dump the results to a json file.
             exp_dir: Root expeirment directory.
@@ -116,12 +123,16 @@ class ImplicitronEvaluator(EvaluatorBase):
                         implicitron_render,
                         bg_color="black",
                         lpips_model=lpips_model,
+                        source_cameras=(  # None will make it use batch’s known cameras
+                            None if self.is_multisequence else all_train_cameras
+                        ),
                     )
                 )
 
         _, category_result = evaluate.summarize_nvs_eval_results(
             per_batch_eval_results,
             self.is_multisequence,
+            self.camera_difficulty_bin_breaks,
         )
 
         results = category_result["results"]
@@ -148,11 +159,14 @@ def _dump_to_json(
 
 def _get_eval_frame_data(frame_data: Any) -> Any:
     """
-    Masks the target image data to make sure we cannot use it at model evaluation
-    time. Assumes the first batch element is target, the rest are source.
+    Masks the unknown image data to make sure we cannot use it at model evaluation time.
     """
     frame_data_for_eval = copy.deepcopy(frame_data)
+    is_known = ds_utils.is_known_frame(frame_data.frame_type).type_as(
+        frame_data.image_rgb
+    )[:, None, None, None]
     for k in ("image_rgb", "depth_map", "fg_probability", "mask_crop"):
         value = getattr(frame_data_for_eval, k)
-        value[0].zero_()
+        value_masked = value.clone() * is_known if value is not None else None
+        setattr(frame_data_for_eval, k, value_masked)
     return frame_data_for_eval

pytorch3d/implicitron/models/__init__.py

@@ -3,8 +3,3 @@
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
-
-# Allows to register the models
-# see: pytorch3d.implicitron.tools.config.registry:register
-from pytorch3d.implicitron.models.generic_model import GenericModel
-from pytorch3d.implicitron.models.overfit_model import OverfitModel

pytorch3d/implicitron/models/base_model.py

@@ -8,11 +8,11 @@ from dataclasses import dataclass, field
 from typing import Any, Dict, List, Optional
 
 import torch
-
-from pytorch3d.implicitron.models.renderer.base import EvaluationMode
 from pytorch3d.implicitron.tools.config import ReplaceableBase
 from pytorch3d.renderer.cameras import CamerasBase
 
+from .renderer.base import EvaluationMode
+
 
 @dataclass
 class ImplicitronRender:
@@ -49,6 +49,9 @@ class ImplicitronModelBase(ReplaceableBase, torch.nn.Module):
     # the training loop.
     log_vars: List[str] = field(default_factory=lambda: ["objective"])
 
+    def __init__(self) -> None:
+        super().__init__()
+
     def forward(
         self,
         *,  # force keyword-only arguments

pytorch3d/implicitron/models/feature_extractor/feature_extractor.py

@@ -15,6 +15,9 @@ class FeatureExtractorBase(ReplaceableBase, torch.nn.Module):
     Base class for an extractor of a set of features from images.
     """
 
+    def __init__(self):
+        super().__init__()
+
     def get_feat_dims(self) -> int:
         """
         Returns:

pytorch3d/implicitron/models/feature_extractor/resnet_feature_extractor.py

@@ -78,6 +78,7 @@ class ResNetFeatureExtractor(FeatureExtractorBase):
     feature_rescale: float = 1.0
 
     def __post_init__(self):
+        super().__init__()
         if self.normalize_image:
             # register buffers needed to normalize the image
             for k, v in (("_resnet_mean", _RESNET_MEAN), ("_resnet_std", _RESNET_STD)):
@@ -143,6 +144,7 @@ class ResNetFeatureExtractor(FeatureExtractorBase):
         return (img - self._resnet_mean) / self._resnet_std
 
     def get_feat_dims(self) -> int:
+        # pyre-fixme[29]
         return sum(self._feat_dim.values())
 
     def forward(
@@ -180,8 +182,13 @@ class ResNetFeatureExtractor(FeatureExtractorBase):
                 imgs_normed = self._resnet_normalize_image(imgs_resized)
             else:
                 imgs_normed = imgs_resized
+            # pyre-fixme[29]: `Union[torch.Tensor, torch.nn.modules.module.Module]`
             #  is not a function.
             feats = self.stem(imgs_normed)
+            # pyre-fixme[6]: For 1st param expected `Iterable[Variable[_T1]]` but
+            #  got `Union[Tensor, Module]`.
+            # pyre-fixme[6]: For 2nd param expected `Iterable[Variable[_T2]]` but
+            #  got `Union[Tensor, Module]`.
             for stage, (layer, proj) in enumerate(zip(self.layers, self.proj_layers)):
                 feats = layer(feats)
                 # just a sanity check below

pytorch3d/implicitron/models/generic_model.py

@@ -9,56 +9,66 @@
 # which are part of implicitron. They ensure that the registry is prepopulated.
 
 import logging
+import warnings
 from dataclasses import field
 from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING, Union
 
 import torch
+import tqdm
 from omegaconf import DictConfig
-
-from pytorch3d.implicitron.models.base_model import (
-    ImplicitronModelBase,
-    ImplicitronRender,
-)
-from pytorch3d.implicitron.models.feature_extractor import FeatureExtractorBase
-from pytorch3d.implicitron.models.global_encoder.global_encoder import GlobalEncoderBase
-from pytorch3d.implicitron.models.implicit_function.base import ImplicitFunctionBase
+from pytorch3d.common.compat import prod
 from pytorch3d.implicitron.models.metrics import (
     RegularizationMetricsBase,
     ViewMetricsBase,
 )
-
-from pytorch3d.implicitron.models.renderer.base import (
-    BaseRenderer,
-    EvaluationMode,
-    ImplicitFunctionWrapper,
-    ImplicitronRayBundle,
-    RendererOutput,
-    RenderSamplingMode,
-)
-from pytorch3d.implicitron.models.renderer.ray_sampler import RaySamplerBase
-
-from pytorch3d.implicitron.models.utils import (
-    apply_chunked,
-    chunk_generator,
-    log_loss_weights,
-    preprocess_input,
-    weighted_sum_losses,
-)
-from pytorch3d.implicitron.models.view_pooler.view_pooler import ViewPooler
-from pytorch3d.implicitron.tools import vis_utils
+from pytorch3d.implicitron.models.renderer.base import ImplicitronRayBundle
+from pytorch3d.implicitron.tools import image_utils, vis_utils
 from pytorch3d.implicitron.tools.config import (
     expand_args_fields,
     registry,
     run_auto_creation,
 )
-
 from pytorch3d.implicitron.tools.rasterize_mc import rasterize_sparse_ray_bundle
+from pytorch3d.implicitron.tools.utils import cat_dataclass
 from pytorch3d.renderer import utils as rend_utils
-from pytorch3d.renderer.cameras import CamerasBase
 
+from pytorch3d.renderer.cameras import CamerasBase
 
 if TYPE_CHECKING:
     from visdom import Visdom
+
+from .base_model import ImplicitronModelBase, ImplicitronRender
+from .feature_extractor import FeatureExtractorBase
+from .feature_extractor.resnet_feature_extractor import ResNetFeatureExtractor  # noqa
+from .global_encoder.global_encoder import GlobalEncoderBase
+from .implicit_function.base import ImplicitFunctionBase
+from .implicit_function.idr_feature_field import IdrFeatureField  # noqa
+from .implicit_function.neural_radiance_field import (  # noqa
+    NeRFormerImplicitFunction,
+    NeuralRadianceFieldImplicitFunction,
+)
+from .implicit_function.scene_representation_networks import (  # noqa
+    SRNHyperNetImplicitFunction,
+    SRNImplicitFunction,
+)
+from .implicit_function.voxel_grid_implicit_function import (  # noqa
+    VoxelGridImplicitFunction,
+)
+
+from .renderer.base import (
+    BaseRenderer,
+    EvaluationMode,
+    ImplicitFunctionWrapper,
+    RendererOutput,
+    RenderSamplingMode,
+)
+from .renderer.lstm_renderer import LSTMRenderer  # noqa
+from .renderer.multipass_ea import MultiPassEmissionAbsorptionRenderer  # noqa
+from .renderer.ray_sampler import RaySamplerBase
+from .renderer.sdf_renderer import SignedDistanceFunctionRenderer  # noqa
+from .view_pooler.view_pooler import ViewPooler
+
+
 logger = logging.getLogger(__name__)
 
 
@@ -293,38 +303,9 @@ class GenericModel(ImplicitronModelBase):  # pyre-ignore: 13
         ]
     )
 
-    @classmethod
-    def pre_expand(cls) -> None:
-        # use try/finally to bypass cinder's lazy imports
-        try:
-            from pytorch3d.implicitron.models.feature_extractor.resnet_feature_extractor import (  # noqa: F401, B950
-                ResNetFeatureExtractor,
-            )
-            from pytorch3d.implicitron.models.implicit_function.idr_feature_field import (  # noqa: F401, B950
-                IdrFeatureField,
-            )
-            from pytorch3d.implicitron.models.implicit_function.neural_radiance_field import (  # noqa: F401, B950
-                NeRFormerImplicitFunction,
-            )
-            from pytorch3d.implicitron.models.implicit_function.scene_representation_networks import (  # noqa: F401, B950
-                SRNHyperNetImplicitFunction,
-            )
-            from pytorch3d.implicitron.models.implicit_function.voxel_grid_implicit_function import (  # noqa: F401, B950
-                VoxelGridImplicitFunction,
-            )
-            from pytorch3d.implicitron.models.renderer.lstm_renderer import (  # noqa: F401
-                LSTMRenderer,
-            )
-            from pytorch3d.implicitron.models.renderer.multipass_ea import (  # noqa
-                MultiPassEmissionAbsorptionRenderer,
-            )
-            from pytorch3d.implicitron.models.renderer.sdf_renderer import (  # noqa: F401
-                SignedDistanceFunctionRenderer,
-            )
-        finally:
-            pass
-
     def __post_init__(self):
+        super().__init__()
+
         if self.view_pooler_enabled:
             if self.image_feature_extractor_class_type is None:
                 raise ValueError(
@@ -334,7 +315,7 @@ class GenericModel(ImplicitronModelBase):  # pyre-ignore: 13
 
         self._implicit_functions = self._construct_implicit_functions()
 
-        log_loss_weights(self.loss_weights, logger)
+        self.log_loss_weights()
 
     def forward(
         self,
@@ -360,7 +341,7 @@ class GenericModel(ImplicitronModelBase):  # pyre-ignore: 13
                 and source images, which will be used for intersecting with target rays.
             fg_probability: A tensor of shape `(B, 1, H, W)` containing a batch of
                 foreground masks.
-            mask_crop: A binary tensor of shape `(B, 1, H, W)` denoting valid
+            mask_crop: A binary tensor of shape `(B, 1, H, W)` deonting valid
                 regions in the input images (i.e. regions that do not correspond
                 to, e.g., zero-padding). When the `RaySampler`'s sampling mode is set to
                 "mask_sample", rays  will be sampled in the non zero regions.
@@ -378,14 +359,8 @@ class GenericModel(ImplicitronModelBase):  # pyre-ignore: 13
             preds: A dictionary containing all outputs of the forward pass including the
                 rendered images, depths, masks, losses and other metrics.
         """
-        image_rgb, fg_probability, depth_map = preprocess_input(
-            image_rgb,
-            fg_probability,
-            depth_map,
-            self.mask_images,
-            self.mask_depths,
-            self.mask_threshold,
-            self.bg_color,
+        image_rgb, fg_probability, depth_map = self._preprocess_input(
+            image_rgb, fg_probability, depth_map
         )
 
         # Obtain the batch size from the camera as this is the only required input.
@@ -463,15 +438,19 @@ class GenericModel(ImplicitronModelBase):  # pyre-ignore: 13
             )
         custom_args["global_code"] = global_code
 
+        # pyre-fixme[29]:
+        #  `Union[BoundMethod[typing.Callable(torch.Tensor.__iter__)[[Named(self,
+        #  torch.Tensor)], typing.Iterator[typing.Any]], torch.Tensor], torch.Tensor,
+        #  torch.nn.Module]` is not a function.
         for func in self._implicit_functions:
             func.bind_args(**custom_args)
 
-        inputs_to_be_chunked = {}
+        chunked_renderer_inputs = {}
         if fg_probability is not None and self.renderer.requires_object_mask():
             sampled_fb_prob = rend_utils.ndc_grid_sample(
                 fg_probability[:n_targets], ray_bundle.xys, mode="nearest"
             )
-            inputs_to_be_chunked["object_mask"] = sampled_fb_prob > 0.5
+            chunked_renderer_inputs["object_mask"] = sampled_fb_prob > 0.5
 
         # (5)-(6) Implicit function evaluation and Rendering
         rendered = self._render(
@@ -479,12 +458,16 @@ class GenericModel(ImplicitronModelBase):  # pyre-ignore: 13
             sampling_mode=sampling_mode,
             evaluation_mode=evaluation_mode,
             implicit_functions=self._implicit_functions,
-            inputs_to_be_chunked=inputs_to_be_chunked,
+            chunked_inputs=chunked_renderer_inputs,
         )
 
         # Unbind the custom arguments to prevent pytorch from storing
         # large buffers of intermediate results due to points in the
         # bound arguments.
+        # pyre-fixme[29]:
+        #  `Union[BoundMethod[typing.Callable(torch.Tensor.__iter__)[[Named(self,
+        #  torch.Tensor)], typing.Iterator[typing.Any]], torch.Tensor], torch.Tensor,
+        #  torch.nn.Module]` is not a function.
         for func in self._implicit_functions:
             func.unbind_args()
 
@@ -539,18 +522,30 @@ class GenericModel(ImplicitronModelBase):  # pyre-ignore: 13
             raise AssertionError("Unreachable state")
 
         # (7) Compute losses
+        # finally get the optimization objective using self.loss_weights
         objective = self._get_objective(preds)
         if objective is not None:
             preds["objective"] = objective
 
         return preds
 
-    def _get_objective(self, preds: Dict[str, torch.Tensor]) -> Optional[torch.Tensor]:
+    def _get_objective(self, preds) -> Optional[torch.Tensor]:
         """
         A helper function to compute the overall loss as the dot product
         of individual loss functions with the corresponding weights.
         """
-        return weighted_sum_losses(preds, self.loss_weights)
+        losses_weighted = [
+            preds[k] * float(w)
+            for k, w in self.loss_weights.items()
+            if (k in preds and w != 0.0)
+        ]
+        if len(losses_weighted) == 0:
+            warnings.warn("No main objective found.")
+            return None
+        loss = sum(losses_weighted)
+        assert torch.is_tensor(loss)
+        # pyre-fixme[7]: Expected `Optional[Tensor]` but got `int`.
+        return loss
 
     def visualize(
         self,
@@ -582,7 +577,7 @@ class GenericModel(ImplicitronModelBase):  # pyre-ignore: 13
         self,
         *,
         ray_bundle: ImplicitronRayBundle,
-        inputs_to_be_chunked: Dict[str, torch.Tensor],
+        chunked_inputs: Dict[str, torch.Tensor],
         sampling_mode: RenderSamplingMode,
         **kwargs,
     ) -> RendererOutput:
@@ -590,7 +585,7 @@ class GenericModel(ImplicitronModelBase):  # pyre-ignore: 13
         Args:
             ray_bundle: A `ImplicitronRayBundle` object containing the parametrizations of the
                 sampled rendering rays.
-            inputs_to_be_chunked: A collection of tensor of shape `(B, _, H, W)`. E.g.
+            chunked_inputs: A collection of tensor of shape `(B, _, H, W)`. E.g.
                 SignedDistanceFunctionRenderer requires "object_mask", shape
                 (B, 1, H, W), the silhouette of the object in the image. When
                 chunking, they are passed to the renderer as shape
@@ -602,27 +597,30 @@ class GenericModel(ImplicitronModelBase):  # pyre-ignore: 13
             An instance of RendererOutput
         """
         if sampling_mode == RenderSamplingMode.FULL_GRID and self.chunk_size_grid > 0:
-            return apply_chunked(
+            return _apply_chunked(
                 self.renderer,
-                chunk_generator(
+                _chunk_generator(
                     self.chunk_size_grid,
                     ray_bundle,
-                    inputs_to_be_chunked,
+                    chunked_inputs,
                     self.tqdm_trigger_threshold,
                     **kwargs,
                 ),
-                lambda batch: torch.cat(batch, dim=1).reshape(
-                    *ray_bundle.lengths.shape[:-1], -1
-                ),
+                lambda batch: _tensor_collator(batch, ray_bundle.lengths.shape[:-1]),
             )
         else:
             # pyre-fixme[29]: `BaseRenderer` is not a function.
             return self.renderer(
                 ray_bundle=ray_bundle,
-                **inputs_to_be_chunked,
+                **chunked_inputs,
                 **kwargs,
             )
 
+    def _get_global_encoder_encoding_dim(self) -> int:
+        if self.global_encoder is None:
+            return 0
+        return self.global_encoder.get_encoding_dim()
+
     def _get_viewpooled_feature_dim(self) -> int:
         if self.view_pooler is None:
             return 0
@@ -714,29 +712,30 @@ class GenericModel(ImplicitronModelBase):  # pyre-ignore: 13
         function(s) are initialized.
         """
         extra_args = {}
-        global_encoder_dim = (
-            0 if self.global_encoder is None else self.global_encoder.get_encoding_dim()
-        )
-        viewpooled_feature_dim = self._get_viewpooled_feature_dim()
-
         if self.implicit_function_class_type in (
             "NeuralRadianceFieldImplicitFunction",
             "NeRFormerImplicitFunction",
         ):
-            extra_args["latent_dim"] = viewpooled_feature_dim + global_encoder_dim
+            extra_args["latent_dim"] = (
+                self._get_viewpooled_feature_dim()
+                + self._get_global_encoder_encoding_dim()
+            )
             extra_args["color_dim"] = self.render_features_dimensions
 
         if self.implicit_function_class_type == "IdrFeatureField":
             extra_args["feature_vector_size"] = self.render_features_dimensions
-            extra_args["encoding_dim"] = global_encoder_dim
+            extra_args["encoding_dim"] = self._get_global_encoder_encoding_dim()
 
         if self.implicit_function_class_type == "SRNImplicitFunction":
-            extra_args["latent_dim"] = viewpooled_feature_dim + global_encoder_dim
+            extra_args["latent_dim"] = (
+                self._get_viewpooled_feature_dim()
+                + self._get_global_encoder_encoding_dim()
+            )
 
         # srn_hypernet preprocessing
         if self.implicit_function_class_type == "SRNHyperNetImplicitFunction":
-            extra_args["latent_dim"] = viewpooled_feature_dim
-            extra_args["latent_dim_hypernet"] = global_encoder_dim
+            extra_args["latent_dim"] = self._get_viewpooled_feature_dim()
+            extra_args["latent_dim_hypernet"] = self._get_global_encoder_encoding_dim()
 
         # check that for srn, srn_hypernet, idr we have self.num_passes=1
         implicit_function_type = registry.get(
@@ -763,3 +762,147 @@ class GenericModel(ImplicitronModelBase):  # pyre-ignore: 13
             for _ in range(self.num_passes)
         ]
         return torch.nn.ModuleList(implicit_functions_list)
+
+    def log_loss_weights(self) -> None:
+        """
+        Print a table of the loss weights.
+        """
+        loss_weights_message = (
+            "-------\nloss_weights:\n"
+            + "\n".join(f"{k:40s}: {w:1.2e}" for k, w in self.loss_weights.items())
+            + "-------"
+        )
+        logger.info(loss_weights_message)
+
+    def _preprocess_input(
+        self,
+        image_rgb: Optional[torch.Tensor],
+        fg_probability: Optional[torch.Tensor],
+        depth_map: Optional[torch.Tensor],
+    ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]:
+        """
+        Helper function to preprocess the input images and optional depth maps
+        to apply masking if required.
+
+        Args:
+            image_rgb: A tensor of shape `(B, 3, H, W)` containing a batch of rgb images
+                corresponding to the source viewpoints from which features will be extracted
+            fg_probability: A tensor of shape `(B, 1, H, W)` containing a batch
+                of foreground masks with values in [0, 1].
+            depth_map: A tensor of shape `(B, 1, H, W)` containing a batch of depth maps.
+
+        Returns:
+            Modified image_rgb, fg_mask, depth_map
+        """
+        if image_rgb is not None and image_rgb.ndim == 3:
+            # The FrameData object is used for both frames and batches of frames,
+            # and a user might get this error if those were confused.
+            # Perhaps a user has a FrameData `fd` representing a single frame and
+            # wrote something like `model(**fd)` instead of
+            # `model(**fd.collate([fd]))`.
+            raise ValueError(
+                "Model received unbatched inputs. "
+                + "Perhaps they came from a FrameData which had not been collated."
+            )
+
+        fg_mask = fg_probability
+        if fg_mask is not None and self.mask_threshold > 0.0:
+            # threshold masks
+            warnings.warn("Thresholding masks!")
+            fg_mask = (fg_mask >= self.mask_threshold).type_as(fg_mask)
+
+        if self.mask_images and fg_mask is not None and image_rgb is not None:
+            # mask the image
+            warnings.warn("Masking images!")
+            image_rgb = image_utils.mask_background(
+                image_rgb, fg_mask, dim_color=1, bg_color=torch.tensor(self.bg_color)
+            )
+
+        if self.mask_depths and fg_mask is not None and depth_map is not None:
+            # mask the depths
+            assert (
+                self.mask_threshold > 0.0
+            ), "Depths should be masked only with thresholded masks"
+            warnings.warn("Masking depths!")
+            depth_map = depth_map * fg_mask
+
+        return image_rgb, fg_mask, depth_map
+
+
+def _apply_chunked(func, chunk_generator, tensor_collator):
+    """
+    Helper function to apply a function on a sequence of
+    chunked inputs yielded by a generator and collate
+    the result.
+    """
+    processed_chunks = [
+        func(*chunk_args, **chunk_kwargs)
+        for chunk_args, chunk_kwargs in chunk_generator
+    ]
+
+    return cat_dataclass(processed_chunks, tensor_collator)
+
+
+def _tensor_collator(batch, new_dims) -> torch.Tensor:
+    """
+    Helper function to reshape the batch to the desired shape
+    """
+    return torch.cat(batch, dim=1).reshape(*new_dims, -1)
+
+
+def _chunk_generator(
+    chunk_size: int,
+    ray_bundle: ImplicitronRayBundle,
+    chunked_inputs: Dict[str, torch.Tensor],
+    tqdm_trigger_threshold: int,
+    *args,
+    **kwargs,
+):
+    """
+    Helper function which yields chunks of rays from the
+    input ray_bundle, to be used when the number of rays is
+    large and will not fit in memory for rendering.
+    """
+    (
+        batch_size,
+        *spatial_dim,
+        n_pts_per_ray,
+    ) = ray_bundle.lengths.shape  # B x ... x n_pts_per_ray
+    if n_pts_per_ray > 0 and chunk_size % n_pts_per_ray != 0:
+        raise ValueError(
+            f"chunk_size_grid ({chunk_size}) should be divisible "
+            f"by n_pts_per_ray ({n_pts_per_ray})"
+        )
+
+    n_rays = prod(spatial_dim)
+    # special handling for raytracing-based methods
+    n_chunks = -(-n_rays * max(n_pts_per_ray, 1) // chunk_size)
+    chunk_size_in_rays = -(-n_rays // n_chunks)
+
+    iter = range(0, n_rays, chunk_size_in_rays)
+    if len(iter) >= tqdm_trigger_threshold:
+        iter = tqdm.tqdm(iter)
+
+    def _safe_slice(
+        tensor: Optional[torch.Tensor], start_idx: int, end_idx: int
+    ) -> Any:
+        return tensor[start_idx:end_idx] if tensor is not None else None
+
+    for start_idx in iter:
+        end_idx = min(start_idx + chunk_size_in_rays, n_rays)
+        ray_bundle_chunk = ImplicitronRayBundle(
+            origins=ray_bundle.origins.reshape(batch_size, -1, 3)[:, start_idx:end_idx],
+            directions=ray_bundle.directions.reshape(batch_size, -1, 3)[
+                :, start_idx:end_idx
+            ],
+            lengths=ray_bundle.lengths.reshape(batch_size, n_rays, n_pts_per_ray)[
+                :, start_idx:end_idx
+            ],
+            xys=ray_bundle.xys.reshape(batch_size, -1, 2)[:, start_idx:end_idx],
+            camera_ids=_safe_slice(ray_bundle.camera_ids, start_idx, end_idx),
+            camera_counts=_safe_slice(ray_bundle.camera_counts, start_idx, end_idx),
+        )
+        extra_args = kwargs.copy()
+        for k, v in chunked_inputs.items():
+            extra_args[k] = v.flatten(2)[:, :, start_idx:end_idx]
+        yield [ray_bundle_chunk, *args], extra_args

pytorch3d/implicitron/models/global_encoder/autodecoder.py

@@ -29,6 +29,8 @@ class Autodecoder(Configurable, torch.nn.Module):
     ignore_input: bool = False
 
     def __post_init__(self):
+        super().__init__()
+
         if self.n_instances <= 0:
             raise ValueError(f"Invalid n_instances {self.n_instances}")
 
@@ -69,7 +71,7 @@ class Autodecoder(Configurable, torch.nn.Module):
         return key_map
 
     def calculate_squared_encoding_norm(self) -> Optional[torch.Tensor]:
-        return (self._autodecoder_codes.weight**2).mean()
+        return (self._autodecoder_codes.weight**2).mean()  # pyre-ignore[16]
 
     def get_encoding_dim(self) -> int:
         return self.encoding_dim
@@ -93,6 +95,7 @@ class Autodecoder(Configurable, torch.nn.Module):
                 # pyre-fixme[9]: x has type `Union[List[str], LongTensor]`; used as
                 #  `Tensor`.
                 x = torch.tensor(
+                    # pyre-ignore[29]
                     [self._key_map[elem] for elem in x],
                     dtype=torch.long,
                     device=next(self.parameters()).device,
@@ -100,6 +103,7 @@ class Autodecoder(Configurable, torch.nn.Module):
             except StopIteration:
                 raise ValueError("Not enough n_instances in the autodecoder") from None
 
+        # pyre-fixme[29]: `Union[torch.Tensor, torch.nn.Module]` is not a function.
         return self._autodecoder_codes(x)
 
     def _load_key_map_hook(

pytorch3d/implicitron/models/global_encoder/global_encoder.py

@@ -26,6 +26,9 @@ class GlobalEncoderBase(ReplaceableBase):
     (`SequenceAutodecoder`).
     """
 
+    def __init__(self) -> None:
+        super().__init__()
+
     def get_encoding_dim(self):
         """
         Returns the dimensionality of the returned encoding.
@@ -66,6 +69,7 @@ class SequenceAutodecoder(GlobalEncoderBase, torch.nn.Module):  # pyre-ignore: 1
     autodecoder: Autodecoder
 
     def __post_init__(self):
+        super().__init__()
         run_auto_creation(self)
 
     def get_encoding_dim(self):
@@ -99,6 +103,7 @@ class HarmonicTimeEncoder(GlobalEncoderBase, torch.nn.Module):
     time_divisor: float = 1.0
 
     def __post_init__(self):
+        super().__init__()
         self._harmonic_embedding = HarmonicEmbedding(
             n_harmonic_functions=self.n_harmonic_functions,
             append_input=self.append_input,
@@ -119,7 +124,7 @@ class HarmonicTimeEncoder(GlobalEncoderBase, torch.nn.Module):
         if frame_timestamp.shape[-1] != 1:
             raise ValueError("Frame timestamp's last dimensions should be one.")
         time = frame_timestamp / self.time_divisor
-        return self._harmonic_embedding(time)
+        return self._harmonic_embedding(time)  # pyre-ignore: 29
 
     def calculate_squared_encoding_norm(self) -> Optional[torch.Tensor]:
         return None

pytorch3d/implicitron/models/implicit_function/base.py

@@ -14,6 +14,9 @@ from pytorch3d.renderer.cameras import CamerasBase
 
 
 class ImplicitFunctionBase(ABC, ReplaceableBase):
+    def __init__(self):
+        super().__init__()
+
     @abstractmethod
     def forward(
         self,

pytorch3d/implicitron/models/implicit_function/decoding_functions.py

@@ -45,6 +45,9 @@ class DecoderFunctionBase(ReplaceableBase, torch.nn.Module):
     space and transforms it into the required quantity (for example density and color).
     """
 
+    def __post_init__(self):
+        super().__init__()
+
     def forward(
         self, features: torch.Tensor, z: Optional[torch.Tensor] = None
     ) -> torch.Tensor:
@@ -80,6 +83,7 @@ class ElementwiseDecoder(DecoderFunctionBase):
     operation: DecoderActivation = DecoderActivation.IDENTITY
 
     def __post_init__(self):
+        super().__post_init__()
         if self.operation not in [
             DecoderActivation.RELU,
             DecoderActivation.SOFTPLUS,
@@ -159,6 +163,8 @@ class MLPWithInputSkips(Configurable, torch.nn.Module):
     use_xavier_init: bool = True
 
     def __post_init__(self):
+        super().__init__()
+
         try:
             last_activation = {
                 DecoderActivation.RELU: torch.nn.ReLU(True),
@@ -230,9 +236,14 @@ class MLPWithInputSkips(Configurable, torch.nn.Module):
             # if the skip tensor is None, we use `x` instead.
             z = x
         skipi = 0
+        # pyre-fixme[6]: For 1st param expected `Iterable[Variable[_T]]` but got
+        #  `Union[Tensor, Module]`.
         for li, layer in enumerate(self.mlp):
+            # pyre-fixme[58]: `in` is not supported for right operand type
+            #  `Union[torch._tensor.Tensor, torch.nn.modules.module.Module]`.
             if li in self._input_skips:
                 if self._skip_affine_trans:
+                    # pyre-fixme[29]: `Union[BoundMethod[typing.Callable(torch._C._Te...
                     y = self._apply_affine_layer(self.skip_affines[skipi], y, z)
                 else:
                     y = torch.cat((y, z), dim=-1)
@@ -273,6 +284,7 @@ class MLPDecoder(DecoderFunctionBase):
     network: MLPWithInputSkips
 
     def __post_init__(self):
+        super().__post_init__()
         run_auto_creation(self)
 
     def forward(

pytorch3d/implicitron/models/implicit_function/idr_feature_field.py

@@ -66,6 +66,8 @@ class IdrFeatureField(ImplicitFunctionBase, torch.nn.Module):
     encoding_dim: int = 0
 
     def __post_init__(self):
+        super().__init__()
+
         dims = [self.d_in] + list(self.dims) + [self.d_out + self.feature_vector_size]
 
         self.embed_fn = None
@@ -141,11 +143,16 @@ class IdrFeatureField(ImplicitFunctionBase, torch.nn.Module):
             self.embed_fn is None and fun_viewpool is None and global_code is None
         ):
             return torch.tensor(
-                [], device=rays_points_world.device, dtype=rays_points_world.dtype
+                [],
+                device=rays_points_world.device,
+                dtype=rays_points_world.dtype
+                # pyre-fixme[6]: For 2nd param expected `int` but got `Union[Module,
+                #  Tensor]`.
             ).view(0, self.out_dim)
 
         embeddings = []
         if self.embed_fn is not None:
+            # pyre-fixme[29]: `Union[torch.Tensor, torch.nn.Module]` is not a function.
             embeddings.append(self.embed_fn(rays_points_world))
 
         if fun_viewpool is not None:
@@ -164,13 +171,17 @@ class IdrFeatureField(ImplicitFunctionBase, torch.nn.Module):
 
         embedding = torch.cat(embeddings, dim=-1)
         x = embedding
+        # pyre-fixme[29]: `Union[BoundMethod[typing.Callable(torch._C._TensorBase.__s...
         for layer_idx in range(self.num_layers - 1):
             if layer_idx in self.skip_in:
                 x = torch.cat([x, embedding], dim=-1) / 2**0.5
 
+            # pyre-fixme[29]: `Union[torch.Tensor, torch.nn.Module]` is not a function.
             x = self.linear_layers[layer_idx](x)
 
+            # pyre-fixme[29]: `Union[BoundMethod[typing.Callable(torch._C._TensorBase...
             if layer_idx < self.num_layers - 2:
+                # pyre-fixme[29]: `Union[torch.Tensor, torch.nn.Module]` is not a function.
                 x = self.softplus(x)
 
         return x

pytorch3d/implicitron/models/implicit_function/neural_radiance_field.py

@@ -9,14 +9,11 @@ from typing import Optional, Tuple
 
 import torch
 from pytorch3d.common.linear_with_repeat import LinearWithRepeat
-from pytorch3d.implicitron.models.renderer.base import (
-    conical_frustum_to_gaussian,
-    ImplicitronRayBundle,
-)
+from pytorch3d.implicitron.models.renderer.base import ImplicitronRayBundle
 from pytorch3d.implicitron.tools.config import expand_args_fields, registry
+from pytorch3d.renderer import ray_bundle_to_ray_points
 from pytorch3d.renderer.cameras import CamerasBase
 from pytorch3d.renderer.implicit import HarmonicEmbedding
-from pytorch3d.renderer.implicit.utils import ray_bundle_to_ray_points
 
 from .base import ImplicitFunctionBase
 
@@ -39,7 +36,6 @@ class NeuralRadianceFieldBase(ImplicitFunctionBase, torch.nn.Module):
     input_xyz: bool = True
     xyz_ray_dir_in_camera_coords: bool = False
     color_dim: int = 3
-    use_integrated_positional_encoding: bool = False
     """
     Args:
         n_harmonic_functions_xyz: The number of harmonic functions
@@ -57,13 +53,10 @@ class NeuralRadianceFieldBase(ImplicitFunctionBase, torch.nn.Module):
         n_layers_xyz: The number of layers of the MLP that outputs the
             occupancy field.
         append_xyz: The list of indices of the skip layers of the occupancy MLP.
-        use_integrated_positional_encoding: If True, use integrated positional enoding
-            as defined in `MIP-NeRF <https://arxiv.org/abs/2103.13415>`_.
-            If False, use the classical harmonic embedding
-            defined in `NeRF <https://arxiv.org/abs/2003.08934>`_.
     """
 
     def __post_init__(self):
+        super().__init__()
         # The harmonic embedding layer converts input 3D coordinates
         # to a representation that is more suitable for
         # processing with a deep neural network.
@@ -121,8 +114,10 @@ class NeuralRadianceFieldBase(ImplicitFunctionBase, torch.nn.Module):
         # Normalize the ray_directions to unit l2 norm.
         rays_directions_normed = torch.nn.functional.normalize(rays_directions, dim=-1)
         # Obtain the harmonic embedding of the normalized ray directions.
+        # pyre-fixme[29]: `Union[torch.Tensor, torch.nn.Module]` is not a function.
         rays_embedding = self.harmonic_embedding_dir(rays_directions_normed)
 
+        # pyre-fixme[29]: `Union[torch.Tensor, torch.nn.Module]` is not a function.
         return self.color_layer((self.intermediate_linear(features), rays_embedding))
 
     @staticmethod
@@ -157,10 +152,6 @@ class NeuralRadianceFieldBase(ImplicitFunctionBase, torch.nn.Module):
                     containing the direction vectors of sampling rays in world coords.
                 lengths: A tensor of shape `(minibatch, ..., num_points_per_ray)`
                     containing the lengths at which the rays are sampled.
-                bins: An optional tensor of shape `(minibatch,..., num_points_per_ray + 1)`
-                    containing the bins at which the rays are sampled. In this case
-                    lengths is equal to the midpoints of bins.
-
             fun_viewpool: an optional callback with the signature
                     fun_fiewpool(points) -> pooled_features
                 where points is a [N_TGT x N x 3] tensor of world coords,
@@ -172,26 +163,16 @@ class NeuralRadianceFieldBase(ImplicitFunctionBase, torch.nn.Module):
                 denoting the opacitiy of each ray point.
             rays_colors: A tensor of shape `(minibatch, ..., num_points_per_ray, 3)`
                 denoting the color of each ray point.
-
-        Raises:
-            ValueError: If `use_integrated_positional_encoding` is True and
-                `ray_bundle.bins` is None.
         """
-        if self.use_integrated_positional_encoding and ray_bundle.bins is None:
-            raise ValueError(
-                "When use_integrated_positional_encoding is True, ray_bundle.bins must be set."
-                "Have you set to True `AbstractMaskRaySampler.use_bins_for_ray_sampling`?"
-            )
-
-        rays_points_world, diag_cov = (
-            conical_frustum_to_gaussian(ray_bundle)
-            if self.use_integrated_positional_encoding
-            else (ray_bundle_to_ray_points(ray_bundle), None)  # pyre-ignore
-        )
+        # We first convert the ray parametrizations to world
+        # coordinates with `ray_bundle_to_ray_points`.
+        # pyre-ignore[6]
+        rays_points_world = ray_bundle_to_ray_points(ray_bundle)
         # rays_points_world.shape = [minibatch x ... x pts_per_ray x 3]
 
         embeds = create_embeddings_for_implicit_function(
             xyz_world=rays_points_world,
+            # pyre-fixme[6]: Expected `Optional[typing.Callable[..., typing.Any]]`
             #  for 2nd param but got `Union[None, torch.Tensor, torch.nn.Module]`.
             xyz_embedding_function=self.harmonic_embedding_xyz
             if self.input_xyz
@@ -200,16 +181,17 @@ class NeuralRadianceFieldBase(ImplicitFunctionBase, torch.nn.Module):
             fun_viewpool=fun_viewpool,
             xyz_in_camera_coords=self.xyz_ray_dir_in_camera_coords,
             camera=camera,
-            diag_cov=diag_cov,
         )
 
         # embeds.shape = [minibatch x n_src x n_rays x n_pts x self.n_harmonic_functions*6+3]
+        # pyre-fixme[29]: `Union[torch.Tensor, torch.nn.Module]` is not a function.
         features = self.xyz_encoder(embeds)
         # features.shape = [minibatch x ... x self.n_hidden_neurons_xyz]
         # NNs operate on the flattenned rays; reshaping to the correct spatial size
         # TODO: maybe make the transformer work on non-flattened tensors to avoid this reshape
         features = features.reshape(*rays_points_world.shape[:-1], -1)
 
+        # pyre-fixme[29]: `Union[torch.Tensor, torch.nn.Module]` is not a function.
         raw_densities = self.density_layer(features)
         # raw_densities.shape = [minibatch x ... x 1] in [0-1]
 

pytorch3d/implicitron/models/implicit_function/scene_representation_networks.py

@@ -44,6 +44,7 @@ class SRNRaymarchFunction(Configurable, torch.nn.Module):
     raymarch_function: Any = None
 
     def __post_init__(self):
+        super().__init__()
         self._harmonic_embedding = HarmonicEmbedding(
             self.n_harmonic_functions, append_input=True
         )
@@ -101,6 +102,8 @@ class SRNRaymarchFunction(Configurable, torch.nn.Module):
 
         embeds = create_embeddings_for_implicit_function(
             xyz_world=rays_points_world,
+            # pyre-fixme[6]: Expected `Optional[typing.Callable[..., typing.Any]]`
+            #  for 2nd param but got `Union[torch.Tensor, torch.nn.Module]`.
             xyz_embedding_function=self._harmonic_embedding,
             global_code=global_code,
             fun_viewpool=fun_viewpool,
@@ -110,6 +113,7 @@ class SRNRaymarchFunction(Configurable, torch.nn.Module):
 
         # Before running the network, we have to resize embeds to ndims=3,
         # otherwise the SRN layers consume huge amounts of memory.
+        # pyre-fixme[29]: `Union[torch.Tensor, torch.nn.Module]` is not a function.
         raymarch_features = self._net(
             embeds.view(embeds.shape[0], -1, embeds.shape[-1])
         )
@@ -131,6 +135,7 @@ class SRNPixelGenerator(Configurable, torch.nn.Module):
     ray_dir_in_camera_coords: bool = False
 
     def __post_init__(self):
+        super().__init__()
         self._harmonic_embedding = HarmonicEmbedding(
             self.n_harmonic_functions, append_input=True
         )
@@ -164,7 +169,9 @@ class SRNPixelGenerator(Configurable, torch.nn.Module):
         # Normalize the ray_directions to unit l2 norm.
         rays_directions_normed = torch.nn.functional.normalize(rays_directions, dim=-1)
         # Obtain the harmonic embedding of the normalized ray directions.
+        # pyre-fixme[29]: `Union[torch.Tensor, torch.nn.Module]` is not a function.
         rays_embedding = self._harmonic_embedding(rays_directions_normed)
+        # pyre-fixme[29]: `Union[torch.Tensor, torch.nn.Module]` is not a function.
         return self._color_layer((features, rays_embedding))
 
     def forward(
@@ -193,6 +200,7 @@ class SRNPixelGenerator(Configurable, torch.nn.Module):
                 denoting the color of each ray point.
         """
         # raymarch_features.shape = [minibatch x ... x pts_per_ray x 3]
+        # pyre-fixme[29]: `Union[torch.Tensor, torch.nn.Module]` is not a function.
         features = self._net(raymarch_features)
         # features.shape = [minibatch x ... x self.n_hidden_units]
 
@@ -207,6 +215,7 @@ class SRNPixelGenerator(Configurable, torch.nn.Module):
         # NNs operate on the flattenned rays; reshaping to the correct spatial size
         features = features.reshape(*raymarch_features.shape[:-1], -1)
 
+        # pyre-fixme[29]: `Union[torch.Tensor, torch.nn.Module]` is not a function.
         raw_densities = self._density_layer(features)
 
         rays_colors = self._get_colors(features, directions)
@@ -240,6 +249,7 @@ class SRNRaymarchHyperNet(Configurable, torch.nn.Module):
     xyz_in_camera_coords: bool = False
 
     def __post_init__(self):
+        super().__init__()
         raymarch_input_embedding_dim = (
             HarmonicEmbedding.get_output_dim_static(
                 self.in_features,
@@ -267,6 +277,7 @@ class SRNRaymarchHyperNet(Configurable, torch.nn.Module):
         srn_raymarch_function.
         """
 
+        # pyre-fixme[29]: `Union[torch.Tensor, torch.nn.Module]` is not a function.
         net = self._hypernet(global_code)
 
         # use the hyper-net generated network to instantiate the raymarch module
@@ -302,6 +313,7 @@ class SRNRaymarchHyperNet(Configurable, torch.nn.Module):
         # across LSTM iterations for the same global_code.
         if self.cached_srn_raymarch_function is None:
             # generate the raymarching network from the hypernet
+            # pyre-fixme[16]: `SRNRaymarchHyperNet` has no attribute
             self.cached_srn_raymarch_function = self._run_hypernet(global_code)
         (srn_raymarch_function,) = cast(
             Tuple[SRNRaymarchFunction], self.cached_srn_raymarch_function
@@ -323,11 +335,13 @@ class SRNImplicitFunction(ImplicitFunctionBase, torch.nn.Module):
     pixel_generator: SRNPixelGenerator
 
     def __post_init__(self):
+        super().__init__()
         run_auto_creation(self)
 
     def create_raymarch_function(self) -> None:
         self.raymarch_function = SRNRaymarchFunction(
             latent_dim=self.latent_dim,
+            # pyre-ignore[32]
             **self.raymarch_function_args,
         )
 
@@ -379,12 +393,14 @@ class SRNHyperNetImplicitFunction(ImplicitFunctionBase, torch.nn.Module):
     pixel_generator: SRNPixelGenerator
 
     def __post_init__(self):
+        super().__init__()
         run_auto_creation(self)
 
     def create_hypernet(self) -> None:
         self.hypernet = SRNRaymarchHyperNet(
             latent_dim=self.latent_dim,
             latent_dim_hypernet=self.latent_dim_hypernet,
+            # pyre-ignore[32]
             **self.hypernet_args,
         )
 

pytorch3d/implicitron/models/implicit_function/utils.py

@@ -36,7 +36,6 @@ def create_embeddings_for_implicit_function(
     camera: Optional[CamerasBase],
     fun_viewpool: Optional[Callable],
     xyz_embedding_function: Optional[Callable],
-    diag_cov: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
 
     bs, *spatial_size, pts_per_ray, _ = xyz_world.shape
@@ -60,11 +59,11 @@ def create_embeddings_for_implicit_function(
             prod(spatial_size),
             pts_per_ray,
             0,
+            dtype=xyz_world.dtype,
+            device=xyz_world.device,
         )
     else:
-
-        embeds = xyz_embedding_function(ray_points_for_embed, diag_cov=diag_cov)
-        embeds = embeds.reshape(
+        embeds = xyz_embedding_function(ray_points_for_embed).reshape(
             bs,
             1,
             prod(spatial_size),

pytorch3d/implicitron/models/implicit_function/voxel_grid.py

@@ -81,6 +81,7 @@ class VoxelGridBase(ReplaceableBase, torch.nn.Module):
     )
 
     def __post_init__(self):
+        super().__init__()
         if 0 not in self.resolution_changes:
             raise ValueError("There has to be key `0` in `resolution_changes`.")
 
@@ -267,6 +268,7 @@ class VoxelGridBase(ReplaceableBase, torch.nn.Module):
                 for name, tensor in vars(grid_values_with_wanted_resolution).items()
             }
 
+        # pyre-ignore[29]
         return self.values_type(**params), True
 
     def get_resolution_change_epochs(self) -> Tuple[int, ...]:
@@ -855,6 +857,7 @@ class VoxelGridModule(Configurable, torch.nn.Module):
     param_groups: Dict[str, str] = field(default_factory=lambda: {})
 
     def __post_init__(self):
+        super().__init__()
         run_auto_creation(self)
         n_grids = 1  # Voxel grid objects are batched. We need only a single grid.
         shapes = self.voxel_grid.get_shapes(epoch=0)
@@ -880,6 +883,8 @@ class VoxelGridModule(Configurable, torch.nn.Module):
             torch.Tensor of shape (..., n_features)
         """
         locator = self._get_volume_locator()
+        # pyre-fixme[29]: `Union[torch._tensor.Tensor,
+        #  torch.nn.modules.module.Module]` is not a function.
         grid_values = self.voxel_grid.values_type(**self.params)
         # voxel grids operate with extra n_grids dimension, which we fix to one
         return self.voxel_grid.evaluate_world(points[None], grid_values, locator)[0]
@@ -893,6 +898,8 @@ class VoxelGridModule(Configurable, torch.nn.Module):
                 replace current parameters
         """
         if self.hold_voxel_grid_as_parameters:
+            # pyre-ignore [16]
+            # Nones are converted to empty tensors by Parameter()
             self.params = torch.nn.ParameterDict(
                 {
                     k: torch.nn.Parameter(val)
@@ -943,6 +950,7 @@ class VoxelGridModule(Configurable, torch.nn.Module):
         Returns:
             True if parameter change has happened else False.
         """
+        # pyre-ignore[29]
         grid_values = self.voxel_grid.values_type(**self.params)
         grid_values, change = self.voxel_grid.change_resolution(
             grid_values, epoch=epoch
@@ -990,16 +998,19 @@ class VoxelGridModule(Configurable, torch.nn.Module):
         """
         '''
         new_params = {}
+        # pyre-ignore[29]
         for name in self.params:
             key = prefix + "params." + name
             if key in state_dict:
                 new_params[name] = torch.zeros_like(state_dict[key])
+        # pyre-ignore[29]
         self.set_voxel_grid_parameters(self.voxel_grid.values_type(**new_params))
 
     def get_device(self) -> torch.device:
         """
         Returns torch.device on which module parameters are located
         """
+        # pyre-ignore[29]
         return next(val for val in self.params.values() if val is not None).device
 
     def crop_self(self, min_point: torch.Tensor, max_point: torch.Tensor) -> None:
@@ -1015,6 +1026,7 @@ class VoxelGridModule(Configurable, torch.nn.Module):
             nothing
         """
         locator = self._get_volume_locator()
+        # pyre-fixme[29]: `Union[torch._tensor.Tensor,
         #  torch.nn.modules.module.Module]` is not a function.
         old_grid_values = self.voxel_grid.values_type(**self.params)
         new_grid_values = self.voxel_grid.crop_world(
@@ -1023,6 +1035,7 @@ class VoxelGridModule(Configurable, torch.nn.Module):
         grid_values, _ = self.voxel_grid.change_resolution(
             new_grid_values, grid_values_with_wanted_resolution=old_grid_values
         )
+        # pyre-ignore [16]
         self.params = torch.nn.ParameterDict(
             {
                 k: torch.nn.Parameter(val)

pytorch3d/implicitron/models/implicit_function/voxel_grid_implicit_function.py

@@ -186,17 +186,23 @@ class VoxelGridImplicitFunction(ImplicitFunctionBase, torch.nn.Module):
     volume_cropping_epochs: Tuple[int, ...] = ()
 
     def __post_init__(self) -> None:
+        super().__init__()
         run_auto_creation(self)
+        # pyre-ignore[16]
         self.voxel_grid_scaffold = self._create_voxel_grid_scaffold()
+        # pyre-ignore[16]
         self.harmonic_embedder_xyz_density = HarmonicEmbedding(
             **self.harmonic_embedder_xyz_density_args
         )
+        # pyre-ignore[16]
         self.harmonic_embedder_xyz_color = HarmonicEmbedding(
             **self.harmonic_embedder_xyz_color_args
         )
+        # pyre-ignore[16]
         self.harmonic_embedder_dir_color = HarmonicEmbedding(
             **self.harmonic_embedder_dir_color_args
         )
+        # pyre-ignore[16]
         self._scaffold_ready = False
 
     def forward(
@@ -247,6 +253,7 @@ class VoxelGridImplicitFunction(ImplicitFunctionBase, torch.nn.Module):
         # ########## filter the points using the scaffold ########## #
         if self._scaffold_ready and self.scaffold_filter_points:
             with torch.no_grad():
+                # pyre-ignore[29]
                 non_empty_points = self.voxel_grid_scaffold(points)[..., 0] > 0
             points = points[non_empty_points]
             if len(points) == 0:
@@ -358,6 +365,7 @@ class VoxelGridImplicitFunction(ImplicitFunctionBase, torch.nn.Module):
                 feature dimensionality which `decoder_density` returns
         """
         embeds_density = self.voxel_grid_density(points)
+        # pyre-ignore[29]
         harmonic_embedding_density = self.harmonic_embedder_xyz_density(embeds_density)
         # shape = [..., density_dim]
         return self.decoder_density(harmonic_embedding_density)
@@ -400,11 +408,13 @@ class VoxelGridImplicitFunction(ImplicitFunctionBase, torch.nn.Module):
 
         # ########## embed with the harmonic function ########## #
         # Obtain the harmonic embedding of the voxel grid output.
+        # pyre-ignore[29]
         harmonic_embedding_color = self.harmonic_embedder_xyz_color(embeds_color)
 
         # Normalize the ray_directions to unit l2 norm.
         rays_directions_normed = torch.nn.functional.normalize(directions, dim=-1)
         # Obtain the harmonic embedding of the normalized ray directions.
+        # pyre-ignore[29]
         harmonic_embedding_dir = self.harmonic_embedder_dir_color(
             rays_directions_normed
         )
@@ -473,8 +483,10 @@ class VoxelGridImplicitFunction(ImplicitFunctionBase, torch.nn.Module):
             an object inside, else False.
         """
         # find bounding box
+        # pyre-ignore[16]
         points = self.voxel_grid_scaffold.get_grid_points(epoch=epoch)
         assert self._scaffold_ready, "Scaffold has to be calculated before cropping."
+        # pyre-ignore[29]
         occupancy = self.voxel_grid_scaffold(points)[..., 0] > 0
         non_zero_idxs = torch.nonzero(occupancy)
         if len(non_zero_idxs) == 0:
@@ -506,6 +518,7 @@ class VoxelGridImplicitFunction(ImplicitFunctionBase, torch.nn.Module):
         """
 
         planes = []
+        # pyre-ignore[16]
         points = self.voxel_grid_scaffold.get_grid_points(epoch=epoch)
 
         chunk_size = (
@@ -525,7 +538,9 @@ class VoxelGridImplicitFunction(ImplicitFunctionBase, torch.nn.Module):
             stride=1,
         )
         occupancy_cube = density_cube > self.scaffold_empty_space_threshold
+        # pyre-ignore[16]
         self.voxel_grid_scaffold.params["voxel_grid"] = occupancy_cube.float()
+        # pyre-ignore[16]
         self._scaffold_ready = True
 
         return False
@@ -542,6 +557,7 @@ class VoxelGridImplicitFunction(ImplicitFunctionBase, torch.nn.Module):
         decoding function to this value.
         """
         grid_args = self.voxel_grid_density_args
+        # pyre-ignore[6]
         grid_output_dim = VoxelGridModule.get_output_dim(grid_args)
 
         embedder_args = self.harmonic_embedder_xyz_density_args
@@ -570,6 +586,7 @@ class VoxelGridImplicitFunction(ImplicitFunctionBase, torch.nn.Module):
         decoding function to this value.
         """
         grid_args = self.voxel_grid_color_args
+        # pyre-ignore[6]
         grid_output_dim = VoxelGridModule.get_output_dim(grid_args)
 
         embedder_args = self.harmonic_embedder_xyz_color_args
@@ -603,7 +620,9 @@ class VoxelGridImplicitFunction(ImplicitFunctionBase, torch.nn.Module):
                     `self.voxel_grid_density`
         """
         return VoxelGridModule(
+            # pyre-ignore[29]
             extents=self.voxel_grid_density_args["extents"],
+            # pyre-ignore[29]
             translation=self.voxel_grid_density_args["translation"],
             voxel_grid_class_type="FullResolutionVoxelGrid",
             hold_voxel_grid_as_parameters=False,

pytorch3d/implicitron/models/metrics.py

@@ -25,6 +25,9 @@ class RegularizationMetricsBase(ReplaceableBase, torch.nn.Module):
     depend on the model's parameters.
     """
 
+    def __post_init__(self) -> None:
+        super().__init__()
+
     def forward(
         self, model: Any, keys_prefix: str = "loss_", **kwargs
     ) -> Dict[str, Any]:
@@ -53,6 +56,9 @@ class ViewMetricsBase(ReplaceableBase, torch.nn.Module):
     `forward()` method produces losses and other metrics.
     """
 
+    def __post_init__(self) -> None:
+        super().__init__()
+
     def forward(
         self,
         raymarched: RendererOutput,

pytorch3d/implicitron/models/model_dbir.py

@@ -41,8 +41,9 @@ class ModelDBIR(ImplicitronModelBase):
     bg_color: Tuple[float, float, float] = (0.0, 0.0, 0.0)
     max_points: int = -1
 
-    # pyre-fixme[14]: `forward` overrides method defined in `ImplicitronModelBase`
-    #  inconsistently.
+    def __post_init__(self):
+        super().__init__()
+
     def forward(
         self,
         *,  # force keyword-only arguments

pytorch3d/implicitron/models/overfit_model.py

@@ -1,664 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the BSD-style license found in the
-# LICENSE file in the root directory of this source tree.
-
-
-# Note: The #noqa comments below are for unused imports of pluggable implementations
-# which are part of implicitron. They ensure that the registry is prepopulated.
-
-import functools
-import logging
-from dataclasses import field
-from typing import Any, Callable, Dict, List, Optional, Tuple, TYPE_CHECKING, Union
-
-import torch
-from omegaconf import DictConfig
-
-from pytorch3d.implicitron.models.base_model import (
-    ImplicitronModelBase,
-    ImplicitronRender,
-)
-from pytorch3d.implicitron.models.global_encoder.global_encoder import GlobalEncoderBase
-from pytorch3d.implicitron.models.implicit_function.base import ImplicitFunctionBase
-from pytorch3d.implicitron.models.metrics import (
-    RegularizationMetricsBase,
-    ViewMetricsBase,
-)
-
-from pytorch3d.implicitron.models.renderer.base import (
-    BaseRenderer,
-    EvaluationMode,
-    ImplicitronRayBundle,
-    RendererOutput,
-    RenderSamplingMode,
-)
-from pytorch3d.implicitron.models.renderer.ray_sampler import RaySamplerBase
-from pytorch3d.implicitron.models.utils import (
-    apply_chunked,
-    chunk_generator,
-    log_loss_weights,
-    preprocess_input,
-    weighted_sum_losses,
-)
-from pytorch3d.implicitron.tools import vis_utils
-from pytorch3d.implicitron.tools.config import (
-    expand_args_fields,
-    registry,
-    run_auto_creation,
-)
-
-from pytorch3d.implicitron.tools.rasterize_mc import rasterize_sparse_ray_bundle
-from pytorch3d.renderer import utils as rend_utils
-from pytorch3d.renderer.cameras import CamerasBase
-
-
-if TYPE_CHECKING:
-    from visdom import Visdom
-logger = logging.getLogger(__name__)
-
-IMPLICIT_FUNCTION_ARGS_TO_REMOVE: List[str] = [
-    "feature_vector_size",
-    "encoding_dim",
-    "latent_dim",
-    "color_dim",
-]
-
-
-@registry.register
-class OverfitModel(ImplicitronModelBase):  # pyre-ignore: 13
-    """
-    OverfitModel is a wrapper for the neural implicit
-    rendering and reconstruction pipeline which consists
-    of the following sequence of 4 steps:
-
-
-        (1) Ray Sampling
-        ------------------
-        Rays are sampled from an image grid based on the target view(s).
-                │
-                ▼
-        (2) Implicit Function Evaluation
-        ------------------
-        Evaluate the implicit function(s) at the sampled ray points
-        (also optionally pass in a global encoding from global_encoder).
-                │
-                ▼
-        (3) Rendering
-        ------------------
-        Render the image into the target cameras by raymarching along
-        the sampled rays and aggregating the colors and densities
-        output by the implicit function in (2).
-                │
-                ▼
-        (4) Loss Computation
-        ------------------
-        Compute losses based on the predicted target image(s).
-
-
-    The `forward` function of OverfitModel executes
-    this sequence of steps. Currently, steps 1, 2, 3
-    can be customized by intializing a subclass of the appropriate
-    base class and adding the newly created module to the registry.
-    Please see https://github.com/facebookresearch/pytorch3d/blob/main/projects/implicitron_trainer/README.md#custom-plugins
-    for more details on how to create and register a custom component.
-
-    In the config .yaml files for experiments, the parameters below are
-    contained in the
-    `model_factory_ImplicitronModelFactory_args.model_OverfitModel_args`
-    node. As OverfitModel derives from ReplaceableBase, the input arguments are
-    parsed by the run_auto_creation function to initialize the
-    necessary member modules. Please see implicitron_trainer/README.md
-    for more details on this process.
-
-    Args:
-        mask_images: Whether or not to mask the RGB image background given the
-            foreground mask (the `fg_probability` argument of `GenericModel.forward`)
-        mask_depths: Whether or not to mask the depth image background given the
-            foreground mask (the `fg_probability` argument of `GenericModel.forward`)
-        render_image_width: Width of the output image to render
-        render_image_height: Height of the output image to render
-        mask_threshold: If greater than 0.0, the foreground mask is
-            thresholded by this value before being applied to the RGB/Depth images
-        output_rasterized_mc: If True, visualize the Monte-Carlo pixel renders by
-            splatting onto an image grid. Default: False.
-        bg_color: RGB values for setting the background color of input image
-            if mask_images=True. Defaults to (0.0, 0.0, 0.0). Each renderer has its own
-            way to determine the background color of its output, unrelated to this.
-        chunk_size_grid: The total number of points which can be rendered
-            per chunk. This is used to compute the number of rays used
-            per chunk when the chunked version of the renderer is used (in order
-            to fit rendering on all rays in memory)
-        render_features_dimensions: The number of output features to render.
-            Defaults to 3, corresponding to RGB images.
-        sampling_mode_training: The sampling method to use during training. Must be
-            a value from the RenderSamplingMode Enum.
-        sampling_mode_evaluation: Same as above but for evaluation.
-        global_encoder_class_type: The name of the class to use for global_encoder,
-            which must be available in the registry. Or `None` to disable global encoder.
-        global_encoder: An instance of `GlobalEncoder`. This is used to generate an encoding
-            of the image (referred to as the global_code) that can be used to model aspects of
-            the scene such as multiple objects or morphing objects. It is up to the implicit
-            function definition how to use it, but the most typical way is to broadcast and
-            concatenate to the other inputs for the implicit function.
-        raysampler_class_type: The name of the raysampler class which is available
-            in the global registry.
-        raysampler: An instance of RaySampler which is used to emit
-            rays from the target view(s).
-        renderer_class_type: The name of the renderer class which is available in the global
-            registry.
-        renderer: A renderer class which inherits from BaseRenderer. This is used to
-            generate the images from the target view(s).
-        share_implicit_function_across_passes: If set to True
-            coarse_implicit_function is automatically set as implicit_function
-            (coarse_implicit_function=implicit_funciton). The
-            implicit_functions are then run sequentially during the rendering.
-        implicit_function_class_type: The type of implicit function to use which
-            is available in the global registry.
-        implicit_function: An instance of ImplicitFunctionBase.
-        coarse_implicit_function_class_type: The type of implicit function to use which
-            is available in the global registry.
-        coarse_implicit_function: An instance of ImplicitFunctionBase.
-            If set and `share_implicit_function_across_passes` is set to False,
-            coarse_implicit_function is instantiated on itself. It
-            is then used as the second pass during the rendering.
-            If set to None, we only do a single pass with implicit_function.
-        view_metrics: An instance of ViewMetricsBase used to compute loss terms which
-            are independent of the model's parameters.
-        view_metrics_class_type: The type of view metrics to use, must be available in
-            the global registry.
-        regularization_metrics: An instance of RegularizationMetricsBase used to compute
-            regularization terms which can depend on the model's parameters.
-        regularization_metrics_class_type: The type of regularization metrics to use,
-            must be available in the global registry.
-        loss_weights: A dictionary with a {loss_name: weight} mapping; see documentation
-            for `ViewMetrics` class for available loss functions.
-        log_vars: A list of variable names which should be logged.
-            The names should correspond to a subset of the keys of the
-            dict `preds` output by the `forward` function.
-    """  # noqa: B950
-
-    mask_images: bool = True
-    mask_depths: bool = True
-    render_image_width: int = 400
-    render_image_height: int = 400
-    mask_threshold: float = 0.5
-    output_rasterized_mc: bool = False
-    bg_color: Tuple[float, float, float] = (0.0, 0.0, 0.0)
-    chunk_size_grid: int = 4096
-    render_features_dimensions: int = 3
-    tqdm_trigger_threshold: int = 16
-
-    n_train_target_views: int = 1
-    sampling_mode_training: str = "mask_sample"
-    sampling_mode_evaluation: str = "full_grid"
-
-    # ---- global encoder settings
-    global_encoder_class_type: Optional[str] = None
-    global_encoder: Optional[GlobalEncoderBase]
-
-    # ---- raysampler
-    raysampler_class_type: str = "AdaptiveRaySampler"
-    raysampler: RaySamplerBase
-
-    # ---- renderer configs
-    renderer_class_type: str = "MultiPassEmissionAbsorptionRenderer"
-    renderer: BaseRenderer
-
-    # ---- implicit function settings
-    share_implicit_function_across_passes: bool = False
-    implicit_function_class_type: str = "NeuralRadianceFieldImplicitFunction"
-    implicit_function: ImplicitFunctionBase
-    coarse_implicit_function_class_type: Optional[str] = None
-    coarse_implicit_function: Optional[ImplicitFunctionBase]
-
-    # ----- metrics
-    view_metrics: ViewMetricsBase
-    view_metrics_class_type: str = "ViewMetrics"
-
-    regularization_metrics: RegularizationMetricsBase
-    regularization_metrics_class_type: str = "RegularizationMetrics"
-
-    # ---- loss weights
-    loss_weights: Dict[str, float] = field(
-        default_factory=lambda: {
-            "loss_rgb_mse": 1.0,
-            "loss_prev_stage_rgb_mse": 1.0,
-            "loss_mask_bce": 0.0,
-            "loss_prev_stage_mask_bce": 0.0,
-        }
-    )
-
-    # ---- variables to be logged (logger automatically ignores if not computed)
-    log_vars: List[str] = field(
-        default_factory=lambda: [
-            "loss_rgb_psnr_fg",
-            "loss_rgb_psnr",
-            "loss_rgb_mse",
-            "loss_rgb_huber",
-            "loss_depth_abs",
-            "loss_depth_abs_fg",
-            "loss_mask_neg_iou",
-            "loss_mask_bce",
-            "loss_mask_beta_prior",
-            "loss_eikonal",
-            "loss_density_tv",
-            "loss_depth_neg_penalty",
-            "loss_autodecoder_norm",
-            # metrics that are only logged in 2+stage renderes
-            "loss_prev_stage_rgb_mse",
-            "loss_prev_stage_rgb_psnr_fg",
-            "loss_prev_stage_rgb_psnr",
-            "loss_prev_stage_mask_bce",
-            # basic metrics
-            "objective",
-            "epoch",
-            "sec/it",
-        ]
-    )
-
-    @classmethod
-    def pre_expand(cls) -> None:
-        # use try/finally to bypass cinder's lazy imports
-        try:
-            from pytorch3d.implicitron.models.implicit_function.idr_feature_field import (  # noqa: F401, B950
-                IdrFeatureField,
-            )
-            from pytorch3d.implicitron.models.implicit_function.neural_radiance_field import (  # noqa: F401, B950
-                NeuralRadianceFieldImplicitFunction,
-            )
-            from pytorch3d.implicitron.models.implicit_function.scene_representation_networks import (  # noqa: F401, B950
-                SRNImplicitFunction,
-            )
-            from pytorch3d.implicitron.models.renderer.lstm_renderer import (  # noqa: F401
-                LSTMRenderer,
-            )
-            from pytorch3d.implicitron.models.renderer.multipass_ea import (  # noqa: F401
-                MultiPassEmissionAbsorptionRenderer,
-            )
-            from pytorch3d.implicitron.models.renderer.sdf_renderer import (  # noqa: F401
-                SignedDistanceFunctionRenderer,
-            )
-        finally:
-            pass
-
-    def __post_init__(self):
-        # The attribute will be filled by run_auto_creation
-        run_auto_creation(self)
-        log_loss_weights(self.loss_weights, logger)
-        # We need to set it here since run_auto_creation
-        # will create coarse_implicit_function before implicit_function
-        if self.share_implicit_function_across_passes:
-            self.coarse_implicit_function = self.implicit_function
-
-    def forward(
-        self,
-        *,  # force keyword-only arguments
-        image_rgb: Optional[torch.Tensor],
-        camera: CamerasBase,
-        fg_probability: Optional[torch.Tensor] = None,
-        mask_crop: Optional[torch.Tensor] = None,
-        depth_map: Optional[torch.Tensor] = None,
-        sequence_name: Optional[List[str]] = None,
-        frame_timestamp: Optional[torch.Tensor] = None,
-        evaluation_mode: EvaluationMode = EvaluationMode.EVALUATION,
-        **kwargs,
-    ) -> Dict[str, Any]:
-        """
-        Args:
-            image_rgb: A tensor of shape `(B, 3, H, W)` containing a batch of rgb images;
-                the first `min(B, n_train_target_views)` images are considered targets and
-                are used to supervise the renders; the rest corresponding to the source
-                viewpoints from which features will be extracted.
-            camera: An instance of CamerasBase containing a batch of `B` cameras corresponding
-                to the viewpoints of target images, from which the rays will be sampled,
-                and source images, which will be used for intersecting with target rays.
-            fg_probability: A tensor of shape `(B, 1, H, W)` containing a batch of
-                foreground masks.
-            mask_crop: A binary tensor of shape `(B, 1, H, W)` deonting valid
-                regions in the input images (i.e. regions that do not correspond
-                to, e.g., zero-padding). When the `RaySampler`'s sampling mode is set to
-                "mask_sample", rays  will be sampled in the non zero regions.
-            depth_map: A tensor of shape `(B, 1, H, W)` containing a batch of depth maps.
-            sequence_name: A list of `B` strings corresponding to the sequence names
-                from which images `image_rgb` were extracted. They are used to match
-                target frames with relevant source frames.
-            frame_timestamp: Optionally a tensor of shape `(B,)` containing a batch
-                of frame timestamps.
-            evaluation_mode: one of EvaluationMode.TRAINING or
-                EvaluationMode.EVALUATION which determines the settings used for
-                rendering.
-
-        Returns:
-            preds: A dictionary containing all outputs of the forward pass including the
-                rendered images, depths, masks, losses and other metrics.
-        """
-        image_rgb, fg_probability, depth_map = preprocess_input(
-            image_rgb,
-            fg_probability,
-            depth_map,
-            self.mask_images,
-            self.mask_depths,
-            self.mask_threshold,
-            self.bg_color,
-        )
-
-        # Determine the used ray sampling mode.
-        sampling_mode = RenderSamplingMode(
-            self.sampling_mode_training
-            if evaluation_mode == EvaluationMode.TRAINING
-            else self.sampling_mode_evaluation
-        )
-
-        # (1) Sample rendering rays with the ray sampler.
-        # pyre-ignore[29]
-        ray_bundle: ImplicitronRayBundle = self.raysampler(
-            camera,
-            evaluation_mode,
-            mask=mask_crop
-            if mask_crop is not None and sampling_mode == RenderSamplingMode.MASK_SAMPLE
-            else None,
-        )
-
-        inputs_to_be_chunked = {}
-        if fg_probability is not None and self.renderer.requires_object_mask():
-            sampled_fb_prob = rend_utils.ndc_grid_sample(
-                fg_probability, ray_bundle.xys, mode="nearest"
-            )
-            inputs_to_be_chunked["object_mask"] = sampled_fb_prob > 0.5
-
-        # (2)-(3) Implicit function evaluation and Rendering
-        implicit_functions: List[Union[Callable, ImplicitFunctionBase]] = [
-            self.implicit_function
-        ]
-        if self.coarse_implicit_function is not None:
-            implicit_functions = [self.coarse_implicit_function, self.implicit_function]
-
-        if self.global_encoder is not None:
-            global_code = self.global_encoder(  # pyre-fixme[29]
-                sequence_name=sequence_name,
-                frame_timestamp=frame_timestamp,
-            )
-            implicit_functions = [
-                functools.partial(implicit_function, global_code=global_code)
-                if isinstance(implicit_function, Callable)
-                else functools.partial(
-                    implicit_function.forward, global_code=global_code
-                )
-                for implicit_function in implicit_functions
-            ]
-        rendered = self._render(
-            ray_bundle=ray_bundle,
-            sampling_mode=sampling_mode,
-            evaluation_mode=evaluation_mode,
-            implicit_functions=implicit_functions,
-            inputs_to_be_chunked=inputs_to_be_chunked,
-        )
-
-        # A dict to store losses as well as rendering results.
-        preds: Dict[str, Any] = self.view_metrics(
-            results={},
-            raymarched=rendered,
-            ray_bundle=ray_bundle,
-            image_rgb=image_rgb,
-            depth_map=depth_map,
-            fg_probability=fg_probability,
-            mask_crop=mask_crop,
-        )
-
-        preds.update(
-            self.regularization_metrics(
-                results=preds,
-                model=self,
-            )
-        )
-
-        if sampling_mode == RenderSamplingMode.MASK_SAMPLE:
-            if self.output_rasterized_mc:
-                # Visualize the monte-carlo pixel renders by splatting onto
-                # an image grid.
-                (
-                    preds["images_render"],
-                    preds["depths_render"],
-                    preds["masks_render"],
-                ) = rasterize_sparse_ray_bundle(
-                    ray_bundle,
-                    rendered.features,
-                    (self.render_image_height, self.render_image_width),
-                    rendered.depths,
-                    masks=rendered.masks,
-                )
-        elif sampling_mode == RenderSamplingMode.FULL_GRID:
-            preds["images_render"] = rendered.features.permute(0, 3, 1, 2)
-            preds["depths_render"] = rendered.depths.permute(0, 3, 1, 2)
-            preds["masks_render"] = rendered.masks.permute(0, 3, 1, 2)
-
-            preds["implicitron_render"] = ImplicitronRender(
-                image_render=preds["images_render"],
-                depth_render=preds["depths_render"],
-                mask_render=preds["masks_render"],
-            )
-        else:
-            raise AssertionError("Unreachable state")
-
-        # (4) Compute losses
-        # finally get the optimization objective using self.loss_weights
-        objective = self._get_objective(preds)
-        if objective is not None:
-            preds["objective"] = objective
-
-        return preds
-
-    def _get_objective(self, preds: Dict[str, torch.Tensor]) -> Optional[torch.Tensor]:
-        """
-        A helper function to compute the overall loss as the dot product
-        of individual loss functions with the corresponding weights.
-        """
-        return weighted_sum_losses(preds, self.loss_weights)
-
-    def visualize(
-        self,
-        viz: Optional["Visdom"],
-        visdom_env_imgs: str,
-        preds: Dict[str, Any],
-        prefix: str,
-    ) -> None:
-        """
-        Helper function to visualize the predictions generated
-        in the forward pass.
-
-        Args:
-            viz: Visdom connection object
-            visdom_env_imgs: name of visdom environment for the images.
-            preds: predictions dict like returned by forward()
-            prefix: prepended to the names of images
-        """
-        if viz is None or not viz.check_connection():
-            logger.info("no visdom server! -> skipping batch vis")
-            return
-
-        idx_image = 0
-        title = f"{prefix}_im{idx_image}"
-
-        vis_utils.visualize_basics(viz, preds, visdom_env_imgs, title=title)
-
-    def _render(
-        self,
-        *,
-        ray_bundle: ImplicitronRayBundle,
-        inputs_to_be_chunked: Dict[str, torch.Tensor],
-        sampling_mode: RenderSamplingMode,
-        **kwargs,
-    ) -> RendererOutput:
-        """
-        Args:
-            ray_bundle: A `ImplicitronRayBundle` object containing the parametrizations of the
-                sampled rendering rays.
-            inputs_to_be_chunked: A collection of tensor of shape `(B, _, H, W)`. E.g.
-                SignedDistanceFunctionRenderer requires "object_mask", shape
-                (B, 1, H, W), the silhouette of the object in the image. When
-                chunking, they are passed to the renderer as shape
-                `(B, _, chunksize)`.
-            sampling_mode: The sampling method to use. Must be a value from the
-                RenderSamplingMode Enum.
-
-        Returns:
-            An instance of RendererOutput
-        """
-        if sampling_mode == RenderSamplingMode.FULL_GRID and self.chunk_size_grid > 0:
-            return apply_chunked(
-                self.renderer,
-                chunk_generator(
-                    self.chunk_size_grid,
-                    ray_bundle,
-                    inputs_to_be_chunked,
-                    self.tqdm_trigger_threshold,
-                    **kwargs,
-                ),
-                lambda batch: torch.cat(batch, dim=1).reshape(
-                    *ray_bundle.lengths.shape[:-1], -1
-                ),
-            )
-        else:
-            # pyre-fixme[29]: `BaseRenderer` is not a function.
-            return self.renderer(
-                ray_bundle=ray_bundle,
-                **inputs_to_be_chunked,
-                **kwargs,
-            )
-
-    @classmethod
-    def raysampler_tweak_args(cls, type, args: DictConfig) -> None:
-        """
-        We don't expose certain fields of the raysampler because we want to set
-        them from our own members.
-        """
-        del args["sampling_mode_training"]
-        del args["sampling_mode_evaluation"]
-        del args["image_width"]
-        del args["image_height"]
-
-    def create_raysampler(self):
-        extra_args = {
-            "sampling_mode_training": self.sampling_mode_training,
-            "sampling_mode_evaluation": self.sampling_mode_evaluation,
-            "image_width": self.render_image_width,
-            "image_height": self.render_image_height,
-        }
-        raysampler_args = getattr(
-            self, "raysampler_" + self.raysampler_class_type + "_args"
-        )
-        self.raysampler = registry.get(RaySamplerBase, self.raysampler_class_type)(
-            **raysampler_args, **extra_args
-        )
-
-    @classmethod
-    def renderer_tweak_args(cls, type, args: DictConfig) -> None:
-        """
-        We don't expose certain fields of the renderer because we want to set
-        them based on other inputs.
-        """
-        args.pop("render_features_dimensions", None)
-        args.pop("object_bounding_sphere", None)
-
-    def create_renderer(self):
-        extra_args = {}
-
-        if self.renderer_class_type == "SignedDistanceFunctionRenderer":
-            extra_args["render_features_dimensions"] = self.render_features_dimensions
-            if not hasattr(self.raysampler, "scene_extent"):
-                raise ValueError(
-                    "SignedDistanceFunctionRenderer requires"
-                    + " a raysampler that defines the 'scene_extent' field"
-                    + " (this field is supported by, e.g., the adaptive raysampler - "
-                    + " self.raysampler_class_type='AdaptiveRaySampler')."
-                )
-            extra_args["object_bounding_sphere"] = self.raysampler.scene_extent
-
-        renderer_args = getattr(self, "renderer_" + self.renderer_class_type + "_args")
-        self.renderer = registry.get(BaseRenderer, self.renderer_class_type)(
-            **renderer_args, **extra_args
-        )
-
-    @classmethod
-    def implicit_function_tweak_args(cls, type, args: DictConfig) -> None:
-        """
-        We don't expose certain implicit_function fields because we want to set
-        them based on other inputs.
-        """
-        for arg in IMPLICIT_FUNCTION_ARGS_TO_REMOVE:
-            args.pop(arg, None)
-
-    @classmethod
-    def coarse_implicit_function_tweak_args(cls, type, args: DictConfig) -> None:
-        """
-        We don't expose certain implicit_function fields because we want to set
-        them based on other inputs.
-        """
-        for arg in IMPLICIT_FUNCTION_ARGS_TO_REMOVE:
-            args.pop(arg, None)
-
-    def _create_extra_args_for_implicit_function(self) -> Dict[str, Any]:
-        extra_args = {}
-        global_encoder_dim = (
-            0 if self.global_encoder is None else self.global_encoder.get_encoding_dim()
-        )
-        if self.implicit_function_class_type in (
-            "NeuralRadianceFieldImplicitFunction",
-            "NeRFormerImplicitFunction",
-        ):
-            extra_args["latent_dim"] = global_encoder_dim
-            extra_args["color_dim"] = self.render_features_dimensions
-
-        if self.implicit_function_class_type == "IdrFeatureField":
-            extra_args["feature_work_size"] = global_encoder_dim
-            extra_args["feature_vector_size"] = self.render_features_dimensions
-
-        if self.implicit_function_class_type == "SRNImplicitFunction":
-            extra_args["latent_dim"] = global_encoder_dim
-        return extra_args
-
-    def create_implicit_function(self) -> None:
-        implicit_function_type = registry.get(
-            ImplicitFunctionBase, self.implicit_function_class_type
-        )
-        expand_args_fields(implicit_function_type)
-
-        config_name = f"implicit_function_{self.implicit_function_class_type}_args"
-        config = getattr(self, config_name, None)
-        if config is None:
-            raise ValueError(f"{config_name} not present")
-
-        extra_args = self._create_extra_args_for_implicit_function()
-        self.implicit_function = implicit_function_type(**config, **extra_args)
-
-    def create_coarse_implicit_function(self) -> None:
-        # If coarse_implicit_function_class_type has been defined
-        # then we init a module based on its arguments
-        if (
-            self.coarse_implicit_function_class_type is not None
-            and not self.share_implicit_function_across_passes
-        ):
-            config_name = "coarse_implicit_function_{0}_args".format(
-                self.coarse_implicit_function_class_type
-            )
-            config = getattr(self, config_name, {})
-
-            implicit_function_type = registry.get(
-                ImplicitFunctionBase,
-                # pyre-ignore: config is None allow to check if this is None.
-                self.coarse_implicit_function_class_type,
-            )
-            expand_args_fields(implicit_function_type)
-
-            extra_args = self._create_extra_args_for_implicit_function()
-            self.coarse_implicit_function = implicit_function_type(
-                **config, **extra_args
-            )
-        elif self.share_implicit_function_across_passes:
-            # Since coarse_implicit_function is initialised before
-            # implicit_function we handle this case in the post_init.
-            pass
-        else:
-            self.coarse_implicit_function = None

pytorch3d/implicitron/models/renderer/base.py

@@ -6,6 +6,8 @@
 
 from __future__ import annotations
 
+import dataclasses
+
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
 from enum import Enum
@@ -14,7 +16,6 @@ from typing import Any, Dict, List, Optional, Tuple
 import torch
 from pytorch3d.implicitron.tools.config import ReplaceableBase
 from pytorch3d.ops import packed_to_padded
-from pytorch3d.renderer.implicit.utils import ray_bundle_variables_to_ray_points
 
 
 class EvaluationMode(Enum):
@@ -27,6 +28,7 @@ class RenderSamplingMode(Enum):
     FULL_GRID = "full_grid"
 
 
+@dataclasses.dataclass
 class ImplicitronRayBundle:
     """
     Parametrizes points along projection rays by storing ray `origins`,
@@ -45,79 +47,14 @@ class ImplicitronRayBundle:
         camera_counts: A tensor of shape (N, ) which how many times the
             coresponding camera in `camera_ids` was sampled.
             `sum(camera_counts) == minibatch`, where `minibatch = origins.shape[0]`.
-
-    Attributes:
-        origins: A tensor of shape `(..., 3)` denoting the
-            origins of the sampling rays in world coords.
-        directions: A tensor of shape `(..., 3)` containing the direction
-            vectors of sampling rays in world coords. They don't have to be normalized;
-            they define unit vectors in the respective 1D coordinate systems; see
-            documentation for :func:`ray_bundle_to_ray_points` for the conversion formula.
-        lengths: A tensor of shape `(..., num_points_per_ray)`
-            containing the lengths at which the rays are sampled.
-        xys: A tensor of shape `(..., 2)`, the xy-locations (`xys`) of the ray pixels
-        camera_ids: An optional tensor of shape (N, ) which indicates which camera
-            was used to sample the rays. `N` is the number of unique sampled cameras.
-        camera_counts: An optional tensor of shape (N, ) indicates how many times the
-            coresponding camera in `camera_ids` was sampled.
-            `sum(camera_counts)==total_number_of_rays`.
-        bins: An optional tensor of shape `(..., num_points_per_ray + 1)`
-            containing the bins at which the rays are sampled. In this case
-            lengths should be equal to the midpoints of bins `(..., num_points_per_ray)`.
-        pixel_radii_2d: An optional tensor of shape `(..., 1)`
-            base radii of the conical frustums.
-
-    Raises:
-        ValueError: If either bins or lengths are not provided.
-        ValueError: If bins is provided and the last dim is inferior or equal to 1.
     """
 
-    def __init__(
-        self,
-        origins: torch.Tensor,
-        directions: torch.Tensor,
-        lengths: Optional[torch.Tensor],
-        xys: torch.Tensor,
-        camera_ids: Optional[torch.LongTensor] = None,
-        camera_counts: Optional[torch.LongTensor] = None,
-        bins: Optional[torch.Tensor] = None,
-        pixel_radii_2d: Optional[torch.Tensor] = None,
-    ):
-        if bins is not None and bins.shape[-1] <= 1:
-            raise ValueError(
-                "The last dim of bins must be at least superior or equal to 2."
-            )
-
-        if bins is None and lengths is None:
-            raise ValueError(
-                "Please set either bins or lengths to initialize an ImplicitronRayBundle."
-            )
-
-        self.origins = origins
-        self.directions = directions
-        self._lengths = lengths if bins is None else None
-        self.xys = xys
-        self.bins = bins
-        self.pixel_radii_2d = pixel_radii_2d
-        self.camera_ids = camera_ids
-        self.camera_counts = camera_counts
-
-    @property
-    def lengths(self) -> torch.Tensor:
-        if self.bins is not None:
-            # equivalent to: 0.5 * (bins[..., 1:] + bins[..., :-1]) but more efficient
-            # pyre-ignore
-            return torch.lerp(self.bins[..., :-1], self.bins[..., 1:], 0.5)
-        return self._lengths
-
-    @lengths.setter
-    def lengths(self, value):
-        if self.bins is not None:
-            raise ValueError(
-                "If the bins attribute is not None you cannot set the lengths attribute."
-            )
-        else:
-            self._lengths = value
+    origins: torch.Tensor
+    directions: torch.Tensor
+    lengths: torch.Tensor
+    xys: torch.Tensor
+    camera_ids: Optional[torch.LongTensor] = None
+    camera_counts: Optional[torch.LongTensor] = None
 
     def is_packed(self) -> bool:
         """
@@ -204,6 +141,9 @@ class BaseRenderer(ABC, ReplaceableBase):
     Base class for all Renderer implementations.
     """
 
+    def __init__(self) -> None:
+        super().__init__()
+
     def requires_object_mask(self) -> bool:
         """
         Whether `forward` needs the object_mask.
@@ -258,154 +198,3 @@ class BaseRenderer(ABC, ReplaceableBase):
             instance of RendererOutput
         """
         pass
-
-
-def compute_3d_diagonal_covariance_gaussian(
-    rays_directions: torch.Tensor,
-    rays_dir_variance: torch.Tensor,
-    radii_variance: torch.Tensor,
-    eps: float = 1e-6,
-) -> torch.Tensor:
-    """
-    Transform the variances (rays_dir_variance, radii_variance) of the gaussians from
-    the coordinate frame of the conical frustum to 3D world coordinates.
-
-    It follows the equation 16 of `MIP-NeRF <https://arxiv.org/abs/2103.13415>`_
-
-    Args:
-        rays_directions: A tensor of shape `(..., 3)`
-        rays_dir_variance: A tensor of shape `(..., num_intervals)` representing
-            the variance of the conical frustum  with respect to the rays direction.
-        radii_variance: A tensor of shape `(..., num_intervals)` representing
-            the variance of the conical frustum with respect to its radius.
-        eps: a small number to prevent division by zero.
-
-    Returns:
-        A tensor of shape `(..., num_intervals, 3)` containing the diagonal
-            of the covariance matrix.
-    """
-    d_outer_diag = torch.pow(rays_directions, 2)
-    dir_mag_sq = torch.clamp(torch.sum(d_outer_diag, dim=-1, keepdim=True), min=eps)
-
-    null_outer_diag = 1 - d_outer_diag / dir_mag_sq
-    ray_dir_cov_diag = rays_dir_variance[..., None] * d_outer_diag[..., None, :]
-    xy_cov_diag = radii_variance[..., None] * null_outer_diag[..., None, :]
-    return ray_dir_cov_diag + xy_cov_diag
-
-
-def approximate_conical_frustum_as_gaussians(
-    bins: torch.Tensor, radii: torch.Tensor
-) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    """
-    Approximates a conical frustum as two Gaussian distributions.
-
-    The Gaussian distributions are characterized by
-    three values:
-
-    - rays_dir_mean: mean along the rays direction
-        (defined as t in the parametric representation of a cone).
-    - rays_dir_variance: the variance of the conical frustum  along the rays direction.
-    - radii_variance: variance of the conical frustum with respect to its radius.
-
-
-    The computation is stable and follows equation 7
-    of `MIP-NeRF <https://arxiv.org/abs/2103.13415>`_.
-
-    For more information on how the mean and variances are computed
-    refers to the appendix of the paper.
-
-    Args:
-        bins: A tensor of shape `(..., num_points_per_ray + 1)`
-            containing the bins at which the rays are sampled.
-            `bin[..., t]` and `bin[..., t+1]` represent respectively
-            the left and right coordinates of the interval.
-        t0: A tensor of shape `(..., num_points_per_ray)`
-            containing the left coordinates of the intervals
-            on which the rays are sampled.
-        t1: A tensor of shape `(..., num_points_per_ray)`
-            containing the rights coordinates of the intervals
-            on which the rays are sampled.
-        radii: A tensor of shape `(..., 1)`
-            base radii of the conical frustums.
-
-    Returns:
-        rays_dir_mean: A tensor of shape `(..., num_intervals)` representing
-            the mean along the rays direction
-            (t in the parametric represention of the cone)
-        rays_dir_variance: A tensor of shape `(...,  num_intervals)` representing
-            the variance of the conical frustum along the rays
-            (t in the parametric represention of the cone).
-        radii_variance: A tensor of shape `(..., num_intervals)` representing
-            the variance of the conical frustum with respect to its radius.
-    """
-    t_mu = torch.lerp(bins[..., 1:], bins[..., :-1], 0.5)
-    t_delta = torch.diff(bins, dim=-1) / 2
-
-    t_mu_pow2 = torch.pow(t_mu, 2)
-    t_delta_pow2 = torch.pow(t_delta, 2)
-    t_delta_pow4 = torch.pow(t_delta, 4)
-
-    den = 3 * t_mu_pow2 + t_delta_pow2
-
-    # mean along the rays direction
-    rays_dir_mean = t_mu + 2 * t_mu * t_delta_pow2 / den
-
-    # Variance of the conical frustum with along the rays directions
-    rays_dir_variance = t_delta_pow2 / 3 - (4 / 15) * (
-        t_delta_pow4 * (12 * t_mu_pow2 - t_delta_pow2) / torch.pow(den, 2)
-    )
-
-    # Variance of the conical frustum with respect to its radius
-    radii_variance = torch.pow(radii, 2) * (
-        t_mu_pow2 / 4 + (5 / 12) * t_delta_pow2 - 4 / 15 * (t_delta_pow4) / den
-    )
-    return rays_dir_mean, rays_dir_variance, radii_variance
-
-
-def conical_frustum_to_gaussian(
-    ray_bundle: ImplicitronRayBundle,
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    """
-    Approximate a conical frustum following a ray bundle as a Gaussian.
-
-    Args:
-        ray_bundle: A `RayBundle` or `HeterogeneousRayBundle` object with fields:
-            origins: A tensor of shape `(..., 3)`
-            directions: A tensor of shape `(..., 3)`
-            lengths: A tensor of shape `(..., num_points_per_ray)`
-            bins: A tensor of shape `(..., num_points_per_ray + 1)`
-                containing the bins at which the rays are sampled. .
-            pixel_radii_2d: A tensor of shape `(..., 1)`
-                base radii of the conical frustums.
-
-    Returns:
-        means: A tensor of shape `(..., num_points_per_ray - 1, 3)`
-            representing the means of the Gaussians
-            approximating the conical frustums.
-        diag_covariances: A tensor of shape `(...,num_points_per_ray -1, 3)`
-            representing the diagonal covariance matrices of our Gaussians.
-    """
-
-    if ray_bundle.pixel_radii_2d is None or ray_bundle.bins is None:
-        raise ValueError(
-            "RayBundle pixel_radii_2d or bins have not been provided."
-            " Look at pytorch3d.renderer.implicit.renderer.ray_sampler::"
-            "AbstractMaskRaySampler to see how to compute them. Have you forgot to set"
-            "`cast_ray_bundle_as_cone` to True?"
-        )
-
-    (
-        rays_dir_mean,
-        rays_dir_variance,
-        radii_variance,
-    ) = approximate_conical_frustum_as_gaussians(
-        ray_bundle.bins,
-        ray_bundle.pixel_radii_2d,
-    )
-    means = ray_bundle_variables_to_ray_points(
-        ray_bundle.origins, ray_bundle.directions, rays_dir_mean
-    )
-    diag_covariances = compute_3d_diagonal_covariance_gaussian(
-        ray_bundle.directions, rays_dir_variance, radii_variance
-    )
-    return means, diag_covariances

pytorch3d/implicitron/models/renderer/lstm_renderer.py

@@ -4,7 +4,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-import copy
+import dataclasses
 import logging
 from typing import List, Optional, Tuple
 
@@ -57,6 +57,7 @@ class LSTMRenderer(BaseRenderer, torch.nn.Module):
     verbose: bool = False
 
     def __post_init__(self):
+        super().__init__()
         self._lstm = torch.nn.LSTMCell(
             input_size=self.n_feature_channels,
             hidden_size=self.hidden_size,
@@ -102,11 +103,12 @@ class LSTMRenderer(BaseRenderer, torch.nn.Module):
             )
 
         # jitter the initial depths
-
-        ray_bundle_t = copy.copy(ray_bundle)
-        ray_bundle_t.lengths = (
-            ray_bundle.lengths
-            + torch.randn_like(ray_bundle.lengths) * self.init_depth_noise_std
+        ray_bundle_t = dataclasses.replace(
+            ray_bundle,
+            lengths=(
+                ray_bundle.lengths
+                + torch.randn_like(ray_bundle.lengths) * self.init_depth_noise_std
+            ),
         )
 
         states: List[Optional[Tuple[torch.Tensor, torch.Tensor]]] = [None]
@@ -133,6 +135,7 @@ class LSTMRenderer(BaseRenderer, torch.nn.Module):
                 break
 
             # run the lstm marcher
+            # pyre-fixme[29]: `Union[torch.Tensor, torch.nn.Module]` is not a function.
             state_h, state_c = self._lstm(
                 raymarch_features.view(-1, raymarch_features.shape[-1]),
                 states[-1],
@@ -140,6 +143,7 @@ class LSTMRenderer(BaseRenderer, torch.nn.Module):
             if state_h.requires_grad:
                 state_h.register_hook(lambda x: x.clamp(min=-10, max=10))
             # predict the next step size
+            # pyre-fixme[29]: `Union[torch.Tensor, torch.nn.Module]` is not a function.
             signed_distance = self._out_layer(state_h).view(ray_bundle_t.lengths.shape)
             # log the lstm states
             states.append((state_h, state_c))