text:
import numpy as np
from ._base_metric import _BaseMetric
from .. import _timing
from functools import partial
from .. import utils
from ..utils import TrackEvalException

class TrackMAP(_BaseMetric):
    """Class which implements the TrackMAP metrics"""

    @staticmethod
    def get_default_metric_config():
        """Default class config values"""
        default_config = {
            'USE_AREA_RANGES': True,  # whether to evaluate for certain area ranges
            'AREA_RANGES': [[0 ** 2, 32 ** 2],  # additional area range sets for which TrackMAP is evaluated
                            [32 ** 2, 96 ** 2],  # (all area range always included), default values for TAO
                            [96 ** 2, 1e5 ** 2]],  # evaluation
            'AREA_RANGE_LABELS': ["area_s", "area_m", "area_l"],  # the labels for the area ranges
            'USE_TIME_RANGES': True,  # whether to evaluate for certain time ranges (length of tracks)
            'TIME_RANGES': [[0, 3], [3, 10], [10, 1e5]],  # additional time range sets for which TrackMAP is evaluated
            # (all time range always included) , default values for TAO evaluation
            'TIME_RANGE_LABELS': ["time_s", "time_m", "time_l"],  # the labels for the time ranges
            'IOU_THRESHOLDS': np.arange(0.5, 0.96, 0.05),  # the IoU thresholds
            'RECALL_THRESHOLDS': np.linspace(0.0, 1.00, int(np.round((1.00 - 0.0) / 0.01) + 1), endpoint=True),
            # recall thresholds at which precision is evaluated
            'MAX_DETECTIONS': 0,  # limit the maximum number of considered tracks per sequence (0 for unlimited)
            'PRINT_CONFIG': True
        }
        return default_config

    def __init__(self, config=None):
        super().__init__()
        self.config = utils.init_config(config, self.get_default_metric_config(), self.get_name())

        self.num_ig_masks = 1
        self.lbls = ['all']
        self.use_area_rngs = self.config['USE_AREA_RANGES']
        if self.use_area_rngs:
            self.area_rngs = self.config['AREA_RANGES']
            self.area_rng_lbls = self.config['AREA_RANGE_LABELS']
            self.num_ig_masks += len(self.area_rng_lbls)
            self.lbls += self.area_rng_lbls

        self.use_time_rngs = self.config['USE_TIME_RANGES']
        if self.use_time_rngs:
            self.time_rngs = self.config['TIME_RANGES']
            self.time_rng_lbls = self.config['TIME_RANGE_LABELS']
            self.num_ig_masks += len(self.time_rng_lbls)
            self.lbls += self.time_rng_lbls

        self.array_labels = self.config['IOU_THRESHOLDS']
        self.rec_thrs = self.config['RECALL_THRESHOLDS']

        self.maxDet = self.config['MAX_DETECTIONS']
        self.float_array_fields = ['AP_' + lbl for lbl in self.lbls] + ['AR_' + lbl for lbl in self.lbls]
        self.fields = self.float_array_fields
        self.summary_fields = self.float_array_fields

    @_timing.time
    def eval_sequence(self, data):
        """Calculates GT and Tracker matches for one sequence for TrackMAP metrics. Adapted from
        https://github.com/TAO-Dataset/"""

        # Initialise results to zero for each sequence as the fields are only defined over the set of all sequences
        res = {}
        for field in self.fields:
            res[field] = [0 for _ in self.array_labels]

        gt_ids, dt_ids = data['gt_track_ids'], data['dt_track_ids']

        if len(gt_ids) == 0 and len(dt_ids) == 0:
            for idx in range(self.num_ig_masks):
                res[idx] = None
            return res

        # get track data
        gt_tr_areas = data.get('gt_track_areas', None) if self.use_area_rngs else None
        gt_tr_lengths = data.get('gt_track_lengths', None) if self.use_time_rngs else None
        gt_tr_iscrowd = data.get('gt_track_iscrowd', None)
        dt_tr_areas = data.get('dt_track_areas', None) if self.use_area_rngs else None
        dt_tr_lengths = data.get('dt_track_lengths', None) if self.use_time_rngs else None
        is_nel = data.get('not_exhaustively_labeled', False)

        # compute ignore masks for different track sets to eval
        gt_ig_masks = self._compute_track_ig_masks(len(gt_ids), track_lengths=gt_tr_lengths, track_areas=gt_tr_areas,
                                                   iscrowd=gt_tr_iscrowd)
        dt_ig_masks = self._compute_track_ig_masks(len(dt_ids), track_lengths=dt_tr_lengths, track_areas=dt_tr_areas,
                                                   is_not_exhaustively_labeled=is_nel, is_gt=False)

        boxformat = data.get('boxformat', 'xywh')
        ious = self._compute_track_ious(data['dt_tracks'], data['gt_tracks'], iou_function=data['iou_type'],
                                        boxformat=boxformat)

        for mask_idx in range(self.num_ig_masks):
            gt_ig_mask = gt_ig_masks[mask_idx]

            # Sort gt ignore last
            gt_idx = np.argsort([g for g in gt_ig_mask], kind="mergesort")
            gt_ids = [gt_ids[i] for i in gt_idx]

            ious_sorted = ious[:, gt_idx] if len(ious) > 0 else ious

            num_thrs = len(self.array_labels)
            num_gt = len(gt_ids)
            num_dt = len(dt_ids)

            # Array to store the "id" of the matched dt/gt
            gt_m = np.zeros((num_thrs, num_gt)) - 1
            dt_m = np.zeros((num_thrs, num_dt)) - 1

            gt_ig = np.array([gt_ig_mask[idx] for idx in gt_idx])
            dt_ig = np.zeros((num_thrs, num_dt))

            for iou_thr_idx, iou_thr in enumerate(self.array_labels):
                if len(ious_sorted) == 0:
                    break

                for dt_idx, _dt in enumerate(dt_ids):
                    iou = min([iou_thr, 1 - 1e-10])
                    # information about best match so far (m=-1 -> unmatched)
                    # store the gt_idx which matched for _dt
                    m = -1
                    for gt_idx, _ in enumerate(gt_ids):
                        # if this gt already matched continue
                        if gt_m[iou_thr_idx, gt_idx] > 0:
                            continue
                        # if _dt matched to reg gt, and on ignore gt, stop
                        if m > -1 and gt_ig[m] == 0 and gt_ig[gt_idx] == 1:
                            break
                        # continue to next gt unless better match made
                        if ious_sorted[dt_idx, gt_idx] < iou - np.finfo('float').eps:
                            continue
                        # if match successful and best so far, store appropriately
                        iou = ious_sorted[dt_idx, gt_idx]
                        m = gt_idx

                    # No match found for _dt, go to next _dt
                    if m == -1:
                        continue

                    # if gt to ignore for some reason update dt_ig.
                    # Should not be used in evaluation.
                    dt_ig[iou_thr_idx, dt_idx] = gt_ig[m]
                    # _dt match found, update gt_m, and dt_m with "id"
                    dt_m[iou_thr_idx, dt_idx] = gt_ids[m]
                    gt_m[iou_thr_idx, m] = _dt

            dt_ig_mask = dt_ig_masks[mask_idx]

            dt_ig_mask = np.array(dt_ig_mask).reshape((1, num_dt))  # 1 X num_dt
            dt_ig_mask = np.repeat(dt_ig_mask, num_thrs, 0)  # num_thrs X num_dt

            # Based on dt_ig_mask ignore any unmatched detection by updating dt_ig
            dt_ig = np.logical_or(dt_ig, np.logical_and(dt_m == -1, dt_ig_mask))
            # store results for given video and category
            res[mask_idx] = {
                "dt_ids": dt_ids,
                "gt_ids": gt_ids,
                "dt_matches": dt_m,
                "gt_matches": gt_m,
                "dt_scores": data['dt_track_scores'],
                "gt_ignore": gt_ig,
                "dt_ignore": dt_ig,
            }

        return res

    def combine_sequences(self, all_res):
        """Combines metrics across all sequences. Computes precision and recall values based on track matches.
        Adapted from https://github.com/TAO-Dataset/
        """
        num_thrs = len(self.array_labels)
        num_recalls = len(self.rec_thrs)

        # -1 for absent categories
        precision = -np.ones(
            (num_thrs, num_recalls, self.num_ig_masks)
        )
        recall = -np.ones((num_thrs, self.num_ig_masks))

        for ig_idx in range(self.num_ig_masks):
            ig_idx_results = [res[ig_idx] for res in all_res.values() if res[ig_idx] is not None]

            # Remove elements which are None
            if len(ig_idx_results) == 0:
                continue

            # Append all scores: shape (N,)
            # limit considered tracks for each sequence if maxDet > 0
            if self.maxDet == 0:
                dt_scores = np.concatenate([res["dt_scores"] for res in ig_idx_results], axis=0)

                dt_idx = np.argsort(-dt_scores, kind="mergesort")

                dt_m = np.concatenate([e["dt_matches"] for e in ig_idx_results],
                                      axis=1)[:, dt_idx]
                dt_ig = np.concatenate([e["dt_ignore"] for e in ig_idx_results],
                                       axis=1)[:, dt_idx]
            elif self.maxDet > 0:
                dt_scores = np.concatenate([res["dt_scores"][0:self.maxDet] for res in ig_idx_results], axis=0)

                dt_idx = np.argsort(-dt_scores, kind="mergesort")

                dt_m = np.concatenate([e["dt_matches"][:, 0:self.maxDet] for e in ig_idx_results],
                                      axis=1)[:, dt_idx]
                dt_ig = np.concatenate([e["dt_ignore"][:, 0:self.maxDet] for e in ig_idx_results],
                                       axis=1)[:, dt_idx]
            else:
                raise Exception("Number of maximum detections must be >= 0, but is set to %i" % self.maxDet)

            gt_ig = np.concatenate([res["gt_ignore"] for res in ig_idx_results])
            # num gt anns to consider
            num_gt = np.count_nonzero(gt_ig == 0)

            if num_gt == 0:
                continue

            tps = np.logical_and(dt_m != -1, np.logical_not(dt_ig))
            fps = np.logical_and(dt_m == -1, np.logical_not(dt_ig))

            tp_sum = np.cumsum(tps, axis=1).astype(dtype=np.float)
            fp_sum = np.cumsum(fps, axis=1).astype(dtype=np.float)

            for iou_thr_idx, (tp, fp) in enumerate(zip(tp_sum, fp_sum)):
                tp = np.array(tp)
                fp = np.array(fp)
                num_tp = len(tp)
                rc = tp / num_gt
                if num_tp:
                    recall[iou_thr_idx, ig_idx] = rc[-1]
                else:
                    recall[iou_thr_idx, ig_idx] = 0

                # np.spacing(1) ~= eps
                pr = tp / (fp + tp + np.spacing(1))
                pr = pr.tolist()

                # Ensure precision values are monotonically decreasing
                for i in range(num_tp - 1, 0, -1):
                    if pr[i] > pr[i - 1]:
                        pr[i - 1] = pr[i]

                # find indices at the predefined recall values
                rec_thrs_insert_idx = np.searchsorted(rc, self.rec_thrs, side="left")

                pr_at_recall = [0.0] * num_recalls

                try:
                    for _idx, pr_idx in enumerate(rec_thrs_insert_idx):
                        pr_at_recall[_idx] = pr[pr_idx]
                except IndexError:
                    pass

                precision[iou_thr_idx, :, ig_idx] = (np.array(pr_at_recall))

        res = {'precision': precision, 'recall': recall}

        # compute the precision and recall averages for the respective alpha thresholds and ignore masks
        for lbl in self.lbls:
            res['AP_' + lbl] = np.zeros((len(self.array_labels)), dtype=np.float)
            res['AR_' + lbl] = np.zeros((len(self.array_labels)), dtype=np.float)

        for a_id, alpha in enumerate(self.array_labels):
            for lbl_idx, lbl in enumerate(self.lbls):
                p = precision[a_id, :, lbl_idx]
                if len(p[p > -1]) == 0:
                    mean_p = -1
                else:
                    mean_p = np.mean(p[p > -1])
                res['AP_' + lbl][a_id] = mean_p
                res['AR_' + lbl][a_id] = recall[a_id, lbl_idx]

        return res

    def combine_classes_class_averaged(self, all_res):
        """Combines metrics across all classes by averaging over the class values"""
        res = {}
        for field in self.fields:
            res[field] = np.zeros((len(self.array_labels)), dtype=np.float)
            field_stacked = np.array([res[field] for res in all_res.values()])

            for a_id, alpha in enumerate(self.array_labels):
                values = field_stacked[:, a_id]
                if len(values[values > -1]) == 0:
                    mean = -1
                else:
                    mean = np.mean(values[values > -1])
                res[field][a_id] = mean
        return res

    def combine_classes_det_averaged(self, all_res):
        """Combines metrics across all classes by averaging over the detection values"""

        res = {}
        for field in self.fields:
            res[field] = np.zeros((len(self.array_labels)), dtype=np.float)
            field_stacked = np.array([res[field] for res in all_res.values()])

            for a_id, alpha in enumerate(self.array_labels):
                values = field_stacked[:, a_id]
                if len(values[values > -1]) == 0:
                    mean = -1
                else:
                    mean = np.mean(values[values > -1])
                res[field][a_id] = mean
        return res

    def _compute_track_ig_masks(self, num_ids, track_lengths=None, track_areas=None, iscrowd=None,
                                is_not_exhaustively_labeled=False, is_gt=True):
        """
        Computes ignore masks for different track sets to evaluate
        :param num_ids: the number of track IDs
        :param track_lengths: the lengths of the tracks (number of timesteps)
        :param track_areas: the average area of a track
        :param iscrowd: whether a track is marked as crowd
        :param is_not_exhaustively_labeled: whether the track category is not exhaustively labeled
        :param is_gt: whether it is gt
        :return: the track ignore masks
        """
        # for TAO tracks for classes which are not exhaustively labeled are not evaluated
        if not is_gt and is_not_exhaustively_labeled:
            track_ig_masks = [[1 for _ in range(num_ids)] for i in range(self.num_ig_masks)]
        else:
            # consider all tracks
            track_ig_masks = [[0 for _ in range(num_ids)]]

            # consider tracks with certain area
            if self.use_area_rngs:
                for rng in self.area_rngs:
                    track_ig_masks.append([0 if rng[0] - np.finfo('float').eps <= area <= rng[1] + np.finfo('float').eps
                                           else 1 for area in track_areas])

            # consider tracks with certain duration
            if self.use_time_rngs:
                for rng in self.time_rngs:
                    track_ig_masks.append([0 if rng[0] - np.finfo('float').eps <= length
                                                <= rng[1] + np.finfo('float').eps else 1 for length in track_lengths])

        # for YouTubeVIS evaluation tracks with crowd tag are not evaluated
        if is_gt and iscrowd:
            track_ig_masks = [np.logical_or(mask, iscrowd) for mask in track_ig_masks]

        return track_ig_masks

    @staticmethod
    def _compute_bb_track_iou(dt_track, gt_track, boxformat='xywh'):
        """
        Calculates the track IoU for one detected track and one ground truth track for bounding boxes
        :param dt_track: the detected track (format: dictionary with frame index as keys and
                            numpy arrays as values)
        :param gt_track: the ground truth track (format: dictionary with frame index as keys and
                        numpy array as values)
        :param boxformat: the format of the boxes
        :return: the track IoU
        """
        intersect = 0
        union = 0
        image_ids = set(gt_track.keys()) | set(dt_track.keys())
        for image in image_ids:
            g = gt_track.get(image, None)
            d = dt_track.get(image, None)
            if boxformat == 'xywh':
                if d is not None and g is not None:
                    dx, dy, dw, dh = d
                    gx, gy, gw, gh = g
                    w = max(min(dx + dw, gx + gw) - max(dx, gx), 0)
                    h = max(min(dy + dh, gy + gh) - max(dy, gy), 0)
                    i = w * h
                    u = dw * dh + gw * gh - i
                    intersect += i
                    union += u
                elif d is None and g is not None:
                    union += g[2] * g[3]
                elif d is not None and g is None:
                    union += d[2] * d[3]
            elif boxformat == 'x0y0x1y1':
                if d is not None and g is not None:
                    dx0, dy0, dx1, dy1 = d
                    gx0, gy0, gx1, gy1 = g
                    w = max(min(dx1, gx1) - max(dx0, gx0), 0)
                    h = max(min(dy1, gy1) - max(dy0, gy0), 0)
                    i = w * h
                    u = (dx1 - dx0) * (dy1 - dy0) + (gx1 - gx0) * (gy1 - gy0) - i
                    intersect += i
                    union += u
                elif d is None and g is not None:
                    union += (g[2] - g[0]) * (g[3] - g[1])
                elif d is not None and g is None:
                    union += (d[2] - d[0]) * (d[3] - d[1])
            else:
                raise TrackEvalException('BoxFormat not implemented')
        if intersect > union:
            raise TrackEvalException("Intersection value > union value. Are the box values corrupted?")
        return intersect / union if union > 0 else 0

    @staticmethod
    def _compute_mask_track_iou(dt_track, gt_track):
        """
        Calculates the track IoU for one detected track and one ground truth track for segmentation masks
        :param dt_track: the detected track (format: dictionary with frame index as keys and
                            pycocotools rle encoded masks as values)
        :param gt_track: the ground truth track (format: dictionary with frame index as keys and
                            pycocotools rle encoded masks as values)
        :return: the track IoU
        """
        # only loaded when needed to reduce minimum requirements
        from pycocotools import mask as mask_utils

        intersect = .0
        union = .0
        image_ids = set(gt_track.keys()) | set(dt_track.keys())
        for image in image_ids:
            g = gt_track.get(image, None)
            d = dt_track.get(image, None)
            if d and g:
                intersect += mask_utils.area(mask_utils.merge([d, g], True))
                union += mask_utils.area(mask_utils.merge([d, g], False))
            elif not d and g:
                union += mask_utils.area(g)
            elif d and not g:
                union += mask_utils.area(d)
        if union < 0.0 - np.finfo('float').eps:
            raise TrackEvalException("Union value < 0. Are the segmentaions corrupted?")
        if intersect > union:
            raise TrackEvalException("Intersection value > union value. Are the segmentations corrupted?")
        iou = intersect / union if union > 0.0 + np.finfo('float').eps else 0.0
        return iou

    @staticmethod
    def _compute_track_ious(dt, gt, iou_function='bbox', boxformat='xywh'):
        """
        Calculate track IoUs for a set of ground truth tracks and a set of detected tracks
        """

        if len(gt) == 0 and len(dt) == 0:
            return []

        if iou_function == 'bbox':
            track_iou_function = partial(TrackMAP._compute_bb_track_iou, boxformat=boxformat)
        elif iou_function == 'mask':
            track_iou_function = partial(TrackMAP._compute_mask_track_iou)
        else:
            raise Exception('IoU function not implemented')

        ious = np.zeros([len(dt), len(gt)])
        for i, j in np.ndindex(ious.shape):
            ious[i, j] = track_iou_function(dt[i], gt[j])
        return ious

    @staticmethod
    def _row_print(*argv):
        """Prints results in an evenly spaced rows, with more space in first row"""
        if len(argv) == 1:
            argv = argv[0]
        to_print = '%-40s' % argv[0]
        for v in argv[1:]:
            to_print += '%-12s' % str(v)
        print(to_print)

