mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-04-29 03:00:14 -04:00
24 lines
979 B
Python
24 lines
979 B
Python
import numpy as np
|
|
from tinygrad import Tensor
|
|
|
|
def nms(boxes, scores, thresh=0.5):
|
|
x1, y1, x2, y2 = np.rollaxis(boxes, 1)
|
|
areas = (x2 - x1 + 1) * (y2 - y1 + 1)
|
|
to_process, keep = scores.argsort()[::-1], []
|
|
while to_process.size > 0:
|
|
cur, to_process = to_process[0], to_process[1:]
|
|
keep.append(cur)
|
|
inter_x1 = np.maximum(x1[cur], x1[to_process])
|
|
inter_y1 = np.maximum(y1[cur], y1[to_process])
|
|
inter_x2 = np.minimum(x2[cur], x2[to_process])
|
|
inter_y2 = np.minimum(y2[cur], y2[to_process])
|
|
inter_area = np.maximum(0, inter_x2 - inter_x1 + 1) * np.maximum(0, inter_y2 - inter_y1 + 1)
|
|
iou = inter_area / (areas[cur] + areas[to_process] - inter_area)
|
|
to_process = to_process[np.where(iou <= thresh)[0]]
|
|
return keep
|
|
|
|
def meshgrid(x, y):
|
|
grid_x = Tensor.cat(*[x[idx:idx+1].expand(y.shape).unsqueeze(0) for idx in range(x.shape[0])])
|
|
grid_y = Tensor.cat(*[y.unsqueeze(0)]*x.shape[0])
|
|
return grid_x.reshape(-1, 1), grid_y.reshape(-1, 1)
|