Ejemplo n.º 1
0
def check_dim(t, dim):
    if check_type(t, [torch.Tensor, np.ndarray]) and check_type(dim, [int]):
        assert t.dim(
        ) == dim, f'tensor should be {dim} dimensional, tensor dimension is {t.dim()}'
    else:
        raise TypeError(
            f'the tensor type is {type(t)}. accepted types are torch.Tensor and numpy.ndarray'
        )
Ejemplo n.º 2
0
def assertDefaultScale(obj=None):
    '''Assert obj has default scale of 1

    Attributes:
        obj -- Object to check. dt.Transform
    '''
    general.check_type(obj, 'obj', [pm.nt.Transform])

    return obj.getScale() == [1, 1, 1]
Ejemplo n.º 3
0
def assertZeroRots(obj=None):
    '''Assert obj has zero rotations

    Attributes:
        obj -- Object to check. dt.Transform
    '''
    general.check_type(obj, 'obj', [pm.nt.Transform])

    return obj.getRotation() == pm.dt.EulerRotation([0, 0, 0])
Ejemplo n.º 4
0
def assertZeroTrans(obj=None):
    '''Assert obj has zero translations

    Attributes:
        obj -- Object to check. dt.Transform
    '''
    general.check_type(obj, 'obj', [pm.nt.Transform])

    return obj.getTranslation() == pm.dt.Vector([0, 0, 0])
Ejemplo n.º 5
0
def assertZeroTrans(obj=None):
    '''Assert obj has zero translations

    Attributes:
        obj -- Object to check. dt.Transform
    '''
    general.check_type(obj, 'obj', [pm.nt.Transform])

    return obj.getTranslation() == pm.dt.Vector([0, 0, 0])
Ejemplo n.º 6
0
def assertZeroRots(obj=None):
    '''Assert obj has zero rotations

    Attributes:
        obj -- Object to check. dt.Transform
    '''
    general.check_type(obj, 'obj', [pm.nt.Transform])

    return obj.getRotation() == pm.dt.EulerRotation([0, 0, 0])
Ejemplo n.º 7
0
def assertDefaultScale(obj=None):
    '''Assert obj has default scale of 1

    Attributes:
        obj -- Object to check. dt.Transform
    '''
    general.check_type(obj, 'obj', [pm.nt.Transform])

    return obj.getScale() == [1, 1, 1]
Ejemplo n.º 8
0
def check_shape(t, last_dim_shape):
    if check_type(t, [torch.Tensor, np.ndarray]) and check_type(
            last_dim_shape, [int]):
        assert t.shape[-1] == last_dim_shape, (
            f'tensor last dimension be {last_dim_shape},' +
            f'tensor last dimension is {t.shape[-1]}')
    else:
        raise TypeError(
            f'the tensor type is {type(t)}. accepted types are torch.Tensor and numpy.ndarray'
        )
Ejemplo n.º 9
0
def assertLocationIs(obj=None, pos=None):
    '''Assert obj is located at position pos

    Attributes:
        obj -- Object to check. (pm.nt.Transform, pm.nt.Joint)
        pos -- pm.dt.Point=(x, y, z)
    '''
    general.check_type(obj, 'obj', [pm.nt.Transform])
    general.check_type(pos, 'pos', [pm.dt.Point])

    return general.assertAlmostEquals(
        pm.dt.Point(pm.xform(obj, q=1, ws=1, t=1)), pos, 3)
Ejemplo n.º 10
0
def assertParentIs(obj=None, prnt=None):
    '''Assert prnt is the paprent of obj

    Attributes:
        obj -- Object that is to be checked
        prnt -- Expected parent  of obj
    '''
    general.check_type(obj, 'obj', [pm.nt.Transform, pm.nt.Joint])
    general.check_type(prnt, 'prnt', [pm.nt.Transform, pm.nt.Joint])

    if obj.getParent() == prnt:
        return True
    return False
Ejemplo n.º 11
0
def assertParentIs(obj=None, prnt=None):
    '''Assert prnt is the paprent of obj

    Attributes:
        obj -- Object that is to be checked
        prnt -- Expected parent  of obj
    '''
    general.check_type(obj, 'obj', [pm.nt.Transform, pm.nt.Joint])
    general.check_type(prnt, 'prnt', [pm.nt.Transform, pm.nt.Joint])

    if obj.getParent() == prnt:
        return True
    return False