max_stars_repo_name: collector-m/UniTrack

max_issues_repo_name: collector-m/UniTrack

max_forks_repo_name: collector-m/UniTrack


text:
import numpy as np
from sklearn.neighbors import NearestNeighbors
import cv2

def estimate_normals(p, k=20):
    neigh = NearestNeighbors(n_neighbors=k)
    neigh.fit(p)
    distances, indices = neigh.kneighbors(return_distance=True)

    dp = (p[indices] - p[:,None])
    U, s, V = np.linalg.svd(dp.transpose(0,2,1))
    nv = U[:, :, -1]
    return nv / np.linalg.norm(nv)

def test_norm():
    np.random.seed(3)
    from matplotlib import pyplot as plt
    pt_map = np.concatenate([get_line(n=33, s=1.0) for _ in range(3)], axis=0)

    pt_map = np.random.normal(loc=pt_map, scale=0.05)
    nv = estimate_normals(pt_map)
    plt.plot(pt_map[:,0], pt_map[:,1], '.', label='map')
    plt.quiver(
                pt_map[:,0], pt_map[:,1],
                nv[:,0], nv[:,1],
                scale_units='xy',
                angles='xy'
                )
    plt.gca().set_aspect('equal', 'datalim')
    plt.legend()
    plt.show()

def stable_subsample(p):
    """
    Geometrically Stable ICP Subsampling with constraints analysis
    Gelfand et al. 2003
    http://www.people.vcu.edu/~dbandyop/pubh8472/Gelfand_SVC.pdf
    https://graphics.stanford.edu/papers/stabicp/stabicp.pdf
    """
    # TODO : implement
    pass

