trafficDetectionHandover/traffic_newest.py

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)