jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def check(self, save_dir, target='string'):
        sub_dirs = os.listdir(save_dir)
        self.assertGreater(len(sub_dirs), 0)

        rospy.sleep(1)

        successfully_saved_count = 0
        for sub_dir in sub_dirs:
            sub_dir = osp.join(save_dir, sub_dir)
            text_filepath = osp.join(sub_dir, 'sample_string.txt')

            # Check if text and yaml or text and image are
            # saved in the subdirectory of `save_dir`.
            if target == 'string':
                yaml_filepath = osp.join(sub_dir, 'sample_topic.yaml')
                if not (osp.exists(text_filepath)
                        and osp.exists(yaml_filepath)):
                    continue
            else:
                img_filepath = osp.join(sub_dir, 'sample_image.png')
                if not (osp.exists(text_filepath)
                        and osp.exists(img_filepath)):
                    continue

            with open(text_filepath) as f:
                self.assertEqual(f.read(), 'spam')
            if target == 'string':
                with open(yaml_filepath) as f:
                    data = yaml.load(f, Loader=yaml.SafeLoader)
                    self.assertEqual(data['data'], 'sample')
            elif target == 'image':
                if (StrictVersion(scipy.version.version) > StrictVersion('1.2.0')):
                    img = imageio.imread(img_filepath)
                else:
                    img = scipy.misc.imread(img_filepath)
                self.assertTrue(np.allclose(img, scipy.misc.face()))
            else:
                raise ValueError('Unexpected target: {}'.format(target))
            successfully_saved_count += 1
        if successfully_saved_count == 0:
            rospy.logerr('Data are not successfully saved in {}'
                         .format(save_dir))
//  [python]
def check(self, save_dir, target='string'):
        sub_dirs = os.listdir(save_dir)
        self.assertGreater(len(sub_dirs), 0)

        rospy.sleep(1)

        successfully_saved_count = 0
        for sub_dir in sub_dirs:
            sub_dir = osp.join(save_dir, sub_dir)
            text_filepath = osp.join(sub_dir, 'sample_string.txt')

            # Check if text and yaml or text and image are
            # saved in the subdirectory of `save_dir`.
            if target == 'string':
                yaml_filepath = osp.join(sub_dir, 'sample_topic.yaml')
                if not (osp.exists(text_filepath)
                        and osp.exists(yaml_filepath)):
                    continue
            else:
                img_filepath = osp.join(sub_dir, 'sample_image.png')
                if not (osp.exists(text_filepath)
                        and osp.exists(img_filepath)):
                    continue

            with open(text_filepath) as f:
                self.assertEqual(f.read(), 'spam')
            if target == 'string':
                with open(yaml_filepath) as f:
                    data = yaml.load(f, Loader=yaml.SafeLoader)
                    self.assertEqual(data['data'], 'sample')
            elif target == 'image':
                if (StrictVersion(scipy.version.version) > StrictVersion('1.2.0')):
                    img = imageio.imread(img_filepath)
                else:
                    img = scipy.misc.imread(img_filepath)
                self.assertTrue(np.allclose(img, scipy.datasets.face()))
            else:
                raise ValueError('Unexpected target: {}'.format(target))
            successfully_saved_count += 1
        if successfully_saved_count == 0:
            rospy.logerr('Data are not successfully saved in {}'
                         .format(save_dir))
//  [python]
def test_pyplot(setup):
    import scipy.misc
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    import numpy as np

    face = scipy.misc.face()
    logger.save_image(face, "face.png")

    fig = plt.figure(figsize=(4, 2))
    xs = np.linspace(0, 5, 1000)
    plt.plot(xs, np.cos(xs))
    logger.savefig("face_02.png", fig=fig)
    plt.close()

    fig = plt.figure(figsize=(4, 2))
    xs = np.linspace(0, 5, 1000)
    plt.plot(xs, np.cos(xs))
    logger.savefig('sine.pdf')
//  [python]
def test_pyplot(setup):
    import scipy.datasets
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    import numpy as np

    face = scipy.datasets.face()
    logger.save_image(face, "face.png")

    fig = plt.figure(figsize=(4, 2))
    xs = np.linspace(0, 5, 1000)
    plt.plot(xs, np.cos(xs))
    logger.savefig("face_02.png", fig=fig)
    plt.close()

    fig = plt.figure(figsize=(4, 2))
    xs = np.linspace(0, 5, 1000)
    plt.plot(xs, np.cos(xs))
    logger.savefig('sine.pdf')
