<commit_before>#!/usr/bin/env python

"""
Registration algorithms and utility functions
"""

from __future__ import print_function
import numpy as np
from .. import BoundingBox, log, force_srs
from .stickscale import get_stick_scale
from pcl.registration import pcl

from patty.utils import (
    log,
    save,
)

from patty.segmentation import (
    boundary_of_center_object,
    boundary_of_drivemap,
    boundary_of_lowest_points,
)

from sklearn.decomposition import PCA

def align_footprints(loose_pc, fixed_pc,
                     allow_scaling=True,
                     allow_rotation=True,
                     allow_translation=True):
    '''
    Align a pointcloud 'loose_pc' by placing it on top of
    'fixed_pc' as good as poosible. Done by aligning the 
    principle axis of both pointclouds.

    NOTE: Both pointclouds are assumed to be the footprint (or projection)
    on the xy plane, with basically zero extent along the z-axis.

    (allow_rotation=True)
        The pointcloud boundary is alinged with the footprint
        by rotating its pricipal axis in the (x,y) plane. 

    (allow_translation=True)
        Then, it is translated so the centers of mass coincide.

    (allow_scaling=True)
        Finally, the pointcloud is scaled to have the same extent.

    Arguments:
        loose_pc          : pcl.PointCloud
        fixed_pc          : pcl.PointCloud

        allow_scaling     : Bolean
        allow_rotation    : Bolean
        allow_translation : Bolean

    Returns:
        rot_matrix, rot_center, scale, translation : np.array()

    '''

    rot_center = loose_pc.center()

    if allow_rotation:
        log(" - Finding rotation")
        rot_matrix = find_rotation_xy(loose_pc, fixed_pc)
        loose_pc.rotate(rot_matrix, origin=rot_center)
    else:
        log(" - Skipping rotation")
        rot_matrix = np.eye(3)

    if allow_scaling:
        log(" - Finding scale")
        loose_bb = BoundingBox(loose_pc)
        fixed_bb = BoundingBox(fixed_pc)
        scale = fixed_bb.size[0:2] / loose_bb.size[0:2]

        # take the average scale factor for the x and y dimensions
        scale = np.mean(scale)
    else:
        log(" - Skipping scale")
        scale = 1.0

    if allow_translation:
        log(" - Finding translation")
        translation = fixed_pc.center() - rot_center
        loose_pc.translate(translation)
    else:
        log(" - Skipping translation")
        translation = np.array([0.0, 0.0, 0.0])

    return rot_matrix, rot_center, scale, translation


def estimate_pancake_up(pointcloud):
    '''
    Assuming a pancake like pointcloud, the up direction is the third PCA.<commit_msg>use the center of the main object for scaling so it stays in place<commit_after>#!/usr/bin/env python

"""
Registration algorithms and utility functions
"""

from __future__ import print_function
import numpy as np
from .. import BoundingBox, log, force_srs
from .stickscale import get_stick_scale
from pcl.registration import pcl

from patty.utils import (
    log,
    save,
)

from patty.segmentation import (
    boundary_of_center_object,
    boundary_of_drivemap,
    boundary_of_lowest_points,
)

from sklearn.decomposition import PCA

def align_footprints(loose_pc, fixed_pc,
                     allow_scaling=True,
                     allow_rotation=True,
                     allow_translation=True):
    '''
    Align a pointcloud 'loose_pc' by placing it on top of
    'fixed_pc' as good as poosible. Done by aligning the 
    principle axis of both pointclouds.

    NOTE: Both pointclouds are assumed to be the footprint (or projection)
    on the xy plane, with basically zero extent along the z-axis.

    (allow_rotation=True)
        The pointcloud boundary is alinged with the footprint
        by rotating its pricipal axis in the (x,y) plane. 

    (allow_translation=True)
        Then, it is translated so the centers of mass coincide.

    (allow_scaling=True)
        Finally, the pointcloud is scaled to have the same extent.

    Arguments:
        loose_pc          : pcl.PointCloud
        fixed_pc          : pcl.PointCloud

        allow_scaling     : Bolean
        allow_rotation    : Bolean
        allow_translation : Bolean

    Returns:
        rot_matrix, rot_center, scale, translation : np.array()

    '''

    rot_center = loose_pc.center()

    if allow_rotation:
        log(" - Finding rotation")
        rot_matrix = find_rotation_xy(loose_pc, fixed_pc)
        loose_pc.rotate(rot_matrix, origin=rot_center)
    else:
        log(" - Skipping rotation")
        rot_matrix = np.eye(3)


    if allow_scaling:
        log(" - Finding scale")
        fixed_bb = BoundingBox(fixed_pc) # used 2x below
        loose_bb = BoundingBox(loose_pc)
        scale = fixed_bb.size[0:2] / loose_bb.size[0:2]

        # take the average scale factor for the x and y dimensions
        scale = np.mean(scale)
        loose_pc.scale(scale, origin=rot_center)
    else:
        log(" - Skipping scale")
        scale = 1.0

    if allow_translation:
        log(" - Finding translation")
        translation = fixed_pc.center() - rot_center
        loose_pc.translate(translation)
    else:
        log(" - Skipping translation")
        translation = np.array([0.0, 0.0, 0.0])

    return rot_matrix, rot_center, scale, translation


def estimate_pancake_up(pointcloud):
    '''
    Assuming a pancake like pointcloud, the up direction is the third PCA.
{'fraction_non_alphanumeric': 0.07408765816981477, 'fraction_numerical': 0.004951402897487621, 'mean_word_length': 2.5323834196891193, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '108679', 'n_tokens_mistral': 1664, 'n_tokens_neox': 1574, 'n_words': 491}