class ICP():
    def __init__(self):
        pass

    @staticmethod
    def best_fit_transform(A, B):
        '''
        Calculates the least-squares best-fit transform that maps corresponding points A to B in m spatial dimensions
        Input:
          A: Naxm numpy array of corresponding points
          B: Nbxm numpy array of corresponding points
        Returns:
          T: (m+1)x(m+1) homogeneous transformation matrix that maps A on to B
          R: mxm rotation matrix
          t: mx1 translation vector
        '''

        # assert A.shape == B.shape

        # get number of dimensions
        m = A.shape[1]

        # translate points to their centroids
        centroid_A = np.mean(A, axis=0)
        centroid_B = np.mean(B, axis=0)
        AA = A - centroid_A
        BB = B - centroid_B

        # rotation matrix
        H = np.dot(AA.T, BB)
        U, S, Vt = np.linalg.svd(H)
        R = np.dot(Vt.T, U.T)

        # special reflection case
        if np.linalg.det(R) < 0:
           Vt[m-1,:] *= -1
           R = np.dot(Vt.T, U.T)

        # translation
        t = centroid_B.T - np.dot(R,centroid_A.T)

        # homogeneous transformation
        T = np.identity(m+1)
        T[:m, :m] = R
        T[:m, m] = t

        return T, R, t

    @staticmethod
    def best_fit_transform_point_to_plane(src, dst):
        # TODO : generalize to N-d cases?
        # TODO : use neighborhood from prior computations?
        nvec = estimate_normals(dst, k=20)

        # construct according to
        # https://gfx.cs.princeton.edu/proj/iccv05_course/iccv05_icp_gr.pdf

        A_lhs = np.cross(src, nvec)[:,None]
        A_rhs = nvec

        Amat = np.concatenate([A_lhs, A_rhs], axis=1) # == Nx3

        # == Experimental : Stability Analysis ==
        #C = Amat.T.dot(Amat) # 3x3 cov-mat
        #w,v = np.linalg.eig(C)
        ##c = w[0] / w[-1] # stability
        #c = w[-1]
        ##print('w[-1] : {}'.format(w[-1]))
        c = None
        # =======================================

        bvec = - ((src - dst)*(nvec)).sum(axis=-1)

        tvec = np.linalg.pinv(Amat).dot(bvec) # == (dh, dx, dy)

        R = Rmat( tvec[0] )
        t = tvec[1:]

        T = np.eye(3, dtype=np.float32)
        T[:2, :2] = R
        T[:2, 2]  = t

        return T, R, t, c

    @staticmethod
    def nearest_neighbor(src, dst):
        '''
        Find the nearest (Euclidean) neighbor in dst for each point in src
        Input:
            src: Naxm array of points
            dst: Nbxm array of points
        Output:
            distances: Euclidean distances of the nearest neighbor
            indices: dst indices of the nearest neighbor
        '''
        assert src.shape[1:] == dst.shape[1:]

        neigh = NearestNeighbors(n_neighbors=1)
        neigh.fit(dst)
        distances, indices = neigh.kneighbors(src, return_distance=True)
        return distances.ravel(), indices.ravel()

    @staticmethod
    def projected_neighbor(src, dst, origin):
        """
        index-matching correspondences, according to Blais and Levine, '95
        https://ieeexplore.ieee.org/stamp/stamp.jsp?tp=&arnumber=400574
        """
        # TODO : implement
        pass

    @staticmethod
    def icp(A, B, init_pose=None, max_iterations=20, tolerance=0.001):
        '''
        The Iterative Closest Point method: finds best-fit transform that maps points A on to points B
        Input:
            A: Naxm numpy array of source mD points
            B: Nbxm numpy array of destination mD point
            init_pose: (m+1)x(m+1) homogeneous transformation
            max_iterations: exit algorithm after max_iterations
            tolerance: convergence criteria
        Output:
            T: final homogeneous transformation that maps A on to B, T.A = B
            distances: Euclidean distances (errors) of the nearest neighbor
            i: number of iterations to converge
        '''

        # assert A.shape == B.shape

        # get number of dimensions

        ndim = A.shape[1]
        nA, nB = A.shape[0], B.shape[0]
        assert(A.shape[1:] == B.shape[1:])

        # make points homogeneous, copy them to maintain the originals
        src = np.ones((nA,ndim+1), dtype=np.float32)
        dst = np.ones((nB,ndim+1), dtype=np.float32)

        src[:, :ndim] = A
        dst[:, :ndim] = B

        # apply the initial pose estimation
        if init_pose is not None:
            src = np.dot(init_pose, src)

        prev_error = 0

        for i in range(max_iterations):
            # find the nearest neighbors between the current source and destination points
            distances, indices = ICP.nearest_neighbor(src[:,:-1], dst[:,:-1])

            # compute the transformation between the current source and nearest destination points
            #T,_,_ = ICP.best_fit_transform(src[:,:-1], dst[indices,:-1])
            T,_,_,_ = ICP.best_fit_transform_point_to_plane(src[:,:-1], dst[indices,:-1])

            src = np.dot(src,T.T) # right-multiply transform matrix

            # check error
            mean_error = np.mean(distances)
            if np.abs(prev_error - mean_error) < tolerance:
                break
            prev_error = mean_error

        # calculate final transformation
        #T,_,_ = ICP.best_fit_transform(A, src[:,:-1])
        T,_,_,c = ICP.best_fit_transform_point_to_plane(A, src[:,:-1])
        #print('stability : {}'.format(c))

        # reprojection-error based filter
        # (i.e. without RANSAC)
        # TODO : this is probably faster, but does performance degrade?
        delta = np.linalg.norm(dst[indices] - src, axis=-1)
        msk   = (delta < 0.2)
        T3    = None

        # alternative - refine transform
        # (mostly, compute the mask and validate truth)
        # opencv version - RANSAC
        # doesn't really do anything
        #T3, msk = cv2.estimateAffine2D(
        #        src[None,:,:-1],
        #        dst[None,indices,:-1],
        #        method=cv2.FM_RANSAC,
        #        ransacReprojThreshold=0.2, # TODO : expose this param
        #        maxIters=2000,
        #        confidence=0.9999,
        #        refineIters=100
        #        )

        ms = msk.sum()
        #if ms != msk.size:
        #    print '{}% = {}/{}'.format(float(ms)/msk.size*100, ms, msk.size)

        inlier_ratio = float(ms) / msk.size
        #inlier_ratio *= np.float32( c > 3e-2 ) # account for stability
        # TODO : is it reasonable to account for stability?
        # I think c==3e-2 equates to ~10 deg rotation / 0.17m translation
        # not 100% sure about this.

        if T3 is not None:
            #print 'T3'
            #print T[:2] - T3
            #if inlier_ratio > 0.75:
            #    T3 = np.concatenate([T3, [[0,0,1]] ], axis=0)
            #    T = T.dot(T3)
            # TODO : revive? idk
            pass

        return T, distances, indices, i, inlier_ratio

def Rmat(x):
    c,s = np.cos(x), np.sin(x)
    R = np.float32([c,-s,s,c]).reshape(2,2)
    return R

def get_line(n, s):
    # return Nx2 line
    t = np.linspace(0, s, n)
    pt = np.stack([t, np.zeros_like(t)], axis=-1)
    h = np.random.uniform(-np.pi, np.pi)
    o = np.random.uniform(-1.0, 1.0)
    return pt.dot(Rmat(h).T) + o

def main():
    np.random.seed(0)
    from matplotlib import pyplot as plt
    # parameters
    n_size = 100
    n_repeat = 100
    s_noise = 0.1

    pt_map = np.concatenate([get_line(n=33, s=1.0) for _ in range(3)], axis=0)

    # get geometric transformation
    # scan_{map} -> scan_{base_link}
    h = np.random.uniform(-np.pi/4, np.pi/4)
    R = Rmat(h)
    t = np.random.uniform(-1.0, 1.0, size=(1,2))
    T0 = np.eye(3)
    T0[:2,:2] = R
    T0[:2,2] = t

    print ('T (x,y,h) : ({}, {}, {})'.format(t[0,0], t[0,1], h))

    errs = []
    #n_sub  = 20
    #n_subs = np.linspace(1, len(pt_map), num=10).astype(np.int32)
    n_subs = [50]
    for n_sub in n_subs:
        print('n_sub = {}'.format(n_sub))
        err_i = []

        for _ in range(n_repeat):
            # subsample
            pt_scan = np.copy(pt_map[np.random.choice(pt_map.shape[0], size=n_sub, replace=False)])

            # apply transform
            pt_scan = pt_scan.dot(R.T) + t

            # apply noise
            pt_scan = np.random.normal(loc=pt_scan, scale=s_noise)
            T, distances, indices, iterations, inl = ICP.icp(pt_scan, pt_map,
                    tolerance=0.000001)

            pt_scan_r = pt_scan.dot(T[:2,:2].T) + T[:2,2].reshape(1,2)

            # scan_{base_link} -> scan_{map}
            h2 = np.arctan2(T[1,0], T[1,1])
            t2 = T[:2,2]
            #print ('T_inv (x,y,h) : ({}, {}, {})'.format(t2[0], t2[1], h2))

            # compute error metric
            #print T0.dot(T)
            err_T = T0.dot(T)
            dt = err_T[:2,2]
            dh = np.arctan2(err_T[1,0], err_T[1,1])
            err_i.append(np.abs([dt[0], dt[1], dh]))
            #print('err (dx,dy,dh) : ({},{},{})'.format(dt[0],dt[1],dh))
        err_i = np.mean(err_i, axis=0)
        errs.append(err_i)

        plt.plot(pt_map[:,0], pt_map[:,1], '.', label='map')
        plt.plot(pt_scan[:,0], pt_scan[:,1], '.', label='scan')
        plt.plot(pt_scan_r[:,0], pt_scan_r[:,1], '.', label='scan-T')
        plt.legend()
        plt.show()

    #plt.plot(np.float32(n_subs) / len(pt_map), errs)
    #plt.legend(['dx','dy','dh'])
    #plt.xlabel('sample ratio')
    #plt.ylabel('error')
    #plt.title('icp characterization')
    #plt.show()

if __name__ == "__main__":
    #main()
    test_norm()

ext: py

reason: import numpy


text:
import numpy as np
from gym.spaces import Dict, Box, Discrete

from noge.data_types import PartialGraphObservation
from .base_graph_env import BaseGraphEnv, StepEvents