//  [python]
def test_make_video(setup):
    import scipy.misc
    from time import sleep

    face = scipy.misc.face()
    h, w, _ = face.shape
    with logger.Sync():
        for i in range(10):
            logger.save_image(face[i:h - 100 + i, i:w - 100 + i], f"face_{i:04d}.png")

    file_list = logger.glob("face_*.png")
    # This one makes the video according to ascending order
    logger.make_video("face_*.png", key="videos/face_move_glob.mp4")
    logger.make_video("face_*.png", order='descending', key="videos/face_move_glob_descending.mp4")
    # This one makes the video according to the random order given by glob
    logger.make_video(file_list, key="videos/face_move_list.mp4")

    # this should not reaise an error
    path = logger.make_video("no_exist_*.png", key="videos/not_exist.mp4")
    assert path is None
//  [python]
def test_make_video(setup):
    import scipy.datasets
    from time import sleep

    face = scipy.datasets.face()
    h, w, _ = face.shape
    with logger.Sync():
        for i in range(10):
            logger.save_image(face[i:h - 100 + i, i:w - 100 + i], f"face_{i:04d}.png")

    file_list = logger.glob("face_*.png")
    # This one makes the video according to ascending order
    logger.make_video("face_*.png", key="videos/face_move_glob.mp4")
    logger.make_video("face_*.png", order='descending', key="videos/face_move_glob_descending.mp4")
    # This one makes the video according to the random order given by glob
    logger.make_video(file_list, key="videos/face_move_list.mp4")

    # this should not reaise an error
    path = logger.make_video("no_exist_*.png", key="videos/not_exist.mp4")
    assert path is None
//  [python]
def test2_8():
    """
    numpy.ix_
    """
    img_face = scipy.misc.face()
    xmax = img_face.shape[0]
    ymax = img_face.shape[1]
    def shuffle_indices(size):
        arr = np.arange(size)
        np.random.shuffle(arr)
        return arr
    xindices = shuffle_indices(xmax)
    # xindices = np.arange(xmax)
    yindices = shuffle_indices(ymax)
    # yindices = np.arange(ymax)
    idxs = np.ix_(xindices, yindices)
    matplotlib.pyplot.imshow(img_face[idxs])
    matplotlib.pyplot.show()
//  [python]
def test2_8():
    """
    numpy.ix_
    """
    img_face = scipy.datasets.face()
    xmax = img_face.shape[0]
    ymax = img_face.shape[1]
    def shuffle_indices(size):
        arr = np.arange(size)
        np.random.shuffle(arr)
        return arr
    xindices = shuffle_indices(xmax)
    # xindices = np.arange(xmax)
    yindices = shuffle_indices(ymax)
    # yindices = np.arange(ymax)
    idxs = np.ix_(xindices, yindices)
    matplotlib.pyplot.imshow(img_face[idxs])
    matplotlib.pyplot.show()
//  [python]
def raccoon(
        self,
    ) -> SpatialData:
        """Raccoon dataset."""
        im_data = s.face()
        im = Image2DModel.parse(im_data, dims=["y", "x", "c"])
        labels_data = slic(im_data, n_segments=100, compactness=10, sigma=1)
        labels = Labels2DModel.parse(labels_data, dims=["y", "x"])
        coords = np.array([[610, 450], [730, 325], [575, 300], [480, 90]])
        circles = ShapesModel.parse(coords, geometry=0, radius=np.array([30, 30, 30, 50]))
        return SpatialData(images={"raccoon": im}, labels={"segmentation": labels}, shapes={"circles": circles})
//  [python]
def raccoon(
        self,
    ) -> SpatialData:
        """Raccoon dataset."""
        im_data = scipy.datasets.face()
        im = Image2DModel.parse(im_data, dims=["y", "x", "c"])
        labels_data = slic(im_data, n_segments=100, compactness=10, sigma=1)
        labels = Labels2DModel.parse(labels_data, dims=["y", "x"])
        coords = np.array([[610, 450], [730, 325], [575, 300], [480, 90]])
        circles = ShapesModel.parse(coords, geometry=0, radius=np.array([30, 30, 30, 50]))
        return SpatialData(images={"raccoon": im}, labels={"segmentation": labels}, shapes={"circles": circles})
