Atss generate_anchors python+c++写法

19 篇文章 1 订阅
14 篇文章 1 订阅

import torch
def gen_single_level_base_anchors(base_size,
                                  scales,
                                  ratios,
                                  center=None):
    """Generate base anchors of a single level.

    Args:
        base_size (int | float): Basic size of an anchor.
        scales (torch.Tensor): Scales of the anchor.
        ratios (torch.Tensor): The ratio between between the height
            and width of anchors in a single level.
        center (tuple[float], optional): The center of the base anchor
            related to a single feature grid. Defaults to None.

    Returns:
        torch.Tensor: Anchors in a single-level feature maps.
    """
    # import pdb;pdb.set_trace()
    w = base_size
    h = base_size
    if center is None:
        x_center = 0 * w
        y_center = 0 * h
    else:
        x_center, y_center = center

    h_ratios = torch.sqrt(ratios)
    w_ratios = 1 / h_ratios
    # if self.scale_major:
    ws = (w * w_ratios[:, None] * scales[None, :]).view(-1)
    hs = (h * h_ratios[:, None] * scales[None, :]).view(-1)
    # else:
    #     ws = (w * scales[:, None] * w_ratios[None, :]).view(-1)
    #     hs = (h * scales[:, None] * h_ratios[None, :]).view(-1)

    # use float anchor and the anchor's center is aligned with the
    # pixel center
    base_anchors = [
        x_center - 0.5 * ws, y_center - 0.5 * hs, x_center + 0.5 * ws,
        y_center + 0.5 * hs
    ]
    base_anchors = torch.stack(base_anchors, dim=-1)
    # tensor([[-8., -8.,  8.,  8.]])
    # print('==base_anchors:', base_anchors)
    return base_anchors

def gen_base_anchors(scales, ratios, base_sizes):
    multi_level_base_anchors = []
    for i, base_size in enumerate(base_sizes):
        multi_level_base_anchors.append(
        gen_single_level_base_anchors(
            base_size,
            scales=scales, #tensor([8.])
            ratios=ratios,#tensor([1.])
            center=None))
    # import pdb;pdb.set_trace()#[tensor([[-32., -32.,  32.,  32.]]),
    # tensor([[-64., -64.,  64.,  64.]]), tensor([[-128., -128.,  128.,  128.]]),
    # tensor([[-256., -256.,  256.,  256.]]), tensor([[-512., -512.,  512.,  512.]])]
    return multi_level_base_anchors

def _meshgrid(x, y, row_major=True):
    """Generate mesh grid of x and y.

    Args:
        x (torch.Tensor): Grids of x dimension.
        y (torch.Tensor): Grids of y dimension.
        row_major (bool, optional): Whether to return y grids first.
            Defaults to True.

    Returns:
        tuple[torch.Tensor]: The mesh grids of x and y.
    """
    xx = x.repeat(len(y))
    yy = y.view(-1, 1).repeat(1, len(x)).view(-1)
    if row_major:
        return xx, yy
    else:
        return yy, xx

def single_level_grid_anchors(
                              base_anchors,
                              featmap_size,
                              stride=(16, 16),
                              device='cpu'):
    """Generate grid anchors of a single level.

    Note:
        This function is usually called by method ``self.grid_anchors``.

    Args:
        base_anchors (torch.Tensor): The base anchors of a feature grid.
        featmap_size (tuple[int]): Size of the feature maps.
        stride (tuple[int], optional): Stride of the feature map.
            Defaults to (16, 16).
        device (str, optional): Device the tensor will be put on.
            Defaults to 'cuda'.

    Returns:
        torch.Tensor: Anchors in the overall feature maps.
    """
    feat_h, feat_w = featmap_size
    # convert Tensor to int, so that we can covert to ONNX correctlly
    feat_h = int(feat_h)
    feat_w = int(feat_w)
    shift_x = torch.arange(0, feat_w, device=device) * stride[0]
    shift_y = torch.arange(0, feat_h, device=device) * stride[1]

    shift_xx, shift_yy = _meshgrid(shift_x, shift_y)
    shifts = torch.stack([shift_xx, shift_yy, shift_xx, shift_yy], dim=-1)
    shifts = shifts.type_as(base_anchors)
    # first feat_w elements correspond to the first row of shifts
    # add A anchors (1, A, 4) to K shifts (K, 1, 4) to get
    # shifted anchors (K, A, 4), reshape to (K*A, 4)

    all_anchors = base_anchors[None, :, :] + shifts[:, None, :]
    all_anchors = all_anchors.view(-1, 4)
    # first A rows correspond to A anchors of (0, 0) in feature map,
    # then (0, 1), (0, 2), ...
    return all_anchors

