import os # os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" # os.environ["CUDA_VISIBLE_DEVICES"] = "1" import asyncio import glob from argparse import ArgumentParser from tqdm import tqdm import matplotlib.pyplot as plt import cv2 import numpy as np import json import torch from mmcv import Config from mmdet.apis import async_inference_detector, inference_detector, show_result_pyplot from mmdet.apis.inference import init_detector from ensemble_boxes import * # tracking model from yolox.tracker.byte_tracker import BYTETracker from track_parse import make_parser import time args = make_parser() ######################## global_fps = 3 SAVED_FRAME = 90 threshold_vehicle = 0.3 threshold_person = 0.3 ####################### # Utils from utils import * # Weighted box fusion weights_vehicle = [2, 1] weights_person = [2, 1] iou_thr = 0.5 skip_box_thr = 0.0001 sigma = 0.1 # write tracking log f = open('log.txt', 'w', encoding='utf8') VIDEO_DIR = "/root/DataSet/experiments/pres_vid/input/" OUT_DIR = "/root/DataSet/experiments/pres_vid/output/" cuda_device = "cuda:1" # VEHICLE DETECTION TRAINED MODELS config_files = ["/root/DataSet/experiments/ours_best_vehicles/tood-r101-balance/tood_r101.py", "/root/DataSet/experiments/ours_best_vehicles/balanced_models/semi1/balance_faster_convnext_l_semi_p1.py"] checkpoints = ["/root/DataSet/experiments/ours_best_vehicles/tood-r101-balance/epoch_4.pth", "/root/DataSet/experiments/ours_best_vehicles/balanced_models/semi1/epoch_10.pth"] # PERSON DETECTION TRAINED MODELS config_files_person = ["/root/DataSet/experiments/signalman_convnext_l/faster-rcnn-convnext-l.py", "/root/DataSet/experiments/signalman_tood_101/tood_101.py"] checkpoints_person = ["/root/DataSet/experiments/signalman_convnext_l/epoch_20.pth", "/root/DataSet/experiments/signalman_tood_101/epoch_30.pth"] # Model init model_1 = init_detector(config_files[0], checkpoints[0], device=cuda_device) model_1.CLASSES = ["Excavator","Dodger","Forklift","Dump Truck","Mixer Truck", "Cargo Truck","Scissor Lift","Crane"] model_2 = init_detector(config_files[1], checkpoints[1], device=cuda_device) model_2.CLASSES = ["Excavator","Dodger","Forklift","Dump Truck","Mixer Truck", "Cargo Truck","Scissor Lift","Crane"] # model_3 = init_detector(config_files[2], checkpoints[2], device=cuda_device) # model_3.CLASSES = ["Excavator","Dodger","Forklift","Dump Truck","Mixer Truck", # "Cargo Truck","Scissor Lift","Crane"] # Model init model_person_1 = init_detector(config_files_person[0], checkpoints_person[0], device=cuda_device) model_person_1.CLASSES = ["Worker", "Signalman"] model_person_2 = init_detector(config_files_person[1], checkpoints_person[1], device=cuda_device) model_person_2.CLASSES = ["Worker", "Signalman"] # model_person_3 = init_detector(config_files_person[2], checkpoints_person[2], device=cuda_device) # model_person_3.CLASSES = ["Worker", "Signalman"] video_list_raw = os.listdir(VIDEO_DIR) video_list = [] video_format = ["mp4","avi"] for video_num in range(len(video_list_raw)): if video_list_raw[video_num].split(".")[-1] in video_format: video_list.append(video_list_raw[video_num]) for video_name in video_list: video_path = os.path.join(VIDEO_DIR, video_name) cap = cv2.VideoCapture(video_path) stop_num = 0 image_array = [] video_length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) video_fps = int(cap.get(cv2.CAP_PROP_FPS)) print('Video ID : ', video_name) print('Video fps : ', video_fps) print('Total len : ', video_length) target_margin = int(video_fps / global_fps) count_margin = 0 with tqdm(total=SAVED_FRAME) as pbar: while True: ret, frame = cap.read() count_margin+=1 if count_margin >= target_margin: count_margin = 0 if frame is not None: # frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) image_size = frame.shape # VEHICLE DETECTIONS result_1 = inference_detector(model_1, frame) result_2 = inference_detector(model_2, frame) # result_3 = inference_detector(model_3, frame) ## Merge list boxes_list1, labels_list1, scores_list1 = convert_wbf(result_1, threshold = threshold_vehicle, image_size = image_size) boxes_list2, labels_list2, scores_list2 = convert_wbf(result_2, threshold = threshold_vehicle, image_size = image_size) # boxes_list3, labels_list3, scores_list3 = convert_wbf(result_3, threshold = threshold_vehicle, image_size = image_size) boxes_list_merges = [] labels_list_merges = [] scores_list_merges = [] boxes_list_merges.append(boxes_list1) boxes_list_merges.append(boxes_list2) # boxes_list_merges.append(boxes_list3) labels_list_merges.append(labels_list1) labels_list_merges.append(labels_list2) # labels_list_merges.append(labels_list3) scores_list_merges.append(scores_list1) scores_list_merges.append(scores_list2) # scores_list_merges.append(scores_list3) boxes, scores, labels = weighted_boxes_fusion(boxes_list_merges, scores_list_merges, labels_list_merges, weights=weights_vehicle, iou_thr=iou_thr, skip_box_thr=skip_box_thr) vehicle_class = {1:'Excavator', 2:'Dodger', 3:'Forklift', 4:'Dump_Truck', 5:'Mixer_Truck', 6:'Cargo_Truck', 7:'Scissor_Lift', 8:'Crane'} img_w = image_size[1] img_h = image_size[0] # in one image -> vehicle vehicle_results=[] for idx,box in enumerate(boxes): if scores[idx]>threshold_vehicle: x = int(box[0]*image_size[1]) y = int(box[1]*image_size[0]) w = int(box[2]*image_size[1]) - x h = int(box[3]*image_size[0]) - y temp={} temp['bbox'] = [x,y,w,h] temp['det_score'] = round(scores[idx],2) temp['cls_name'] = vehicle_class[labels[idx]] vehicle_results.append(temp) # WORKER DETECTION result_person_1 = inference_detector(model_person_1, frame) result_person_2 = inference_detector(model_person_2, frame) # result_person_3 = inference_detector(model_person_3, frame) boxes_list1, labels_list1, scores_list1 = convert_wbf(result_person_1, threshold = threshold_person, image_size = image_size) boxes_list2, labels_list2, scores_list2 = convert_wbf(result_person_2, threshold = threshold_person, image_size = image_size) # boxes_list3, labels_list3, scores_list3 = convert_wbf(result_person_3, threshold = threshold_person, image_size = image_size) boxes_list_merges = [] labels_list_merges = [] scores_list_merges = [] boxes_list_merges.append(boxes_list1) boxes_list_merges.append(boxes_list2) # boxes_list_merges.append(boxes_list3) labels_list_merges.append(labels_list1) labels_list_merges.append(labels_list2) # labels_list_merges.append(labels_list3) scores_list_merges.append(scores_list1) scores_list_merges.append(scores_list2) # scores_list_merges.append(scores_list3) boxes, scores, labels = weighted_boxes_fusion(boxes_list_merges, scores_list_merges, labels_list_merges, weights=weights_person, iou_thr=iou_thr, skip_box_thr=skip_box_thr) # check threshold also worker_bboxes = [] # x,y,w,h for idx,box in enumerate(boxes): if scores[idx]>threshold_person: x = int(box[0]*image_size[1]) y = int(box[1]*image_size[0]) w = int(box[2]*image_size[1] - x) h = int(box[3]*image_size[0] - y) worker_bboxes.append([x,y,w,h]) # check signalman # worker_bboxes = check_signalman(worker_bboxes, frame) worker_class = {1:'worker', 2:'signalman'} ################################################################# update # in one image -> worker worker_results=[] for i in range(len(worker_bboxes)): temp={} temp['bbox'] = worker_bboxes[i][:4] temp['det_score'] = round(scores[i],2) temp['cls_name'] = worker_class[int(labels[i])] temp['risk1'] = 'safe' worker_results.append(temp) # ------------------- DRAW --------------------- # # risk assessment - 1 # 'risk1' : 'danger' or 'safe' vehicle_results, worker_results = risk_algorithm_1(vehicle_results, worker_results) # draw vehicle # img_result = draw_vehicle(frame, vehicle_results) img_result = draw_vehicle2(frame, vehicle_results) # draw worker img_result = draw_worker(img_result, worker_results) image_array.append(img_result) h,w,l = img_result.shape pbar.update(1) # break if stop_num == SAVED_FRAME: break stop_num += 1 else: break size = (w,h) out = cv2.VideoWriter(OUT_DIR+"inferenced_{}.mp4".format(video_name), cv2.VideoWriter_fourcc(*'DIVX'), global_fps, size) print('Done') print() for i in range(len(image_array)): out.write(image_array[i]) out.release() # break ##########################################################################################