//  [python]
def test_conv_2d(wavelet: str, level: int, size: tuple, mode: str, dtype: jnp.dtype):
    """Runs tests of the two-dimensional fwt code."""
    if dtype == jnp.float32:
        atol = 1e-3
    else:
        atol = 1e-8

    wavelet = pywt.Wavelet(wavelet)
    face = jnp.transpose(scipy.misc.face(), [2, 0, 1]).astype(dtype)
    face = face[:, 128 : (128 + size[0]), 256 : (256 + size[1])]

    coeff2d_pywt = pywt.wavedec2(face, wavelet, mode=mode, level=level)
    coeff2d = wavedec2(face, wavelet, level=level, mode=mode)
    # test pywt compatability
    pywt_flat_list = jnp.concatenate(flatten_2d_coeff_lst(coeff2d_pywt), -1)
    jwt_flat_list = jnp.concatenate(flatten_2d_coeff_lst(coeff2d), -1)
    assert jnp.allclose(pywt_flat_list, jwt_flat_list, atol=atol)

    # test invertability
    reconstruction_2d = waverec2(coeff2d, wavelet)[..., : size[0], : size[1]]
    assert jnp.allclose(reconstruction_2d, face, atol=atol)
//  [python]
def test_conv_2d(wavelet: str, level: int, size: tuple, mode: str, dtype: jnp.dtype):
    """Runs tests of the two-dimensional fwt code."""
    if dtype == jnp.float32:
        atol = 1e-3
    else:
        atol = 1e-8

    wavelet = pywt.Wavelet(wavelet)
    face = jnp.transpose(scipy.datasets.face(), [2, 0, 1]).astype(dtype)
    face = face[:, 128 : (128 + size[0]), 256 : (256 + size[1])]

    coeff2d_pywt = pywt.wavedec2(face, wavelet, mode=mode, level=level)
    coeff2d = wavedec2(face, wavelet, level=level, mode=mode)
    # test pywt compatability
    pywt_flat_list = jnp.concatenate(flatten_2d_coeff_lst(coeff2d_pywt), -1)
    jwt_flat_list = jnp.concatenate(flatten_2d_coeff_lst(coeff2d), -1)
    assert jnp.allclose(pywt_flat_list, jwt_flat_list, atol=atol)

    # test invertability
    reconstruction_2d = waverec2(coeff2d, wavelet)[..., : size[0], : size[1]]
    assert jnp.allclose(reconstruction_2d, face, atol=atol)
//  [python]
def test_matrix_analysis_fwt_2d_haar(size: Tuple[int, int], level: int) -> None:
    """Test the fwt-2d matrix-haar transform, should be equal to the pywt."""
    face = np.mean(
        scipy.misc.face()[256 : (256 + size[0]), 256 : (256 + size[1])], -1
    ).astype(np.float64)
    wavelet = pywt.Wavelet("haar")
    matrixfwt = MatrixWavedec2(wavelet, level=level, separable=False)
    mat_coeff = matrixfwt(torch.from_numpy(face))
    conv_coeff = pywt.wavedec2(face, wavelet, level=level, mode="zero")
    flat_mat_coeff = torch.cat(_flatten_2d_coeff_lst(mat_coeff), -1)
    flat_conv_coeff = np.concatenate(_flatten_2d_coeff_lst(conv_coeff), -1)
    test = np.allclose(flat_mat_coeff.numpy(), flat_conv_coeff)
    test2 = np.allclose(mat_coeff[0].squeeze(0).numpy(), conv_coeff[0])
    test3 = np.allclose(mat_coeff[1][0].squeeze(0).numpy(), conv_coeff[1][0])
    assert test and test2 and test3