def grid_anchors(featmap_sizes, base_anchors, strides, device='cuda'):
    """Generate grid anchors in multiple feature levels.

    Args:
        featmap_sizes (list[tuple]): List of feature map sizes in
            multiple feature levels.
        device (str): Device where the anchors will be put on.

    Return:
        list[torch.Tensor]: Anchors in multiple feature levels. \
            The sizes of each tensor should be [N, 4], where \
            N = width * height * num_base_anchors, width and height \
            are the sizes of the corresponding feature lavel, \
            num_base_anchors is the number of anchors for that level.
    """

    multi_level_anchors = []
    num_levels = 5
    for i in range(num_levels):
        anchors = single_level_grid_anchors(
            base_anchors[i],
            featmap_sizes[i],
            strides[i],
            device=device)
        # import pdb;pdb.set_trace()
        multi_level_anchors.append(anchors)

    return multi_level_anchors

scales = torch.tensor([8.])
ratios = torch.tensor([1.])
base_sizes = [8, 16, 32, 64, 128]
base_anchors = gen_base_anchors(scales, ratios, base_sizes)
print('===base_anchors:', base_anchors)
strides = [(8, 8), (16, 16), (32, 32), (64, 64), (128, 128)]
#((h,w),..)
featmap_sizes = [(80, 48), (40, 24), (20, 12), (10, 6), (5, 3)]
multi_level_anchors = grid_anchors(featmap_sizes, base_anchors, strides, device='cuda')
for multi_level_anchor in multi_level_anchors:
    print('==multi_level_anchor.shape:', multi_level_anchor.shape)
    print('=multi_level_anchor:\n', multi_level_anchor)

c++写法:

//
// Created by fzh on 2021/6/22.
//

//#include <pybind11/numpy.h>
//#include <a.h>
//#include <b.h>
#include <iostream>
#include <opencv2/opencv.hpp>
#include <opencv2/core/core.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <string>
#include <algorithm>
#include <vector>
#include <memory>

typedef cv::Vec<double, 5> Vec5d;
using namespace std;

typedef struct MRect2f {
    float left;   ///< X-coordinate of the upper-left corner.
    float top;    ///< Y-coordinate of the upper-left corner.
    float right;  ///< X-coordinate of the bottom-right corner.
    float bottom; ///< Y-coordinate of the bottom-right corner.
    friend std::ostream& operator<<(std::ostream& out, const MRect2f& t) {
        out << "l:" << t.left << ", t:" << t.top
            << ", r:" << t.right << ", b:" << t.bottom;
        return out;
    }
} MRect2f;

static std::vector<std::vector<MRect2f>> get_anchors(
        std::vector<int>& featmap_heights,
        std::vector<int> featmap_widths,
        std::vector<int> strides,
        const std::vector<float>& anchor_base,
        const std::vector<float>& anchor_scales,
        const std::vector<float>& anchor_ratios) {
    size_t stage_num = strides.size();
    std::vector<std::vector<MRect2f>> anchors(stage_num);

    for (int r = 0; r < anchor_ratios.size(); r++) {
        float h_ratios = std::sqrt(anchor_ratios[r]);
        float w_ratios = 1 / h_ratios;

        for (size_t stage = 0; stage < strides.size(); stage++) {
            int h = featmap_heights[stage];
            int w = featmap_widths[stage];
            int stride = strides[stage];
            int scale = anchor_scales[0];
            int base_size = anchor_base[stage];
            float x_ctr = 0;
            float y_ctr = 0;
            float ws = base_size * w_ratios * scale;
            float hs = base_size * h_ratios * scale;
            float base_left = x_ctr - 0.5f * (ws);
            float base_top = y_ctr - 0.5f * (hs);
            float base_right = x_ctr + 0.5f * (ws);
            float base_bottom = y_ctr + 0.5f * (hs);
            for (int i = 0; i < h; i++) {
                for (int j = 0; j < w; j++) {
                    MRect2f rect = {.left = base_left + float(j * stride),
                            .top = base_top + float(i * stride),
                            .right = base_right + float(j * stride),
                            .bottom = base_bottom + float(i * stride)};

                    anchors[stage].push_back(rect);
                }
            }
        }
    }
    return anchors;
}

int main(int argc, char** argv)
{

    vector<int> strides_ = {8, 16, 32, 64, 128};

    vector<int> featmap_heights_ = {80, 40, 20, 10, 5};
    vector<int> featmap_widths_ = {48, 24, 12, 6, 3};
    vector<float> anchor_base_ = {8, 16, 32, 64, 128};
    vector<float> anchor_scales_ = {8};
    vector<float> anchor_ratios_ = {1.0};
    std::vector<std::vector<MRect2f>> mlvl_anchors_ = get_anchors(featmap_heights_, featmap_widths_, strides_,
                anchor_base_, anchor_scales_, anchor_ratios_);
    cout<<mlvl_anchors_.size()<<endl;
    for(int i=0; i<mlvl_anchors_.size(); i++){
        for(int j=0; j < mlvl_anchors_[i].size(); j++){
            cout<<mlvl_anchors_[i][j];
        }
        cout<<endl;
    }
}

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值