class OnlineGraphEnv(BaseGraphEnv):

    def __init__(self, max_episode_steps, reward_type='path_length', data_generator=None,
                 max_nodes=None, max_edges=None, nn_feat=False):
        super().__init__(max_episode_steps=max_episode_steps, reward_type=reward_type, data_generator=data_generator)

        self.max_nodes = max_nodes
        self.max_edges = max_edges
        self.nn_feat = nn_feat

        # online state (for more efficient feeding to torch geometric)
        self._idx_to_node = np.empty(shape=(max_nodes,), dtype=np.int64)  # map online nodes to their original names
        self._node_to_idx = None  # map discovered nodes to online indices {0, 1, ..., N-1}
        self._online_frontier = None  # the frontier nodes (online ids)

        # initialize these arrays only once, then overwrite in each episode
        max_visits = min(max_episode_steps + 1, max_nodes)
        self._node_counts = np.empty(shape=(max_visits,), dtype=np.int64)
        self._edge_counts = np.empty(shape=(max_visits,), dtype=np.int64)
        self._online_visit_sequence = np.empty(shape=(max_visits,), dtype=np.int64)
        self._online_discovery_sequence = np.empty(shape=(max_nodes,), dtype=np.int64)
        self._known_edge_index = np.empty(shape=(2, max_edges), dtype=np.int64)

        # observations
        num_output_meas = 9
        self.observation_space = Dict(dict(
            t=Discrete(n=max_episode_steps),
            new_nodes=Box(low=0, high=max_nodes, shape=(max_nodes,), dtype=np.int64),
            new_edges=Box(low=0, high=max_edges, shape=(2, max_edges), dtype=np.int64),
            visited_node=Discrete(n=max_nodes),
            nn=Discrete(n=max_nodes),
            bf=Discrete(n=max_nodes),
            df=Discrete(n=max_nodes),
            meas=Box(low=-np.inf, high=np.inf, shape=(num_output_meas,), dtype=np.float32)
        ))

        # actions
        self.action_space = Discrete(n=max_nodes)

    def reset(self, **kwargs):

        events: StepEvents = super().reset(**kwargs)

        m = len(events['new_edges'])

        # track the discovered data renaming nodes as {0,1,...,N}
        self._idx_to_node[0] = self._source
        self._idx_to_node[1:m+1] = events['new_nodes']
        self._node_to_idx = {x: i for i, x in enumerate(self._idx_to_node[:m+1])}

        online_frontier_list = list(range(1, m+1))
        self._online_frontier = online_frontier_list

        # track the explored nodes as visited_seq of shape [N]
        online_source = 0
        self._online_visit_sequence[0] = online_source

        # track discovery sequence (of sets)
        self._online_discovery_sequence[0] = online_source
        self._online_discovery_sequence[1:1+m] = online_frontier_list

        # track the discovered edges as edge_index of shape [2, E]
        self._known_edge_index[0, :m] = online_source
        self._known_edge_index[1, :m] = online_frontier_list
        self._known_edge_index[0, m:2*m] = online_frontier_list
        self._known_edge_index[1, m:2*m] = online_source

        # reset counts
        self._node_counts[0] = m + 1
        self._edge_counts[0] = 2*m

        obs = self._get_partial_obs(events)
        return obs

    def _get_partial_obs(self, events: StepEvents) -> PartialGraphObservation:

        t = self._time_step

        if t > 0:
            n_old = self._node_counts[t-1]
            m_old = self._edge_counts[t-1]
        else:
            n_old = 0
            m_old = 0

        # new nodes
        n_new = self._node_counts[t]
        new_nodes = self._online_discovery_sequence[n_old:n_new].copy()

        # new edges
        m_new = self._edge_counts[t]
        new_edges = self._known_edge_index[:, m_old:m_new].copy()

        # visited
        visited_node = self._online_visit_sequence[t]

        if len(self._online_frontier):
            if self.nn_feat:
                nn = self.nearest_neighbor_policy()
            else:
                nn = -1

            bf = self._online_frontier[0]
            df = self._online_frontier[-1]
        else:
            nn = bf = df = -1

        path_nodes = np.array([self._node_to_idx[v] for v in events['path']], dtype=np.int64)
        delta_obs = PartialGraphObservation(t=t,
                                            new_nodes=new_nodes,
                                            new_edges=new_edges,
                                            visited_node=visited_node,
                                            path_nodes=path_nodes,
                                            path_cost=events['cost'],
                                            nn=nn,
                                            bf=bf,
                                            df=df,
                                            meas=self.get_current_measurements()
                                            )

        return delta_obs

    def nearest_neighbor_policy(self):
        v = super().nearest_neighbor_policy()
        return self._node_to_idx[v]

    def bfs_policy(self):
        return self._online_frontier[0]

    def dfs_policy(self):
        return self._online_frontier[-1]

    def random_policy(self):
        i = self.rng.randint(len(self._online_frontier))
        return self._online_frontier[i]

    def step(self, action):
        offline_subgoal = self._idx_to_node[action]
        events, reward, done, info = super().step(offline_subgoal)

        t = self._time_step  # time step > 0 (has been incremented in BaseGraphEnv)

        online_subgoal = action  # action = online node index = integer in [0, ..., N-1]
        self._online_frontier.remove(online_subgoal)
        self._online_visit_sequence[t] = online_subgoal

        # track discovered edges online indices
        num_discovered_edges = len(events['new_edges'])
        self._edge_counts[t] = self._edge_counts[t-1] + num_discovered_edges * 2

        new_nodes = events['new_nodes']
        num_discovered_nodes = len(new_nodes)
        self._node_counts[t] = self._node_counts[t-1] + num_discovered_nodes

        if num_discovered_edges > 0:
            # NOTE: Remember that events.discovered_frontier_neighbors is a set, not a list
            discovered_frontier_nodes_oidx = [self._node_to_idx[v] for v in events['found_frontier_nodes']]

            if num_discovered_nodes > 0:
                # track discovered nodes online indices
                n_prev = self._node_counts[t-1]
                n_curr = self._node_counts[t]
                new_nodes_oidx = list(range(n_prev, n_curr))
                self._node_to_idx.update(dict(zip(new_nodes, new_nodes_oidx)))
                self._idx_to_node[n_prev:n_curr] = new_nodes

                # update online node discovery sequence
                self._online_discovery_sequence[n_prev:n_curr] = new_nodes_oidx
                discovered_ends_oidx = new_nodes_oidx + discovered_frontier_nodes_oidx

                # update online frontier
                self._online_frontier.extend(new_nodes_oidx)

            else:
                discovered_ends_oidx = discovered_frontier_nodes_oidx

            # add both directions for undirected graph (symmetric adjacency)
            a = self._edge_counts[t-1]
            b = a + num_discovered_edges
            c = b + num_discovered_edges
            self._known_edge_index[0, a:b] = online_subgoal
            self._known_edge_index[1, a:b] = discovered_ends_oidx
            self._known_edge_index[0, b:c] = discovered_ends_oidx
            self._known_edge_index[1, b:c] = online_subgoal

        obs = self._get_partial_obs(events)

        return obs, reward, done, info

    def get_current_measurements(self):
        """Always return the cumulative path length.
           Conversion to other measurements will happen at some later stage of the algorithm"""

        t = self._time_step
        C = t + 1
        L = self._total_path_cost
        V = self._node_counts[t]
        E = self._edge_counts[t]

        exp_rate = self._exploration_rates[t]  # t / L_t
        plot = 0 if t == 0 else L / t  # path length over time
        node_discovery = C / V
        edge_discovery = C / E
        sparsity = (V - 1) / E

        meas = np.array([t, L, V, E, exp_rate, node_discovery, edge_discovery, sparsity, plot], dtype=np.float32)

        return meas

max_stars_repo_name: johny-c/noge

max_issues_repo_name: johny-c/noge

max_forks_repo_name: johny-c/noge


text:
import os.path as osp
import pickle

import numpy as np
from pyturbo import Stage

class LoadFeature(Stage):

    """
    Input: video_id
    Output: features [N x D]
    """

    def allocate_resource(self, resources, *, feature_dir, file_suffix='pkl',
                          worker_per_cpu=1):
        self.feature_dir = feature_dir
        self.file_suffix = file_suffix
        return resources.split(len(resources.get('cpu'))) * worker_per_cpu

    @staticmethod
    def load_features(feature_path):
        features = []
        with open(feature_path, 'rb') as f:
            while True:
                try:
                    _, frame_feature = pickle.load(f)
                    features.append(frame_feature)
                except EOFError:
                    break
        return features

    def process(self, task):
        task.start(self)
        video_id = task.content
        feature_path = osp.join(
            self.feature_dir, f'{video_id}.{self.file_suffix}')
        features = self.load_features(feature_path)
        return task.finish(features)


text:
"""

test_pq_gen.py

Author: Jordan Mirocha
Affiliation: UCLA
Created on: Thu Dec 15 14:31:54 PST 2016

Description:

"""

import ares
import numpy as np

def test():
    # Make sure we can parameterize a bunch of things

    pars = ares.util.ParameterBundle('pop:sfe-dpl')

    pop = ares.populations.GalaxyPopulation(**pars)

    base_kwargs = \
    {
     'pq_func': 'dpl',
     'pq_func_var': 'Mh',
     'pq_func_par0': 0.1,
     'pq_func_par1': 3e11,
     'pq_func_par2': 0.6,
     'pq_func_par3': -0.6,
     'pq_func_par4': 1e10,
    }

    val = []
    parameterizable_things = \
        ['fstar', 'fshock', 'fpoll', 'fstall', 'rad_yield', 'fduty', 'fesc_LW']

    for par in parameterizable_things:

        pars = base_kwargs.copy()
        pars['pop_{!s}'.format(par)] = 'pq'

        pop = ares.populations.GalaxyPopulation(**pars)

        func = pop.__getattr__(par)
        val.append(func(z=6, Mh=1e12))

    print('{!s}'.format(val))
    assert np.unique(val).size == 1, "Error in building ParameterizedQuantity!"

if __name__ == '__main__':
    test()

max_stars_repo_name: jlashner/ares

max_issues_repo_name: jlashner/ares

max_forks_repo_name: jlashner/ares


text:
import numpy as np
import scipy.stats as ss
import tensorflow as tf
from tensorflow.keras import backend as K

rating_weights = K.expand_dims(tf.constant(np.arange(1, 11), dtype='float32'), -1)

def earth_movers_distance(y_true, y_pred):
    cdf_true = K.cumsum(y_true, axis=-1)
    cdf_pred = K.cumsum(y_pred, axis=-1)
    emd = K.sqrt(K.mean(K.square(cdf_true - cdf_pred), axis=-1))
    return K.mean(emd)

def pearson_correlation(y_true, y_pred):
    # y_true = y_true * rating_weights
    # y_pred = y_pred * rating_weights
    #
    # xm = y_true - K.mean(y_true)
    # ym = y_pred - K.mean(y_pred)
    # print(xm, ym)
    # pearson_correlation = K.sum(xm * ym) / K.sqrt(K.sum(K.square(xm) * K.square(ym)))
    # print(pearson_correlation)
    # return K.square(pearson_correlation)  # is is actually R-squared from regression
    means_true = y_true - K.mean(y_true)
    means_pred = y_pred - K.mean(y_pred)

    # normalizing stage - setting a 1 variance
    means_true = K.l2_normalize(means_true, axis=0)
    means_pred = K.l2_normalize(means_pred, axis=0)

    # final result
    pearson_correlation = K.sum(means_true * means_pred)

    return pearson_correlation

def pearson_correlation_ava(y_true, y_pred):
    x = K.cumsum(K.dot(y_true, rating_weights))
    y = K.cumsum(K.dot(y_pred, rating_weights))
    return pearson_correlation(x, y)

def spearman_correlation(y_true, y_pred):
    return ss.spearmanr(y_true.numpy(), y_pred.numpy())[0]

def two_class_quality(y_true, y_pred):
    x = K.dot(y_true, rating_weights)
    y = K.dot(y_pred, rating_weights)
    score = K.equal(tf.floor(x / 5), tf.floor(y / 5))
    return K.mean(score)

def two_class_quality_acc(y_true, y_pred):
    x = np.array(y_true)
    y = np.array(y_pred)
    score = np.equal(np.floor(x / 5), np.floor(y / 5))
    return np.mean(score)

def mean_abs_percentage(y_true, y_pred):
    abs_diff = K.abs(y_pred - y_true) / y_true
    return K.mean(1 - abs_diff)

def mean_abs_percentage_acc(y_true, y_pred):
    x, y = np.array(y_true), np.array(y_pred)
    abs_diff = np.abs(x - y) / x
    return np.mean(1 - abs_diff)

def mean_abs_percentage_ava(y_true, y_pred):
    x = K.dot(y_true, rating_weights)
    y = K.dot(y_pred, rating_weights)
    return mean_abs_percentage(x, y)

if __name__ == '__main__':
    a = np.array([1.62, 1.83, 1.89, 1.55, 1.74, 1.6, 1.6, 1.72, 1.54, 1.82])
    b = np.array([57.15, 91.69, 95.27, 56.16, 78.52, 66.09, 63.71, 79.58, 50.22, 93.39])

    print(a)
    print(b)
    print(pearson_correlation(a, b))
    print(ss.pearsonr(a, b))
    print(spearman_correlation(a, b))
    print(ss.spearmanr(a, b))

    a = np.array([1.62, 4.83, 5.89, 8.55, 8.74, 6.6, ])
    b = np.array([2.62, 3.83, 1.89, 6.55, 5.74, 4.6, ])
    print(two_class_quality(a, b))
    print(mean_abs_percentage(a, b))

max_stars_repo_name: maindolaamit/nima

max_issues_repo_name: maindolaamit/nima

max_forks_repo_name: maindolaamit/nima

reason: import numpy,import scipy


text:
'''getting started with diffusions, continuous time stochastic processes

Author: josef-pktd
License: BSD

References
----------

An Algorithmic Introduction to Numerical Simulation of Stochastic Differential
Equations
Author(s): Desmond J. Higham
Source: SIAM Review, Vol. 43, No. 3 (Sep., 2001), pp. 525-546
Published by: Society for Industrial and Applied Mathematics
Stable URL: http://www.jstor.org/stable/3649798

http://www.sitmo.com/  especially the formula collection

Notes
-----

OU process: use same trick for ARMA with constant (non-zero mean) and drift
some of the processes have easy multivariate extensions

*Open Issues*

include xzero in returned sample or not? currently not

*TODOS*

* Milstein from Higham paper, for which processes does it apply
* Maximum Likelihood estimation
* more statistical properties (useful for tests)
* helper functions for display and MonteCarlo summaries (also for testing/checking)
* more processes for the menagerie (e.g. from empirical papers)
* characteristic functions
* transformations, non-linear e.g. log
* special estimators, e.g. Ait Sahalia, empirical characteristic functions
* fft examples
* check naming of methods, "simulate", "sample", "simexact", ... ?

stochastic volatility models: estimation unclear

finance applications ? option pricing, interest rate models

'''
from __future__ import print_function
import numpy as np
from scipy import stats, signal
import matplotlib.pyplot as plt

