[python]numba安装后测试代码

from numba import jit
import numpy as np
from scipy.spatial import KDTree
import time
from tqdm import tqdm

# 计算iou
@jit(nopython=True)
def bbox_iou(box1, box2):
    x1, y1, x2, y2 = box1
    x1_, y1_, x2_, y2_ = box2

    inter_x1 = max(x1, x1_)
    inter_y1 = max(y1, y1_)
    inter_x2 = min(x2, x2_)
    inter_y2 = min(y2, y2_)
    inter_w = max(0, inter_x2 - inter_x1)
    inter_h = max(0, inter_y2 - inter_y1)

    union = (x2 - x1) * (y2 - y1) + (x2_ - x1_) * (y2_ - y1_) - inter_w * inter_h
    iou = inter_w * inter_h / union
    return iou

# 定义两组bbox
bbox1 = np.random.rand(1000, 4)
bbox2 = np.random.rand(500, 4)

# 利用for循环计算匹配
start_time = time.time()
matches = []
for i in tqdm(range(bbox1.shape[0])):
    max_iou = 0
    match_idx = -1
    for j in range(bbox2.shape[0]):
        iou = bbox_iou(bbox1[i], bbox2[j])
        if iou > max_iou:
            max_iou = iou
            match_idx = j
    matches.append(match_idx)
print("Matching with for loop takes: %.3f s" % (time.time() - start_time))
@jit(nopython=True)   # 
def match_loop(bbox1,bbox2):
    matches = []
    for i in range(bbox1.shape[0]):
        max_iou = 0
        match_idx = -1
        for j in range(bbox2.shape[0]):
            iou = bbox_iou(bbox1[i], bbox2[j])
            if iou > max_iou:
                max_iou = iou
                match_idx = j
        matches.append(match_idx)

start_time = time.time()
match_loop(bbox1,bbox2)
print("Matching with for loop takes: %.3f s" % (time.time() - start_time))