Ejemplo n.º 12
0
def assertLocationsMatch(a=None, b=None):
    '''Assert a and b locations match

    Attributes:
        a -- Object a. [pm.nt.Transform, pm.nt.Joint]
        b -- Object b. [pm.nt.Transform, pm.nt.Joint
    '''
    general.check_type(a, 'a', [pm.nt.Transform, pm.nt.Joint])
    general.check_type(b, 'b', [pm.nt.Transform, pm.nt.Joint])

    pos_a = pm.xform(a, q=1, ws=1, rp=1)
    pos_b = pm.xform(b, q=1, ws=1, rp=1)

    return general.assertAlmostEquals(pos_a, pos_b, 2)
Ejemplo n.º 13
0
def assertLocationsMatch(a=None, b=None):
    '''Assert a and b locations match

    Attributes:
        a -- Object a. [pm.nt.Transform, pm.nt.Joint]
        b -- Object b. [pm.nt.Transform, pm.nt.Joint
    '''
    general.check_type(a, 'a', [pm.nt.Transform, pm.nt.Joint])
    general.check_type(b, 'b', [pm.nt.Transform, pm.nt.Joint])

    pos_a = pm.xform(a, q=1, ws=1, rp=1)
    pos_b = pm.xform(b, q=1, ws=1, rp=1)

    return general.assertAlmostEquals(pos_a, pos_b, 2)
Ejemplo n.º 14
0
def assertLocationIs(obj=None, pos=None):
    '''Assert obj is located at position pos

    Attributes:
        obj -- Object to check. (pm.nt.Transform, pm.nt.Joint)
        pos -- pm.dt.Point=(x, y, z)
    '''
    general.check_type(obj, 'obj', [pm.nt.Transform])
    general.check_type(pos, 'pos', [pm.dt.Point])

    return general.assertAlmostEquals(pm.dt.Point(pm.xform(obj,
                                                           q=1,
                                                           ws=1,
                                                           t=1)), pos, 3)
Ejemplo n.º 15
0
def assertAllZero(obj=None):
    '''Assert obj has zero translations/rotations/scale

    Attributes:
        obj -- Object to check. dt.Transform
    '''
    general.check_type(obj, 'obj', [pm.nt.Transform])

    expected = pm.dt.TransformationMatrix([[1.0, 0.0, 0.0, 0.0],
                                           [0.0, 1.0, 0.0, 0.0],
                                           [0.0, 0.0, 1.0, 0.0],
                                           [0.0, 0.0, 0.0, 1.0]])

    if obj.getTransformation() != expected:
        return False

    return True
Ejemplo n.º 16
0
def assertAllZero(obj=None):
    '''Assert obj has zero translations/rotations/scale

    Attributes:
        obj -- Object to check. dt.Transform
    '''
    general.check_type(obj, 'obj', [pm.nt.Transform])

    expected = pm.dt.TransformationMatrix([[1.0, 0.0, 0.0, 0.0],
                                           [0.0, 1.0, 0.0, 0.0],
                                           [0.0, 0.0, 1.0, 0.0],
                                           [0.0, 0.0, 0.0, 1.0]])

    if obj.getTransformation() != expected:
        return False

    return True
Ejemplo n.º 17
0
def assertOrientationMatches(a=None, b=None):
    '''Check that axises and position match for a and b

    Attributes:
        a -- Object to check against b. [pm.nt.Transform, pm.nt.Joint]
        a -- Object to check against a. [pm.nt.Transform, pm.nt.Joint]

    '''

    def makeLoc(obj=None):
        # Place locator as child to jnt and zero it
        loc = pm.spaceLocator()
        pm.parent(loc, obj)
        loc.setTranslation(0)
        loc.setRotation([0, 0, 0])
        return loc

    def moveLoc(loc=None, axis=None, amount=None):
        # Move loc along axis by amount, return position
        if axis == 'x':
            loc.tx.set(amount)
        if axis == 'y':
            loc.ty.set(amount)
        if axis == 'z':
            loc.tz.set(amount)
        return pm.xform(loc, q=1, ws=1, rp=1)

    general.check_type(a, 'a', [pm.nt.Transform, pm.nt.Joint])
    general.check_type(b, 'b', [pm.nt.Transform, pm.nt.Joint])

    a_pos = []
    b_pos = []
    for axis in ['x', 'y', 'z']:
        # Make locators
        loc_a = makeLoc(a)
        loc_b = makeLoc(b)

        a_pos.append(moveLoc(loc_a, axis, 5))
        b_pos.append(moveLoc(loc_b, axis, 5))

        pm.delete(loc_a)
        pm.delete(loc_b)

    return general.assertAlmostEquals(a_pos, b_pos)