#np.random.seed(987656789)

class Diffusion(object):
    '''Wiener Process, Brownian Motion with mu=0 and sigma=1
    '''
    def __init__(self):
        pass

    def simulateW(self, nobs=100, T=1, dt=None, nrepl=1):
        '''generate sample of Wiener Process
        '''
        dt = T*1.0/nobs
        t = np.linspace(dt, 1, nobs)
        dW = np.sqrt(dt)*np.random.normal(size=(nrepl, nobs))
        W = np.cumsum(dW,1)
        self.dW = dW
        return W, t

    def expectedsim(self, func, nobs=100, T=1, dt=None, nrepl=1):
        '''get expectation of a function of a Wiener Process by simulation

        initially test example from
        '''
        W, t = self.simulateW(nobs=nobs, T=T, dt=dt, nrepl=nrepl)
        U = func(t, W)
        Umean = U.mean(0)
        return U, Umean, t

class AffineDiffusion(Diffusion):
    '''

    differential equation:

    :math::
    dx_t = f(t,x)dt + \sigma(t,x)dW_t

    integral:

    :math::
    x_T = x_0 + \\int_{0}^{T}f(t,S)dt + \\int_0^T  \\sigma(t,S)dW_t

    TODO: check definition, affine, what about jump diffusion?

    '''

    def __init__(self):
        pass

    def sim(self, nobs=100, T=1, dt=None, nrepl=1):
        # this doesn't look correct if drift or sig depend on x
        # see arithmetic BM
        W, t = self.simulateW(nobs=nobs, T=T, dt=dt, nrepl=nrepl)
        dx =  self._drift() + self._sig() * W
        x  = np.cumsum(dx,1)
        xmean = x.mean(0)
        return x, xmean, t

    def simEM(self, xzero=None, nobs=100, T=1, dt=None, nrepl=1, Tratio=4):
        '''

        from Higham 2001

        TODO: reverse parameterization to start with final nobs and DT
        TODO: check if I can skip the loop using my way from exactprocess
              problem might be Winc (reshape into 3d and sum)
        TODO: (later) check memory efficiency for large simulations
        '''
        #TODO: reverse parameterization to start with final nobs and DT
        nobs = nobs * Tratio  # simple way to change parameter
        # maybe wrong parameterization,
        # drift too large, variance too small ? which dt/Dt
        # _drift, _sig independent of dt is wrong
        if xzero is None:
            xzero = self.xzero
        if dt is None:
            dt = T*1.0/nobs
        W, t = self.simulateW(nobs=nobs, T=T, dt=dt, nrepl=nrepl)
        dW = self.dW
        t = np.linspace(dt, 1, nobs)
        Dt = Tratio*dt;
        L = nobs/Tratio;        # L EM steps of size Dt = R*dt
        Xem = np.zeros((nrepl,L));    # preallocate for efficiency
        Xtemp = xzero
        Xem[:,0] = xzero
        for j in np.arange(1,L):
           #Winc = np.sum(dW[:,Tratio*(j-1)+1:Tratio*j],1)
           Winc = np.sum(dW[:,np.arange(Tratio*(j-1)+1,Tratio*j)],1)
           #Xtemp = Xtemp + Dt*lamda*Xtemp + mu*Xtemp*Winc;
           Xtemp = Xtemp + self._drift(x=Xtemp) + self._sig(x=Xtemp) * Winc
           #Dt*lamda*Xtemp + mu*Xtemp*Winc;
           Xem[:,j] = Xtemp
        return Xem

'''
    R = 4; Dt = R*dt; L = N/R;        % L EM steps of size Dt = R*dt
    Xem = zeros(1,L);                 % preallocate for efficiency
    Xtemp = Xzero;
    for j = 1:L
       Winc = sum(dW(R*(j-1)+1:R*j));
       Xtemp = Xtemp + Dt*lambda*Xtemp + mu*Xtemp*Winc;
       Xem(j) = Xtemp;
    end
'''

class ExactDiffusion(AffineDiffusion):
    '''Diffusion that has an exact integral representation

    this is currently mainly for geometric, log processes

    '''

    def __init__(self):
        pass

    def exactprocess(self, xzero, nobs, ddt=1., nrepl=2):
        '''ddt : discrete delta t

        should be the same as an AR(1)
        not tested yet
        '''
        t = np.linspace(ddt, nobs*ddt, nobs)
        #expnt = np.exp(-self.lambd * t)
        expddt = np.exp(-self.lambd * ddt)
        normrvs = np.random.normal(size=(nrepl,nobs))
        #do I need lfilter here AR(1) ? if mean reverting lag-coeff<1
        #lfilter doesn't handle 2d arrays, it does?
        inc = self._exactconst(expddt) + self._exactstd(expddt) * normrvs
        return signal.lfilter([1.], [1.,-expddt], inc)

    def exactdist(self, xzero, t):
        expnt = np.exp(-self.lambd * t)
        meant = xzero * expnt + self._exactconst(expnt)
        stdt = self._exactstd(expnt)
        return stats.norm(loc=meant, scale=stdt)

class ArithmeticBrownian(AffineDiffusion):
    '''
    :math::
    dx_t &= \\mu dt + \\sigma dW_t
    '''

    def __init__(self, xzero, mu, sigma):
        self.xzero = xzero
        self.mu = mu
        self.sigma = sigma

    def _drift(self, *args, **kwds):
        return self.mu
    def _sig(self, *args, **kwds):
        return self.sigma
    def exactprocess(self, nobs, xzero=None, ddt=1., nrepl=2):
        '''ddt : discrete delta t

        not tested yet
        '''
        if xzero is None:
            xzero = self.xzero
        t = np.linspace(ddt, nobs*ddt, nobs)
        normrvs = np.random.normal(size=(nrepl,nobs))
        inc = self._drift + self._sigma * np.sqrt(ddt) * normrvs
        #return signal.lfilter([1.], [1.,-1], inc)
        return xzero + np.cumsum(inc,1)

    def exactdist(self, xzero, t):
        expnt = np.exp(-self.lambd * t)
        meant = self._drift * t
        stdt = self._sigma * np.sqrt(t)
        return stats.norm(loc=meant, scale=stdt)

class GeometricBrownian(AffineDiffusion):
    '''Geometric Brownian Motion

    :math::
    dx_t &= \\mu x_t dt + \\sigma x_t dW_t

    $x_t $ stochastic process of Geometric Brownian motion,
    $\mu $ is the drift,
    $\sigma $ is the Volatility,
    $W$ is the Wiener process (Brownian motion).

    '''
    def __init__(self, xzero, mu, sigma):
        self.xzero = xzero
        self.mu = mu
        self.sigma = sigma

    def _drift(self, *args, **kwds):
        x = kwds['x']
        return self.mu * x
    def _sig(self, *args, **kwds):
        x = kwds['x']
        return self.sigma * x

class OUprocess(AffineDiffusion):
    '''Ornstein-Uhlenbeck

    :math::
      dx_t&=\\lambda(\\mu - x_t)dt+\\sigma dW_t

    mean reverting process

    TODO: move exact higher up in class hierarchy
    '''
    def __init__(self, xzero, mu, lambd, sigma):
        self.xzero = xzero
        self.lambd = lambd
        self.mu = mu
        self.sigma = sigma

    def _drift(self, *args, **kwds):
        x = kwds['x']
        return self.lambd * (self.mu - x)
    def _sig(self, *args, **kwds):
        x = kwds['x']
        return self.sigma * x
    def exact(self, xzero, t, normrvs):
        #TODO: aggregate over time for process with observations for all t
        #      i.e. exact conditional distribution for discrete time increment
        #      -> exactprocess
        #TODO: for single t, return stats.norm -> exactdist
        expnt = np.exp(-self.lambd * t)
        return (xzero * expnt + self.mu * (1-expnt) +
                self.sigma * np.sqrt((1-expnt*expnt)/2./self.lambd) * normrvs)

    def exactprocess(self, xzero, nobs, ddt=1., nrepl=2):
        '''ddt : discrete delta t

        should be the same as an AR(1)
        not tested yet
        # after writing this I saw the same use of lfilter in sitmo
        '''
        t = np.linspace(ddt, nobs*ddt, nobs)
        expnt = np.exp(-self.lambd * t)
        expddt = np.exp(-self.lambd * ddt)
        normrvs = np.random.normal(size=(nrepl,nobs))
        #do I need lfilter here AR(1) ? lfilter doesn't handle 2d arrays, it does?
        from scipy import signal
        #xzero * expnt
        inc = ( self.mu * (1-expddt) +
                self.sigma * np.sqrt((1-expddt*expddt)/2./self.lambd) * normrvs )

        return signal.lfilter([1.], [1.,-expddt], inc)

    def exactdist(self, xzero, t):
        #TODO: aggregate over time for process with observations for all t
        #TODO: for single t, return stats.norm
        expnt = np.exp(-self.lambd * t)
        meant = xzero * expnt + self.mu * (1-expnt)
        stdt = self.sigma * np.sqrt((1-expnt*expnt)/2./self.lambd)
        from scipy import stats
        return stats.norm(loc=meant, scale=stdt)

    def fitls(self, data, dt):
        '''assumes data is 1d, univariate time series
        formula from sitmo
        '''
        # brute force, no parameter estimation errors
        nobs = len(data)-1
        exog = np.column_stack((np.ones(nobs), data[:-1]))
        parest, res, rank, sing = np.linalg.lstsq(exog, data[1:], rcond=-1)
        const, slope = parest
        errvar = res/(nobs-2.)
        lambd = -np.log(slope)/dt
        sigma = np.sqrt(-errvar * 2.*np.log(slope)/ (1-slope**2)/dt)
        mu = const / (1-slope)
        return mu, lambd, sigma

class SchwartzOne(ExactDiffusion):
    '''the Schwartz type 1 stochastic process

    :math::
    dx_t = \\kappa (\\mu - \\ln x_t) x_t dt + \\sigma x_tdW \\

    The Schwartz type 1 process is a log of the Ornstein-Uhlenbeck stochastic
    process.

    '''

    def __init__(self, xzero, mu, kappa, sigma):
        self.xzero = xzero
        self.mu = mu
        self.kappa = kappa
        self.lambd = kappa #alias until I fix exact
        self.sigma = sigma

    def _exactconst(self, expnt):
        return (1-expnt) * (self.mu - self.sigma**2 / 2. /self.kappa)

    def _exactstd(self, expnt):
        return self.sigma * np.sqrt((1-expnt*expnt)/2./self.kappa)

    def exactprocess(self, xzero, nobs, ddt=1., nrepl=2):
        '''uses exact solution for log of process
        '''
        lnxzero = np.log(xzero)
        lnx = super(self.__class__, self).exactprocess(xzero, nobs, ddt=ddt, nrepl=nrepl)
        return np.exp(lnx)

    def exactdist(self, xzero, t):
        expnt = np.exp(-self.lambd * t)
        #TODO: check this is still wrong, just guessing
        meant = np.log(xzero) * expnt + self._exactconst(expnt)
        stdt = self._exactstd(expnt)
        return stats.lognorm(loc=meant, scale=stdt)

    def fitls(self, data, dt):
        '''assumes data is 1d, univariate time series
        formula from sitmo
        '''
        # brute force, no parameter estimation errors
        nobs = len(data)-1
        exog = np.column_stack((np.ones(nobs),np.log(data[:-1])))
        parest, res, rank, sing = np.linalg.lstsq(exog, np.log(data[1:]), rcond=-1)
        const, slope = parest
        errvar = res/(nobs-2.)  #check denominator estimate, of sigma too low
        kappa = -np.log(slope)/dt
        sigma = np.sqrt(errvar * kappa / (1-np.exp(-2*kappa*dt)))
        mu = const / (1-np.exp(-kappa*dt)) + sigma**2/2./kappa
        if np.shape(mu)== (1,): mu = mu[0]   # how to remove scalar array ?
        if np.shape(sigma)== (1,): sigma = sigma[0]
        #mu, kappa are good, sigma too small
        return mu, kappa, sigma

