99 lines
2.7 KiB
Python
99 lines
2.7 KiB
Python
import sys
|
|
import os
|
|
import time
|
|
import math
|
|
import torch
|
|
import numpy as np
|
|
from torch.autograd import Variable
|
|
|
|
|
|
def bbox_ious(boxes1, boxes2, x1y1x2y2=True):
|
|
if x1y1x2y2:
|
|
mx = torch.min(boxes1[0], boxes2[0])
|
|
Mx = torch.max(boxes1[2], boxes2[2])
|
|
my = torch.min(boxes1[1], boxes2[1])
|
|
My = torch.max(boxes1[3], boxes2[3])
|
|
w1 = boxes1[2] - boxes1[0]
|
|
h1 = boxes1[3] - boxes1[1]
|
|
w2 = boxes2[2] - boxes2[0]
|
|
h2 = boxes2[3] - boxes2[1]
|
|
else:
|
|
mx = torch.min(boxes1[0] - boxes1[2] / 2.0, boxes2[0] - boxes2[2] / 2.0)
|
|
Mx = torch.max(boxes1[0] + boxes1[2] / 2.0, boxes2[0] + boxes2[2] / 2.0)
|
|
my = torch.min(boxes1[1] - boxes1[3] / 2.0, boxes2[1] - boxes2[3] / 2.0)
|
|
My = torch.max(boxes1[1] + boxes1[3] / 2.0, boxes2[1] + boxes2[3] / 2.0)
|
|
w1 = boxes1[2]
|
|
h1 = boxes1[3]
|
|
w2 = boxes2[2]
|
|
h2 = boxes2[3]
|
|
uw = Mx - mx
|
|
uh = My - my
|
|
cw = w1 + w2 - uw
|
|
ch = h1 + h2 - uh
|
|
mask = ((cw <= 0) + (ch <= 0) > 0)
|
|
area1 = w1 * h1
|
|
area2 = w2 * h2
|
|
carea = cw * ch
|
|
carea[mask] = 0
|
|
uarea = area1 + area2 - carea
|
|
return carea / uarea
|
|
|
|
|
|
def get_region_boxes(boxes_and_confs):
|
|
|
|
# print('Getting boxes from boxes and confs ...')
|
|
|
|
boxes_list = []
|
|
confs_list = []
|
|
|
|
for item in boxes_and_confs:
|
|
boxes_list.append(item[0])
|
|
confs_list.append(item[1])
|
|
|
|
# boxes: [batch, num1 + num2 + num3, 1, 4]
|
|
# confs: [batch, num1 + num2 + num3, num_classes]
|
|
boxes = torch.cat(boxes_list, dim=1)
|
|
confs = torch.cat(confs_list, dim=1)
|
|
|
|
return [boxes, confs]
|
|
|
|
|
|
def convert2cpu(gpu_matrix):
|
|
return torch.FloatTensor(gpu_matrix.size()).copy_(gpu_matrix)
|
|
|
|
|
|
def convert2cpu_long(gpu_matrix):
|
|
return torch.LongTensor(gpu_matrix.size()).copy_(gpu_matrix)
|
|
|
|
|
|
|
|
def do_detect(model, img, conf_thresh, nms_thresh, use_cuda=1):
|
|
model.eval()
|
|
t0 = time.time()
|
|
|
|
if type(img) == np.ndarray and len(img.shape) == 3: # cv2 image
|
|
img = torch.from_numpy(img.transpose(2, 0, 1)).float().div(255.0).unsqueeze(0)
|
|
elif type(img) == np.ndarray and len(img.shape) == 4:
|
|
img = torch.from_numpy(img.transpose(0, 3, 1, 2)).float().div(255.0)
|
|
else:
|
|
print("unknow image type")
|
|
exit(-1)
|
|
|
|
if use_cuda:
|
|
img = img.cuda()
|
|
img = torch.autograd.Variable(img)
|
|
|
|
t1 = time.time()
|
|
|
|
output = model(img)
|
|
|
|
t2 = time.time()
|
|
|
|
print('-----------------------------------')
|
|
print(' Preprocess : %f' % (t1 - t0))
|
|
print(' Model Inference : %f' % (t2 - t1))
|
|
print('-----------------------------------')
|
|
|
|
return utils.post_processing(img, conf_thresh, nms_thresh, output)
|
|
|