import math
import megengine as mge
import megengine.functional as F
import numpy as np
from megengine import Tensor
import pdb

def restore_bbox(rois, deltas, unnormalize=True, config = None):

    assert deltas.ndim == 3
    if unnormalize:
        std_opr = mge.tensor(config.bbox_normalize_stds.reshape(1, 1, -1))
        mean_opr = mge.tensor(config.bbox_normalize_means.reshape(1, 1, -1))
        deltas = deltas * std_opr
        deltas = deltas + mean_opr

    # n = deltas.shape[1]
    n, c = deltas.shape[0], deltas.shape[1]
    all_rois = F.broadcast_to(F.expand_dims(rois, 1), (n, c, rois.shape[1])).reshape(-1, rois.shape[1])
    deltas = deltas.reshape(-1, deltas.shape[2])
    pred_bbox = bbox_transform_inv_opr(all_rois, deltas)
    pred_bbox = pred_bbox.reshape(-1, c, pred_bbox.shape[1])
    return pred_bbox

def filter_boxes_opr(boxes, min_size):

    """Remove all boxes with any side smaller than min_size."""
    wh = boxes[:, 2:4] - boxes[:, 0:2] + 1
    keep_mask = 
F.prod(wh >= min_size, axis = 