23 lines
728 B
Python
23 lines
728 B
Python
# coding: utf-8
|
|
import torch
|
|
import numpy as np
|
|
|
|
def make_grid(nx, ny):
|
|
yv, xv = torch.meshgrid([torch.arange(ny), torch.arange(nx)])
|
|
grids = torch.stack((xv, yv), 2).view((1, 1, ny, nx, 2)).float()
|
|
return grids
|
|
|
|
def Yolov5_generate_npy(input_h, input_w) :
|
|
name = ['80','40','20']
|
|
ny, nx = int(input_h // pow(2,3)), int(input_w // pow(2,3))
|
|
for i in range(3):
|
|
ny_i, nx_i = int(ny // pow(2,i)), int(nx // pow(2,i))
|
|
grid = make_grid(nx_i, ny_i)
|
|
grid = grid.numpy()
|
|
file_name = name[i]+'_'+str(input_w)+'x'+str(input_h)+'.npy'
|
|
np.save(file_name,grid)
|
|
|
|
if __name__ == '__main__':
|
|
input_h, input_w = 256, 480
|
|
Yolov5_generate_npy(input_h, input_w)
|