Ejemplo n.º 18
0
def assertOrientationMatches(a=None, b=None):
    '''Check that axises and position match for a and b

    Attributes:
        a -- Object to check against b. [pm.nt.Transform, pm.nt.Joint]
        a -- Object to check against a. [pm.nt.Transform, pm.nt.Joint]

    '''
    def makeLoc(obj=None):
        # Place locator as child to jnt and zero it
        loc = pm.spaceLocator()
        pm.parent(loc, obj)
        loc.setTranslation(0)
        loc.setRotation([0, 0, 0])
        return loc

    def moveLoc(loc=None, axis=None, amount=None):
        # Move loc along axis by amount, return position
        if axis == 'x':
            loc.tx.set(amount)
        if axis == 'y':
            loc.ty.set(amount)
        if axis == 'z':
            loc.tz.set(amount)
        return pm.xform(loc, q=1, ws=1, rp=1)

    general.check_type(a, 'a', [pm.nt.Transform, pm.nt.Joint])
    general.check_type(b, 'b', [pm.nt.Transform, pm.nt.Joint])

    a_pos = []
    b_pos = []
    for axis in ['x', 'y', 'z']:
        # Make locators
        loc_a = makeLoc(a)
        loc_b = makeLoc(b)

        a_pos.append(moveLoc(loc_a, axis, 5))
        b_pos.append(moveLoc(loc_b, axis, 5))

        pm.delete(loc_a)
        pm.delete(loc_b)

    return general.assertAlmostEquals(a_pos, b_pos)
Ejemplo n.º 19
0
def iou(b1, b2, f='c'):
    """calculate the `IOU` of 2 boxes namely b1 and b2.

        Args:
            b1: first box
            b2: second box
            f: box format. `c` for corner, `m` for midpoint.

        both b1 and b2 should be of type `torch.Tensor` or `numpy.ndarray`.
    """

    # ~~~~~~~~~~~~~~~~~~~~~ type checking ~~~~~~~~~~~~~~~~~~~~~ #
    known_types = [torch.Tensor, np.ndarray]
    if not check_type(b1, known_types):
        raise TypeError(
            f'the b1 has a type of {type(b1)}. accepted types are torch.Tensor and numpy.ndarray'
        )

    if not check_type(b2, known_types):
        raise TypeError(
            f'the b1 has a type of {type(b2)}. accepted types are torch.Tensor and numpy.ndarray'
        )

    if isinstance(b1, np.ndarray):
        b1 = torch.from_numpy(b1)
    if isinstance(b2, np.ndarray):
        b2 = torch.from_numpy(b2)

# ~~~~~~~~~~~~~~~~~~~~~ coordinate conversion ~~~~~~~~~~~~~~~~~~~~~ #

    if f == 'm':
        b1 = midpoint_to_corner(b1)
        b2 = midpoint_to_corner(b2)

# ~~~~~~~~~~~~~~~~~~~~~ IOU calculation ~~~~~~~~~~~~~~~~~~~~~ #

    b1_area = box_area(b1)
    b2_area = box_area(b2)
    intersection = calculate_intersection(b1, b2)
    return intersection / (b1_area + b2_area - intersection + 1e-6)
Ejemplo n.º 20
0
def nms(boxes, iou_threshold, prob_threshold):
    # input format [[class,probability,x1,y1,x2,y2]] shape is (N,6)

    # ~~~~~~~~~~~~~~~~~~~~~ type checking ~~~~~~~~~~~~~~~~~~~~~ #
    known_types = [torch.Tensor, np.ndarray]
    if not check_type(boxes, known_types):
        raise TypeError(
            f'the b1 has a type of {type(boxes)}. accepted types are torch.Tensor and numpy.ndarray'
        )
    if isinstance(boxes, np.ndarray):
        boxes = torch.from_numpy(boxes)

    check_shape(boxes, 6)
    # ~~~~~~~~~~~~~~~~~~~~~ nms ~~~~~~~~~~~~~~~~~~~~~ #
    # remove all the boxes below prob_threshold
    boxes = boxes[boxes[..., 1] > prob_threshold]
    # get the all unique classes
    classes = torch.unique(boxes[..., 0])
    check_dim(classes, 1)
    bounding_boxes = []
    # for each classes calculate nms
    for c in classes:

        b = boxes[boxes[..., 0] == c]
        check_dim(b, 2)
        if b.shape[0] > 1:

            # sort the boxes based on the probability
            torch.sort(b, 0)

            # taking the max box out
            max_box = b[0].unsqueeze(0)
            b = b[1:]
            check_dim(max_box, 2)
            bounding_boxes.append(max_box)

            # calculate the iou
            ious = iou(max_box[..., 2:], b[..., 2:]).squeeze(-1)
            check_dim(ious, 1)

            # apply nms
            b = b[ious > iou_threshold]
            if b.shape[0] != 0:
                bounding_boxes.append(b)

        else:
            bounding_boxes.append(b)
    bounding_boxes = torch.cat(bounding_boxes, dim=0)
    check_dim(bounding_boxes, 2)
    check_shape(bounding_boxes, 6)
    return bounding_boxes
