mirror of
https://github.com/facebookresearch/pytorch3d.git
synced 2026-02-27 00:36:02 +08:00
Compare commits
1 Commits
cbcae096a0
...
bottler/un
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
62a2031dd4 |
@@ -10,7 +10,7 @@
|
||||
DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
|
||||
DIR=$(dirname "${DIR}")
|
||||
|
||||
if [[ -f "${DIR}/BUCK" ]]
|
||||
if [[ -f "${DIR}/TARGETS" ]]
|
||||
then
|
||||
pyfmt "${DIR}"
|
||||
else
|
||||
|
||||
@@ -19,6 +19,7 @@
|
||||
#
|
||||
import os
|
||||
import sys
|
||||
|
||||
import unittest.mock as mock
|
||||
|
||||
from recommonmark.parser import CommonMarkParser
|
||||
|
||||
@@ -48,18 +48,22 @@ The outputs of the experiment are saved and logged in multiple ways:
|
||||
import logging
|
||||
import os
|
||||
import warnings
|
||||
|
||||
from dataclasses import field
|
||||
|
||||
import hydra
|
||||
|
||||
import torch
|
||||
from accelerate import Accelerator
|
||||
from omegaconf import DictConfig, OmegaConf
|
||||
from packaging import version
|
||||
|
||||
from pytorch3d.implicitron.dataset.data_source import (
|
||||
DataSourceBase,
|
||||
ImplicitronDataSource,
|
||||
)
|
||||
from pytorch3d.implicitron.models.base_model import ImplicitronModelBase
|
||||
|
||||
from pytorch3d.implicitron.models.renderer.multipass_ea import (
|
||||
MultiPassEmissionAbsorptionRenderer,
|
||||
)
|
||||
|
||||
@@ -11,6 +11,7 @@ import os
|
||||
from typing import Optional
|
||||
|
||||
import torch.optim
|
||||
|
||||
from accelerate import Accelerator
|
||||
from pytorch3d.implicitron.models.base_model import ImplicitronModelBase
|
||||
from pytorch3d.implicitron.tools import model_io
|
||||
|
||||
@@ -14,7 +14,9 @@ from dataclasses import field
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
import torch.optim
|
||||
|
||||
from accelerate import Accelerator
|
||||
|
||||
from pytorch3d.implicitron.models.base_model import ImplicitronModelBase
|
||||
from pytorch3d.implicitron.tools import model_io
|
||||
from pytorch3d.implicitron.tools.config import (
|
||||
|
||||
@@ -12,6 +12,7 @@ import unittest
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
from hydra import compose, initialize_config_dir
|
||||
from omegaconf import OmegaConf
|
||||
from projects.implicitron_trainer.impl.optimizer_factory import (
|
||||
|
||||
@@ -6,4 +6,4 @@
|
||||
|
||||
# pyre-unsafe
|
||||
|
||||
__version__ = "0.7.9"
|
||||
__version__ = "0.7.8"
|
||||
|
||||
@@ -32,9 +32,7 @@ __global__ void BallQueryKernel(
|
||||
at::PackedTensorAccessor64<int64_t, 3, at::RestrictPtrTraits> idxs,
|
||||
at::PackedTensorAccessor64<scalar_t, 3, at::RestrictPtrTraits> dists,
|
||||
const int64_t K,
|
||||
const float radius,
|
||||
const float radius2,
|
||||
const bool skip_points_outside_cube) {
|
||||
const float radius2) {
|
||||
const int64_t N = p1.size(0);
|
||||
const int64_t chunks_per_cloud = (1 + (p1.size(1) - 1) / blockDim.x);
|
||||
const int64_t chunks_to_do = N * chunks_per_cloud;
|
||||
@@ -53,19 +51,7 @@ __global__ void BallQueryKernel(
|
||||
// Iterate over points in p2 until desired count is reached or
|
||||
// all points have been considered
|
||||
for (int64_t j = 0, count = 0; j < lengths2[n] && count < K; ++j) {
|
||||
if (skip_points_outside_cube) {
|
||||
bool is_within_radius = true;
|
||||
// Filter when any one coordinate is already outside the radius
|
||||
for (int d = 0; is_within_radius && d < D; ++d) {
|
||||
scalar_t abs_diff = fabs(p1[n][i][d] - p2[n][j][d]);
|
||||
is_within_radius = (abs_diff <= radius);
|
||||
}
|
||||
if (!is_within_radius) {
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
// Else, calculate the distance between the points and compare
|
||||
// Calculate the distance between the points
|
||||
scalar_t dist2 = 0.0;
|
||||
for (int d = 0; d < D; ++d) {
|
||||
scalar_t diff = p1[n][i][d] - p2[n][j][d];
|
||||
@@ -91,8 +77,7 @@ std::tuple<at::Tensor, at::Tensor> BallQueryCuda(
|
||||
const at::Tensor& lengths1, // (N,)
|
||||
const at::Tensor& lengths2, // (N,)
|
||||
int K,
|
||||
float radius,
|
||||
bool skip_points_outside_cube) {
|
||||
float radius) {
|
||||
// Check inputs are on the same device
|
||||
at::TensorArg p1_t{p1, "p1", 1}, p2_t{p2, "p2", 2},
|
||||
lengths1_t{lengths1, "lengths1", 3}, lengths2_t{lengths2, "lengths2", 4};
|
||||
@@ -135,9 +120,7 @@ std::tuple<at::Tensor, at::Tensor> BallQueryCuda(
|
||||
idxs.packed_accessor64<int64_t, 3, at::RestrictPtrTraits>(),
|
||||
dists.packed_accessor64<float, 3, at::RestrictPtrTraits>(),
|
||||
K_64,
|
||||
radius,
|
||||
radius2,
|
||||
skip_points_outside_cube);
|
||||
radius2);
|
||||
}));
|
||||
|
||||
AT_CUDA_CHECK(cudaGetLastError());
|
||||
|
||||
@@ -25,9 +25,6 @@
|
||||
// within the radius
|
||||
// radius: the radius around each point within which the neighbors need to be
|
||||
// located
|
||||
// skip_points_outside_cube: If true, reduce multiplications of float values
|
||||
// by not explicitly calculating distances to points that fall outside the
|
||||
// D-cube with side length (2*radius) centered at each point in p1.
|
||||
//
|
||||
// Returns:
|
||||
// p1_neighbor_idx: LongTensor of shape (N, P1, K), where
|
||||
@@ -49,8 +46,7 @@ std::tuple<at::Tensor, at::Tensor> BallQueryCpu(
|
||||
const at::Tensor& lengths1,
|
||||
const at::Tensor& lengths2,
|
||||
const int K,
|
||||
const float radius,
|
||||
const bool skip_points_outside_cube);
|
||||
const float radius);
|
||||
|
||||
// CUDA implementation
|
||||
std::tuple<at::Tensor, at::Tensor> BallQueryCuda(
|
||||
@@ -59,8 +55,7 @@ std::tuple<at::Tensor, at::Tensor> BallQueryCuda(
|
||||
const at::Tensor& lengths1,
|
||||
const at::Tensor& lengths2,
|
||||
const int K,
|
||||
const float radius,
|
||||
const bool skip_points_outside_cube);
|
||||
const float radius);
|
||||
|
||||
// Implementation which is exposed
|
||||
// Note: the backward pass reuses the KNearestNeighborBackward kernel
|
||||
@@ -70,8 +65,7 @@ inline std::tuple<at::Tensor, at::Tensor> BallQuery(
|
||||
const at::Tensor& lengths1,
|
||||
const at::Tensor& lengths2,
|
||||
int K,
|
||||
float radius,
|
||||
bool skip_points_outside_cube) {
|
||||
float radius) {
|
||||
if (p1.is_cuda() || p2.is_cuda()) {
|
||||
#ifdef WITH_CUDA
|
||||
CHECK_CUDA(p1);
|
||||
@@ -82,20 +76,16 @@ inline std::tuple<at::Tensor, at::Tensor> BallQuery(
|
||||
lengths1.contiguous(),
|
||||
lengths2.contiguous(),
|
||||
K,
|
||||
radius,
|
||||
skip_points_outside_cube);
|
||||
radius);
|
||||
#else
|
||||
AT_ERROR("Not compiled with GPU support.");
|
||||
#endif
|
||||
}
|
||||
CHECK_CPU(p1);
|
||||
CHECK_CPU(p2);
|
||||
return BallQueryCpu(
|
||||
p1.contiguous(),
|
||||
p2.contiguous(),
|
||||
lengths1.contiguous(),
|
||||
lengths2.contiguous(),
|
||||
K,
|
||||
radius,
|
||||
skip_points_outside_cube);
|
||||
radius);
|
||||
}
|
||||
|
||||
@@ -6,7 +6,6 @@
|
||||
* LICENSE file in the root directory of this source tree.
|
||||
*/
|
||||
|
||||
#include <math.h>
|
||||
#include <torch/extension.h>
|
||||
#include <tuple>
|
||||
|
||||
@@ -16,8 +15,7 @@ std::tuple<at::Tensor, at::Tensor> BallQueryCpu(
|
||||
const at::Tensor& lengths1,
|
||||
const at::Tensor& lengths2,
|
||||
int K,
|
||||
float radius,
|
||||
bool skip_points_outside_cube) {
|
||||
float radius) {
|
||||
const int N = p1.size(0);
|
||||
const int P1 = p1.size(1);
|
||||
const int D = p1.size(2);
|
||||
@@ -39,16 +37,6 @@ std::tuple<at::Tensor, at::Tensor> BallQueryCpu(
|
||||
const int64_t length2 = lengths2_a[n];
|
||||
for (int64_t i = 0; i < length1; ++i) {
|
||||
for (int64_t j = 0, count = 0; j < length2 && count < K; ++j) {
|
||||
if (skip_points_outside_cube) {
|
||||
bool is_within_radius = true;
|
||||
for (int d = 0; is_within_radius && d < D; ++d) {
|
||||
float abs_diff = fabs(p1_a[n][i][d] - p2_a[n][j][d]);
|
||||
is_within_radius = (abs_diff <= radius);
|
||||
}
|
||||
if (!is_within_radius) {
|
||||
continue;
|
||||
}
|
||||
}
|
||||
float dist2 = 0;
|
||||
for (int d = 0; d < D; ++d) {
|
||||
float diff = p1_a[n][i][d] - p2_a[n][j][d];
|
||||
|
||||
@@ -98,11 +98,6 @@ at::Tensor SigmoidAlphaBlendBackward(
|
||||
AT_ERROR("Not compiled with GPU support.");
|
||||
#endif
|
||||
}
|
||||
CHECK_CPU(distances);
|
||||
CHECK_CPU(pix_to_face);
|
||||
CHECK_CPU(alphas);
|
||||
CHECK_CPU(grad_alphas);
|
||||
|
||||
return SigmoidAlphaBlendBackwardCpu(
|
||||
grad_alphas, alphas, distances, pix_to_face, sigma);
|
||||
}
|
||||
|
||||
@@ -33,11 +33,11 @@ __global__ void alphaCompositeCudaForwardKernel(
|
||||
const int64_t W = points_idx.size(3);
|
||||
|
||||
// Get the batch and index
|
||||
const auto batch = blockIdx.x;
|
||||
const int batch = blockIdx.x;
|
||||
|
||||
const int num_pixels = C * H * W;
|
||||
const auto num_threads = gridDim.y * blockDim.x;
|
||||
const auto tid = blockIdx.y * blockDim.x + threadIdx.x;
|
||||
const int num_threads = gridDim.y * blockDim.x;
|
||||
const int tid = blockIdx.y * blockDim.x + threadIdx.x;
|
||||
|
||||
// Iterate over each feature in each pixel
|
||||
for (int pid = tid; pid < num_pixels; pid += num_threads) {
|
||||
@@ -83,11 +83,11 @@ __global__ void alphaCompositeCudaBackwardKernel(
|
||||
const int64_t W = points_idx.size(3);
|
||||
|
||||
// Get the batch and index
|
||||
const auto batch = blockIdx.x;
|
||||
const int batch = blockIdx.x;
|
||||
|
||||
const int num_pixels = C * H * W;
|
||||
const auto num_threads = gridDim.y * blockDim.x;
|
||||
const auto tid = blockIdx.y * blockDim.x + threadIdx.x;
|
||||
const int num_threads = gridDim.y * blockDim.x;
|
||||
const int tid = blockIdx.y * blockDim.x + threadIdx.x;
|
||||
|
||||
// Parallelize over each feature in each pixel in images of size H * W,
|
||||
// for each image in the batch of size batch_size
|
||||
|
||||
@@ -74,9 +74,6 @@ torch::Tensor alphaCompositeForward(
|
||||
AT_ERROR("Not compiled with GPU support");
|
||||
#endif
|
||||
} else {
|
||||
CHECK_CPU(features);
|
||||
CHECK_CPU(alphas);
|
||||
CHECK_CPU(points_idx);
|
||||
return alphaCompositeCpuForward(features, alphas, points_idx);
|
||||
}
|
||||
}
|
||||
@@ -104,11 +101,6 @@ std::tuple<torch::Tensor, torch::Tensor> alphaCompositeBackward(
|
||||
AT_ERROR("Not compiled with GPU support");
|
||||
#endif
|
||||
} else {
|
||||
CHECK_CPU(grad_outputs);
|
||||
CHECK_CPU(features);
|
||||
CHECK_CPU(alphas);
|
||||
CHECK_CPU(points_idx);
|
||||
|
||||
return alphaCompositeCpuBackward(
|
||||
grad_outputs, features, alphas, points_idx);
|
||||
}
|
||||
|
||||
@@ -33,11 +33,11 @@ __global__ void weightedSumNormCudaForwardKernel(
|
||||
const int64_t W = points_idx.size(3);
|
||||
|
||||
// Get the batch and index
|
||||
const auto batch = blockIdx.x;
|
||||
const int batch = blockIdx.x;
|
||||
|
||||
const int num_pixels = C * H * W;
|
||||
const auto num_threads = gridDim.y * blockDim.x;
|
||||
const auto tid = blockIdx.y * blockDim.x + threadIdx.x;
|
||||
const int num_threads = gridDim.y * blockDim.x;
|
||||
const int tid = blockIdx.y * blockDim.x + threadIdx.x;
|
||||
|
||||
// Parallelize over each feature in each pixel in images of size H * W,
|
||||
// for each image in the batch of size batch_size
|
||||
@@ -96,11 +96,11 @@ __global__ void weightedSumNormCudaBackwardKernel(
|
||||
const int64_t W = points_idx.size(3);
|
||||
|
||||
// Get the batch and index
|
||||
const auto batch = blockIdx.x;
|
||||
const int batch = blockIdx.x;
|
||||
|
||||
const int num_pixels = C * W * H;
|
||||
const auto num_threads = gridDim.y * blockDim.x;
|
||||
const auto tid = blockIdx.y * blockDim.x + threadIdx.x;
|
||||
const int num_threads = gridDim.y * blockDim.x;
|
||||
const int tid = blockIdx.y * blockDim.x + threadIdx.x;
|
||||
|
||||
// Parallelize over each feature in each pixel in images of size H * W,
|
||||
// for each image in the batch of size batch_size
|
||||
|
||||
@@ -73,10 +73,6 @@ torch::Tensor weightedSumNormForward(
|
||||
AT_ERROR("Not compiled with GPU support");
|
||||
#endif
|
||||
} else {
|
||||
CHECK_CPU(features);
|
||||
CHECK_CPU(alphas);
|
||||
CHECK_CPU(points_idx);
|
||||
|
||||
return weightedSumNormCpuForward(features, alphas, points_idx);
|
||||
}
|
||||
}
|
||||
@@ -104,11 +100,6 @@ std::tuple<torch::Tensor, torch::Tensor> weightedSumNormBackward(
|
||||
AT_ERROR("Not compiled with GPU support");
|
||||
#endif
|
||||
} else {
|
||||
CHECK_CPU(grad_outputs);
|
||||
CHECK_CPU(features);
|
||||
CHECK_CPU(alphas);
|
||||
CHECK_CPU(points_idx);
|
||||
|
||||
return weightedSumNormCpuBackward(
|
||||
grad_outputs, features, alphas, points_idx);
|
||||
}
|
||||
|
||||
@@ -31,11 +31,11 @@ __global__ void weightedSumCudaForwardKernel(
|
||||
const int64_t W = points_idx.size(3);
|
||||
|
||||
// Get the batch and index
|
||||
const auto batch = blockIdx.x;
|
||||
const int batch = blockIdx.x;
|
||||
|
||||
const int num_pixels = C * H * W;
|
||||
const auto num_threads = gridDim.y * blockDim.x;
|
||||
const auto tid = blockIdx.y * blockDim.x + threadIdx.x;
|
||||
const int num_threads = gridDim.y * blockDim.x;
|
||||
const int tid = blockIdx.y * blockDim.x + threadIdx.x;
|
||||
|
||||
// Parallelize over each feature in each pixel in images of size H * W,
|
||||
// for each image in the batch of size batch_size
|
||||
@@ -78,11 +78,11 @@ __global__ void weightedSumCudaBackwardKernel(
|
||||
const int64_t W = points_idx.size(3);
|
||||
|
||||
// Get the batch and index
|
||||
const auto batch = blockIdx.x;
|
||||
const int batch = blockIdx.x;
|
||||
|
||||
const int num_pixels = C * H * W;
|
||||
const auto num_threads = gridDim.y * blockDim.x;
|
||||
const auto tid = blockIdx.y * blockDim.x + threadIdx.x;
|
||||
const int num_threads = gridDim.y * blockDim.x;
|
||||
const int tid = blockIdx.y * blockDim.x + threadIdx.x;
|
||||
|
||||
// Iterate over each pixel to compute the contribution to the
|
||||
// gradient for the features and weights
|
||||
|
||||
@@ -72,9 +72,6 @@ torch::Tensor weightedSumForward(
|
||||
AT_ERROR("Not compiled with GPU support");
|
||||
#endif
|
||||
} else {
|
||||
CHECK_CPU(features);
|
||||
CHECK_CPU(alphas);
|
||||
CHECK_CPU(points_idx);
|
||||
return weightedSumCpuForward(features, alphas, points_idx);
|
||||
}
|
||||
}
|
||||
@@ -101,11 +98,6 @@ std::tuple<torch::Tensor, torch::Tensor> weightedSumBackward(
|
||||
AT_ERROR("Not compiled with GPU support");
|
||||
#endif
|
||||
} else {
|
||||
CHECK_CPU(grad_outputs);
|
||||
CHECK_CPU(features);
|
||||
CHECK_CPU(alphas);
|
||||
CHECK_CPU(points_idx);
|
||||
|
||||
return weightedSumCpuBackward(grad_outputs, features, alphas, points_idx);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -8,6 +8,7 @@
|
||||
|
||||
// clang-format off
|
||||
#include "./pulsar/global.h" // Include before <torch/extension.h>.
|
||||
#include <torch/extension.h>
|
||||
// clang-format on
|
||||
#include "./pulsar/pytorch/renderer.h"
|
||||
#include "./pulsar/pytorch/tensor_util.h"
|
||||
@@ -105,16 +106,15 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
py::class_<
|
||||
pulsar::pytorch::Renderer,
|
||||
std::shared_ptr<pulsar::pytorch::Renderer>>(m, "PulsarRenderer")
|
||||
.def(
|
||||
py::init<
|
||||
const uint&,
|
||||
const uint&,
|
||||
const uint&,
|
||||
const bool&,
|
||||
const bool&,
|
||||
const float&,
|
||||
const uint&,
|
||||
const uint&>())
|
||||
.def(py::init<
|
||||
const uint&,
|
||||
const uint&,
|
||||
const uint&,
|
||||
const bool&,
|
||||
const bool&,
|
||||
const float&,
|
||||
const uint&,
|
||||
const uint&>())
|
||||
.def(
|
||||
"__eq__",
|
||||
[](const pulsar::pytorch::Renderer& a,
|
||||
|
||||
@@ -60,8 +60,6 @@ std::tuple<at::Tensor, at::Tensor> FaceAreasNormalsForward(
|
||||
AT_ERROR("Not compiled with GPU support.");
|
||||
#endif
|
||||
}
|
||||
CHECK_CPU(verts);
|
||||
CHECK_CPU(faces);
|
||||
return FaceAreasNormalsForwardCpu(verts, faces);
|
||||
}
|
||||
|
||||
@@ -82,9 +80,5 @@ at::Tensor FaceAreasNormalsBackward(
|
||||
AT_ERROR("Not compiled with GPU support.");
|
||||
#endif
|
||||
}
|
||||
CHECK_CPU(grad_areas);
|
||||
CHECK_CPU(grad_normals);
|
||||
CHECK_CPU(verts);
|
||||
CHECK_CPU(faces);
|
||||
return FaceAreasNormalsBackwardCpu(grad_areas, grad_normals, verts, faces);
|
||||
}
|
||||
|
||||
@@ -20,14 +20,14 @@ __global__ void GatherScatterCudaKernel(
|
||||
const size_t V,
|
||||
const size_t D,
|
||||
const size_t E) {
|
||||
const auto tid = threadIdx.x;
|
||||
const int tid = threadIdx.x;
|
||||
|
||||
// Reverse the vertex order if backward.
|
||||
const int v0_idx = backward ? 1 : 0;
|
||||
const int v1_idx = backward ? 0 : 1;
|
||||
|
||||
// Edges are split evenly across the blocks.
|
||||
for (auto e = blockIdx.x; e < E; e += gridDim.x) {
|
||||
for (int e = blockIdx.x; e < E; e += gridDim.x) {
|
||||
// Get indices of vertices which form the edge.
|
||||
const int64_t v0 = edges[2 * e + v0_idx];
|
||||
const int64_t v1 = edges[2 * e + v1_idx];
|
||||
@@ -35,7 +35,7 @@ __global__ void GatherScatterCudaKernel(
|
||||
// Split vertex features evenly across threads.
|
||||
// This implementation will be quite wasteful when D<128 since there will be
|
||||
// a lot of threads doing nothing.
|
||||
for (auto d = tid; d < D; d += blockDim.x) {
|
||||
for (int d = tid; d < D; d += blockDim.x) {
|
||||
const float val = input[v1 * D + d];
|
||||
float* address = output + v0 * D + d;
|
||||
atomicAdd(address, val);
|
||||
|
||||
@@ -53,7 +53,5 @@ at::Tensor GatherScatter(
|
||||
AT_ERROR("Not compiled with GPU support.");
|
||||
#endif
|
||||
}
|
||||
CHECK_CPU(input);
|
||||
CHECK_CPU(edges);
|
||||
return GatherScatterCpu(input, edges, directed, backward);
|
||||
}
|
||||
|
||||
@@ -20,8 +20,8 @@ __global__ void InterpFaceAttrsForwardKernel(
|
||||
const size_t P,
|
||||
const size_t F,
|
||||
const size_t D) {
|
||||
const auto tid = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
const auto num_threads = blockDim.x * gridDim.x;
|
||||
const int tid = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
const int num_threads = blockDim.x * gridDim.x;
|
||||
for (int pd = tid; pd < P * D; pd += num_threads) {
|
||||
const int p = pd / D;
|
||||
const int d = pd % D;
|
||||
@@ -93,8 +93,8 @@ __global__ void InterpFaceAttrsBackwardKernel(
|
||||
const size_t P,
|
||||
const size_t F,
|
||||
const size_t D) {
|
||||
const auto tid = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
const auto num_threads = blockDim.x * gridDim.x;
|
||||
const int tid = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
const int num_threads = blockDim.x * gridDim.x;
|
||||
for (int pd = tid; pd < P * D; pd += num_threads) {
|
||||
const int p = pd / D;
|
||||
const int d = pd % D;
|
||||
|
||||
@@ -57,8 +57,6 @@ at::Tensor InterpFaceAttrsForward(
|
||||
AT_ERROR("Not compiled with GPU support.");
|
||||
#endif
|
||||
}
|
||||
CHECK_CPU(face_attrs);
|
||||
CHECK_CPU(barycentric_coords);
|
||||
return InterpFaceAttrsForwardCpu(pix_to_face, barycentric_coords, face_attrs);
|
||||
}
|
||||
|
||||
@@ -108,9 +106,6 @@ std::tuple<at::Tensor, at::Tensor> InterpFaceAttrsBackward(
|
||||
AT_ERROR("Not compiled with GPU support.");
|
||||
#endif
|
||||
}
|
||||
CHECK_CPU(face_attrs);
|
||||
CHECK_CPU(barycentric_coords);
|
||||
CHECK_CPU(grad_pix_attrs);
|
||||
return InterpFaceAttrsBackwardCpu(
|
||||
pix_to_face, barycentric_coords, face_attrs, grad_pix_attrs);
|
||||
}
|
||||
|
||||
@@ -44,7 +44,5 @@ inline std::tuple<at::Tensor, at::Tensor> IoUBox3D(
|
||||
AT_ERROR("Not compiled with GPU support.");
|
||||
#endif
|
||||
}
|
||||
CHECK_CPU(boxes1);
|
||||
CHECK_CPU(boxes2);
|
||||
return IoUBox3DCpu(boxes1.contiguous(), boxes2.contiguous());
|
||||
}
|
||||
|
||||
@@ -74,8 +74,6 @@ std::tuple<at::Tensor, at::Tensor> KNearestNeighborIdx(
|
||||
AT_ERROR("Not compiled with GPU support.");
|
||||
#endif
|
||||
}
|
||||
CHECK_CPU(p1);
|
||||
CHECK_CPU(p2);
|
||||
return KNearestNeighborIdxCpu(p1, p2, lengths1, lengths2, norm, K);
|
||||
}
|
||||
|
||||
@@ -142,8 +140,6 @@ std::tuple<at::Tensor, at::Tensor> KNearestNeighborBackward(
|
||||
AT_ERROR("Not compiled with GPU support.");
|
||||
#endif
|
||||
}
|
||||
CHECK_CPU(p1);
|
||||
CHECK_CPU(p2);
|
||||
return KNearestNeighborBackwardCpu(
|
||||
p1, p2, lengths1, lengths2, idxs, norm, grad_dists);
|
||||
}
|
||||
|
||||
@@ -58,6 +58,5 @@ inline std::tuple<at::Tensor, at::Tensor, at::Tensor> MarchingCubes(
|
||||
AT_ERROR("Not compiled with GPU support.");
|
||||
#endif
|
||||
}
|
||||
CHECK_CPU(vol);
|
||||
return MarchingCubesCpu(vol.contiguous(), isolevel);
|
||||
}
|
||||
|
||||
@@ -88,8 +88,6 @@ at::Tensor PackedToPadded(
|
||||
AT_ERROR("Not compiled with GPU support.");
|
||||
#endif
|
||||
}
|
||||
CHECK_CPU(inputs_packed);
|
||||
CHECK_CPU(first_idxs);
|
||||
return PackedToPaddedCpu(inputs_packed, first_idxs, max_size);
|
||||
}
|
||||
|
||||
@@ -107,7 +105,5 @@ at::Tensor PaddedToPacked(
|
||||
AT_ERROR("Not compiled with GPU support.");
|
||||
#endif
|
||||
}
|
||||
CHECK_CPU(inputs_padded);
|
||||
CHECK_CPU(first_idxs);
|
||||
return PaddedToPackedCpu(inputs_padded, first_idxs, num_inputs);
|
||||
}
|
||||
|
||||
@@ -174,8 +174,8 @@ std::tuple<at::Tensor, at::Tensor> HullHullDistanceForwardCpu(
|
||||
at::Tensor idxs = at::zeros({A_N,}, as_first_idx.options());
|
||||
// clang-format on
|
||||
|
||||
auto as_a = as.accessor<float, H1 == 1 ? 2 : 3>();
|
||||
auto bs_a = bs.accessor<float, H2 == 1 ? 2 : 3>();
|
||||
auto as_a = as.accessor < float, H1 == 1 ? 2 : 3 > ();
|
||||
auto bs_a = bs.accessor < float, H2 == 1 ? 2 : 3 > ();
|
||||
auto as_first_idx_a = as_first_idx.accessor<int64_t, 1>();
|
||||
auto bs_first_idx_a = bs_first_idx.accessor<int64_t, 1>();
|
||||
auto dists_a = dists.accessor<float, 1>();
|
||||
@@ -230,10 +230,10 @@ std::tuple<at::Tensor, at::Tensor> HullHullDistanceBackwardCpu(
|
||||
at::Tensor grad_as = at::zeros_like(as);
|
||||
at::Tensor grad_bs = at::zeros_like(bs);
|
||||
|
||||
auto as_a = as.accessor<float, H1 == 1 ? 2 : 3>();
|
||||
auto bs_a = bs.accessor<float, H2 == 1 ? 2 : 3>();
|
||||
auto grad_as_a = grad_as.accessor<float, H1 == 1 ? 2 : 3>();
|
||||
auto grad_bs_a = grad_bs.accessor<float, H2 == 1 ? 2 : 3>();
|
||||
auto as_a = as.accessor < float, H1 == 1 ? 2 : 3 > ();
|
||||
auto bs_a = bs.accessor < float, H2 == 1 ? 2 : 3 > ();
|
||||
auto grad_as_a = grad_as.accessor < float, H1 == 1 ? 2 : 3 > ();
|
||||
auto grad_bs_a = grad_bs.accessor < float, H2 == 1 ? 2 : 3 > ();
|
||||
auto idx_bs_a = idx_bs.accessor<int64_t, 1>();
|
||||
auto grad_dists_a = grad_dists.accessor<float, 1>();
|
||||
|
||||
|
||||
@@ -110,7 +110,7 @@ __global__ void DistanceForwardKernel(
|
||||
__syncthreads();
|
||||
|
||||
// Perform reduction in shared memory.
|
||||
for (auto s = blockDim.x / 2; s > 32; s >>= 1) {
|
||||
for (int s = blockDim.x / 2; s > 32; s >>= 1) {
|
||||
if (tid < s) {
|
||||
if (min_dists[tid] > min_dists[tid + s]) {
|
||||
min_dists[tid] = min_dists[tid + s];
|
||||
@@ -502,8 +502,8 @@ __global__ void PointFaceArrayForwardKernel(
|
||||
const float3* tris_f3 = (float3*)tris;
|
||||
|
||||
// Parallelize over P * S computations
|
||||
const auto num_threads = gridDim.x * blockDim.x;
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int num_threads = gridDim.x * blockDim.x;
|
||||
const int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
for (int t_i = tid; t_i < P * T; t_i += num_threads) {
|
||||
const int t = t_i / P; // segment index.
|
||||
@@ -576,8 +576,8 @@ __global__ void PointFaceArrayBackwardKernel(
|
||||
const float3* tris_f3 = (float3*)tris;
|
||||
|
||||
// Parallelize over P * S computations
|
||||
const auto num_threads = gridDim.x * blockDim.x;
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int num_threads = gridDim.x * blockDim.x;
|
||||
const int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
for (int t_i = tid; t_i < P * T; t_i += num_threads) {
|
||||
const int t = t_i / P; // triangle index.
|
||||
@@ -683,8 +683,8 @@ __global__ void PointEdgeArrayForwardKernel(
|
||||
float3* segms_f3 = (float3*)segms;
|
||||
|
||||
// Parallelize over P * S computations
|
||||
const auto num_threads = gridDim.x * blockDim.x;
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int num_threads = gridDim.x * blockDim.x;
|
||||
const int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
for (int t_i = tid; t_i < P * S; t_i += num_threads) {
|
||||
const int s = t_i / P; // segment index.
|
||||
@@ -752,8 +752,8 @@ __global__ void PointEdgeArrayBackwardKernel(
|
||||
float3* segms_f3 = (float3*)segms;
|
||||
|
||||
// Parallelize over P * S computations
|
||||
const auto num_threads = gridDim.x * blockDim.x;
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int num_threads = gridDim.x * blockDim.x;
|
||||
const int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
for (int t_i = tid; t_i < P * S; t_i += num_threads) {
|
||||
const int s = t_i / P; // segment index.
|
||||
|
||||
@@ -88,10 +88,6 @@ std::tuple<torch::Tensor, torch::Tensor> PointFaceDistanceForward(
|
||||
AT_ERROR("Not compiled with GPU support.");
|
||||
#endif
|
||||
}
|
||||
CHECK_CPU(points);
|
||||
CHECK_CPU(points_first_idx);
|
||||
CHECK_CPU(tris);
|
||||
CHECK_CPU(tris_first_idx);
|
||||
return PointFaceDistanceForwardCpu(
|
||||
points, points_first_idx, tris, tris_first_idx, min_triangle_area);
|
||||
}
|
||||
@@ -147,10 +143,6 @@ std::tuple<torch::Tensor, torch::Tensor> PointFaceDistanceBackward(
|
||||
AT_ERROR("Not compiled with GPU support.");
|
||||
#endif
|
||||
}
|
||||
CHECK_CPU(points);
|
||||
CHECK_CPU(tris);
|
||||
CHECK_CPU(idx_points);
|
||||
CHECK_CPU(grad_dists);
|
||||
return PointFaceDistanceBackwardCpu(
|
||||
points, tris, idx_points, grad_dists, min_triangle_area);
|
||||
}
|
||||
@@ -229,10 +221,6 @@ std::tuple<torch::Tensor, torch::Tensor> FacePointDistanceForward(
|
||||
AT_ERROR("Not compiled with GPU support.");
|
||||
#endif
|
||||
}
|
||||
CHECK_CPU(points);
|
||||
CHECK_CPU(points_first_idx);
|
||||
CHECK_CPU(tris);
|
||||
CHECK_CPU(tris_first_idx);
|
||||
return FacePointDistanceForwardCpu(
|
||||
points, points_first_idx, tris, tris_first_idx, min_triangle_area);
|
||||
}
|
||||
@@ -289,10 +277,6 @@ std::tuple<torch::Tensor, torch::Tensor> FacePointDistanceBackward(
|
||||
AT_ERROR("Not compiled with GPU support.");
|
||||
#endif
|
||||
}
|
||||
CHECK_CPU(points);
|
||||
CHECK_CPU(tris);
|
||||
CHECK_CPU(idx_tris);
|
||||
CHECK_CPU(grad_dists);
|
||||
return FacePointDistanceBackwardCpu(
|
||||
points, tris, idx_tris, grad_dists, min_triangle_area);
|
||||
}
|
||||
@@ -362,10 +346,6 @@ std::tuple<torch::Tensor, torch::Tensor> PointEdgeDistanceForward(
|
||||
AT_ERROR("Not compiled with GPU support.");
|
||||
#endif
|
||||
}
|
||||
CHECK_CPU(points);
|
||||
CHECK_CPU(points_first_idx);
|
||||
CHECK_CPU(segms);
|
||||
CHECK_CPU(segms_first_idx);
|
||||
return PointEdgeDistanceForwardCpu(
|
||||
points, points_first_idx, segms, segms_first_idx, max_points);
|
||||
}
|
||||
@@ -416,10 +396,6 @@ std::tuple<torch::Tensor, torch::Tensor> PointEdgeDistanceBackward(
|
||||
AT_ERROR("Not compiled with GPU support.");
|
||||
#endif
|
||||
}
|
||||
CHECK_CPU(points);
|
||||
CHECK_CPU(segms);
|
||||
CHECK_CPU(idx_points);
|
||||
CHECK_CPU(grad_dists);
|
||||
return PointEdgeDistanceBackwardCpu(points, segms, idx_points, grad_dists);
|
||||
}
|
||||
|
||||
@@ -488,10 +464,6 @@ std::tuple<torch::Tensor, torch::Tensor> EdgePointDistanceForward(
|
||||
AT_ERROR("Not compiled with GPU support.");
|
||||
#endif
|
||||
}
|
||||
CHECK_CPU(points);
|
||||
CHECK_CPU(points_first_idx);
|
||||
CHECK_CPU(segms);
|
||||
CHECK_CPU(segms_first_idx);
|
||||
return EdgePointDistanceForwardCpu(
|
||||
points, points_first_idx, segms, segms_first_idx, max_segms);
|
||||
}
|
||||
@@ -542,10 +514,6 @@ std::tuple<torch::Tensor, torch::Tensor> EdgePointDistanceBackward(
|
||||
AT_ERROR("Not compiled with GPU support.");
|
||||
#endif
|
||||
}
|
||||
CHECK_CPU(points);
|
||||
CHECK_CPU(segms);
|
||||
CHECK_CPU(idx_segms);
|
||||
CHECK_CPU(grad_dists);
|
||||
return EdgePointDistanceBackwardCpu(points, segms, idx_segms, grad_dists);
|
||||
}
|
||||
|
||||
@@ -599,8 +567,6 @@ torch::Tensor PointFaceArrayDistanceForward(
|
||||
AT_ERROR("Not compiled with GPU support.");
|
||||
#endif
|
||||
}
|
||||
CHECK_CPU(points);
|
||||
CHECK_CPU(tris);
|
||||
return PointFaceArrayDistanceForwardCpu(points, tris, min_triangle_area);
|
||||
}
|
||||
|
||||
@@ -647,9 +613,6 @@ std::tuple<torch::Tensor, torch::Tensor> PointFaceArrayDistanceBackward(
|
||||
AT_ERROR("Not compiled with GPU support.");
|
||||
#endif
|
||||
}
|
||||
CHECK_CPU(points);
|
||||
CHECK_CPU(tris);
|
||||
CHECK_CPU(grad_dists);
|
||||
return PointFaceArrayDistanceBackwardCpu(
|
||||
points, tris, grad_dists, min_triangle_area);
|
||||
}
|
||||
@@ -698,8 +661,6 @@ torch::Tensor PointEdgeArrayDistanceForward(
|
||||
AT_ERROR("Not compiled with GPU support.");
|
||||
#endif
|
||||
}
|
||||
CHECK_CPU(points);
|
||||
CHECK_CPU(segms);
|
||||
return PointEdgeArrayDistanceForwardCpu(points, segms);
|
||||
}
|
||||
|
||||
@@ -742,8 +703,5 @@ std::tuple<torch::Tensor, torch::Tensor> PointEdgeArrayDistanceBackward(
|
||||
AT_ERROR("Not compiled with GPU support.");
|
||||
#endif
|
||||
}
|
||||
CHECK_CPU(points);
|
||||
CHECK_CPU(segms);
|
||||
CHECK_CPU(grad_dists);
|
||||
return PointEdgeArrayDistanceBackwardCpu(points, segms, grad_dists);
|
||||
}
|
||||
|
||||
@@ -104,12 +104,6 @@ inline void PointsToVolumesForward(
|
||||
AT_ERROR("Not compiled with GPU support.");
|
||||
#endif
|
||||
}
|
||||
CHECK_CPU(points_3d);
|
||||
CHECK_CPU(points_features);
|
||||
CHECK_CPU(volume_densities);
|
||||
CHECK_CPU(volume_features);
|
||||
CHECK_CPU(grid_sizes);
|
||||
CHECK_CPU(mask);
|
||||
PointsToVolumesForwardCpu(
|
||||
points_3d,
|
||||
points_features,
|
||||
@@ -189,14 +183,6 @@ inline void PointsToVolumesBackward(
|
||||
AT_ERROR("Not compiled with GPU support.");
|
||||
#endif
|
||||
}
|
||||
CHECK_CPU(points_3d);
|
||||
CHECK_CPU(points_features);
|
||||
CHECK_CPU(grid_sizes);
|
||||
CHECK_CPU(mask);
|
||||
CHECK_CPU(grad_volume_densities);
|
||||
CHECK_CPU(grad_volume_features);
|
||||
CHECK_CPU(grad_points_3d);
|
||||
CHECK_CPU(grad_points_features);
|
||||
PointsToVolumesBackwardCpu(
|
||||
points_3d,
|
||||
points_features,
|
||||
|
||||
@@ -15,8 +15,8 @@
|
||||
#endif
|
||||
|
||||
#if defined(_WIN64) || defined(_WIN32)
|
||||
using uint = unsigned int;
|
||||
using ushort = unsigned short;
|
||||
#define uint unsigned int
|
||||
#define ushort unsigned short
|
||||
#endif
|
||||
|
||||
#include "./logging.h" // <- include before torch/extension.h
|
||||
|
||||
@@ -417,7 +417,7 @@ __device__ static float atomicMin(float* address, float val) {
|
||||
(OUT_PTR), \
|
||||
(NUM_SELECTED_PTR), \
|
||||
(NUM_ITEMS), \
|
||||
(STREAM));
|
||||
stream = (STREAM));
|
||||
|
||||
#define COPY_HOST_DEV(PTR_D, PTR_H, TYPE, SIZE) \
|
||||
HANDLECUDA(cudaMemcpy( \
|
||||
|
||||
@@ -357,11 +357,11 @@ void MAX_WS(
|
||||
//
|
||||
//
|
||||
#define END_PARALLEL() \
|
||||
end_parallel:; \
|
||||
end_parallel :; \
|
||||
}
|
||||
#define END_PARALLEL_NORET() }
|
||||
#define END_PARALLEL_2D() \
|
||||
end_parallel:; \
|
||||
end_parallel :; \
|
||||
} \
|
||||
}
|
||||
#define END_PARALLEL_2D_NORET() \
|
||||
|
||||
@@ -70,6 +70,11 @@ struct CamGradInfo {
|
||||
float3 pixel_dir_y;
|
||||
};
|
||||
|
||||
// TODO: remove once https://github.com/NVlabs/cub/issues/172 is resolved.
|
||||
struct IntWrapper {
|
||||
int val;
|
||||
};
|
||||
|
||||
} // namespace pulsar
|
||||
|
||||
#endif
|
||||
|
||||
@@ -149,6 +149,11 @@ IHD CamGradInfo operator*(const CamGradInfo& a, const float& b) {
|
||||
return res;
|
||||
}
|
||||
|
||||
IHD IntWrapper operator+(const IntWrapper& a, const IntWrapper& b) {
|
||||
IntWrapper res;
|
||||
res.val = a.val + b.val;
|
||||
return res;
|
||||
}
|
||||
} // namespace pulsar
|
||||
|
||||
#endif
|
||||
|
||||
@@ -155,8 +155,8 @@ void backward(
|
||||
stream);
|
||||
CHECKLAUNCH();
|
||||
SUM_WS(
|
||||
self->ids_sorted_d,
|
||||
self->n_grad_contributions_d,
|
||||
(IntWrapper*)(self->ids_sorted_d),
|
||||
(IntWrapper*)(self->n_grad_contributions_d),
|
||||
static_cast<int>(num_balls),
|
||||
self->workspace_d,
|
||||
self->workspace_size,
|
||||
|
||||
@@ -52,7 +52,7 @@ HOST void construct(
|
||||
self->cam.film_width = width;
|
||||
self->cam.film_height = height;
|
||||
self->max_num_balls = max_num_balls;
|
||||
MALLOC(self->result_d, float, width * height * n_channels);
|
||||
MALLOC(self->result_d, float, width* height* n_channels);
|
||||
self->cam.orthogonal_projection = orthogonal_projection;
|
||||
self->cam.right_handed = right_handed_system;
|
||||
self->cam.background_normalization_depth = background_normalization_depth;
|
||||
@@ -93,7 +93,7 @@ HOST void construct(
|
||||
MALLOC(self->di_sorted_d, DrawInfo, max_num_balls);
|
||||
MALLOC(self->region_flags_d, char, max_num_balls);
|
||||
MALLOC(self->num_selected_d, size_t, 1);
|
||||
MALLOC(self->forw_info_d, float, width * height * (3 + 2 * n_track));
|
||||
MALLOC(self->forw_info_d, float, width* height * (3 + 2 * n_track));
|
||||
MALLOC(self->min_max_pixels_d, IntersectInfo, 1);
|
||||
MALLOC(self->grad_pos_d, float3, max_num_balls);
|
||||
MALLOC(self->grad_col_d, float, max_num_balls* n_channels);
|
||||
|
||||
@@ -255,7 +255,7 @@ GLOBAL void calc_signature(
|
||||
* for every iteration through the loading loop every thread could add a
|
||||
* 'hit' to the buffer.
|
||||
*/
|
||||
#define RENDER_BUFFER_SIZE RENDER_BLOCK_SIZE * RENDER_BLOCK_SIZE * 2
|
||||
#define RENDER_BUFFER_SIZE RENDER_BLOCK_SIZE* RENDER_BLOCK_SIZE * 2
|
||||
/**
|
||||
* The threshold after which the spheres that are in the render buffer
|
||||
* are rendered and the buffer is flushed.
|
||||
|
||||
@@ -6,6 +6,9 @@
|
||||
* LICENSE file in the root directory of this source tree.
|
||||
*/
|
||||
|
||||
#include "./global.h"
|
||||
#include "./logging.h"
|
||||
|
||||
/**
|
||||
* A compilation unit to provide warnings about the code and avoid
|
||||
* repeated messages.
|
||||
|
||||
@@ -25,7 +25,7 @@ class BitMask {
|
||||
|
||||
// Use all threads in the current block to clear all bits of this BitMask
|
||||
__device__ void block_clear() {
|
||||
for (auto i = threadIdx.x; i < H * W * D; i += blockDim.x) {
|
||||
for (int i = threadIdx.x; i < H * W * D; i += blockDim.x) {
|
||||
data[i] = 0;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
@@ -23,8 +23,8 @@ __global__ void TriangleBoundingBoxKernel(
|
||||
const float blur_radius,
|
||||
float* bboxes, // (4, F)
|
||||
bool* skip_face) { // (F,)
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const auto num_threads = blockDim.x * gridDim.x;
|
||||
const int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int num_threads = blockDim.x * gridDim.x;
|
||||
const float sqrt_radius = sqrt(blur_radius);
|
||||
for (int f = tid; f < F; f += num_threads) {
|
||||
const float v0x = face_verts[f * 9 + 0 * 3 + 0];
|
||||
@@ -56,8 +56,8 @@ __global__ void PointBoundingBoxKernel(
|
||||
const int P,
|
||||
float* bboxes, // (4, P)
|
||||
bool* skip_points) {
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const auto num_threads = blockDim.x * gridDim.x;
|
||||
const int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int num_threads = blockDim.x * gridDim.x;
|
||||
for (int p = tid; p < P; p += num_threads) {
|
||||
const float x = points[p * 3 + 0];
|
||||
const float y = points[p * 3 + 1];
|
||||
@@ -113,7 +113,7 @@ __global__ void RasterizeCoarseCudaKernel(
|
||||
const int chunks_per_batch = 1 + (E - 1) / chunk_size;
|
||||
const int num_chunks = N * chunks_per_batch;
|
||||
|
||||
for (auto chunk = blockIdx.x; chunk < num_chunks; chunk += gridDim.x) {
|
||||
for (int chunk = blockIdx.x; chunk < num_chunks; chunk += gridDim.x) {
|
||||
const int batch_idx = chunk / chunks_per_batch; // batch index
|
||||
const int chunk_idx = chunk % chunks_per_batch;
|
||||
const int elem_chunk_start_idx = chunk_idx * chunk_size;
|
||||
@@ -123,7 +123,7 @@ __global__ void RasterizeCoarseCudaKernel(
|
||||
const int64_t elem_stop_idx = elem_start_idx + elems_per_batch[batch_idx];
|
||||
|
||||
// Have each thread handle a different face within the chunk
|
||||
for (auto e = threadIdx.x; e < chunk_size; e += blockDim.x) {
|
||||
for (int e = threadIdx.x; e < chunk_size; e += blockDim.x) {
|
||||
const int e_idx = elem_chunk_start_idx + e;
|
||||
|
||||
// Check that we are still within the same element of the batch
|
||||
@@ -170,7 +170,7 @@ __global__ void RasterizeCoarseCudaKernel(
|
||||
// Now we have processed every elem in the current chunk. We need to
|
||||
// count the number of elems in each bin so we can write the indices
|
||||
// out to global memory. We have each thread handle a different bin.
|
||||
for (auto byx = threadIdx.x; byx < num_bins_y * num_bins_x;
|
||||
for (int byx = threadIdx.x; byx < num_bins_y * num_bins_x;
|
||||
byx += blockDim.x) {
|
||||
const int by = byx / num_bins_x;
|
||||
const int bx = byx % num_bins_x;
|
||||
|
||||
@@ -260,8 +260,8 @@ __global__ void RasterizeMeshesNaiveCudaKernel(
|
||||
float* pix_dists,
|
||||
float* bary) {
|
||||
// Simple version: One thread per output pixel
|
||||
auto num_threads = gridDim.x * blockDim.x;
|
||||
auto tid = blockDim.x * blockIdx.x + threadIdx.x;
|
||||
int num_threads = gridDim.x * blockDim.x;
|
||||
int tid = blockDim.x * blockIdx.x + threadIdx.x;
|
||||
|
||||
for (int i = tid; i < N * H * W; i += num_threads) {
|
||||
// Convert linear index to 3D index
|
||||
@@ -446,8 +446,8 @@ __global__ void RasterizeMeshesBackwardCudaKernel(
|
||||
|
||||
// Parallelize over each pixel in images of
|
||||
// size H * W, for each image in the batch of size N.
|
||||
const auto num_threads = gridDim.x * blockDim.x;
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int num_threads = gridDim.x * blockDim.x;
|
||||
const int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
for (int t_i = tid; t_i < N * H * W; t_i += num_threads) {
|
||||
// Convert linear index to 3D index
|
||||
@@ -650,8 +650,8 @@ __global__ void RasterizeMeshesFineCudaKernel(
|
||||
) {
|
||||
// This can be more than H * W if H or W are not divisible by bin_size.
|
||||
int num_pixels = N * BH * BW * bin_size * bin_size;
|
||||
auto num_threads = gridDim.x * blockDim.x;
|
||||
auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int num_threads = gridDim.x * blockDim.x;
|
||||
int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
for (int pid = tid; pid < num_pixels; pid += num_threads) {
|
||||
// Convert linear index into bin and pixel indices. We make the within
|
||||
|
||||
@@ -138,9 +138,6 @@ RasterizeMeshesNaive(
|
||||
AT_ERROR("Not compiled with GPU support");
|
||||
#endif
|
||||
} else {
|
||||
CHECK_CPU(face_verts);
|
||||
CHECK_CPU(mesh_to_face_first_idx);
|
||||
CHECK_CPU(num_faces_per_mesh);
|
||||
return RasterizeMeshesNaiveCpu(
|
||||
face_verts,
|
||||
mesh_to_face_first_idx,
|
||||
@@ -235,11 +232,6 @@ torch::Tensor RasterizeMeshesBackward(
|
||||
AT_ERROR("Not compiled with GPU support");
|
||||
#endif
|
||||
} else {
|
||||
CHECK_CPU(face_verts);
|
||||
CHECK_CPU(pix_to_face);
|
||||
CHECK_CPU(grad_zbuf);
|
||||
CHECK_CPU(grad_bary);
|
||||
CHECK_CPU(grad_dists);
|
||||
return RasterizeMeshesBackwardCpu(
|
||||
face_verts,
|
||||
pix_to_face,
|
||||
@@ -314,9 +306,6 @@ torch::Tensor RasterizeMeshesCoarse(
|
||||
AT_ERROR("Not compiled with GPU support");
|
||||
#endif
|
||||
} else {
|
||||
CHECK_CPU(face_verts);
|
||||
CHECK_CPU(mesh_to_face_first_idx);
|
||||
CHECK_CPU(num_faces_per_mesh);
|
||||
return RasterizeMeshesCoarseCpu(
|
||||
face_verts,
|
||||
mesh_to_face_first_idx,
|
||||
@@ -434,8 +423,6 @@ RasterizeMeshesFine(
|
||||
AT_ERROR("Not compiled with GPU support");
|
||||
#endif
|
||||
} else {
|
||||
CHECK_CPU(face_verts);
|
||||
CHECK_CPU(bin_faces);
|
||||
AT_ERROR("NOT IMPLEMENTED");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -97,8 +97,8 @@ __global__ void RasterizePointsNaiveCudaKernel(
|
||||
float* zbuf, // (N, H, W, K)
|
||||
float* pix_dists) { // (N, H, W, K)
|
||||
// Simple version: One thread per output pixel
|
||||
const auto num_threads = gridDim.x * blockDim.x;
|
||||
const auto tid = blockDim.x * blockIdx.x + threadIdx.x;
|
||||
const int num_threads = gridDim.x * blockDim.x;
|
||||
const int tid = blockDim.x * blockIdx.x + threadIdx.x;
|
||||
for (int i = tid; i < N * H * W; i += num_threads) {
|
||||
// Convert linear index to 3D index
|
||||
const int n = i / (H * W); // Batch index
|
||||
@@ -237,8 +237,8 @@ __global__ void RasterizePointsFineCudaKernel(
|
||||
float* pix_dists) { // (N, H, W, K)
|
||||
// This can be more than H * W if H or W are not divisible by bin_size.
|
||||
const int num_pixels = N * BH * BW * bin_size * bin_size;
|
||||
const auto num_threads = gridDim.x * blockDim.x;
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int num_threads = gridDim.x * blockDim.x;
|
||||
const int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
for (int pid = tid; pid < num_pixels; pid += num_threads) {
|
||||
// Convert linear index into bin and pixel indices. We make the within
|
||||
@@ -376,8 +376,8 @@ __global__ void RasterizePointsBackwardCudaKernel(
|
||||
float* grad_points) { // (P, 3)
|
||||
// Parallelized over each of K points per pixel, for each pixel in images of
|
||||
// size H * W, for each image in the batch of size N.
|
||||
auto num_threads = gridDim.x * blockDim.x;
|
||||
auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int num_threads = gridDim.x * blockDim.x;
|
||||
int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
for (int i = tid; i < N * H * W * K; i += num_threads) {
|
||||
// const int n = i / (H * W * K); // batch index (not needed).
|
||||
const int yxk = i % (H * W * K);
|
||||
|
||||
@@ -91,10 +91,6 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> RasterizePointsNaive(
|
||||
AT_ERROR("Not compiled with GPU support");
|
||||
#endif
|
||||
} else {
|
||||
CHECK_CPU(points);
|
||||
CHECK_CPU(cloud_to_packed_first_idx);
|
||||
CHECK_CPU(num_points_per_cloud);
|
||||
CHECK_CPU(radius);
|
||||
return RasterizePointsNaiveCpu(
|
||||
points,
|
||||
cloud_to_packed_first_idx,
|
||||
@@ -170,10 +166,6 @@ torch::Tensor RasterizePointsCoarse(
|
||||
AT_ERROR("Not compiled with GPU support");
|
||||
#endif
|
||||
} else {
|
||||
CHECK_CPU(points);
|
||||
CHECK_CPU(cloud_to_packed_first_idx);
|
||||
CHECK_CPU(num_points_per_cloud);
|
||||
CHECK_CPU(radius);
|
||||
return RasterizePointsCoarseCpu(
|
||||
points,
|
||||
cloud_to_packed_first_idx,
|
||||
@@ -240,8 +232,6 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> RasterizePointsFine(
|
||||
AT_ERROR("Not compiled with GPU support");
|
||||
#endif
|
||||
} else {
|
||||
CHECK_CPU(points);
|
||||
CHECK_CPU(bin_points);
|
||||
AT_ERROR("NOT IMPLEMENTED");
|
||||
}
|
||||
}
|
||||
@@ -294,10 +284,6 @@ torch::Tensor RasterizePointsBackward(
|
||||
AT_ERROR("Not compiled with GPU support");
|
||||
#endif
|
||||
} else {
|
||||
CHECK_CPU(points);
|
||||
CHECK_CPU(idxs);
|
||||
CHECK_CPU(grad_zbuf);
|
||||
CHECK_CPU(grad_dists);
|
||||
return RasterizePointsBackwardCpu(points, idxs, grad_zbuf, grad_dists);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -107,8 +107,7 @@ at::Tensor FarthestPointSamplingCuda(
|
||||
const at::Tensor& points, // (N, P, 3)
|
||||
const at::Tensor& lengths, // (N,)
|
||||
const at::Tensor& K, // (N,)
|
||||
const at::Tensor& start_idxs,
|
||||
const int64_t max_K_known = -1) {
|
||||
const at::Tensor& start_idxs) {
|
||||
// Check inputs are on the same device
|
||||
at::TensorArg p_t{points, "points", 1}, lengths_t{lengths, "lengths", 2},
|
||||
k_t{K, "K", 3}, start_idxs_t{start_idxs, "start_idxs", 4};
|
||||
@@ -130,12 +129,7 @@ at::Tensor FarthestPointSamplingCuda(
|
||||
|
||||
const int64_t N = points.size(0);
|
||||
const int64_t P = points.size(1);
|
||||
int64_t max_K;
|
||||
if (max_K_known > 0) {
|
||||
max_K = max_K_known;
|
||||
} else {
|
||||
max_K = at::max(K).item<int64_t>();
|
||||
}
|
||||
const int64_t max_K = at::max(K).item<int64_t>();
|
||||
|
||||
// Initialize the output tensor with the sampled indices
|
||||
auto idxs = at::full({N, max_K}, -1, lengths.options());
|
||||
|
||||
@@ -43,8 +43,7 @@ at::Tensor FarthestPointSamplingCuda(
|
||||
const at::Tensor& points,
|
||||
const at::Tensor& lengths,
|
||||
const at::Tensor& K,
|
||||
const at::Tensor& start_idxs,
|
||||
const int64_t max_K_known = -1);
|
||||
const at::Tensor& start_idxs);
|
||||
|
||||
at::Tensor FarthestPointSamplingCpu(
|
||||
const at::Tensor& points,
|
||||
@@ -57,23 +56,17 @@ at::Tensor FarthestPointSampling(
|
||||
const at::Tensor& points,
|
||||
const at::Tensor& lengths,
|
||||
const at::Tensor& K,
|
||||
const at::Tensor& start_idxs,
|
||||
const int64_t max_K_known = -1) {
|
||||
const at::Tensor& start_idxs) {
|
||||
if (points.is_cuda() || lengths.is_cuda() || K.is_cuda()) {
|
||||
#ifdef WITH_CUDA
|
||||
CHECK_CUDA(points);
|
||||
CHECK_CUDA(lengths);
|
||||
CHECK_CUDA(K);
|
||||
CHECK_CUDA(start_idxs);
|
||||
return FarthestPointSamplingCuda(
|
||||
points, lengths, K, start_idxs, max_K_known);
|
||||
return FarthestPointSamplingCuda(points, lengths, K, start_idxs);
|
||||
#else
|
||||
AT_ERROR("Not compiled with GPU support.");
|
||||
#endif
|
||||
}
|
||||
CHECK_CPU(points);
|
||||
CHECK_CPU(lengths);
|
||||
CHECK_CPU(K);
|
||||
CHECK_CPU(start_idxs);
|
||||
return FarthestPointSamplingCpu(points, lengths, K, start_idxs);
|
||||
}
|
||||
|
||||
@@ -71,8 +71,6 @@ inline void SamplePdf(
|
||||
AT_ERROR("Not compiled with GPU support.");
|
||||
#endif
|
||||
}
|
||||
CHECK_CPU(weights);
|
||||
CHECK_CPU(outputs);
|
||||
CHECK_CONTIGUOUS(outputs);
|
||||
SamplePdfCpu(bins, weights, outputs, eps);
|
||||
}
|
||||
|
||||
@@ -99,7 +99,8 @@ namespace {
|
||||
// and increment it via template recursion until it is equal to the run-time
|
||||
// argument N.
|
||||
template <
|
||||
template <typename, int64_t> class Kernel,
|
||||
template <typename, int64_t>
|
||||
class Kernel,
|
||||
typename T,
|
||||
int64_t minN,
|
||||
int64_t maxN,
|
||||
@@ -123,7 +124,8 @@ struct DispatchKernelHelper1D {
|
||||
// 1D dispatch: Specialization when curN == maxN
|
||||
// We need this base case to avoid infinite template recursion.
|
||||
template <
|
||||
template <typename, int64_t> class Kernel,
|
||||
template <typename, int64_t>
|
||||
class Kernel,
|
||||
typename T,
|
||||
int64_t minN,
|
||||
int64_t maxN,
|
||||
@@ -143,7 +145,8 @@ struct DispatchKernelHelper1D<Kernel, T, minN, maxN, maxN, Args...> {
|
||||
// the run-time values of N and M, at which point we dispatch to the run
|
||||
// method of the kernel.
|
||||
template <
|
||||
template <typename, int64_t, int64_t> class Kernel,
|
||||
template <typename, int64_t, int64_t>
|
||||
class Kernel,
|
||||
typename T,
|
||||
int64_t minN,
|
||||
int64_t maxN,
|
||||
@@ -200,7 +203,8 @@ struct DispatchKernelHelper2D {
|
||||
|
||||
// 2D dispatch, specialization for curN == maxN
|
||||
template <
|
||||
template <typename, int64_t, int64_t> class Kernel,
|
||||
template <typename, int64_t, int64_t>
|
||||
class Kernel,
|
||||
typename T,
|
||||
int64_t minN,
|
||||
int64_t maxN,
|
||||
@@ -239,7 +243,8 @@ struct DispatchKernelHelper2D<
|
||||
|
||||
// 2D dispatch, specialization for curM == maxM
|
||||
template <
|
||||
template <typename, int64_t, int64_t> class Kernel,
|
||||
template <typename, int64_t, int64_t>
|
||||
class Kernel,
|
||||
typename T,
|
||||
int64_t minN,
|
||||
int64_t maxN,
|
||||
@@ -278,7 +283,8 @@ struct DispatchKernelHelper2D<
|
||||
|
||||
// 2D dispatch, specialization for curN == maxN, curM == maxM
|
||||
template <
|
||||
template <typename, int64_t, int64_t> class Kernel,
|
||||
template <typename, int64_t, int64_t>
|
||||
class Kernel,
|
||||
typename T,
|
||||
int64_t minN,
|
||||
int64_t maxN,
|
||||
@@ -307,7 +313,8 @@ struct DispatchKernelHelper2D<
|
||||
|
||||
// This is the function we expect users to call to dispatch to 1D functions
|
||||
template <
|
||||
template <typename, int64_t> class Kernel,
|
||||
template <typename, int64_t>
|
||||
class Kernel,
|
||||
typename T,
|
||||
int64_t minN,
|
||||
int64_t maxN,
|
||||
@@ -323,7 +330,8 @@ void DispatchKernel1D(const int64_t N, Args... args) {
|
||||
|
||||
// This is the function we expect users to call to dispatch to 2D functions
|
||||
template <
|
||||
template <typename, int64_t, int64_t> class Kernel,
|
||||
template <typename, int64_t, int64_t>
|
||||
class Kernel,
|
||||
typename T,
|
||||
int64_t minN,
|
||||
int64_t maxN,
|
||||
|
||||
@@ -15,7 +15,3 @@
|
||||
#define CHECK_CONTIGUOUS_CUDA(x) \
|
||||
CHECK_CUDA(x); \
|
||||
CHECK_CONTIGUOUS(x)
|
||||
#define CHECK_CPU(x) \
|
||||
TORCH_CHECK( \
|
||||
x.device().type() == torch::kCPU, \
|
||||
"Cannot use CPU implementation: " #x " not on CPU.")
|
||||
|
||||
@@ -21,6 +21,7 @@ from typing import (
|
||||
)
|
||||
|
||||
import torch
|
||||
|
||||
from pytorch3d.implicitron.dataset.frame_data import FrameData
|
||||
from pytorch3d.implicitron.dataset.utils import GenericWorkaround
|
||||
|
||||
|
||||
@@ -25,6 +25,7 @@ from typing import (
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from pytorch3d.implicitron.dataset import orm_types, types
|
||||
from pytorch3d.implicitron.dataset.utils import (
|
||||
adjust_camera_to_bbox_crop_,
|
||||
|
||||
@@ -38,6 +38,7 @@ from pytorch3d.implicitron.dataset.utils import is_known_frame_scalar
|
||||
from pytorch3d.implicitron.tools.config import registry, ReplaceableBase
|
||||
from pytorch3d.renderer.camera_utils import join_cameras_as_batch
|
||||
from pytorch3d.renderer.cameras import CamerasBase
|
||||
|
||||
from tqdm import tqdm
|
||||
|
||||
|
||||
@@ -326,9 +327,9 @@ class JsonIndexDataset(DatasetBase, ReplaceableBase):
|
||||
assert os.path.normpath(
|
||||
# pyre-ignore[16]
|
||||
self.frame_annots[idx]["frame_annotation"].image.path
|
||||
) == os.path.normpath(path), (
|
||||
f"Inconsistent frame indices {seq_name, frame_no, path}."
|
||||
)
|
||||
) == os.path.normpath(
|
||||
path
|
||||
), f"Inconsistent frame indices {seq_name, frame_no, path}."
|
||||
return idx
|
||||
|
||||
dataset_idx = [
|
||||
|
||||
@@ -21,6 +21,7 @@ from pytorch3d.renderer.cameras import CamerasBase
|
||||
|
||||
from .dataset_map_provider import DatasetMap, DatasetMapProviderBase, PathManagerFactory
|
||||
from .json_index_dataset import JsonIndexDataset
|
||||
|
||||
from .utils import (
|
||||
DATASET_TYPE_KNOWN,
|
||||
DATASET_TYPE_TEST,
|
||||
|
||||
@@ -18,6 +18,7 @@ from typing import Dict, List, Optional, Tuple, Type, Union
|
||||
|
||||
import numpy as np
|
||||
from iopath.common.file_io import PathManager
|
||||
|
||||
from omegaconf import DictConfig
|
||||
from pytorch3d.implicitron.dataset.dataset_map_provider import (
|
||||
DatasetMap,
|
||||
@@ -30,6 +31,7 @@ from pytorch3d.implicitron.tools.config import (
|
||||
registry,
|
||||
run_auto_creation,
|
||||
)
|
||||
|
||||
from pytorch3d.renderer.cameras import CamerasBase
|
||||
from tqdm import tqdm
|
||||
|
||||
|
||||
@@ -12,6 +12,7 @@ import torch
|
||||
from pytorch3d.implicitron.tools.config import registry
|
||||
|
||||
from .load_llff import load_llff_data
|
||||
|
||||
from .single_sequence_dataset import (
|
||||
_interpret_blender_cameras,
|
||||
SingleSceneDatasetMapProviderBase,
|
||||
|
||||
@@ -8,6 +8,7 @@ import os
|
||||
import warnings
|
||||
|
||||
import numpy as np
|
||||
|
||||
from PIL import Image
|
||||
|
||||
|
||||
|
||||
@@ -13,6 +13,7 @@ import struct
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
|
||||
from pytorch3d.implicitron.dataset.types import (
|
||||
DepthAnnotation,
|
||||
ImageAnnotation,
|
||||
@@ -21,6 +22,7 @@ from pytorch3d.implicitron.dataset.types import (
|
||||
VideoAnnotation,
|
||||
ViewpointAnnotation,
|
||||
)
|
||||
|
||||
from sqlalchemy import LargeBinary
|
||||
from sqlalchemy.orm import (
|
||||
composite,
|
||||
|
||||
@@ -10,6 +10,7 @@ import hashlib
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
|
||||
import urllib
|
||||
from dataclasses import dataclass, Field, field
|
||||
from typing import (
|
||||
@@ -31,11 +32,13 @@ 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 (
|
||||
FrameData,
|
||||
FrameDataBuilder, # noqa
|
||||
FrameDataBuilderBase,
|
||||
)
|
||||
|
||||
from pytorch3d.implicitron.tools.config import (
|
||||
registry,
|
||||
ReplaceableBase,
|
||||
@@ -483,10 +486,9 @@ class SqlIndexDataset(DatasetBase, ReplaceableBase):
|
||||
*self._get_pick_filters(),
|
||||
*self._get_exclude_filters(),
|
||||
]
|
||||
if pick_sequences_sql_clause := self.pick_sequences_sql_clause:
|
||||
if self.pick_sequences_sql_clause:
|
||||
print("Applying the custom SQL clause.")
|
||||
# pyre-ignore[6]: TextClause is compatible with where conditions
|
||||
where_conditions.append(sa.text(pick_sequences_sql_clause))
|
||||
where_conditions.append(sa.text(self.pick_sequences_sql_clause))
|
||||
|
||||
def add_where(stmt):
|
||||
return stmt.where(*where_conditions) if where_conditions else stmt
|
||||
@@ -506,7 +508,6 @@ class SqlIndexDataset(DatasetBase, ReplaceableBase):
|
||||
|
||||
subquery = add_where(subquery).subquery()
|
||||
stmt = sa.select(subquery.c.sequence_name).where(
|
||||
# pyre-ignore[6]: SQLAlchemy column comparison returns ColumnElement, not bool
|
||||
subquery.c.row_number <= self.limit_sequences_per_category_to
|
||||
)
|
||||
|
||||
@@ -635,10 +636,9 @@ class SqlIndexDataset(DatasetBase, ReplaceableBase):
|
||||
)
|
||||
)
|
||||
|
||||
if pick_frames_sql_clause := self.pick_frames_sql_clause:
|
||||
if self.pick_frames_sql_clause:
|
||||
logger.info("Applying the custom SQL clause.")
|
||||
# pyre-ignore[6]: TextClause is compatible with where conditions
|
||||
pick_frames_criteria.append(sa.text(pick_frames_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)
|
||||
@@ -701,10 +701,9 @@ class SqlIndexDataset(DatasetBase, ReplaceableBase):
|
||||
)
|
||||
)
|
||||
|
||||
if pick_frames_sql_clause := self.pick_frames_sql_clause:
|
||||
if self.pick_frames_sql_clause:
|
||||
logger.info(" applying custom SQL clause")
|
||||
# pyre-ignore[6]: TextClause is compatible with where conditions
|
||||
where_conditions.append(sa.text(pick_frames_sql_clause))
|
||||
where_conditions.append(sa.text(self.pick_frames_sql_clause))
|
||||
|
||||
if where_conditions:
|
||||
stmt = stmt.where(*where_conditions)
|
||||
@@ -756,7 +755,7 @@ class SqlIndexDataset(DatasetBase, ReplaceableBase):
|
||||
if pick_sequences:
|
||||
old_len = len(eval_batches)
|
||||
eval_batches = [b for b in eval_batches if b[0][0] in pick_sequences]
|
||||
logger.warning(
|
||||
logger.warn(
|
||||
f"Picked eval batches by sequence/cat: {old_len} -> {len(eval_batches)}"
|
||||
)
|
||||
|
||||
@@ -764,7 +763,7 @@ class SqlIndexDataset(DatasetBase, ReplaceableBase):
|
||||
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.warning(
|
||||
logger.warn(
|
||||
f"Excluded eval batches by sequence: {old_len} -> {len(eval_batches)}"
|
||||
)
|
||||
|
||||
|
||||
@@ -12,7 +12,9 @@ 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,
|
||||
|
||||
@@ -18,6 +18,7 @@ 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__)
|
||||
|
||||
@@ -15,6 +15,7 @@ from typing import List, Optional, Tuple, TypeVar, Union
|
||||
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
|
||||
|
||||
@@ -14,6 +14,7 @@ import warnings
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
import tqdm
|
||||
from pytorch3d.implicitron.evaluation import evaluate_new_view_synthesis as evaluate
|
||||
from pytorch3d.implicitron.models.base_model import EvaluationMode, ImplicitronModelBase
|
||||
|
||||
@@ -10,6 +10,7 @@ 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
|
||||
|
||||
@@ -16,6 +16,7 @@ from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING, Union
|
||||
|
||||
import torch
|
||||
from omegaconf import DictConfig
|
||||
|
||||
from pytorch3d.implicitron.models.base_model import (
|
||||
ImplicitronModelBase,
|
||||
ImplicitronRender,
|
||||
@@ -27,6 +28,7 @@ from pytorch3d.implicitron.models.metrics import (
|
||||
RegularizationMetricsBase,
|
||||
ViewMetricsBase,
|
||||
)
|
||||
|
||||
from pytorch3d.implicitron.models.renderer.base import (
|
||||
BaseRenderer,
|
||||
EvaluationMode,
|
||||
@@ -36,6 +38,7 @@ from pytorch3d.implicitron.models.renderer.base import (
|
||||
RenderSamplingMode,
|
||||
)
|
||||
from pytorch3d.implicitron.models.renderer.ray_sampler import RaySamplerBase
|
||||
|
||||
from pytorch3d.implicitron.models.utils import (
|
||||
apply_chunked,
|
||||
chunk_generator,
|
||||
@@ -50,6 +53,7 @@ from pytorch3d.implicitron.tools.config import (
|
||||
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
|
||||
|
||||
@@ -10,6 +10,7 @@ from abc import ABC, abstractmethod
|
||||
from typing import Optional
|
||||
|
||||
from pytorch3d.implicitron.models.renderer.base import ImplicitronRayBundle
|
||||
|
||||
from pytorch3d.implicitron.tools.config import ReplaceableBase
|
||||
from pytorch3d.renderer.cameras import CamerasBase
|
||||
|
||||
|
||||
@@ -16,11 +16,14 @@ This file contains
|
||||
|
||||
import logging
|
||||
from dataclasses import field
|
||||
|
||||
from enum import Enum
|
||||
from typing import Dict, Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
from omegaconf import DictConfig
|
||||
|
||||
from pytorch3d.implicitron.tools.config import (
|
||||
Configurable,
|
||||
registry,
|
||||
|
||||
@@ -11,6 +11,7 @@ import torch
|
||||
from pytorch3d.implicitron.models.renderer.base import ImplicitronRayBundle
|
||||
from pytorch3d.implicitron.tools.config import registry
|
||||
from pytorch3d.renderer.implicit import HarmonicEmbedding
|
||||
|
||||
from torch import nn
|
||||
|
||||
from .base import ImplicitFunctionBase
|
||||
|
||||
@@ -21,6 +21,7 @@ from pytorch3d.renderer.implicit import HarmonicEmbedding
|
||||
from pytorch3d.renderer.implicit.utils import ray_bundle_to_ray_points
|
||||
|
||||
from .base import ImplicitFunctionBase
|
||||
|
||||
from .decoding_functions import ( # noqa
|
||||
_xavier_init,
|
||||
MLPWithInputSkips,
|
||||
|
||||
@@ -9,6 +9,7 @@
|
||||
from typing import Callable, Optional
|
||||
|
||||
import torch
|
||||
|
||||
import torch.nn.functional as F
|
||||
from pytorch3d.common.compat import prod
|
||||
from pytorch3d.implicitron.models.renderer.base import ImplicitronRayBundle
|
||||
|
||||
@@ -21,6 +21,8 @@ import logging
|
||||
import warnings
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from distutils.version import LooseVersion
|
||||
from typing import Any, Callable, ClassVar, Dict, Iterator, List, Optional, Tuple, Type
|
||||
|
||||
import torch
|
||||
@@ -220,8 +222,7 @@ class VoxelGridBase(ReplaceableBase, torch.nn.Module):
|
||||
+ "| 'bicubic' | 'linear' | 'area' | 'nearest-exact'"
|
||||
)
|
||||
|
||||
# We assume PyTorch 1.11 and newer.
|
||||
interpolate_has_antialias = True
|
||||
interpolate_has_antialias = LooseVersion(torch.__version__) >= "1.11"
|
||||
|
||||
if antialias and not interpolate_has_antialias:
|
||||
warnings.warn("Antialiased interpolation requires PyTorch 1.11+; ignoring")
|
||||
|
||||
@@ -13,7 +13,9 @@ from dataclasses import fields
|
||||
from typing import Callable, Dict, Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
from omegaconf import DictConfig
|
||||
|
||||
from pytorch3d.implicitron.models.implicit_function.base import ImplicitFunctionBase
|
||||
from pytorch3d.implicitron.models.implicit_function.decoding_functions import (
|
||||
DecoderFunctionBase,
|
||||
|
||||
@@ -17,6 +17,7 @@ from typing import Any, Callable, Dict, List, Optional, Tuple, TYPE_CHECKING, Un
|
||||
|
||||
import torch
|
||||
from omegaconf import DictConfig
|
||||
|
||||
from pytorch3d.implicitron.models.base_model import (
|
||||
ImplicitronModelBase,
|
||||
ImplicitronRender,
|
||||
@@ -27,6 +28,7 @@ from pytorch3d.implicitron.models.metrics import (
|
||||
RegularizationMetricsBase,
|
||||
ViewMetricsBase,
|
||||
)
|
||||
|
||||
from pytorch3d.implicitron.models.renderer.base import (
|
||||
BaseRenderer,
|
||||
EvaluationMode,
|
||||
@@ -48,6 +50,7 @@ from pytorch3d.implicitron.tools.config import (
|
||||
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
|
||||
|
||||
@@ -11,6 +11,7 @@ import copy
|
||||
import torch
|
||||
from pytorch3d.implicitron.models.renderer.base import ImplicitronRayBundle
|
||||
from pytorch3d.implicitron.tools.config import Configurable, expand_args_fields
|
||||
|
||||
from pytorch3d.renderer.implicit.sample_pdf import sample_pdf
|
||||
|
||||
|
||||
|
||||
@@ -12,6 +12,7 @@ import torch
|
||||
from pytorch3d.implicitron.models.renderer.base import ImplicitronRayBundle
|
||||
from pytorch3d.implicitron.tools.config import enable_get_default_args
|
||||
from pytorch3d.renderer.implicit import HarmonicEmbedding
|
||||
|
||||
from torch import nn
|
||||
|
||||
|
||||
|
||||
@@ -17,8 +17,11 @@ from typing import Any, Dict, Optional, Tuple
|
||||
import torch
|
||||
import tqdm
|
||||
from pytorch3d.common.compat import prod
|
||||
|
||||
from pytorch3d.implicitron.models.renderer.base import ImplicitronRayBundle
|
||||
|
||||
from pytorch3d.implicitron.tools import image_utils
|
||||
|
||||
from pytorch3d.implicitron.tools.utils import cat_dataclass
|
||||
|
||||
|
||||
@@ -80,9 +83,9 @@ def preprocess_input(
|
||||
|
||||
if mask_depths and fg_mask is not None and depth_map is not None:
|
||||
# mask the depths
|
||||
assert mask_threshold > 0.0, (
|
||||
"Depths should be masked only with thresholded masks"
|
||||
)
|
||||
assert (
|
||||
mask_threshold > 0.0
|
||||
), "Depths should be masked only with thresholded masks"
|
||||
warnings.warn("Masking depths!")
|
||||
depth_map = depth_map * fg_mask
|
||||
|
||||
|
||||
@@ -304,11 +304,11 @@ def _show_predictions(
|
||||
assert isinstance(preds, list)
|
||||
|
||||
pred_all = []
|
||||
# Randomly choose a subset of the rendered images, sort by order in the sequence
|
||||
# Randomly choose a subset of the rendered images, sort by ordr in the sequence
|
||||
n_samples = min(n_samples, len(preds))
|
||||
pred_idx = sorted(random.sample(list(range(len(preds))), n_samples))
|
||||
for predi in pred_idx:
|
||||
# Make the concatenation for the same camera vertically
|
||||
# Make the concatentation for the same camera vertically
|
||||
pred_all.append(
|
||||
torch.cat(
|
||||
[
|
||||
@@ -359,7 +359,7 @@ def _generate_prediction_videos(
|
||||
vws = {}
|
||||
for k in predicted_keys:
|
||||
if k not in preds[0]:
|
||||
logger.warning(f"Cannot generate video for prediction key '{k}'")
|
||||
logger.warn(f"Cannot generate video for prediction key '{k}'")
|
||||
continue
|
||||
cache_dir = (
|
||||
None
|
||||
|
||||
@@ -10,6 +10,7 @@ import math
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import pytorch3d
|
||||
|
||||
import torch
|
||||
from pytorch3d.ops import packed_to_padded
|
||||
from pytorch3d.renderer import PerspectiveCameras
|
||||
|
||||
@@ -499,7 +499,7 @@ class StatsJSONEncoder(json.JSONEncoder):
|
||||
return enc
|
||||
else:
|
||||
raise TypeError(
|
||||
f"Object of type {o.__class__.__name__} is not JSON serializable"
|
||||
f"Object of type {o.__class__.__name__} " f"is not JSON serializable"
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -17,6 +17,7 @@ import matplotlib
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from PIL import Image
|
||||
|
||||
_NO_TORCHVISION = False
|
||||
|
||||
@@ -796,7 +796,7 @@ def save_obj(
|
||||
# Create .mtl file with the material name and texture map filename
|
||||
# TODO: enable material properties to also be saved.
|
||||
with _open_file(mtl_path, path_manager, "w") as f_mtl:
|
||||
lines = f"newmtl mesh\nmap_Kd {output_path.stem}.png\n"
|
||||
lines = f"newmtl mesh\n" f"map_Kd {output_path.stem}.png\n"
|
||||
f_mtl.write(lines)
|
||||
|
||||
|
||||
|
||||
@@ -8,8 +8,11 @@
|
||||
|
||||
|
||||
from .chamfer import chamfer_distance
|
||||
|
||||
from .mesh_edge_loss import mesh_edge_loss
|
||||
|
||||
from .mesh_laplacian_smoothing import mesh_laplacian_smoothing
|
||||
|
||||
from .mesh_normal_consistency import mesh_normal_consistency
|
||||
from .point_mesh_distance import point_mesh_edge_distance, point_mesh_face_distance
|
||||
|
||||
|
||||
@@ -8,14 +8,17 @@
|
||||
|
||||
from .ball_query import ball_query
|
||||
from .cameras_alignment import corresponding_cameras_alignment
|
||||
|
||||
from .cubify import cubify
|
||||
from .graph_conv import GraphConv
|
||||
from .interp_face_attrs import interpolate_face_attributes
|
||||
from .iou_box3d import box3d_overlap
|
||||
from .knn import knn_gather, knn_points
|
||||
from .laplacian_matrices import cot_laplacian, laplacian, norm_laplacian
|
||||
|
||||
from .mesh_face_areas_normals import mesh_face_areas_normals
|
||||
from .mesh_filtering import taubin_smoothing
|
||||
|
||||
from .packed_to_padded import packed_to_padded, padded_to_packed
|
||||
from .perspective_n_points import efficient_pnp
|
||||
from .points_alignment import corresponding_points_alignment, iterative_closest_point
|
||||
@@ -27,7 +30,9 @@ from .points_to_volumes import (
|
||||
add_pointclouds_to_volumes,
|
||||
add_points_features_to_volume_densities_features,
|
||||
)
|
||||
|
||||
from .sample_farthest_points import sample_farthest_points
|
||||
|
||||
from .sample_points_from_meshes import sample_points_from_meshes
|
||||
from .subdivide_meshes import SubdivideMeshes
|
||||
from .utils import (
|
||||
@@ -37,6 +42,7 @@ from .utils import (
|
||||
is_pointclouds,
|
||||
wmean,
|
||||
)
|
||||
|
||||
from .vert_align import vert_align
|
||||
|
||||
|
||||
|
||||
@@ -23,13 +23,11 @@ class _ball_query(Function):
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, p1, p2, lengths1, lengths2, K, radius, skip_points_outside_cube):
|
||||
def forward(ctx, p1, p2, lengths1, lengths2, K, radius):
|
||||
"""
|
||||
Arguments defintions the same as in the ball_query function
|
||||
"""
|
||||
idx, dists = _C.ball_query(
|
||||
p1, p2, lengths1, lengths2, K, radius, skip_points_outside_cube
|
||||
)
|
||||
idx, dists = _C.ball_query(p1, p2, lengths1, lengths2, K, radius)
|
||||
ctx.save_for_backward(p1, p2, lengths1, lengths2, idx)
|
||||
ctx.mark_non_differentiable(idx)
|
||||
return dists, idx
|
||||
@@ -51,7 +49,7 @@ class _ball_query(Function):
|
||||
grad_p1, grad_p2 = _C.knn_points_backward(
|
||||
p1, p2, lengths1, lengths2, idx, 2, grad_dists
|
||||
)
|
||||
return grad_p1, grad_p2, None, None, None, None, None
|
||||
return grad_p1, grad_p2, None, None, None, None
|
||||
|
||||
|
||||
def ball_query(
|
||||
@@ -62,7 +60,6 @@ def ball_query(
|
||||
K: int = 500,
|
||||
radius: float = 0.2,
|
||||
return_nn: bool = True,
|
||||
skip_points_outside_cube: bool = False,
|
||||
):
|
||||
"""
|
||||
Ball Query is an alternative to KNN. It can be
|
||||
@@ -101,9 +98,6 @@ def ball_query(
|
||||
within the radius
|
||||
radius: the radius around each point within which the neighbors need to be located
|
||||
return_nn: If set to True returns the K neighbor points in p2 for each point in p1.
|
||||
skip_points_outside_cube: If set to True, reduce multiplications of float values
|
||||
by not explicitly calculating distances to points that fall outside the
|
||||
D-cube with side length (2*radius) centered at each point in p1.
|
||||
|
||||
Returns:
|
||||
dists: Tensor of shape (N, P1, K) giving the squared distances to
|
||||
@@ -140,9 +134,7 @@ def ball_query(
|
||||
if lengths2 is None:
|
||||
lengths2 = torch.full((N,), P2, dtype=torch.int64, device=p1.device)
|
||||
|
||||
dists, idx = _ball_query.apply(
|
||||
p1, p2, lengths1, lengths2, K, radius, skip_points_outside_cube
|
||||
)
|
||||
dists, idx = _ball_query.apply(p1, p2, lengths1, lengths2, K, radius)
|
||||
|
||||
# Gather the neighbors if needed
|
||||
points_nn = masked_gather(p2, idx) if return_nn else None
|
||||
|
||||
@@ -11,7 +11,9 @@ from typing import Optional
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from pytorch3d.common.compat import meshgrid_ij
|
||||
|
||||
from pytorch3d.structures import Meshes
|
||||
|
||||
|
||||
|
||||
@@ -47,7 +47,8 @@ def laplacian(verts: torch.Tensor, edges: torch.Tensor) -> torch.Tensor:
|
||||
# i.e. A[i, j] = 1 if (i,j) is an edge, or
|
||||
# A[e0, e1] = 1 & A[e1, e0] = 1
|
||||
ones = torch.ones(idx.shape[1], dtype=torch.float32, device=verts.device)
|
||||
A = torch.sparse_coo_tensor(idx, ones, (V, V), dtype=torch.float32)
|
||||
# pyre-fixme[16]: Module `sparse` has no attribute `FloatTensor`.
|
||||
A = torch.sparse.FloatTensor(idx, ones, (V, V))
|
||||
|
||||
# the sum of i-th row of A gives the degree of the i-th vertex
|
||||
deg = torch.sparse.sum(A, dim=1).to_dense()
|
||||
@@ -61,13 +62,15 @@ def laplacian(verts: torch.Tensor, edges: torch.Tensor) -> torch.Tensor:
|
||||
# pyre-fixme[58]: `/` is not supported for operand types `float` and `Tensor`.
|
||||
deg1 = torch.where(deg1 > 0.0, 1.0 / deg1, deg1)
|
||||
val = torch.cat([deg0, deg1])
|
||||
L = torch.sparse_coo_tensor(idx, val, (V, V), dtype=torch.float32)
|
||||
# pyre-fixme[16]: Module `sparse` has no attribute `FloatTensor`.
|
||||
L = torch.sparse.FloatTensor(idx, val, (V, V))
|
||||
|
||||
# Then we add the diagonal values L[i, i] = -1.
|
||||
idx = torch.arange(V, device=verts.device)
|
||||
idx = torch.stack([idx, idx], dim=0)
|
||||
ones = torch.ones(idx.shape[1], dtype=torch.float32, device=verts.device)
|
||||
L -= torch.sparse_coo_tensor(idx, ones, (V, V), dtype=torch.float32)
|
||||
# pyre-fixme[16]: Module `sparse` has no attribute `FloatTensor`.
|
||||
L -= torch.sparse.FloatTensor(idx, ones, (V, V))
|
||||
|
||||
return L
|
||||
|
||||
@@ -123,7 +126,8 @@ def cot_laplacian(
|
||||
ii = faces[:, [1, 2, 0]]
|
||||
jj = faces[:, [2, 0, 1]]
|
||||
idx = torch.stack([ii, jj], dim=0).view(2, F * 3)
|
||||
L = torch.sparse_coo_tensor(idx, cot.view(-1), (V, V), dtype=torch.float32)
|
||||
# pyre-fixme[16]: Module `sparse` has no attribute `FloatTensor`.
|
||||
L = torch.sparse.FloatTensor(idx, cot.view(-1), (V, V))
|
||||
|
||||
# Make it symmetric; this means we are also setting
|
||||
# L[v2, v1] = cota
|
||||
@@ -163,7 +167,7 @@ def norm_laplacian(
|
||||
v0, v1 = edge_verts[:, 0], edge_verts[:, 1]
|
||||
|
||||
# Side lengths of each edge, of shape (E,)
|
||||
w01 = torch.reciprocal((v0 - v1).norm(dim=1) + eps)
|
||||
w01 = 1.0 / ((v0 - v1).norm(dim=1) + eps)
|
||||
|
||||
# Construct a sparse matrix by basically doing:
|
||||
# L[v0, v1] = w01
|
||||
@@ -171,7 +175,8 @@ def norm_laplacian(
|
||||
e01 = edges.t() # (2, E)
|
||||
|
||||
V = verts.shape[0]
|
||||
L = torch.sparse_coo_tensor(e01, w01, (V, V), dtype=torch.float32)
|
||||
# pyre-fixme[16]: Module `sparse` has no attribute `FloatTensor`.
|
||||
L = torch.sparse.FloatTensor(e01, w01, (V, V))
|
||||
L = L + L.t()
|
||||
|
||||
return L
|
||||
|
||||
@@ -55,7 +55,6 @@ def sample_farthest_points(
|
||||
N, P, D = points.shape
|
||||
device = points.device
|
||||
|
||||
constant_length = lengths is None
|
||||
# Validate inputs
|
||||
if lengths is None:
|
||||
lengths = torch.full((N,), P, dtype=torch.int64, device=device)
|
||||
@@ -66,9 +65,7 @@ def sample_farthest_points(
|
||||
raise ValueError("A value in lengths was too large.")
|
||||
|
||||
# TODO: support providing K as a ratio of the total number of points instead of as an int
|
||||
max_K = -1
|
||||
if isinstance(K, int):
|
||||
max_K = K
|
||||
K = torch.full((N,), K, dtype=torch.int64, device=device)
|
||||
elif isinstance(K, list):
|
||||
K = torch.tensor(K, dtype=torch.int64, device=device)
|
||||
@@ -85,19 +82,15 @@ def sample_farthest_points(
|
||||
K = K.to(torch.int64)
|
||||
|
||||
# Generate the starting indices for sampling
|
||||
start_idxs = torch.zeros_like(lengths)
|
||||
if random_start_point:
|
||||
if constant_length:
|
||||
start_idxs = torch.randint(high=P, size=(N,), device=device)
|
||||
else:
|
||||
start_idxs = (lengths * torch.rand(lengths.size(), device=device)).to(
|
||||
torch.int64
|
||||
)
|
||||
else:
|
||||
start_idxs = torch.zeros_like(lengths)
|
||||
for n in range(N):
|
||||
# pyre-fixme[6]: For 1st param expected `int` but got `Tensor`.
|
||||
start_idxs[n] = torch.randint(high=lengths[n], size=(1,)).item()
|
||||
|
||||
with torch.no_grad():
|
||||
# pyre-fixme[16]: `pytorch3d_._C` has no attribute `sample_farthest_points`.
|
||||
idx = _C.sample_farthest_points(points, lengths, K, start_idxs, max_K)
|
||||
idx = _C.sample_farthest_points(points, lengths, K, start_idxs)
|
||||
sampled_points = masked_gather(points, idx)
|
||||
|
||||
return sampled_points, idx
|
||||
|
||||
@@ -16,7 +16,9 @@ import sys
|
||||
from typing import Tuple, Union
|
||||
|
||||
import torch
|
||||
|
||||
from pytorch3d.ops.mesh_face_areas_normals import mesh_face_areas_normals
|
||||
|
||||
from pytorch3d.ops.packed_to_padded import packed_to_padded
|
||||
from pytorch3d.renderer.mesh.rasterizer import Fragments as MeshFragments
|
||||
|
||||
|
||||
@@ -69,6 +69,7 @@ from .mesh import (
|
||||
TexturesUV,
|
||||
TexturesVertex,
|
||||
)
|
||||
|
||||
from .points import (
|
||||
AlphaCompositor,
|
||||
NormWeightedCompositor,
|
||||
|
||||
@@ -153,12 +153,12 @@ def _pulsar_from_opencv_projection(
|
||||
# Check image sizes.
|
||||
image_w = image_size_wh[0, 0]
|
||||
image_h = image_size_wh[0, 1]
|
||||
assert torch.all(image_size_wh[:, 0] == image_w), (
|
||||
"All images in a batch must have the same width!"
|
||||
)
|
||||
assert torch.all(image_size_wh[:, 1] == image_h), (
|
||||
"All images in a batch must have the same height!"
|
||||
)
|
||||
assert torch.all(
|
||||
image_size_wh[:, 0] == image_w
|
||||
), "All images in a batch must have the same width!"
|
||||
assert torch.all(
|
||||
image_size_wh[:, 1] == image_h
|
||||
), "All images in a batch must have the same height!"
|
||||
# Focal length.
|
||||
fx = camera_matrix[:, 0, 0].unsqueeze(1)
|
||||
fy = camera_matrix[:, 1, 1].unsqueeze(1)
|
||||
|
||||
@@ -12,6 +12,7 @@ from .clip import (
|
||||
ClippedFaces,
|
||||
convert_clipped_rasterization_to_original_faces,
|
||||
)
|
||||
|
||||
from .rasterize_meshes import rasterize_meshes
|
||||
from .rasterizer import MeshRasterizer, RasterizationSettings
|
||||
from .renderer import MeshRenderer, MeshRendererWithFragments
|
||||
|
||||
@@ -14,6 +14,7 @@ import torch
|
||||
from pytorch3d import _C
|
||||
|
||||
from ..utils import parse_image_size
|
||||
|
||||
from .clip import (
|
||||
clip_faces,
|
||||
ClipFrustum,
|
||||
|
||||
@@ -625,7 +625,9 @@ class TexturesAtlas(TexturesBase):
|
||||
of length `k`.
|
||||
"""
|
||||
if len(faces_ids_list) != len(self.atlas_list()):
|
||||
raise IndexError("faces_ids_list must be of the same length as atlas_list.")
|
||||
raise IndexError(
|
||||
"faces_ids_list must be of " "the same length as atlas_list."
|
||||
)
|
||||
|
||||
sub_features = []
|
||||
for atlas, faces_ids in zip(self.atlas_list(), faces_ids_list):
|
||||
@@ -1655,7 +1657,7 @@ class TexturesUV(TexturesBase):
|
||||
raise NotImplementedError("This function does not support multiple maps.")
|
||||
if len(faces_ids_list) != len(self.faces_uvs_padded()):
|
||||
raise IndexError(
|
||||
"faces_uvs_padded must be of the same length as face_ids_list."
|
||||
"faces_uvs_padded must be of " "the same length as face_ids_list."
|
||||
)
|
||||
|
||||
sub_faces_uvs, sub_verts_uvs, sub_maps = [], [], []
|
||||
@@ -1869,7 +1871,7 @@ class TexturesVertex(TexturesBase):
|
||||
"""
|
||||
if len(vertex_ids_list) != len(self.verts_features_list()):
|
||||
raise IndexError(
|
||||
"verts_features_list must be of the same length as vertex_ids_list."
|
||||
"verts_features_list must be of " "the same length as vertex_ids_list."
|
||||
)
|
||||
|
||||
sub_features = []
|
||||
|
||||
@@ -24,6 +24,7 @@ from typing import Any, Dict
|
||||
|
||||
os.environ["PYOPENGL_PLATFORM"] = "egl"
|
||||
import OpenGL.EGL as egl # noqa
|
||||
|
||||
import pycuda.driver as cuda # noqa
|
||||
from OpenGL._opaque import opaque_pointer_cls # noqa
|
||||
from OpenGL.raw.EGL._errors import EGLError # noqa
|
||||
|
||||
@@ -17,12 +17,15 @@ import numpy as np
|
||||
import OpenGL.GL as gl
|
||||
import pycuda.gl
|
||||
import torch
|
||||
|
||||
import torch.nn as nn
|
||||
|
||||
from pytorch3d.structures.meshes import Meshes
|
||||
|
||||
from ..cameras import FoVOrthographicCameras, FoVPerspectiveCameras
|
||||
from ..mesh.rasterizer import Fragments, RasterizationSettings
|
||||
from ..utils import parse_image_size
|
||||
|
||||
from .opengl_utils import _torch_to_opengl, global_device_context_store
|
||||
|
||||
# Shader strings, used below to compile an OpenGL program.
|
||||
|
||||
@@ -9,7 +9,9 @@
|
||||
import torch
|
||||
|
||||
from .compositor import AlphaCompositor, NormWeightedCompositor
|
||||
|
||||
from .pulsar.unified import PulsarPointsRenderer
|
||||
|
||||
from .rasterize_points import rasterize_points
|
||||
from .rasterizer import PointsRasterizationSettings, PointsRasterizer
|
||||
from .renderer import PointsRenderer
|
||||
|
||||
@@ -11,6 +11,7 @@ from typing import List, Optional, Tuple, Union
|
||||
import numpy as np
|
||||
import torch
|
||||
from pytorch3d import _C
|
||||
|
||||
from pytorch3d.renderer.mesh.rasterize_meshes import pix_to_non_square_ndc
|
||||
|
||||
from ..utils import parse_image_size
|
||||
|
||||
@@ -531,9 +531,9 @@ class Meshes:
|
||||
list of tensors of vertices of shape (V_n, 3).
|
||||
"""
|
||||
if self._verts_list is None:
|
||||
assert self._verts_padded is not None, (
|
||||
"verts_padded is required to compute verts_list."
|
||||
)
|
||||
assert (
|
||||
self._verts_padded is not None
|
||||
), "verts_padded is required to compute verts_list."
|
||||
self._verts_list = struct_utils.padded_to_list(
|
||||
self._verts_padded, self.num_verts_per_mesh().tolist()
|
||||
)
|
||||
@@ -547,9 +547,9 @@ class Meshes:
|
||||
list of tensors of faces of shape (F_n, 3).
|
||||
"""
|
||||
if self._faces_list is None:
|
||||
assert self._faces_padded is not None, (
|
||||
"faces_padded is required to compute faces_list."
|
||||
)
|
||||
assert (
|
||||
self._faces_padded is not None
|
||||
), "faces_padded is required to compute faces_list."
|
||||
self._faces_list = struct_utils.padded_to_list(
|
||||
self._faces_padded, self.num_faces_per_mesh().tolist()
|
||||
)
|
||||
@@ -925,9 +925,9 @@ class Meshes:
|
||||
|
||||
verts_list = self.verts_list()
|
||||
faces_list = self.faces_list()
|
||||
assert faces_list is not None and verts_list is not None, (
|
||||
"faces_list and verts_list arguments are required"
|
||||
)
|
||||
assert (
|
||||
faces_list is not None and verts_list is not None
|
||||
), "faces_list and verts_list arguments are required"
|
||||
|
||||
if self.isempty():
|
||||
self._faces_padded = torch.zeros(
|
||||
|
||||
@@ -433,9 +433,9 @@ class Pointclouds:
|
||||
list of tensors of points of shape (P_n, 3).
|
||||
"""
|
||||
if self._points_list is None:
|
||||
assert self._points_padded is not None, (
|
||||
"points_padded is required to compute points_list."
|
||||
)
|
||||
assert (
|
||||
self._points_padded is not None
|
||||
), "points_padded is required to compute points_list."
|
||||
points_list = []
|
||||
for i in range(self._N):
|
||||
points_list.append(
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user