//  [python]
def test_matrix_analysis_fwt_2d_haar(size: Tuple[int, int], level: int) -> None:
    """Test the fwt-2d matrix-haar transform, should be equal to the pywt."""
    face = np.mean(
        scipy.datasets.face()[256 : (256 + size[0]), 256 : (256 + size[1])], -1
    ).astype(np.float64)
    wavelet = pywt.Wavelet("haar")
    matrixfwt = MatrixWavedec2(wavelet, level=level, separable=False)
    mat_coeff = matrixfwt(torch.from_numpy(face))
    conv_coeff = pywt.wavedec2(face, wavelet, level=level, mode="zero")
    flat_mat_coeff = torch.cat(_flatten_2d_coeff_lst(mat_coeff), -1)
    flat_conv_coeff = np.concatenate(_flatten_2d_coeff_lst(conv_coeff), -1)
    test = np.allclose(flat_mat_coeff.numpy(), flat_conv_coeff)
    test2 = np.allclose(mat_coeff[0].squeeze(0).numpy(), conv_coeff[0])
    test3 = np.allclose(mat_coeff[1][0].squeeze(0).numpy(), conv_coeff[1][0])
    assert test and test2 and test3
//  [python]
def test_boundary_matrix_fwt_2d(
    wavelet_str: str, size: Tuple[int, int], level: int, separable: bool
) -> None:
    """Ensure the boundary matrix fwt is invertable."""
    face = np.mean(
        scipy.misc.face()[256 : (256 + size[0]), 256 : (256 + size[1])], -1
    ).astype(np.float64)
    wavelet = pywt.Wavelet(wavelet_str)
    matrixfwt = MatrixWavedec2(wavelet, level=level, separable=separable)
    mat_coeff = matrixfwt(torch.from_numpy(face))
    matrixifwt = MatrixWaverec2(wavelet, separable=separable)
    reconstruction = matrixifwt(mat_coeff).squeeze(0)
    # remove the padding
    if size[0] % 2 != 0:
        reconstruction = reconstruction[:-1, :]
    if size[1] % 2 != 0:
        reconstruction = reconstruction[:, :-1]
    assert np.allclose(reconstruction.numpy(), face)
    # test the operator matrices
    if not separable and not matrixfwt.padded and not matrixifwt.padded:
        # padding happens outside of the matrix structure.
        # our matrices therefore only describe pad-free cases.
        test_mat = torch.sparse.mm(
            matrixifwt.sparse_ifwt_operator, matrixfwt.sparse_fwt_operator
        )
        assert np.allclose(test_mat.to_dense().numpy(), np.eye(test_mat.shape[0]))
//  [python]
def test_boundary_matrix_fwt_2d(
    wavelet_str: str, size: Tuple[int, int], level: int, separable: bool
) -> None:
    """Ensure the boundary matrix fwt is invertable."""
    face = np.mean(
        scipy.datasets.face()[256 : (256 + size[0]), 256 : (256 + size[1])], -1
    ).astype(np.float64)
    wavelet = pywt.Wavelet(wavelet_str)
    matrixfwt = MatrixWavedec2(wavelet, level=level, separable=separable)
    mat_coeff = matrixfwt(torch.from_numpy(face))
    matrixifwt = MatrixWaverec2(wavelet, separable=separable)
    reconstruction = matrixifwt(mat_coeff).squeeze(0)
    # remove the padding
    if size[0] % 2 != 0:
        reconstruction = reconstruction[:-1, :]
    if size[1] % 2 != 0:
        reconstruction = reconstruction[:, :-1]
    assert np.allclose(reconstruction.numpy(), face)
    # test the operator matrices
    if not separable and not matrixfwt.padded and not matrixifwt.padded:
        # padding happens outside of the matrix structure.
        # our matrices therefore only describe pad-free cases.
        test_mat = torch.sparse.mm(
            matrixifwt.sparse_ifwt_operator, matrixfwt.sparse_fwt_operator
        )
        assert np.allclose(test_mat.to_dense().numpy(), np.eye(test_mat.shape[0]))
//  [python]
def test_batched_2d_matrix_fwt_ifwt(
    wavelet_str: str, level: int, size: Tuple[int, int], separable: bool
) -> None:
    """Ensure the batched matrix fwt works properly."""
    face = s.face(gray=False)[256 : (256 + size[0]), 256 : (256 + size[1])].astype(
        np.float64
    )
    pt_face = torch.from_numpy(face).permute([2, 0, 1])
    wavelet = pywt.Wavelet(wavelet_str)
    matrixfwt = MatrixWavedec2(wavelet, level=level, separable=separable)
    mat_coeff = matrixfwt(pt_face)
    matrixifwt = MatrixWaverec2(wavelet, separable=separable)
    reconstruction = matrixifwt(mat_coeff)
    assert (
        np.allclose(reconstruction[0].numpy(), face[:, :, 0])
        and np.allclose(reconstruction[1].numpy(), face[:, :, 1])
        and np.allclose(reconstruction[2].numpy(), face[:, :, 2])
    )