class BrownianBridge(object):
    def __init__(self):
        pass

    def simulate(self, x0, x1, nobs, nrepl=1, ddt=1., sigma=1.):
        nobs=nobs+1
        dt = ddt*1./nobs
        t = np.linspace(dt, ddt-dt, nobs)
        t = np.linspace(dt, ddt, nobs)
        wm = [t/ddt, 1-t/ddt]
        #wmi = wm[1]
        #wm1 = x1*wm[0]
        wmi = 1-dt/(ddt-t)
        wm1 = x1*(dt/(ddt-t))
        su = sigma* np.sqrt(t*(1-t)/ddt)
        s = sigma* np.sqrt(dt*(ddt-t-dt)/(ddt-t))
        x = np.zeros((nrepl, nobs))
        x[:,0] = x0
        rvs = s*np.random.normal(size=(nrepl,nobs))
        for i in range(1,nobs):
            x[:,i] = x[:,i-1]*wmi[i] + wm1[i] + rvs[:,i]
        return x, t, su

class CompoundPoisson(object):
    '''nobs iid compound poisson distributions, not a process in time
    '''
    def __init__(self, lambd, randfn=np.random.normal):
        if len(lambd) != len(randfn):
            raise ValueError('lambd and randfn need to have the same number of elements')

        self.nobj = len(lambd)
        self.randfn = randfn
        self.lambd = np.asarray(lambd)

    def simulate(self, nobs, nrepl=1):
        nobj = self.nobj
        x = np.zeros((nrepl, nobs, nobj))
        N = np.random.poisson(self.lambd[None,None,:], size=(nrepl,nobs,nobj))
        for io in range(nobj):
            randfnc = self.randfn[io]

            nc = N[:,:,io]
            #print nrepl,nobs,nc
            #xio = randfnc(size=(nrepl,nobs,np.max(nc))).cumsum(-1)[np.arange(nrepl)[:,None],np.arange(nobs),nc-1]
            rvs = randfnc(size=(nrepl,nobs,np.max(nc)))
            print('rvs.sum()', rvs.sum(), rvs.shape)
            xio = rvs.cumsum(-1)[np.arange(nrepl)[:,None],np.arange(nobs),nc-1]
            #print xio.shape
            x[:,:,io] = xio
        x[N==0] = 0
        return x, N

'''
randn('state',100)                                % set the state of randn
T = 1; N = 500; dt = T/N; t = [dt:dt:1];

M = 1000;                                         % M paths simultaneously
dW = sqrt(dt)*randn(M,N);                         % increments
W = cumsum(dW,2);                                 % cumulative sum
U = exp(repmat(t,[M 1]) + 0.5*W);
Umean = mean(U);
plot([0,t],[1,Umean],'b-'), hold on               % plot mean over M paths
plot([0,t],[ones(5,1),U(1:5,:)],'r--'), hold off  % plot 5 individual paths
xlabel('t','FontSize',16)
ylabel('U(t)','FontSize',16,'Rotation',0,'HorizontalAlignment','right')
legend('mean of 1000 paths','5 individual paths',2)

averr = norm((Umean - exp(9*t/8)),'inf')          % sample error
'''

if __name__ == '__main__':
    doplot = 1
    nrepl = 1000
    examples = []#['all']

    if 'all' in examples:
        w = Diffusion()

        # Wiener Process
        # ^^^^^^^^^^^^^^

        ws = w.simulateW(1000, nrepl=nrepl)
        if doplot:
            plt.figure()
            tmp = plt.plot(ws[0].T)
            tmp = plt.plot(ws[0].mean(0), linewidth=2)
            plt.title('Standard Brownian Motion (Wiener Process)')

        func = lambda t, W: np.exp(t + 0.5*W)
        us = w.expectedsim(func, nobs=500, nrepl=nrepl)
        if doplot:
            plt.figure()
            tmp = plt.plot(us[0].T)
            tmp = plt.plot(us[1], linewidth=2)
            plt.title('Brownian Motion - exp')
        #plt.show()
        averr = np.linalg.norm(us[1] - np.exp(9*us[2]/8.), np.inf)
        print(averr)
        #print us[1][:10]
        #print np.exp(9.*us[2][:10]/8.)

        # Geometric Brownian
        # ^^^^^^^^^^^^^^^^^^

        gb = GeometricBrownian(xzero=1., mu=0.01, sigma=0.5)
        gbs = gb.simEM(nobs=100, nrepl=100)
        if doplot:
            plt.figure()
            tmp = plt.plot(gbs.T)
            tmp = plt.plot(gbs.mean(0), linewidth=2)
            plt.title('Geometric Brownian')
            plt.figure()
            tmp = plt.plot(np.log(gbs).T)
            tmp = plt.plot(np.log(gbs.mean(0)), linewidth=2)
            plt.title('Geometric Brownian - log-transformed')

        ab = ArithmeticBrownian(xzero=1, mu=0.05, sigma=1)
        abs = ab.simEM(nobs=100, nrepl=100)
        if doplot:
            plt.figure()
            tmp = plt.plot(abs.T)
            tmp = plt.plot(abs.mean(0), linewidth=2)
            plt.title('Arithmetic Brownian')

        # Ornstein-Uhlenbeck
        # ^^^^^^^^^^^^^^^^^^

        ou = OUprocess(xzero=2, mu=1, lambd=0.5, sigma=0.1)
        ous = ou.simEM()
        oue = ou.exact(1, 1, np.random.normal(size=(5,10)))
        ou.exact(0, np.linspace(0,10,10/0.1), 0)
        ou.exactprocess(0,10)
        print(ou.exactprocess(0,10, ddt=0.1,nrepl=10).mean(0))
        #the following looks good, approaches mu
        oues = ou.exactprocess(0,100, ddt=0.1,nrepl=100)
        if doplot:
            plt.figure()
            tmp = plt.plot(oues.T)
            tmp = plt.plot(oues.mean(0), linewidth=2)
            plt.title('Ornstein-Uhlenbeck')

        # SchwartsOne
        # ^^^^^^^^^^^

        so = SchwartzOne(xzero=0, mu=1, kappa=0.5, sigma=0.1)
        sos = so.exactprocess(0,50, ddt=0.1,nrepl=100)
        print(sos.mean(0))
        print(np.log(sos.mean(0)))
        doplot = 1
        if doplot:
            plt.figure()
            tmp = plt.plot(sos.T)
            tmp = plt.plot(sos.mean(0), linewidth=2)
            plt.title('Schwartz One')
        print(so.fitls(sos[0,:],dt=0.1))
        sos2 = so.exactprocess(0,500, ddt=0.1,nrepl=5)
        print('true: mu=1, kappa=0.5, sigma=0.1')
        for i in range(5):
            print(so.fitls(sos2[i],dt=0.1))

        # Brownian Bridge
        # ^^^^^^^^^^^^^^^

        bb = BrownianBridge()
        #bbs = bb.sample(x0, x1, nobs, nrepl=1, ddt=1., sigma=1.)
        bbs, t, wm = bb.simulate(0, 0.5, 99, nrepl=500, ddt=1., sigma=0.1)
        if doplot:
            plt.figure()
            tmp = plt.plot(bbs.T)
            tmp = plt.plot(bbs.mean(0), linewidth=2)
            plt.title('Brownian Bridge')
            plt.figure()
            plt.plot(wm,'r', label='theoretical')
            plt.plot(bbs.std(0), label='simulated')
            plt.title('Brownian Bridge - Variance')
            plt.legend()

    # Compound Poisson
    # ^^^^^^^^^^^^^^^^
    cp = CompoundPoisson([1,1], [np.random.normal,np.random.normal])
    cps = cp.simulate(nobs=20000,nrepl=3)
    print(cps[0].sum(-1).sum(-1))
    print(cps[0].sum())
    print(cps[0].mean(-1).mean(-1))
    print(cps[0].mean())
    print(cps[1].size)
    print(cps[1].sum())
    #Note Y = sum^{N} X is compound poisson of iid x, then
    #E(Y) = E(N)*E(X)   eg. eq. (6.37) page 385 in http://ee.stanford.edu/~gray/sp.html

    #plt.show()

reason: import numpy,from scipy


text:
# -*- coding: utf-8 -*-

"""Source code for simple environmental map used in multi-agent world.

Author: Yoshinari Motokawa <yoshinari.moto@fuji.waseda.jp>
"""
import numpy as np

from .abstract_map import AbstractMap

class FourRoomsMap(AbstractMap):
    def locate_walls(self):
        self.wall_matrix[np.array([0, self.SIZE_X - 1]), :] = 1
        self.wall_matrix[:, np.array([0, self.SIZE_Y - 1])] = 1

        self.wall_matrix[
            np.array([1, 2, 3, 7, 8, 9, 10, 13, 14, 15, 16, 20, 21, 22]), 10
        ] = 1
        self.wall_matrix[
            np.array([1, 2, 3, 7, 8, 9, 10, 13, 14, 15, 16, 20, 21, 22]), 13
        ] = 1
        self.wall_matrix[
            10, np.array([1, 2, 3, 7, 8, 9, 10, 13, 14, 15, 16, 20, 21, 22])
        ] = 1
        self.wall_matrix[
            13, np.array([1, 2, 3, 7, 8, 9, 10, 13, 14, 15, 16, 20, 21, 22])
        ] = 1

    def set_objects_area(self):
        if self.config.type_objects == 1:
            # Set objects area for object 0
            self.objects_area_matrix[0, 1: 10, 1: 10] = 1
            self.objects_area_matrix[0, 1: 10, 14: self.SIZE_Y - 1] = 1
            self.objects_area_matrix[0, 14: self.SIZE_X - 1, 1: 10] = 1
            self.objects_area_matrix[0, 14: self.SIZE_X - 1, 14: self.SIZE_Y - 1] = 1

        elif self.config.type_objects == 2:
            # Set objects area for object 0
            self.objects_area_matrix[0, 1: 10, 1: 10] = 1
            self.objects_area_matrix[0, 1: 10, 14: self.SIZE_Y - 1] = 1
            # Set objects area for object 1
            self.objects_area_matrix[1, 14: self.SIZE_X - 1, 1: 10] = 1
            self.objects_area_matrix[1, 14: self.SIZE_X - 1, 14: self.SIZE_Y - 1] = 1


text:
# Copyright (c) OpenMMLab. All rights reserved.
import mmcv
import numpy as np
from torch.nn.modules.utils import _pair

from registry import PIPELINES
from .utils import random_choose_unknown

