110 lines
3.7 KiB
Python
110 lines
3.7 KiB
Python
import os
|
|
import time
|
|
|
|
import cv2
|
|
import torch
|
|
|
|
from DMPRUtils.DMPR_process import DMPR_process, plot_points
|
|
from DMPRUtils.model.detector import DirectionalPointDetector
|
|
from DMPR_YOLO.jointUtil import dmpr_yolo
|
|
from conf import config
|
|
from models.experimental import attempt_load
|
|
from models.yolo_process import yolo_process
|
|
from utils.plots import plot_one_box
|
|
from utils.torch_utils import select_device
|
|
|
|
|
|
def main():
|
|
##预先设置的参数
|
|
device_ = '0' ##选定模型，可选 cpu,'0','1'
|
|
|
|
##以下参数目前不可改
|
|
Detweights = 'weights/urbanManagement/yolo/best.pt'
|
|
seg_nclass = 2
|
|
DMPRweights = "weights/urbanManagement/DMPR/dp_detector_499.pth"
|
|
conf_thres, iou_thres, classes = 0.25, 0.45, 3
|
|
labelnames = "weights/yolov5/class5/labelnames.json"
|
|
rainbows = [[0, 0, 255], [0, 255, 0], [255, 0, 0], [255, 0, 255], [255, 255, 0], [255, 129, 0], [255, 0, 127],
|
|
[127, 255, 0], [0, 255, 127], [0, 127, 255], [127, 0, 255], [255, 127, 255], [255, 255, 127],
|
|
[127, 255, 255], [0, 255, 255], [255, 127, 255], [127, 255, 255], [0, 127, 0], [0, 0, 127],
|
|
[0, 255, 255]]
|
|
allowedList = [0, 1, 2, 3]
|
|
|
|
##加载模型，准备好显示字符
|
|
device = select_device(device_)
|
|
|
|
half = device.type != 'cpu' # half precision only supported on CUDA
|
|
# yolov5 model
|
|
model = attempt_load(Detweights, map_location=device)
|
|
if half:
|
|
model.half()
|
|
|
|
# DMPR model
|
|
args = config.get_parser_for_inference().parse_args()
|
|
DMPRmodel = DirectionalPointDetector(3, args.depth_factor, config.NUM_FEATURE_MAP_CHANNEL).to(device)
|
|
DMPRmodel.load_state_dict(torch.load(DMPRweights))
|
|
|
|
# 图像测试
|
|
# impth = 'images/input'
|
|
impth = 'images/debug'
|
|
# outpth = 'images/output'
|
|
outpth = 'images/debug_out'
|
|
folders = os.listdir(impth)
|
|
for file in folders:
|
|
imgpath = os.path.join(impth, file)
|
|
img0 = cv2.imread(imgpath)
|
|
assert img0 is not None, 'Image Not Found ' + imgpath
|
|
|
|
# t_start = time.time()
|
|
# yolo process
|
|
det0 = yolo_process(img0, model, device, args, half)
|
|
det0 = det0.cpu().detach().numpy()
|
|
# t_yolo = time.time()
|
|
# print(f't_yolo. ({t_yolo - t_start:.3f}s)')
|
|
|
|
# plot所有box
|
|
# for *xyxy, conf, cls in reversed(det0):
|
|
# label = f'{int(cls)} {conf:.2f}'
|
|
# plot_one_box(xyxy, img0, label=label, color=rainbows[int(cls)], line_thickness=2)
|
|
|
|
# DMPR process
|
|
det1 = DMPR_process(img0, DMPRmodel, device, args)
|
|
det1 = det1.cpu().detach().numpy()
|
|
|
|
# t_dmpr = time.time()
|
|
# print(f't_dmpr. ({t_dmpr - t_yolo:.3f}s)')
|
|
|
|
# 绘制角点
|
|
plot_points(img0, det1)
|
|
|
|
# save
|
|
# cv2.imwrite(file, img0)
|
|
|
|
# yolo joint DMPR
|
|
cls = 0 #需要过滤的box类别
|
|
joint_det, dilate_box = dmpr_yolo(det1, det0, img0.shape, cls)
|
|
|
|
# t_joint = time.time()
|
|
# print(f't_joint. ({t_joint - t_dmpr:.3f}s)')
|
|
|
|
# t_end = time.time()
|
|
# print(f'Done. ({t_end - t_start:.3f}s)')
|
|
# 绘制膨胀box
|
|
for *xyxy, flag in dilate_box:
|
|
plot_one_box(xyxy, img0, color=rainbows[int(cls)], line_thickness=2)
|
|
#
|
|
# # 绘制删除满足 在膨胀框内 && 角度差小于90度 的box
|
|
for *xyxy, conf, cls, flag in reversed(joint_det):
|
|
if flag == 0:
|
|
# label = f'{int(cls)} {conf:.2f}'
|
|
label = None
|
|
plot_one_box(xyxy, img0, label=label, color=rainbows[int(cls)], line_thickness=2)
|
|
|
|
# save
|
|
save_path = os.path.join(outpth, file)
|
|
cv2.imwrite(save_path, img0)
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main() |