from models.model_stages import BiSeNet from predict_city.heliushuju import Heliushuju from torch.utils.data import DataLoader import pandas as pd import numpy as np import os import argparse import cv2 import torch import torchvision.transforms as transforms from trafficDetectionUtils import PostProcessing os.environ['CUDA_VISIBLE_DEVICES'] = '0' class MscEvalV0(object): def __init__(self, scaleH=1/3, scaleW=1/3, ignore_label=255): self.ignore_label = ignore_label self.scaleH = scaleH self.scaleW = scaleW self.to_tensor = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)), ]) def __call__(self, net, dl, n_classes): # evaluate label_info = get_label_info('./class_dict.csv') maskPath = '../trafficDetectionTestData/trafficAccidentTest/masks' testImagePath = '../trafficDetectionTestData/trafficAccidentTest/images' File1 = os.listdir(testImagePath) for file in File1: print('####beg to :', file) txtRowContent = [] VehicleCoordinate = [] singleTxtContent = [] txtPath = '../trafficDetectionTestData/trafficAccidentTest/detections' + os.sep + file[:-4] + '.txt' testImage = testImagePath + os.sep + file testImageArray = cv2.imread(testImage) txtContent = open(txtPath, 'r', encoding='utf-8') content = txtContent.readlines() for line in content: if line[0].isnumeric(): e = line.splitlines(False)[0] f = e.split(',', -1) if float(f[0]) == 0: for i in range(len(f)): txtRowContent.append(float(f[i])) VehicleCoordinate.append((int(txtRowContent[1]), int(txtRowContent[2]))) VehicleCoordinate.append((int(txtRowContent[3]), int(txtRowContent[4]))) singleTxtContent.append(txtRowContent) txtRowContent = [] # print("line53, singleTxtContent: ", singleTxtContent) txtContent.close() # 字典形式传参数 traffic_dict = {'label_info': label_info, 'RoadArea': 16000, 'roadVehicleAngle': 15, 'vehicleCOOR': VehicleCoordinate, 'roundness': 0.7, 'cls': 0, 'vehicleFactor': 0.6, 'det': singleTxtContent, 'modelSize': (640, 360), 'testImageName': file, 'radius': 50, 'distanceFlag': False, 'vehicleFlag': False} save_path = './demo/' + file mask = cv2.imread(maskPath + os.sep + file[:-4] + '_mask.png') mask = mask[:, :, 0] targetList, time_infos = PostProcessing(mask, testImageArray, traffic_dict, file) # print("line64", time_infos, mask.shape, traffic_dict) print("line64", time_infos) # print("line65", targetList) # 在测试图片上画出检测框 for i in range(len(targetList)): if targetList[i][10] != 3: X1 = targetList[i][1] Y1 = targetList[i][2] X2 = targetList[i][3] Y2 = targetList[i][4] cv2.rectangle(testImageArray, (int(X1), int(Y1)), (int(X2), int(Y2)), (0, 0, 255), thickness=3, lineType=cv2.LINE_AA) font = cv2.FONT_HERSHEY_SIMPLEX cv2.putText(testImageArray, str(format(targetList[i][9], ".2f")), (int(X1) + 4, int(Y1 - 1)), font, 1, (0, 255, 0), 2, cv2.LINE_AA) cv2.imwrite(save_path, testImageArray) def get_label_info(csv_path): ann = pd.read_csv(csv_path) label = {} for iter, row in ann.iterrows(): label_name = row['name'] r = row['r'] g = row['g'] b = row['b'] label[label_name] = [int(r), int(g), int(b)] return label def evaluatev0(respth='', dspth='', backbone='', scaleH=1/3, scaleW=1/3,use_boundary_2=False, use_boundary_4=False, use_boundary_8=False, use_boundary_16=False, use_conv_last=False): # dataset batchsize = 1 n_workers = 0 dsval = Heliushuju(dspth, mode='test') dl = DataLoader(dsval, batch_size=batchsize, shuffle=False, num_workers=n_workers, drop_last=False) n_classes = 3 # print("backbone:", backbone) net = BiSeNet(backbone=backbone, n_classes=n_classes, use_boundary_2=use_boundary_2, use_boundary_4=use_boundary_4, use_boundary_8=use_boundary_8, use_boundary_16=use_boundary_16, use_conv_last=use_conv_last) net.load_state_dict(torch.load(respth)) net.cuda() net.eval() with torch.no_grad(): single_scale = MscEvalV0(scaleH=scaleH, scaleW=scaleW) single_scale(net, dl, 3) if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument('--weights', nargs='+', type=str, default='./model_save/pths/best.pt', help='model.pt path(s)') parser.add_argument('--source', type=str, default='./data/test/images', help='source') # file/folder, 0 for webcam parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)') parser.add_argument('--conf-thres', type=float, default=0.25, help='object confidence threshold') parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold for NMS') parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') parser.add_argument('--view-img', action='store_true', help='display results') parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels') parser.add_argument('--nosave', action='store_true', help='do not save images/videos') parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3') parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS') parser.add_argument('--augment', action='store_true', help='augmented inference') parser.add_argument('--update', action='store_true', help='update all models') opt = parser.parse_args() # stdc_360X640_highWay.pth model_final.pth evaluatev0(respth='./model_save/pths/stdc_360X640_highWay.pth', dspth='../trafficDetectionTestData/trafficAccidentTest/masks', backbone='STDCNet813', scaleH=1 / 3, scaleW=1 / 3, use_boundary_2=False, use_boundary_4=False, use_boundary_8=False, use_boundary_16=False, use_conv_last=False)