@PIPELINES.register_module()
class Crop:
    """Crop data to specific size for training.

    Args:
        keys (Sequence[str]): The images to be cropped.
        crop_size (Tuple[int]): Target spatial size (h, w).
        random_crop (bool): If set to True, it will random crop
            image. Otherwise, it will work as center crop.
    """

    def __init__(self, keys, crop_size, random_crop=True):
        if not mmcv.is_tuple_of(crop_size, int):
            raise TypeError(
                'Elements of crop_size must be int and crop_size must be'
                f' tuple, but got {type(crop_size[0])} in {type(crop_size)}')

        self.keys = keys
        self.crop_size = crop_size
        self.random_crop = random_crop

    def _crop(self, data):
        if not isinstance(data, list):
            data_list = [data]
        else:
            data_list = data

        crop_bbox_list = []
        data_list_ = []

        for item in data_list:
            data_h, data_w = item.shape[:2]
            crop_h, crop_w = self.crop_size
            crop_h = min(data_h, crop_h)
            crop_w = min(data_w, crop_w)

            if self.random_crop:
                x_offset = np.random.randint(0, data_w - crop_w + 1)
                y_offset = np.random.randint(0, data_h - crop_h + 1)
            else:
                x_offset = max(0, (data_w - crop_w)) // 2
                y_offset = max(0, (data_h - crop_h)) // 2

            crop_bbox = [x_offset, y_offset, crop_w, crop_h]
            item_ = item[y_offset:y_offset + crop_h,
                         x_offset:x_offset + crop_w, ...]
            crop_bbox_list.append(crop_bbox)
            data_list_.append(item_)

        if not isinstance(data, list):
            return data_list_[0], crop_bbox_list[0]
        return data_list_, crop_bbox_list

    def __call__(self, results):
        """Call function.

        Args:
            results (dict): A dict containing the necessary information and
                data for augmentation.

        Returns:
            dict: A dict containing the processed data and information.
        """
        for k in self.keys:
            data_, crop_bbox = self._crop(results[k])
            results[k] = data_
            results[k + '_crop_bbox'] = crop_bbox
        results['crop_size'] = self.crop_size
        return results

    def __repr__(self):
        repr_str = self.__class__.__name__
        repr_str += (f'keys={self.keys}, crop_size={self.crop_size}, '
                     f'random_crop={self.random_crop}')

        return repr_str

@PIPELINES.register_module()
class CropSequence(Crop):
    """Crop a sequence to specific size for training.

    The main difference to 'Crop' is that the region to be cropped is the same
    for every images in the sequence.

    Args:
        keys (Sequence[str]): The images to be cropped.
        crop_size (Tuple[int]): Target spatial size (h, w).
        random_crop (bool): If set to True, it will random crop
            image. Otherwise, it will work as center crop.
    """

    def _crop(self, data):
        if not isinstance(data, list):
            raise TypeError(f'Input must be a list, but got {type(data)}.')

        # determine crop location. Must be the same for all images
        data_h, data_w = data[0].shape[:2]
        crop_h, crop_w = self.crop_size
        crop_h = min(data_h, crop_h)
        crop_w = min(data_w, crop_w)

        if self.random_crop:
            x_offset = np.random.randint(0, data_w - crop_w + 1)
            y_offset = np.random.randint(0, data_h - crop_h + 1)
        else:
            x_offset = max(0, (data_w - crop_w)) // 2
            y_offset = max(0, (data_h - crop_h)) // 2
        crop_bbox = [x_offset, y_offset, crop_w, crop_h]

        data_list = []
        for item in data:
            item = item[y_offset:y_offset + crop_h, x_offset:x_offset + crop_w,
                        ...]
            data_list.append(item)

        return data_list, crop_bbox

@PIPELINES.register_module()
class FixedCrop:
    """Crop paired data (at a specific position) to specific size for training.

    Args:
        keys (Sequence[str]): The images to be cropped.
        crop_size (Tuple[int]): Target spatial size (h, w).
        crop_pos (Tuple[int]): Specific position (x, y). If set to None,
            random initialize the position to crop paired data batch.
    """

    def __init__(self, keys, crop_size, crop_pos=None):
        if not mmcv.is_tuple_of(crop_size, int):
            raise TypeError(
                'Elements of crop_size must be int and crop_size must be'
                f' tuple, but got {type(crop_size[0])} in {type(crop_size)}')
        if not mmcv.is_tuple_of(crop_pos, int) and (crop_pos is not None):
            raise TypeError(
                'Elements of crop_pos must be int and crop_pos must be'
                f' tuple or None, but got {type(crop_pos[0])} in '
                f'{type(crop_pos)}')

        self.keys = keys
        self.crop_size = crop_size
        self.crop_pos = crop_pos

    def _crop(self, data, x_offset, y_offset, crop_w, crop_h):
        crop_bbox = [x_offset, y_offset, crop_w, crop_h]
        data_ = data[y_offset:y_offset + crop_h, x_offset:x_offset + crop_w,
                     ...]
        return data_, crop_bbox

    def __call__(self, results):
        """Call function.

        Args:
            results (dict): A dict containing the necessary information and
                data for augmentation.

        Returns:
            dict: A dict containing the processed data and information.
        """

        if isinstance(results[self.keys[0]], list):
            data_h, data_w = results[self.keys[0]][0].shape[:2]
        else:
            data_h, data_w = results[self.keys[0]].shape[:2]
        crop_h, crop_w = self.crop_size
        crop_h = min(data_h, crop_h)
        crop_w = min(data_w, crop_w)

        if self.crop_pos is None:
            x_offset = np.random.randint(0, data_w - crop_w + 1)
            y_offset = np.random.randint(0, data_h - crop_h + 1)
        else:
            x_offset, y_offset = self.crop_pos
            crop_w = min(data_w - x_offset, crop_w)
            crop_h = min(data_h - y_offset, crop_h)

        for k in self.keys:
            images = results[k]
            is_list = isinstance(images, list)
            if not is_list:
                images = [images]
            cropped_images = []
            crop_bbox = None
            for image in images:
                # In fixed crop for paired images, sizes should be the same
                if (image.shape[0] != data_h or image.shape[1] != data_w):
                    raise ValueError(
                        'The sizes of paired images should be the same. '
                        f'Expected ({data_h}, {data_w}), '
                        f'but got ({image.shape[0]}, '
                        f'{image.shape[1]}).')
                data_, crop_bbox = self._crop(image, x_offset, y_offset,
                                              crop_w, crop_h)
                cropped_images.append(data_)
            results[k + '_crop_bbox'] = crop_bbox
            if not is_list:
                cropped_images = cropped_images[0]
            results[k] = cropped_images
        results['crop_size'] = self.crop_size
        results['crop_pos'] = self.crop_pos
        return results

    def __repr__(self):
        repr_str = self.__class__.__name__
        repr_str += (f'keys={self.keys}, crop_size={self.crop_size}, '
                     f'crop_pos={self.crop_pos}')
        return repr_str

@PIPELINES.register_module()
class TripleRandomCrop:
    """Paried random crop.

    It crops a pair of lq and gt images with corresponding locations.
    It also supports accepting lq list and gt list.
    Required keys are "scale", "lq", and "gt",
    added or modified keys are "lq" and "gt".

    Args:
        gt_patch_size (int): cropped gt patch size.
    """

    def __init__(self, gt_patch_size):
        self.gt_patch_size = gt_patch_size

    def __call__(self, results):
        """Call function.

        Args:
            results (dict): A dict containing the necessary information and
                data for augmentation.

        Returns:
            dict: A dict containing the processed data and information.
        """
        scale = results['scale']
        lq_patch_size = self.gt_patch_size // scale

        lq_is_list = isinstance(results['lq'], list)
        if not lq_is_list:
            results['lq'] = [results['lq']]
        gt_is_list = isinstance(results['gt'], list)
        if not gt_is_list:
            results['gt'] = [results['gt']]

        h_lq, w_lq, _ = results['lq'][0].shape
        h_gt, w_gt, _ = results['gt'][0].shape

        if h_gt != h_lq * scale or w_gt != w_lq * scale:
            raise ValueError(
                f'Scale mismatches. GT ({h_gt}, {w_gt}) is not {scale}x ',
                f'multiplication of LQ ({h_lq}, {w_lq}).')
        if h_lq < lq_patch_size or w_lq < lq_patch_size:
            raise ValueError(
                f'LQ ({h_lq}, {w_lq}) is smaller than patch size ',
                f'({lq_patch_size}, {lq_patch_size}). Please check '
                f'{results["lq_path"][0]} and {results["gt_path"][0]}.')

        # randomly choose top and left coordinates for lq patch
        top = np.random.randint(h_lq - lq_patch_size + 1)
        left = np.random.randint(w_lq - lq_patch_size + 1)
        # crop lq patch
        results['lq'] = [
            v[top:top + lq_patch_size, left:left + lq_patch_size, ...]
            for v in results['lq']
        ]
        # crop corresponding gt patch
        top_gt, left_gt = int(top * scale), int(left * scale)
        results['gt'] = [
            v[top_gt:top_gt + self.gt_patch_size,
              left_gt:left_gt + self.gt_patch_size, ...] for v in results['gt']
        ]

        if not lq_is_list:
            results['lq'] = results['lq'][0]
        if not gt_is_list:
            results['gt'] = results['gt'][0]
        return results

    def __repr__(self):
        repr_str = self.__class__.__name__
        repr_str += f'(gt_patch_size={self.gt_patch_size})'
        return repr_str

@PIPELINES.register_module()
class PairedRandomCrop:
    """Paried random crop.

    It crops a pair of lq and gt images with corresponding locations.
    It also supports accepting lq list and gt list.
    Required keys are "scale", "lq", and "gt",
    added or modified keys are "lq" and "gt".

    Args:
        gt_patch_size (int): cropped gt patch size.
    """

    def __init__(self, gt_patch_size):
        self.gt_patch_size = gt_patch_size

    def __call__(self, results):
        """Call function.

        Args:
            results (dict): A dict containing the necessary information and
                data for augmentation.

        Returns:
            dict: A dict containing the processed data and information.
        """
        scale = results['scale']
        lq_patch_size = self.gt_patch_size // scale

        lq_is_list = isinstance(results['lq'], list)
        if not lq_is_list:
            results['lq'] = [results['lq']]
        gt_is_list = isinstance(results['gt'], list)
        if not gt_is_list:
            results['gt'] = [results['gt']]

        h_lq, w_lq, _ = results['lq'][0].shape
        h_gt, w_gt, _ = results['gt'][0].shape

        if h_gt != h_lq * scale or w_gt != w_lq * scale:
            raise ValueError(
                f'Scale mismatches. GT ({h_gt}, {w_gt}) is not {scale}x ',
                f'multiplication of LQ ({h_lq}, {w_lq}).')
        if h_lq < lq_patch_size or w_lq < lq_patch_size:
            raise ValueError(
                f'LQ ({h_lq}, {w_lq}) is smaller than patch size ',
                f'({lq_patch_size}, {lq_patch_size}). Please check '
                f'{results["lq_path"][0]} and {results["gt_path"][0]}.')

        # randomly choose top and left coordinates for lq patch
        top = np.random.randint(h_lq - lq_patch_size + 1)
        left = np.random.randint(w_lq - lq_patch_size + 1)
        # crop lq patch
        results['lq'] = [
            v[top:top + lq_patch_size, left:left + lq_patch_size, ...]
            for v in results['lq']
        ]
        # crop corresponding gt patch
        top_gt, left_gt = int(top * scale), int(left * scale)
        results['gt'] = [
            v[top_gt:top_gt + self.gt_patch_size,
              left_gt:left_gt + self.gt_patch_size, ...] for v in results['gt']
        ]

        if not lq_is_list:
            results['lq'] = results['lq'][0]
        if not gt_is_list:
            results['gt'] = results['gt'][0]
        return results

    def __repr__(self):
        repr_str = self.__class__.__name__
        repr_str += f'(gt_patch_size={self.gt_patch_size})'
        return repr_str