Ejemplo n.º 21
0
def assertAimingAt(a=None, b=None, axis=None, log=False):
    '''Assert a aims at b along axis

    Attributes:
        a -- Object a. [pm.nt.Transform, pm.nt.Joint]
        b -- Object b. [pm.nt.Transform, pm.nt.Joint]
        axis -- axis aiming at b. 'x' 'y' or 'z'
    '''
    general.check_type(a, 'a', [pm.nt.Transform, pm.nt.Joint])
    general.check_type(b, 'b', [pm.nt.Transform, pm.nt.Joint])
    general.check_type(axis, 'axis', [str])

    if axis not in ['x', 'y', 'z']:
        raise errors.InputError(axis, 'axis', ['x', 'y', 'z'])

    pos_a = pm.dt.Vector(pm.xform(a, q=1, ws=1, rp=1))
    pos_b = pm.dt.Vector(pm.xform(b, q=1, ws=1, rp=1))
    vec = pos_b - pos_a

    if not vec.length() > 0:
        raise errors.ObjectError(a, 'Different position than b',
                                 'Same position as b')

    loc = pm.spaceLocator()
    pm.parent(loc, a)
    loc.setTranslation(0)
    loc.setRotation([0, 0, 0])
    if axis == 'x':
        loc.tx.set(vec.length())
    elif axis == 'y':
        loc.ty.set(vec.length())
    elif axis == 'z':
        loc.tz.set(vec.length())

    pos_loc = pm.dt.Vector(pm.xform(loc, q=1, ws=1, rp=1))
    pm.delete(loc)

    if log:
        str_1 = 'pos_a: ', pos_a
        str_2 = 'pos_b: ', pos_b
        str_3 = 'pos_loc: ', pos_loc
        general.logging.debug(str_1)
        general.logging.debug(str_2)
        general.logging.debug(str_3)

    return general.assertAlmostEquals(pos_loc, pos_b, 3)
Ejemplo n.º 22
0
def assertAimingAt(a=None, b=None, axis=None, log=False):
    '''Assert a aims at b along axis

    Attributes:
        a -- Object a. [pm.nt.Transform, pm.nt.Joint]
        b -- Object b. [pm.nt.Transform, pm.nt.Joint]
        axis -- axis aiming at b. 'x' 'y' or 'z'
    '''
    general.check_type(a, 'a', [pm.nt.Transform, pm.nt.Joint])
    general.check_type(b, 'b', [pm.nt.Transform, pm.nt.Joint])
    general.check_type(axis, 'axis', [str])

    if axis not in ['x', 'y', 'z']:
        raise errors.InputError(axis, 'axis', ['x', 'y', 'z'])

    pos_a = pm.dt.Vector(pm.xform(a, q=1, ws=1, rp=1))
    pos_b = pm.dt.Vector(pm.xform(b, q=1, ws=1, rp=1))
    vec = pos_b - pos_a

    if not vec.length() > 0:
        raise errors.ObjectError(a, 'Different position than b',
                                 'Same position as b')

    loc = pm.spaceLocator()
    pm.parent(loc, a)
    loc.setTranslation(0)
    loc.setRotation([0, 0, 0])
    if axis == 'x':
        loc.tx.set(vec.length())
    elif axis == 'y':
        loc.ty.set(vec.length())
    elif axis == 'z':
        loc.tz.set(vec.length())

    pos_loc = pm.dt.Vector(pm.xform(loc, q=1, ws=1, rp=1))
    pm.delete(loc)

    if log:
        str_1 = 'pos_a: ', pos_a
        str_2 = 'pos_b: ', pos_b
        str_3 = 'pos_loc: ', pos_loc
        general.logging.debug(str_1)
        general.logging.debug(str_2)
        general.logging.debug(str_3)

    return general.assertAlmostEquals(pos_loc, pos_b, 3)