mirror of
https://github.com/facebookresearch/pytorch3d.git
synced 2025-07-31 10:52:50 +08:00
83bacda8fb
Summary: Fix recent flake complaints Reviewed By: MichaelRamamonjisoa Differential Revision: D51811912 fbshipit-source-id: 65183f5bc7058da910e4d5a63b2250ce8637f1cc
264 lines
10 KiB
Python
264 lines
10 KiB
Python
# 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 contextlib
|
|
import gzip
|
|
import os
|
|
import unittest
|
|
from typing import List
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from pytorch3d.implicitron.dataset import types
|
|
from pytorch3d.implicitron.dataset.dataset_base import FrameData
|
|
from pytorch3d.implicitron.dataset.frame_data import FrameDataBuilder
|
|
from pytorch3d.implicitron.dataset.utils import (
|
|
get_bbox_from_mask,
|
|
load_16big_png_depth,
|
|
load_1bit_png_mask,
|
|
load_depth,
|
|
load_depth_mask,
|
|
load_image,
|
|
load_mask,
|
|
safe_as_tensor,
|
|
transpose_normalize_image,
|
|
)
|
|
from pytorch3d.implicitron.tools.config import get_default_args
|
|
from pytorch3d.renderer.cameras import PerspectiveCameras
|
|
|
|
from tests.common_testing import TestCaseMixin
|
|
from tests.implicitron.common_resources import get_skateboard_data
|
|
|
|
|
|
class TestFrameDataBuilder(TestCaseMixin, unittest.TestCase):
|
|
def setUp(self):
|
|
torch.manual_seed(42)
|
|
|
|
category = "skateboard"
|
|
stack = contextlib.ExitStack()
|
|
self.dataset_root, self.path_manager = stack.enter_context(
|
|
get_skateboard_data()
|
|
)
|
|
self.addCleanup(stack.close)
|
|
self.image_height = 768
|
|
self.image_width = 512
|
|
|
|
self.frame_data_builder = FrameDataBuilder(
|
|
image_height=self.image_height,
|
|
image_width=self.image_width,
|
|
dataset_root=self.dataset_root,
|
|
path_manager=self.path_manager,
|
|
)
|
|
|
|
# loading single frame annotation of dataset (see JsonIndexDataset._load_frames())
|
|
frame_file = os.path.join(self.dataset_root, category, "frame_annotations.jgz")
|
|
local_file = self.path_manager.get_local_path(frame_file)
|
|
with gzip.open(local_file, "rt", encoding="utf8") as zipfile:
|
|
frame_annots_list = types.load_dataclass(
|
|
zipfile, List[types.FrameAnnotation]
|
|
)
|
|
self.frame_annotation = frame_annots_list[0]
|
|
|
|
sequence_annotations_file = os.path.join(
|
|
self.dataset_root, category, "sequence_annotations.jgz"
|
|
)
|
|
local_file = self.path_manager.get_local_path(sequence_annotations_file)
|
|
with gzip.open(local_file, "rt", encoding="utf8") as zipfile:
|
|
seq_annots_list = types.load_dataclass(
|
|
zipfile, List[types.SequenceAnnotation]
|
|
)
|
|
seq_annots = {entry.sequence_name: entry for entry in seq_annots_list}
|
|
self.seq_annotation = seq_annots[self.frame_annotation.sequence_name]
|
|
|
|
point_cloud = self.seq_annotation.point_cloud
|
|
self.frame_data = FrameData(
|
|
frame_number=safe_as_tensor(self.frame_annotation.frame_number, torch.long),
|
|
frame_timestamp=safe_as_tensor(
|
|
self.frame_annotation.frame_timestamp, torch.float
|
|
),
|
|
sequence_name=self.frame_annotation.sequence_name,
|
|
sequence_category=self.seq_annotation.category,
|
|
camera_quality_score=safe_as_tensor(
|
|
self.seq_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,
|
|
)
|
|
|
|
def test_frame_data_builder_args(self):
|
|
# test that FrameDataBuilder works with get_default_args
|
|
get_default_args(FrameDataBuilder)
|
|
|
|
def test_fix_point_cloud_path(self):
|
|
"""Some files in Co3Dv2 have an accidental absolute path stored."""
|
|
original_path = "some_file_path"
|
|
modified_path = self.frame_data_builder._fix_point_cloud_path(original_path)
|
|
self.assertIn(original_path, modified_path)
|
|
self.assertIn(self.frame_data_builder.dataset_root, modified_path)
|
|
|
|
def test_load_and_adjust_frame_data(self):
|
|
self.frame_data.image_size_hw = safe_as_tensor(
|
|
self.frame_annotation.image.size, torch.long
|
|
)
|
|
self.frame_data.effective_image_size_hw = self.frame_data.image_size_hw
|
|
|
|
fg_mask_np, mask_path = self.frame_data_builder._load_fg_probability(
|
|
self.frame_annotation
|
|
)
|
|
self.frame_data.mask_path = mask_path
|
|
self.frame_data.fg_probability = safe_as_tensor(fg_mask_np, torch.float)
|
|
mask_thr = self.frame_data_builder.box_crop_mask_thr
|
|
bbox_xywh = get_bbox_from_mask(fg_mask_np, mask_thr)
|
|
self.frame_data.bbox_xywh = safe_as_tensor(bbox_xywh, torch.long)
|
|
|
|
self.assertIsNotNone(self.frame_data.mask_path)
|
|
self.assertTrue(torch.is_tensor(self.frame_data.fg_probability))
|
|
self.assertTrue(torch.is_tensor(self.frame_data.bbox_xywh))
|
|
# assert bboxes shape
|
|
self.assertEqual(self.frame_data.bbox_xywh.shape, torch.Size([4]))
|
|
|
|
image_path = os.path.join(
|
|
self.frame_data_builder.dataset_root, self.frame_annotation.image.path
|
|
)
|
|
image_np = load_image(self.frame_data_builder._local_path(image_path))
|
|
self.assertIsInstance(image_np, np.ndarray)
|
|
self.frame_data.image_rgb = self.frame_data_builder._postprocess_image(
|
|
image_np, self.frame_annotation.image.size, self.frame_data.fg_probability
|
|
)
|
|
self.assertIsInstance(self.frame_data.image_rgb, torch.Tensor)
|
|
|
|
(
|
|
self.frame_data.depth_map,
|
|
depth_path,
|
|
self.frame_data.depth_mask,
|
|
) = self.frame_data_builder._load_mask_depth(
|
|
self.frame_annotation,
|
|
self.frame_data.fg_probability,
|
|
)
|
|
self.assertTrue(torch.is_tensor(self.frame_data.depth_map))
|
|
self.assertIsNotNone(depth_path)
|
|
self.assertTrue(torch.is_tensor(self.frame_data.depth_mask))
|
|
|
|
new_size = (self.image_height, self.image_width)
|
|
|
|
if self.frame_data_builder.box_crop:
|
|
self.frame_data.crop_by_metadata_bbox_(
|
|
self.frame_data_builder.box_crop_context,
|
|
)
|
|
|
|
# assert image and mask shapes after resize
|
|
self.frame_data.resize_frame_(
|
|
new_size_hw=torch.tensor(new_size, dtype=torch.long),
|
|
)
|
|
self.assertEqual(
|
|
self.frame_data.mask_crop.shape,
|
|
torch.Size([1, self.image_height, self.image_width]),
|
|
)
|
|
self.assertEqual(
|
|
self.frame_data.image_rgb.shape,
|
|
torch.Size([3, self.image_height, self.image_width]),
|
|
)
|
|
self.assertEqual(
|
|
self.frame_data.mask_crop.shape,
|
|
torch.Size([1, self.image_height, self.image_width]),
|
|
)
|
|
self.assertEqual(
|
|
self.frame_data.fg_probability.shape,
|
|
torch.Size([1, self.image_height, self.image_width]),
|
|
)
|
|
self.assertEqual(
|
|
self.frame_data.depth_map.shape,
|
|
torch.Size([1, self.image_height, self.image_width]),
|
|
)
|
|
self.assertEqual(
|
|
self.frame_data.depth_mask.shape,
|
|
torch.Size([1, self.image_height, self.image_width]),
|
|
)
|
|
self.frame_data.camera = self.frame_data_builder._get_pytorch3d_camera(
|
|
self.frame_annotation,
|
|
)
|
|
self.assertEqual(type(self.frame_data.camera), PerspectiveCameras)
|
|
|
|
def test_transpose_normalize_image(self):
|
|
def inverse_transpose_normalize_image(image: np.ndarray) -> np.ndarray:
|
|
im = image * 255.0
|
|
return im.transpose((1, 2, 0)).astype(np.uint8)
|
|
|
|
# Test 2D input
|
|
input_image = np.array(
|
|
[[10, 20, 30], [40, 50, 60], [70, 80, 90]], dtype=np.uint8
|
|
)
|
|
expected_input = inverse_transpose_normalize_image(
|
|
transpose_normalize_image(input_image)
|
|
)
|
|
self.assertClose(input_image[..., None], expected_input)
|
|
|
|
# Test 3D input
|
|
input_image = np.array(
|
|
[
|
|
[[10, 20, 30], [40, 50, 60], [70, 80, 90]],
|
|
[[100, 110, 120], [130, 140, 150], [160, 170, 180]],
|
|
[[190, 200, 210], [220, 230, 240], [250, 255, 255]],
|
|
],
|
|
dtype=np.uint8,
|
|
)
|
|
expected_input = inverse_transpose_normalize_image(
|
|
transpose_normalize_image(input_image)
|
|
)
|
|
self.assertClose(input_image, expected_input)
|
|
|
|
def test_load_image(self):
|
|
path = os.path.join(self.dataset_root, self.frame_annotation.image.path)
|
|
local_path = self.path_manager.get_local_path(path)
|
|
image = load_image(local_path)
|
|
self.assertEqual(image.dtype, np.float32)
|
|
self.assertLessEqual(np.max(image), 1.0)
|
|
self.assertGreaterEqual(np.min(image), 0.0)
|
|
|
|
def test_load_mask(self):
|
|
path = os.path.join(self.dataset_root, self.frame_annotation.mask.path)
|
|
path = self.path_manager.get_local_path(path)
|
|
mask = load_mask(path)
|
|
self.assertEqual(mask.dtype, np.float32)
|
|
self.assertLessEqual(np.max(mask), 1.0)
|
|
self.assertGreaterEqual(np.min(mask), 0.0)
|
|
|
|
def test_load_depth(self):
|
|
path = os.path.join(self.dataset_root, self.frame_annotation.depth.path)
|
|
path = self.path_manager.get_local_path(path)
|
|
depth_map = load_depth(path, self.frame_annotation.depth.scale_adjustment)
|
|
self.assertEqual(depth_map.dtype, np.float32)
|
|
self.assertEqual(len(depth_map.shape), 3)
|
|
|
|
def test_load_16big_png_depth(self):
|
|
path = os.path.join(self.dataset_root, self.frame_annotation.depth.path)
|
|
path = self.path_manager.get_local_path(path)
|
|
depth_map = load_16big_png_depth(path)
|
|
self.assertEqual(depth_map.dtype, np.float32)
|
|
self.assertEqual(len(depth_map.shape), 2)
|
|
|
|
def test_load_1bit_png_mask(self):
|
|
mask_path = os.path.join(
|
|
self.dataset_root, self.frame_annotation.depth.mask_path
|
|
)
|
|
mask_path = self.path_manager.get_local_path(mask_path)
|
|
mask = load_1bit_png_mask(mask_path)
|
|
self.assertEqual(mask.dtype, np.float32)
|
|
self.assertEqual(len(mask.shape), 2)
|
|
|
|
def test_load_depth_mask(self):
|
|
mask_path = os.path.join(
|
|
self.dataset_root, self.frame_annotation.depth.mask_path
|
|
)
|
|
mask_path = self.path_manager.get_local_path(mask_path)
|
|
mask = load_depth_mask(mask_path)
|
|
self.assertEqual(mask.dtype, np.float32)
|
|
self.assertEqual(len(mask.shape), 3)
|