本文代码转自如下链接：
https://zhuanlan.zhihu.com/p/85194783
该代码剥离原来的官方代码，更清晰。同时，数据预处理部分是先固定一边为512，然后另外一边根据比例缩放，然后再填充为512，很值得看下python对这块的操作！
然后后续为了还原到原来的坐标对上述操作再逆操作还原到原图坐标！更值得研究细看！
还有batch操作。
用到的一些技巧函数：
cv2.copyMakeBorder
loaded_iminfos = list(map(prepare_img, image_frames)) ##prepare_img是函数
batch构建：im_batches = [torch.cat(img_tensors[i * opt.batch_size: np.min([(i + 1) * opt.batch_size, len(img_tensors)])]) for i in range(num_batches)]
idx.float() obj_info[3:5].int()
ratios_select = torch.index_select(torch.Tensor(ratios), 0, output[:, 0].long())
output[:, [1, 3]] -= lefts_select.unsqueeze(1)
加了中间变量信息的注释代码如下：
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os
import cv2
import numpy as np
import torch
import time
image_ext = ['jpg', 'jpeg', 'png']
import sys
CENTERNET_PATH = './src/'
sys.path.insert(0, CENTERNET_PATH)
CENTERNET_PATH = './src/lib/'
sys.path.insert(0, CENTERNET_PATH)
from models.model import create_model, load_model
from models.decode import ctdet_decode
class DefaultConfig(object):
def __init__(self):
# self.demo = '../images/17790319373_bd19b24cfc_k.jpg'
self.demo = '/data_2/2019biaozhushuju/0000-2020/20200514_refdet_small_test_data/sample'
self.reso = 512
#self.mean = np.array([0.40789654, 0.44719302, 0.47026115],
# dtype=np.float32).reshape(1, 1, 3)
#self.std = np.array([0.28863828, 0.27408164, 0.27809835],
# dtype=np.float32).reshape(1, 1, 3)
self.mean = np.array([0.408, 0.447, 0.47],
dtype=np.float32).reshape(1, 1, 3)
self.std = np.array([0.289, 0.274, 0.278],
dtype=np.float32).reshape(1, 1, 3)
self.batch_size = 4
self.use_gpu = torch.cuda.is_available()
self.device = torch.device("cuda" if self.use_gpu else "cpu")
if self.use_gpu:
torch.set_default_tensor_type('torch.cuda.FloatTensor')
else:
torch.set_default_tensor_type('torch.FloatTensor')
self.arch = 'dla_34'
self.heads = {'hm': 24, 'wh': 2, 'reg': 2}
self.head_conv = 256
self.load_model = '/data_2/project_202009/pytorch_project/CenterNet/centernet_src_202010/CenterNet-master/myfile/model_last-3epoch-multi.pth'
self.K = 100
self.conf_thresh = 0.3
self.downsample_rate = 4
def parse(self, kwargs):
'''
根据字典kwargs 更新 config参数
'''
for k, v in kwargs.items():
if not hasattr(self, k):
assert ValueError
setattr(self, k, v)
opt = DefaultConfig()
# load model
print('Creating model...')
model = create_model(opt.arch, opt.heads, opt.head_conv)
model = load_model(model, opt.load_model)
model = model.to(opt.device)
model.eval()
# load imgdir
if os.path.isdir(opt.demo):
image_names = []
ls = os.listdir(opt.demo)
for file_name in sorted(ls):
ext = file_name[file_name.rfind('.') + 1:].lower()
image_names.append(os.path.join(opt.demo, file_name))
#if ext in image_ext:
# image_names.append(os.path.join(opt.demo, file_name))
else:
image_names = [opt.demo]
image_frames = [cv2.imread(image_name) for image_name in image_names]
def prepare_img(frame, color=(0, 0, 0)):
h, w, c = frame.shape
## h = 1152 w = 885 c= 3
ratio = np.min([opt.reso / h, opt.reso / w]) #0.44444
new_h = int(h * ratio) #512
new_w = int(w * ratio) #393
if (w, h) != (new_w, new_h):
frame = cv2.resize(frame, (new_w, new_h), interpolation=cv2.INTER_CUBIC) # w,h
# frame = cv2.resize(frame, (new_w, new_h), interpolation=cv2.INTER_LINEAR) # w,h
# frame = cv2.resize(frame, (new_w, new_h)) # w,h
#59.5
dw = (opt.reso - new_w) / 2 # width padding
#0
dh = (opt.reso - new_h) / 2 # height padding
## 0 0
top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
left, right = int(round(dw - 0.1)), int(round(dw + 0.1)) ##59 60
frame = cv2.copyMakeBorder(frame, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) # add border
#[512,512,3]
frame = ((frame / 255. - opt.mean) / opt.std).astype(np.float32)
frame = frame.transpose(2, 0, 1)
return torch.Tensor(frame).unsqueeze(0), ratio, left, top #[1,3,512,512]
loaded_iminfos = list(map(prepare_img, image_frames))
img_tensors = [im_info[0] for im_info in loaded_iminfos]
ratios = [im_info[1] for im_info in loaded_iminfos]
lefts = [im_info[2] for im_info in loaded_iminfos]
tops = [im_info[3] for im_info in loaded_iminfos]
leftover = 0 if len(img_tensors) % opt.batch_size == 0 else 1 # len(img_tensors) = 85 batch_size=4 1
num_batches = len(img_tensors) // opt.batch_size + leftover #22
## [4,3,512,512]*21 + [1,3,512,512]*1
im_batches = [torch.cat(
img_tensors[i * opt.batch_size: np.min([(i + 1) * opt.batch_size, len(img_tensors)])]) for i in
range(num_batches)]
output = []
for i, batch in enumerate(im_batches):
batch = batch.to(opt.device)
batch_size = batch.shape[0]
pre_process_time = time.time()
with torch.no_grad():
mode_output = model(batch)[-1]
hm = mode_output['hm'].sigmoid_() #[4,24,128,128]
wh = mode_output['wh'] #[4,2,128,128]
reg = mode_output['reg']#[4,2,128,128]
# torch.cuda.synchronize()
# forward_time = time.time()
#[4,100,6]
prediction = ctdet_decode(hm, wh, reg=reg, cat_spec_wh=False, K=opt.K) # batch_size K=100 box+scores+clses
# torch.cuda.synchronize()
decode_time = time.time()
# print('net: {}'.format((forward_time-pre_process_time)/batch_size))
# print('dec: {}'.format((decode_time-forward_time)/batch_size))
print('dec: {}'.format((decode_time-pre_process_time)/batch_size))
#[400,6]
prediction = prediction.reshape(-1, prediction.shape[2]) # N*K 6
# aa = torch.arange(batch_size) #[0,1,2,3]
# bb = torch.arange(batch_size).unsqueeze(1)
# cc = torch.arange(batch_size).unsqueeze(1).repeat(1, opt.K) #0*100 1*100 2*100 3*100
#[4] [4,1] [4,100] [400,1]
idx = torch.arange(batch_size).unsqueeze(1).repeat(1, opt.K).view(-1, 1)
# x = prediction.new(prediction.shape[0], 1).fill_(i)
prediction = torch.cat((idx.float(), prediction), 1) #[400,7]
prediction[:, 0] += i * opt.batch_size # 第0为是图片的索引 一张图片有100个一样的索引
output.append(prediction)
# _ * idx x y x y conf class
try:#before: list(22个) [400,7]*21 + [100,7]*1
output = torch.cat(output, 0) #after [8500,7]
except:
output = torch.empty((0, 7), dtype=torch.get_default_dtype())
a1 = torch.Tensor(ratios) #[85]
a2 = output[:, 0] #[8500]
ratios_select = torch.index_select(torch.Tensor(ratios), 0, output[:, 0].long()) #[8500]
lefts_select = torch.index_select(torch.Tensor(lefts), 0, output[:, 0].long()) #[8500]
tops_select = torch.index_select(torch.Tensor(tops), 0, output[:, 0].long()) #[8500]
output[:, 1:5] *= opt.downsample_rate
a3 = output[:, [1, 3]] #[8500,2]
a4 = lefts_select.unsqueeze(1) #[8500,1]
# a5 = torch.Tensor([[2,3]])
# a6 = torch.Tensor([[5]])
# print(a5)
# print(a6)
# a7 = a5 - a6
## 增广
output[:, [1, 3]] -= lefts_select.unsqueeze(1)
output[:, [2, 4]] -= tops_select.unsqueeze(1)
output[:, 1:5] /= ratios_select.unsqueeze(1)
frames_shape = [im.shape[:2] for im in image_frames]
frames_shape_select = torch.index_select(torch.Tensor(frames_shape), 0, output[:, 0].long())
for i in range(output.shape[0]):
output[i, [3, 5]] = torch.clamp(output[i, [3, 5]], 0.0, frames_shape_select[i, 1])
output[i, [4, 6]] = torch.clamp(output[i, [4, 6]], 0.0, frames_shape_select[i, 0])
# print(output)
def write(obj_info, loaded_imgs):
obj_img = loaded_imgs[int(obj_info[0])]
if obj_info[5] < opt.conf_thresh:
return
# obj_class = int(obj_info[1]) idx x y x y conf class
label = '{:.2f}_{}'.format(obj_info[5], obj_info[6])
x1y1 = tuple(obj_info[1:3].int())
x2y2 = tuple(obj_info[3:5].int())
color = (255, 0, 0)
cv2.rectangle(obj_img, x1y1, x2y2, color, thickness=1)
t_size = cv2.getTextSize(label, fontFace=cv2.FONT_HERSHEY_PLAIN, fontScale=1, thickness=1)[0]
x2y2 = (x1y1[0]+t_size[0]+3, x1y1[1]+t_size[1]+4)
cv2.rectangle(obj_img, x1y1, x2y2, color, thickness=-1)
cv2.putText(obj_img, label, org=(x1y1[0], x1y1[1]+t_size[1] + 4),
fontFace=cv2.FONT_HERSHEY_PLAIN, fontScale=1, color=[255,255,255], thickness=1)
list(map(lambda x: write(x, image_frames), output))
os.makedirs('mubai', exist_ok=True)
det_names = ['mubai/det_{}'.format(impath.split('/')[-1]) for impath in image_names]
list(map(cv2.imwrite, det_names, image_frames))
我数据预处理部分直接resize，然后再映射到原图，下面是我写的代码，但是好像哪里有问题啊！！！框不对，我仔细看了一下还是不知道哪里有问题。。。呃呃呃！！
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os
import cv2
import numpy as np
import torch
import time
image_ext = ['jpg', 'jpeg', 'png']
import sys
CENTERNET_PATH = './src/'
sys.path.insert(0, CENTERNET_PATH)
CENTERNET_PATH = './src/lib/'
sys.path.insert(0, CENTERNET_PATH)
from models.model import create_model, load_model
from models.decode import ctdet_decode
class DefaultConfig(object):
def __init__(self):
# self.demo = '../images/17790319373_bd19b24cfc_k.jpg'
self.demo = '/data_2/2019biaozhushuju/0000-2020/20200514_refdet_small_test_data/sample'
self.reso = 512
#self.mean = np.array([0.40789654, 0.44719302, 0.47026115],
# dtype=np.float32).reshape(1, 1, 3)
#self.std = np.array([0.28863828, 0.27408164, 0.27809835],
# dtype=np.float32).reshape(1, 1, 3)
self.mean = np.array([0.408, 0.447, 0.47],
dtype=np.float32).reshape(1, 1, 3)
self.std = np.array([0.289, 0.274, 0.278],
dtype=np.float32).reshape(1, 1, 3)
self.batch_size = 4
self.use_gpu = torch.cuda.is_available()
self.device = torch.device("cuda" if self.use_gpu else "cpu")
if self.use_gpu:
torch.set_default_tensor_type('torch.cuda.FloatTensor')
else:
torch.set_default_tensor_type('torch.FloatTensor')
self.class_name = [
'chedeng', 'chebiao', 'chepai', 'person', 'car_light_detector',
'pingmu', 'luntaiguige', 'qizhifu', 'abs', 'chechuanshui', 'shouxieqianzi',
'seat', 'fuzhuzhidong', 'chemenpentu', 'taban', 'zuoyianquandai', 'xsz_rq', 'shuiyin_sj',
'shuiyin_rq', 'chejiahao', 'wbchepai', 'xsz_fuyezhang', '3cbz', 'button']
self.arch = 'dla_34'
self.heads = {'hm': 24, 'wh': 2, 'reg': 2}
self.head_conv = 256
self.load_model = '/data_2/project_202009/pytorch_project/CenterNet/centernet_src_202010/CenterNet-master/myfile/model_last-3epoch-multi.pth'
self.K = 100
self.conf_thresh = 0.3
self.downsample_rate = 4
def parse(self, kwargs):
'''
根据字典kwargs 更新 config参数
'''
for k, v in kwargs.items():
if not hasattr(self, k):
assert ValueError
setattr(self, k, v)
def init_model(opt):
# load model
print('Creating model...')
model = create_model(opt.arch, opt.heads, opt.head_conv)
model = load_model(model, opt.load_model)
model = model.to(opt.device)
model.eval()
return model
def prepare_img(img_path,opt):
image = cv2.imread(img_path)
h, w, c = image.shape
resized_image = cv2.resize(image, (opt.reso, opt.reso))
inp_image = ((resized_image / 255. - opt.mean) / opt.std).astype(np.float32)
inp_image = inp_image.transpose(2, 0, 1).reshape(1, 3, opt.reso, opt.reso)
inp_image = torch.from_numpy(inp_image)
inp_image = inp_image.to(opt.device)
return inp_image,h,w
def process(model,images,opt,h,w,img_path_):
output = model(images)[-1]
hm = output['hm'].sigmoid_()
wh = output['wh']
reg = output['reg']
dets = ctdet_decode(hm, wh, reg=reg, cat_spec_wh=False, K=opt.K) #[1,100,6] bboxes, scores, clses
dets[...,:4] *= opt.downsample_rate
ratio_h = h * 1.0 / opt.reso
ratio_w = w * 1.0 / opt.reso
dets[..., 0] *= ratio_w #x1
dets[..., 2] *= ratio_w #x2
dets[..., 1] *= ratio_h #y1
dets[..., 3] *= ratio_h #y2
score = dets[..., 4]
score = score >= opt.conf_thresh
score.squeeze_() #[100]
dets = dets[:,score,:]
if 0 == dets.numel():
return
dets = dets.detach().cpu().numpy()
img = cv2.imread(img_path_)
for i in range(dets.shape[1]):
box = dets[0, i, :4]
score = dets[0, i, 4]
if score < opt.conf_thresh:
continue
cls = dets[0, i, 5].astype(np.uint8)
x1y1 = (box[0].astype(np.uint8),box[1].astype(np.uint8))
x2y2 = (box[2].astype(np.uint8),box[3].astype(np.uint8))
color = (255, 0, 0)
cv2.rectangle(img, x1y1, x2y2, color, thickness=1)
t_size = cv2.getTextSize(opt.class_name[cls], fontFace=cv2.FONT_HERSHEY_PLAIN, fontScale=1, thickness=1)[0]
x2y2 = (x1y1[0] + t_size[0] + 3, x1y1[1] + t_size[1] + 4)
cv2.rectangle(img, x1y1, x2y2, color, thickness=-1)
cv2.putText(img, opt.class_name[cls], org=(x1y1[0], x1y1[1] + t_size[1] + 4),
fontFace=cv2.FONT_HERSHEY_PLAIN, fontScale=1, color=[255, 255, 255], thickness=1)
cv2.namedWindow("show",0)
cv2.imshow("show",img)
cv2.waitKey()
if __name__ == '__main__':
opt = DefaultConfig()
model = init_model(opt)
if os.path.isdir(opt.demo):
img_path = []
ls = os.listdir(opt.demo)
for file_name in sorted(ls):
ext = file_name[file_name.rfind('.') + 1:].lower()
img_path.append(os.path.join(opt.demo, file_name))
#if ext in image_ext:
# image_names.append(os.path.join(opt.demo, file_name))
else:
img_path = [opt.demo]
for cnt,img_path_ in enumerate(img_path):
print(cnt,img_path_)
inp_image, h, w = prepare_img(img_path_, opt)
process(model, inp_image, opt, h, w, img_path_)
哎，没有办法，又按照官网代码，用透视变换写了一份代码，该脚本放在CenterNet-master根目录下面运行：
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os
import cv2
import numpy as np
import torch
import time
#import torch.backends.cudnn as cudnn
torch.backends.cudnn.enabled = True
image_ext = ['jpg', 'jpeg', 'png']
import sys
CENTERNET_PATH = './src/'
sys.path.insert(0, CENTERNET_PATH)
CENTERNET_PATH = './src/lib/'
sys.path.insert(0, CENTERNET_PATH)
from models.model import create_model, load_model
from models.decode import ctdet_decode
from utils.image import get_affine_transform
from utils.post_process import ctdet_post_process
def pre_process(opt, image, scale=1.0, meta=None):
height, width = image.shape[0:2]
new_height = int(height * scale)
new_width = int(width * scale)
if 1:# if self.opt.fix_res:
inp_height, inp_width = opt.input_h,opt.input_w
c = np.array([new_width / 2., new_height / 2.], dtype=np.float32)
s = max(height, width) * 1.0
else:
inp_height = (new_height | self.opt.pad) + 1
inp_width = (new_width | self.opt.pad) + 1
c = np.array([new_width // 2, new_height // 2], dtype=np.float32)
s = np.array([inp_width, inp_height], dtype=np.float32)
trans_input = get_affine_transform(c, s, 0, [inp_width, inp_height])
resized_image = cv2.resize(image, (new_width, new_height))
inp_image = cv2.warpAffine(
resized_image, trans_input, (inp_width, inp_height),
flags=cv2.INTER_LINEAR)
inp_image = ((inp_image / 255. - opt.mean) / opt.std).astype(np.float32)
images = inp_image.transpose(2, 0, 1).reshape(1, 3, inp_height, inp_width)
images = torch.from_numpy(images)
meta = {'c': c, 's': s,
'out_height': inp_height // opt.down_ratio,
'out_width': inp_width // opt.down_ratio}
return images, meta
def process_ctdet(model,images,opt):
with torch.no_grad():
output = model(images)[-1]
hm = output['hm'].sigmoid_()
wh = output['wh']
reg = output['reg']
dets = ctdet_decode(hm, wh, reg=reg, cat_spec_wh=False, K=opt.K)
return output, dets
def post_process_ctdet(dets, meta, opt, scale=1.0):
dets = dets.detach().cpu().numpy()
dets = dets.reshape(1, -1, dets.shape[2])
dets = ctdet_post_process(
dets.copy(), [meta['c']], [meta['s']],
meta['out_height'], meta['out_width'], opt.num_classes)
for j in range(1, opt.num_classes + 1):
dets[0][j] = np.array(dets[0][j], dtype=np.float32).reshape(-1, 5)
dets[0][j][:, :4] /= scale
return dets[0]
def merge_outputs(opt, detections):
results = {}
for j in range(1, opt.num_classes + 1):
results[j] = np.concatenate([detection[j] for detection in detections], axis=0).astype(np.float32)
scores = np.hstack([results[j][:, 4] for j in range(1, opt.num_classes + 1)])
if len(scores) > opt.max_per_image:
kth = len(scores) - self.max_per_image
thresh = np.partition(scores, kth)[kth]
for j in range(1, self.num_classes + 1):
keep_inds = (results[j][:, 4] >= thresh)
results[j] = results[j][keep_inds]
return results
def add_coco_bbox(opt, bbox, cat, image, conf=1, show_txt=True):
bbox = np.array(bbox, dtype=np.int32)
# cat = (int(cat) + 1) % 80
cat = int(cat)
# print('cat', cat, self.names[cat])
# c = self.colors[cat][0][0].tolist()
# if self.theme == 'white':
# c = (255 - np.array(c)).tolist()
txt = '{}{:.1f}'.format(opt.class_name[cat], conf)
font = cv2.FONT_HERSHEY_SIMPLEX
cat_size = cv2.getTextSize(txt, font, 0.5, 2)[0]
cv2.rectangle(
image, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (255,0,0), 2)
if show_txt:
cv2.rectangle(image,
(bbox[0], bbox[1] - cat_size[1] - 2),
(bbox[0] + cat_size[0], bbox[1] - 2), (0,0,255), -1)
cv2.putText(image, txt, (bbox[0], bbox[1] - 2),
font, 0.5, (0, 0, 0), thickness=1, lineType=cv2.LINE_AA)
def save_txt(opt,results,path_txt):
with open(path_txt, "a") as fw:
for j in range(1, opt.num_classes + 1):
for bbox_ in results[j]:
if bbox_[4] >= opt.conf_thresh:
bbox = bbox_[:4]
bbox = np.array(bbox, dtype=np.int32)
cat = int(j-1)
str_1 = opt.class_name[cat] + " " + str(bbox_[4]) + " " + str(int(bbox[0])) + " " + str(
int(bbox[1])) + " " + str(int(bbox[2])) + " " + str(int(bbox[3])) + "
"
fw.write(str_1)
def show_results(img_path_,opt, results):
image = cv2.imread(img_path_)
for j in range(1, opt.num_classes + 1):
for bbox in results[j]:
if bbox[4] >= opt.conf_thresh:
add_coco_bbox(opt, bbox[:4], j - 1, image, bbox[4])
cv2.namedWindow("show",0)
cv2.imshow("show",image)
cv2.waitKey(0)
class DefaultConfig(object):
def __init__(self):
# self.demo = '../images/17790319373_bd19b24cfc_k.jpg'
self.demo = '/data_1/test_data/images-optional'
self.reso = 512
self.input_h = 512
self.input_w = 512
self.show_result = 1
self.save_txt = 0
# self.mean = np.array([0.40789654, 0.44719302, 0.47026115],
# dtype=np.float32).reshape(1, 1, 3)
# self.std = np.array([0.28863828, 0.27408164, 0.27809835],
# dtype=np.float32).reshape(1, 1, 3)
self.mean = np.array([0.408, 0.447, 0.47],
dtype=np.float32).reshape(1, 1, 3)
self.std = np.array([0.289, 0.274, 0.278],
dtype=np.float32).reshape(1, 1, 3)
self.batch_size = 4
self.use_gpu = torch.cuda.is_available()
self.device = torch.device("cuda" if self.use_gpu else "cpu")
if self.use_gpu:
torch.set_default_tensor_type('torch.cuda.FloatTensor')
else:
torch.set_default_tensor_type('torch.FloatTensor')
self.class_name = [
'cheliang', 'chewei', 'chelian', 'dibiao_20', 'sanjiaojia',
'qizhibiaozhi', 'motorbike', 'dibiao_0', 'dibiao_qd', 'xiaochebiaozhipai', 'tingchebiaozhipai',
'fanguangbeixin', 'dibiao_10']
self.arch = 'dla_34'
self.num_classes = 13
self.heads = {'hm':self.num_classes, 'wh': 2, 'reg': 2}
self.head_conv = 256
self.load_model = '/data_2/CenterNet-master/myfile/model_last.pth'
self.K = 100
self.conf_thresh = 0.2
self.downsample_rate = 4
self.down_ratio = 4
self.max_per_image = 100
def parse(self, kwargs):
'''
根据字典kwargs 更新 config参数
'''
for k, v in kwargs.items():
if not hasattr(self, k):
assert ValueError
setattr(self, k, v)
def init_model(opt):
# load model
print('Creating model...')
model = create_model(opt.arch, opt.heads, opt.head_conv)
model = load_model(model, opt.load_model)
model = model.to(opt.device)
model.eval()
#cudnn.benchmark = True
return model
if __name__ == '__main__':
opt = DefaultConfig()
model = init_model(opt)
if os.path.isdir(opt.demo):
img_path = []
ls = os.listdir(opt.demo)
for file_name in sorted(ls):
ext = file_name[file_name.rfind('.') + 1:].lower()
img_path.append(os.path.join(opt.demo, file_name))
#if ext in image_ext:
# image_names.append(os.path.join(opt.demo, file_name))
else:
img_path = [opt.demo]
torch.cuda.synchronize()
start = time.time()
for cnt,img_path_ in enumerate(img_path):
#print(cnt,img_path_)
detections = []
image = cv2.imread(img_path_)
images, meta = pre_process(opt,image)
images = images.to(opt.device)
output, dets = process_ctdet(model, images, opt)
dets = post_process_ctdet(dets, meta, opt)
detections.append(dets)
results = merge_outputs(opt,detections)
if opt.show_result:
show_results(img_path_, opt, results)
if opt.save_txt:
path_save_dir = 'save_txt'
os.makedirs(path_save_dir, exist_ok=True)
pos_1 = img_path_.rfind('/')
name_txt_1 = img_path_[pos_1 + 1 : -4] + ".txt"
save_txt(opt, results, path_save_dir + "/" + name_txt_1)
torch.cuda.synchronize()
end = time.time()
time_1 = end - start
print("ave time=",time_1 * 1000.0 / (cnt+1))
print("all=",cnt+1)