@PIPELINES.register_module()
class CropAroundCenter:
    """Randomly crop the images around unknown area in the center 1/4 images.

    This cropping strategy is adopted in GCA matting. The `unknown area` is the
    same as `semi-transparent area`.
    https://arxiv.org/pdf/2001.04069.pdf

    It retains the center 1/4 images and resizes the images to 'crop_size'.
    Required keys are "fg", "bg", "trimap" and "alpha", added or modified keys
    are "crop_bbox", "fg", "bg", "trimap" and "alpha".

    Args:
        crop_size (int | tuple): Desired output size. If int, square crop is
            applied.
    """

    def __init__(self, crop_size):
        if mmcv.is_tuple_of(crop_size, int):
            assert len(crop_size) == 2, 'length of crop_size must be 2.'
        elif not isinstance(crop_size, int):
            raise TypeError('crop_size must be int or a tuple of int, but got '
                            f'{type(crop_size)}')
        self.crop_size = _pair(crop_size)

    def __call__(self, results):
        """Call function.

        Args:
            results (dict): A dict containing the necessary information and
                data for augmentation.

        Returns:
            dict: A dict containing the processed data and information.
        """
        fg = results['fg']
        alpha = results['alpha']
        trimap = results['trimap']
        bg = results['bg']
        h, w = fg.shape[:2]
        assert bg.shape == fg.shape, (f'shape of bg {bg.shape} should be the '
                                      f'same as fg {fg.shape}.')

        crop_h, crop_w = self.crop_size
        # Make sure h >= crop_h, w >= crop_w. If not, rescale imgs
        rescale_ratio = max(crop_h / h, crop_w / w)
        if rescale_ratio > 1:
            new_h = max(int(h * rescale_ratio), crop_h)
            new_w = max(int(w * rescale_ratio), crop_w)
            fg = mmcv.imresize(fg, (new_w, new_h), interpolation='nearest')
            alpha = mmcv.imresize(
                alpha, (new_w, new_h), interpolation='nearest')
            trimap = mmcv.imresize(
                trimap, (new_w, new_h), interpolation='nearest')
            bg = mmcv.imresize(bg, (new_w, new_h), interpolation='bicubic')
            h, w = new_h, new_w

        # resize to 1/4 to ignore small unknown patches
        small_trimap = mmcv.imresize(
            trimap, (w // 4, h // 4), interpolation='nearest')
        # find unknown area in center 1/4 region
        margin_h, margin_w = crop_h // 2, crop_w // 2
        sample_area = small_trimap[margin_h // 4:(h - margin_h) // 4,
                                   margin_w // 4:(w - margin_w) // 4]
        unknown_xs, unknown_ys = np.where(sample_area == 128)
        unknown_num = len(unknown_xs)
        if unknown_num < 10:
            # too few unknown area in the center, crop from the whole image
            top = np.random.randint(0, h - crop_h + 1)
            left = np.random.randint(0, w - crop_w + 1)
        else:
            idx = np.random.randint(unknown_num)
            top = unknown_xs[idx] * 4
            left = unknown_ys[idx] * 4
        bottom = top + crop_h
        right = left + crop_w

        results['fg'] = fg[top:bottom, left:right]
        results['alpha'] = alpha[top:bottom, left:right]
        results['trimap'] = trimap[top:bottom, left:right]
        results['bg'] = bg[top:bottom, left:right]
        results['crop_bbox'] = (left, top, right, bottom)
        return results

    def __repr__(self):
        return self.__class__.__name__ + f'(crop_size={self.crop_size})'

@PIPELINES.register_module()
class CropAroundUnknown:
    """Crop around unknown area with a randomly selected scale.

    Randomly select the w and h from a list of (w, h).
    Required keys are the keys in argument `keys`, added or
    modified keys are "crop_bbox" and the keys in argument `keys`.
    This class assumes value of "alpha" ranges from 0 to 255.

    Args:
        keys (Sequence[str]): The images to be cropped. It must contain
            'alpha'. If unknown_source is set to 'trimap', then it must also
            contain 'trimap'.
        crop_sizes (list[int | tuple[int]]): List of (w, h) to be selected.
        unknown_source (str, optional): Unknown area to select from. It must be
            'alpha' or 'tirmap'. Default to 'alpha'.
        interpolations (str | list[str], optional): Interpolation method of
            mmcv.imresize. The interpolation operation will be applied when
            image size is smaller than the crop_size. If given as a list of
            str, it should have the same length as `keys`. Or if given as a
            str all the keys will be resized with the same method.
            Default to 'bilinear'.
    """

    def __init__(self,
                 keys,
                 crop_sizes,
                 unknown_source='alpha',
                 interpolations='bilinear'):
        if 'alpha' not in keys:
            raise ValueError(f'"alpha" must be in keys, but got {keys}')
        self.keys = keys

        if not isinstance(crop_sizes, list):
            raise TypeError(
                f'Crop sizes must be list, but got {type(crop_sizes)}.')
        self.crop_sizes = [_pair(crop_size) for crop_size in crop_sizes]
        if not mmcv.is_tuple_of(self.crop_sizes[0], int):
            raise TypeError('Elements of crop_sizes must be int or tuple of '
                            f'int, but got {type(self.crop_sizes[0][0])}.')

        if unknown_source not in ['alpha', 'trimap']:
            raise ValueError('unknown_source must be "alpha" or "trimap", '
                             f'but got {unknown_source}')
        if unknown_source not in keys:
            # it could only be trimap, since alpha is checked before
            raise ValueError(
                'if unknown_source is "trimap", it must also be set in keys')
        self.unknown_source = unknown_source

        if isinstance(interpolations, str):
            self.interpolations = [interpolations] * len(self.keys)
        elif mmcv.is_list_of(interpolations,
                             str) and len(interpolations) == len(self.keys):
            self.interpolations = interpolations
        else:
            raise TypeError(
                'interpolations must be a str or list of str with '
                f'the same length as keys, but got {interpolations}')

    def __call__(self, results):
        """Call function.

        Args:
            results (dict): A dict containing the necessary information and
                data for augmentation.

        Returns:
            dict: A dict containing the processed data and information.
        """
        h, w = results[self.keys[0]].shape[:2]

        rand_ind = np.random.randint(len(self.crop_sizes))
        crop_h, crop_w = self.crop_sizes[rand_ind]

        # Make sure h >= crop_h, w >= crop_w. If not, rescale imgs
        rescale_ratio = max(crop_h / h, crop_w / w)
        if rescale_ratio > 1:
            h = max(int(h * rescale_ratio), crop_h)
            w = max(int(w * rescale_ratio), crop_w)
            for key, interpolation in zip(self.keys, self.interpolations):
                results[key] = mmcv.imresize(
                    results[key], (w, h), interpolation=interpolation)

        # Select the cropping top-left point which is an unknown pixel
        if self.unknown_source == 'alpha':
            unknown = (results['alpha'] > 0) & (results['alpha'] < 255)
        else:
            unknown = results['trimap'] == 128
        top, left = random_choose_unknown(unknown.squeeze(), (crop_h, crop_w))

        bottom = top + crop_h
        right = left + crop_w

        for key in self.keys:
            results[key] = results[key][top:bottom, left:right]
        results['crop_bbox'] = (left, top, right, bottom)
        return results

    def __repr__(self):
        repr_str = self.__class__.__name__
        repr_str += (f'(keys={self.keys}, crop_sizes={self.crop_sizes}, '
                     f"unknown_source='{self.unknown_source}', "
                     f'interpolations={self.interpolations})')
        return repr_str

@PIPELINES.register_module()
class CropAroundFg:
    """Crop around the whole foreground in the segmentation mask.

    Required keys are "seg" and the keys in argument `keys`.
    Meanwhile, "seg" must be in argument `keys`. Added or modified keys are
    "crop_bbox" and the keys in argument `keys`.

    Args:
        keys (Sequence[str]): The images to be cropped. It must contain
            'seg'.
        bd_ratio_range (tuple, optional): The range of the boundary (bd) ratio
            to select from. The boundary ratio is the ratio of the boundary to
            the minimal bbox that contains the whole foreground given by
            segmentation. Default to (0.1, 0.4).
        test_mode (bool): Whether use test mode. In test mode, the tight crop
            area of foreground will be extended to the a square.
            Default to False.
    """

    def __init__(self, keys, bd_ratio_range=(0.1, 0.4), test_mode=False):
        if 'seg' not in keys:
            raise ValueError(f'"seg" must be in keys, but got {keys}')
        if (not mmcv.is_tuple_of(bd_ratio_range, float)
                or len(bd_ratio_range) != 2):
            raise TypeError('bd_ratio_range must be a tuple of 2 int, but got '
                            f'{bd_ratio_range}')
        self.keys = keys
        self.bd_ratio_range = bd_ratio_range
        self.test_mode = test_mode

    def __call__(self, results):
        """Call function.

        Args:
            results (dict): A dict containing the necessary information and
                data for augmentation.

        Returns:
            dict: A dict containing the processed data and information.
        """
        seg = results['seg']
        height, width = seg.shape[:2]

        # get foreground bbox
        fg_coor = np.array(np.where(seg))
        top, left = np.amin(fg_coor, axis=1)
        bottom, right = np.amax(fg_coor, axis=1)

        # enlarge bbox
        long_side = np.maximum(bottom - top, right - left)
        if self.test_mode:
            bottom = top + long_side
            right = left + long_side
        boundary_ratio = np.random.uniform(*self.bd_ratio_range)
        boundary = int(np.round(boundary_ratio * long_side))
        # NOTE: Different from the original repo, we keep track of the four
        # corners of the bbox (left, top, right, bottom) while the original
        # repo use (top, left, height, width) to represent bbox. This may
        # introduce an difference of 1 pixel.
        top = max(top - boundary, 0)
        left = max(left - boundary, 0)
        bottom = min(bottom + boundary, height)
        right = min(right + boundary, width)

        for key in self.keys:
            results[key] = results[key][top:bottom, left:right]
        results['crop_bbox'] = (left, top, right, bottom)
        return results

@PIPELINES.register_module()
class ModCrop:
    """Mod crop gt images, used during testing.

    Required keys are "scale" and "gt",
    added or modified keys are "gt".
    """

    def __call__(self, results):
        """Call function.

        Args:
            results (dict): A dict containing the necessary information and
                data for augmentation.

        Returns:
            dict: A dict containing the processed data and information.
        """
        img = results['gt'].copy()
        scale = results['scale']
        if img.ndim in [2, 3]:
            h, w = img.shape[0], img.shape[1]
            h_remainder, w_remainder = h % scale, w % scale
            img = img[:h - h_remainder, :w - w_remainder, ...]
        else:
            raise ValueError(f'Wrong img ndim: {img.ndim}.')
        results['gt'] = img
        return results

@PIPELINES.register_module()
class CropLike:
    """Crop/pad the image in the target_key according to the size of image
        in the reference_key .

    Args:
        target_key (str): The key needs to be cropped.
        reference_key (str | None): The reference key, need its size.
            Default: None.
    """

    def __init__(self, target_key, reference_key=None):

        assert reference_key and target_key
        self.target_key = target_key
        self.reference_key = reference_key

    def __call__(self, results):
        """Call function.

        Args:
            results (dict): A dict containing the necessary information and
                data for augmentation.
                Require self.target_key and self.reference_key.

        Returns:
            dict: A dict containing the processed data and information.
                Modify self.target_key.
        """
        size = results[self.reference_key].shape
        old_image = results[self.target_key]
        old_size = old_image.shape
        h, w = old_size[:2]
        new_size = size[:2] + old_size[2:]
        h_cover, w_cover = min(h, size[0]), min(w, size[1])

        format_image = np.zeros(new_size, dtype=old_image.dtype)
        format_image[:h_cover, :w_cover] = old_image[:h_cover, :w_cover]
        results[self.target_key] = format_image

        return results

    def __repr__(self):
        return (self.__class__.__name__ + f' target_key={self.target_key}, ' +
                f'reference_key={self.reference_key}')

max_stars_repo_name: Guominyingxiongququ/RAFT

max_issues_repo_name: Guominyingxiongququ/RAFT

max_forks_repo_name: Guominyingxiongququ/RAFT