//  [python]
def test_batched_2d_matrix_fwt_ifwt(
    wavelet_str: str, level: int, size: Tuple[int, int], separable: bool
) -> None:
    """Ensure the batched matrix fwt works properly."""
    face = scipy.datasets.face()[256 : (256 + size[0]), 256 : (256 + size[1])].astype(
        np.float64
    )
    pt_face = torch.from_numpy(face).permute([2, 0, 1])
    wavelet = pywt.Wavelet(wavelet_str)
    matrixfwt = MatrixWavedec2(wavelet, level=level, separable=separable)
    mat_coeff = matrixfwt(pt_face)
    matrixifwt = MatrixWaverec2(wavelet, separable=separable)
    reconstruction = matrixifwt(mat_coeff)
    assert (
        np.allclose(reconstruction[0].numpy(), face[:, :, 0])
        and np.allclose(reconstruction[1].numpy(), face[:, :, 1])
        and np.allclose(reconstruction[2].numpy(), face[:, :, 2])
    )
//  [python]
def test_matrix_transform_2d_rebuild(wavelet_str: str, separable: bool) -> None:
    """Ensure the boundary matrix fwt is invertable."""
    wavelet = pywt.Wavelet(wavelet_str)
    matrixifwt = MatrixWaverec2(wavelet, separable=separable)
    for level in [4, 1, None]:
        matrixfwt = MatrixWavedec2(wavelet, level=level, separable=separable)
        for size in [[16, 16], [17, 17]]:
            face = np.mean(
                scipy.misc.face()[256 : (256 + size[0]), 256 : (256 + size[1])], -1
            ).astype(np.float64)
            mat_coeff = matrixfwt(torch.from_numpy(face))
            reconstruction = matrixifwt(mat_coeff).squeeze(0)
            # remove the padding
            if size[0] % 2 != 0:
                reconstruction = reconstruction[:-1, :]
            if size[1] % 2 != 0:
                reconstruction = reconstruction[:, :-1]
            assert np.allclose(reconstruction.numpy(), face)
            # test the operator matrices
            if not separable and not matrixfwt.padded and not matrixifwt.padded:
                test_mat = torch.sparse.mm(
                    matrixifwt.sparse_ifwt_operator, matrixfwt.sparse_fwt_operator
                )
                assert np.allclose(
                    test_mat.to_dense().numpy(), np.eye(test_mat.shape[0])
                )
//  [python]
def test_matrix_transform_2d_rebuild(wavelet_str: str, separable: bool) -> None:
    """Ensure the boundary matrix fwt is invertable."""
    wavelet = pywt.Wavelet(wavelet_str)
    matrixifwt = MatrixWaverec2(wavelet, separable=separable)
    for level in [4, 1, None]:
        matrixfwt = MatrixWavedec2(wavelet, level=level, separable=separable)
        for size in [[16, 16], [17, 17]]:
            face = np.mean(
                scipy.datasets.face()[256 : (256 + size[0]), 256 : (256 + size[1])], -1
            ).astype(np.float64)
            mat_coeff = matrixfwt(torch.from_numpy(face))
            reconstruction = matrixifwt(mat_coeff).squeeze(0)
            # remove the padding
            if size[0] % 2 != 0:
                reconstruction = reconstruction[:-1, :]
            if size[1] % 2 != 0:
                reconstruction = reconstruction[:, :-1]
            assert np.allclose(reconstruction.numpy(), face)
            # test the operator matrices
            if not separable and not matrixfwt.padded and not matrixifwt.padded:
                test_mat = torch.sparse.mm(
                    matrixifwt.sparse_ifwt_operator, matrixfwt.sparse_fwt_operator
                )
                assert np.allclose(
                    test_mat.to_dense().numpy(), np.eye(test_mat.shape[0])
                )