pipeline.py 55 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316
  1. # Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. import os
  15. import yaml
  16. import glob
  17. import cv2
  18. import numpy as np
  19. import math
  20. import paddle
  21. import sys
  22. import copy
  23. import threading
  24. import queue
  25. import time
  26. from collections import Sequence, defaultdict
  27. from datacollector import DataCollector, Result
  28. # add deploy path of PadleDetection to sys.path
  29. parent_path = os.path.abspath(os.path.join(__file__, *(['..'] * 2)))
  30. sys.path.insert(0, parent_path)
  31. from cfg_utils import argsparser, print_arguments, merge_cfg
  32. from pipe_utils import PipeTimer
  33. from pipe_utils import get_test_images, crop_image_with_det, crop_image_with_mot, parse_mot_res, parse_mot_keypoint
  34. from pipe_utils import PushStream
  35. from python.infer import Detector, DetectorPicoDet
  36. from python.keypoint_infer import KeyPointDetector
  37. from python.keypoint_postprocess import translate_to_ori_images
  38. from python.preprocess import decode_image, ShortSizeScale
  39. from python.visualize import visualize_box_mask, visualize_attr, visualize_pose, visualize_action, visualize_vehicleplate, visualize_vehiclepress, visualize_lane, visualize_vehicle_retrograde
  40. from pptracking.python.mot_sde_infer import SDE_Detector
  41. from pptracking.python.mot.visualize import plot_tracking_dict
  42. from pptracking.python.mot.utils import flow_statistic, update_object_info
  43. from pphuman.attr_infer import AttrDetector
  44. from pphuman.video_action_infer import VideoActionRecognizer
  45. from pphuman.action_infer import SkeletonActionRecognizer, DetActionRecognizer, ClsActionRecognizer
  46. from pphuman.action_utils import KeyPointBuff, ActionVisualHelper
  47. from pphuman.reid import ReID
  48. from pphuman.mtmct import mtmct_process
  49. from ppvehicle.vehicle_plate import PlateRecognizer
  50. from ppvehicle.vehicle_attr import VehicleAttr
  51. from ppvehicle.vehicle_pressing import VehiclePressingRecognizer
  52. from ppvehicle.vehicle_retrograde import VehicleRetrogradeRecognizer
  53. from ppvehicle.lane_seg_infer import LaneSegPredictor
  54. from download import auto_download_model
  55. class Pipeline(object):
  56. """
  57. Pipeline
  58. Args:
  59. args (argparse.Namespace): arguments in pipeline, which contains environment and runtime settings
  60. cfg (dict): config of models in pipeline
  61. """
  62. def __init__(self, args, cfg):
  63. self.multi_camera = False
  64. reid_cfg = cfg.get('REID', False)
  65. self.enable_mtmct = reid_cfg['enable'] if reid_cfg else False
  66. self.is_video = False
  67. self.output_dir = args.output_dir
  68. self.vis_result = cfg['visual']
  69. self.input = self._parse_input(args.image_file, args.image_dir,
  70. args.video_file, args.video_dir,
  71. args.camera_id, args.rtsp)
  72. if self.multi_camera:
  73. self.predictor = []
  74. for name in self.input:
  75. predictor_item = PipePredictor(
  76. args, cfg, is_video=True, multi_camera=True)
  77. predictor_item.set_file_name(name)
  78. self.predictor.append(predictor_item)
  79. else:
  80. self.predictor = PipePredictor(args, cfg, self.is_video)
  81. if self.is_video:
  82. self.predictor.set_file_name(self.input)
  83. def _parse_input(self, image_file, image_dir, video_file, video_dir,
  84. camera_id, rtsp):
  85. # parse input as is_video and multi_camera
  86. if image_file is not None or image_dir is not None:
  87. input = get_test_images(image_dir, image_file)
  88. self.is_video = False
  89. self.multi_camera = False
  90. elif video_file is not None:
  91. assert os.path.exists(
  92. video_file
  93. ) or 'rtsp' in video_file, "video_file not exists and not an rtsp site."
  94. self.multi_camera = False
  95. input = video_file
  96. self.is_video = True
  97. elif video_dir is not None:
  98. videof = [os.path.join(video_dir, x) for x in os.listdir(video_dir)]
  99. if len(videof) > 1:
  100. self.multi_camera = True
  101. videof.sort()
  102. input = videof
  103. else:
  104. input = videof[0]
  105. self.is_video = True
  106. elif rtsp is not None:
  107. if len(rtsp) > 1:
  108. rtsp = [rtsp_item for rtsp_item in rtsp if 'rtsp' in rtsp_item]
  109. self.multi_camera = True
  110. input = rtsp
  111. else:
  112. self.multi_camera = False
  113. input = rtsp[0]
  114. self.is_video = True
  115. elif camera_id != -1:
  116. self.multi_camera = False
  117. input = camera_id
  118. self.is_video = True
  119. else:
  120. raise ValueError(
  121. "Illegal Input, please set one of ['video_file', 'camera_id', 'image_file', 'image_dir']"
  122. )
  123. return input
  124. def run_multithreads(self):
  125. if self.multi_camera:
  126. multi_res = []
  127. threads = []
  128. for idx, (predictor,
  129. input) in enumerate(zip(self.predictor, self.input)):
  130. thread = threading.Thread(
  131. name=str(idx).zfill(3),
  132. target=predictor.run,
  133. args=(input, idx))
  134. threads.append(thread)
  135. for thread in threads:
  136. thread.start()
  137. for predictor, thread in zip(self.predictor, threads):
  138. thread.join()
  139. collector_data = predictor.get_result()
  140. multi_res.append(collector_data)
  141. if self.enable_mtmct:
  142. mtmct_process(
  143. multi_res,
  144. self.input,
  145. mtmct_vis=self.vis_result,
  146. output_dir=self.output_dir)
  147. else:
  148. self.predictor.run(self.input)
  149. def run(self):
  150. if self.multi_camera:
  151. multi_res = []
  152. for predictor, input in zip(self.predictor, self.input):
  153. predictor.run(input)
  154. collector_data = predictor.get_result()
  155. multi_res.append(collector_data)
  156. if self.enable_mtmct:
  157. mtmct_process(
  158. multi_res,
  159. self.input,
  160. mtmct_vis=self.vis_result,
  161. output_dir=self.output_dir)
  162. else:
  163. self.predictor.run(self.input)
  164. def get_model_dir(cfg):
  165. """
  166. Auto download inference model if the model_path is a url link.
  167. Otherwise it will use the model_path directly.
  168. """
  169. for key in cfg.keys():
  170. if type(cfg[key]) == dict and \
  171. ("enable" in cfg[key].keys() and cfg[key]['enable']
  172. or "enable" not in cfg[key].keys()):
  173. if "model_dir" in cfg[key].keys():
  174. model_dir = cfg[key]["model_dir"]
  175. downloaded_model_dir = auto_download_model(model_dir)
  176. if downloaded_model_dir:
  177. model_dir = downloaded_model_dir
  178. cfg[key]["model_dir"] = model_dir
  179. print(key, " model dir: ", model_dir)
  180. elif key == "VEHICLE_PLATE":
  181. det_model_dir = cfg[key]["det_model_dir"]
  182. downloaded_det_model_dir = auto_download_model(det_model_dir)
  183. if downloaded_det_model_dir:
  184. det_model_dir = downloaded_det_model_dir
  185. cfg[key]["det_model_dir"] = det_model_dir
  186. print("det_model_dir model dir: ", det_model_dir)
  187. rec_model_dir = cfg[key]["rec_model_dir"]
  188. downloaded_rec_model_dir = auto_download_model(rec_model_dir)
  189. if downloaded_rec_model_dir:
  190. rec_model_dir = downloaded_rec_model_dir
  191. cfg[key]["rec_model_dir"] = rec_model_dir
  192. print("rec_model_dir model dir: ", rec_model_dir)
  193. elif key == "MOT": # for idbased and skeletonbased actions
  194. model_dir = cfg[key]["model_dir"]
  195. downloaded_model_dir = auto_download_model(model_dir)
  196. if downloaded_model_dir:
  197. model_dir = downloaded_model_dir
  198. cfg[key]["model_dir"] = model_dir
  199. print("mot_model_dir model_dir: ", model_dir)
  200. class PipePredictor(object):
  201. """
  202. Predictor in single camera
  203. The pipeline for image input:
  204. 1. Detection
  205. 2. Detection -> Attribute
  206. The pipeline for video input:
  207. 1. Tracking
  208. 2. Tracking -> Attribute
  209. 3. Tracking -> KeyPoint -> SkeletonAction Recognition
  210. 4. VideoAction Recognition
  211. Args:
  212. args (argparse.Namespace): arguments in pipeline, which contains environment and runtime settings
  213. cfg (dict): config of models in pipeline
  214. is_video (bool): whether the input is video, default as False
  215. multi_camera (bool): whether to use multi camera in pipeline,
  216. default as False
  217. """
  218. def __init__(self, args, cfg, is_video=True, multi_camera=False):
  219. # general module for pphuman and ppvehicle
  220. self.with_mot = cfg.get('MOT', False)['enable'] if cfg.get(
  221. 'MOT', False) else False
  222. self.with_human_attr = cfg.get('ATTR', False)['enable'] if cfg.get(
  223. 'ATTR', False) else False
  224. if self.with_mot:
  225. print('Multi-Object Tracking enabled')
  226. if self.with_human_attr:
  227. print('Human Attribute Recognition enabled')
  228. # only for pphuman
  229. self.with_skeleton_action = cfg.get(
  230. 'SKELETON_ACTION', False)['enable'] if cfg.get('SKELETON_ACTION',
  231. False) else False
  232. self.with_video_action = cfg.get(
  233. 'VIDEO_ACTION', False)['enable'] if cfg.get('VIDEO_ACTION',
  234. False) else False
  235. self.with_idbased_detaction = cfg.get(
  236. 'ID_BASED_DETACTION', False)['enable'] if cfg.get(
  237. 'ID_BASED_DETACTION', False) else False
  238. self.with_idbased_clsaction = cfg.get(
  239. 'ID_BASED_CLSACTION', False)['enable'] if cfg.get(
  240. 'ID_BASED_CLSACTION', False) else False
  241. self.with_mtmct = cfg.get('REID', False)['enable'] if cfg.get(
  242. 'REID', False) else False
  243. if self.with_skeleton_action:
  244. print('SkeletonAction Recognition enabled')
  245. if self.with_video_action:
  246. print('VideoAction Recognition enabled')
  247. if self.with_idbased_detaction:
  248. print('IDBASED Detection Action Recognition enabled')
  249. if self.with_idbased_clsaction:
  250. print('IDBASED Classification Action Recognition enabled')
  251. if self.with_mtmct:
  252. print("MTMCT enabled")
  253. # only for ppvehicle
  254. self.with_vehicleplate = cfg.get(
  255. 'VEHICLE_PLATE', False)['enable'] if cfg.get('VEHICLE_PLATE',
  256. False) else False
  257. if self.with_vehicleplate:
  258. print('Vehicle Plate Recognition enabled')
  259. self.with_vehicle_attr = cfg.get(
  260. 'VEHICLE_ATTR', False)['enable'] if cfg.get('VEHICLE_ATTR',
  261. False) else False
  262. if self.with_vehicle_attr:
  263. print('Vehicle Attribute Recognition enabled')
  264. self.with_vehicle_press = cfg.get(
  265. 'VEHICLE_PRESSING', False)['enable'] if cfg.get('VEHICLE_PRESSING',
  266. False) else False
  267. if self.with_vehicle_press:
  268. print('Vehicle Pressing Recognition enabled')
  269. self.with_vehicle_retrograde = cfg.get(
  270. 'VEHICLE_RETROGRADE', False)['enable'] if cfg.get(
  271. 'VEHICLE_RETROGRADE', False) else False
  272. if self.with_vehicle_retrograde:
  273. print('Vehicle Retrograde Recognition enabled')
  274. self.modebase = {
  275. "framebased": False,
  276. "videobased": False,
  277. "idbased": False,
  278. "skeletonbased": False
  279. }
  280. self.basemode = {
  281. "MOT": "idbased",
  282. "ATTR": "idbased",
  283. "VIDEO_ACTION": "videobased",
  284. "SKELETON_ACTION": "skeletonbased",
  285. "ID_BASED_DETACTION": "idbased",
  286. "ID_BASED_CLSACTION": "idbased",
  287. "REID": "idbased",
  288. "VEHICLE_PLATE": "idbased",
  289. "VEHICLE_ATTR": "idbased",
  290. "VEHICLE_PRESSING": "idbased",
  291. "VEHICLE_RETROGRADE": "idbased",
  292. }
  293. self.is_video = is_video
  294. self.multi_camera = multi_camera
  295. self.cfg = cfg
  296. self.output_dir = args.output_dir
  297. self.draw_center_traj = args.draw_center_traj
  298. self.secs_interval = args.secs_interval
  299. self.do_entrance_counting = args.do_entrance_counting
  300. self.do_break_in_counting = args.do_break_in_counting
  301. self.region_type = args.region_type
  302. self.region_polygon = args.region_polygon
  303. self.illegal_parking_time = args.illegal_parking_time
  304. self.warmup_frame = self.cfg['warmup_frame']
  305. self.pipeline_res = Result()
  306. self.pipe_timer = PipeTimer()
  307. self.file_name = None
  308. self.collector = DataCollector()
  309. self.pushurl = args.pushurl
  310. # auto download inference model
  311. get_model_dir(self.cfg)
  312. if self.with_vehicleplate:
  313. vehicleplate_cfg = self.cfg['VEHICLE_PLATE']
  314. self.vehicleplate_detector = PlateRecognizer(args, vehicleplate_cfg)
  315. basemode = self.basemode['VEHICLE_PLATE']
  316. self.modebase[basemode] = True
  317. if self.with_human_attr:
  318. attr_cfg = self.cfg['ATTR']
  319. basemode = self.basemode['ATTR']
  320. self.modebase[basemode] = True
  321. self.attr_predictor = AttrDetector.init_with_cfg(args, attr_cfg)
  322. if self.with_vehicle_attr:
  323. vehicleattr_cfg = self.cfg['VEHICLE_ATTR']
  324. basemode = self.basemode['VEHICLE_ATTR']
  325. self.modebase[basemode] = True
  326. self.vehicle_attr_predictor = VehicleAttr.init_with_cfg(
  327. args, vehicleattr_cfg)
  328. if self.with_vehicle_press:
  329. vehiclepress_cfg = self.cfg['VEHICLE_PRESSING']
  330. basemode = self.basemode['VEHICLE_PRESSING']
  331. self.modebase[basemode] = True
  332. self.vehicle_press_predictor = VehiclePressingRecognizer(
  333. vehiclepress_cfg)
  334. if self.with_vehicle_press or self.with_vehicle_retrograde:
  335. laneseg_cfg = self.cfg['LANE_SEG']
  336. self.laneseg_predictor = LaneSegPredictor(
  337. laneseg_cfg['lane_seg_config'], laneseg_cfg['model_dir'])
  338. if not is_video:
  339. det_cfg = self.cfg['DET']
  340. model_dir = det_cfg['model_dir']
  341. batch_size = det_cfg['batch_size']
  342. self.det_predictor = Detector(
  343. model_dir, args.device, args.run_mode, batch_size,
  344. args.trt_min_shape, args.trt_max_shape, args.trt_opt_shape,
  345. args.trt_calib_mode, args.cpu_threads, args.enable_mkldnn)
  346. else:
  347. if self.with_idbased_detaction:
  348. idbased_detaction_cfg = self.cfg['ID_BASED_DETACTION']
  349. basemode = self.basemode['ID_BASED_DETACTION']
  350. self.modebase[basemode] = True
  351. self.det_action_predictor = DetActionRecognizer.init_with_cfg(
  352. args, idbased_detaction_cfg)
  353. self.det_action_visual_helper = ActionVisualHelper(1)
  354. if self.with_idbased_clsaction:
  355. idbased_clsaction_cfg = self.cfg['ID_BASED_CLSACTION']
  356. basemode = self.basemode['ID_BASED_CLSACTION']
  357. self.modebase[basemode] = True
  358. self.cls_action_predictor = ClsActionRecognizer.init_with_cfg(
  359. args, idbased_clsaction_cfg)
  360. self.cls_action_visual_helper = ActionVisualHelper(1)
  361. if self.with_skeleton_action:
  362. skeleton_action_cfg = self.cfg['SKELETON_ACTION']
  363. display_frames = skeleton_action_cfg['display_frames']
  364. self.coord_size = skeleton_action_cfg['coord_size']
  365. basemode = self.basemode['SKELETON_ACTION']
  366. self.modebase[basemode] = True
  367. skeleton_action_frames = skeleton_action_cfg['max_frames']
  368. self.skeleton_action_predictor = SkeletonActionRecognizer.init_with_cfg(
  369. args, skeleton_action_cfg)
  370. self.skeleton_action_visual_helper = ActionVisualHelper(
  371. display_frames)
  372. kpt_cfg = self.cfg['KPT']
  373. kpt_model_dir = kpt_cfg['model_dir']
  374. kpt_batch_size = kpt_cfg['batch_size']
  375. self.kpt_predictor = KeyPointDetector(
  376. kpt_model_dir,
  377. args.device,
  378. args.run_mode,
  379. kpt_batch_size,
  380. args.trt_min_shape,
  381. args.trt_max_shape,
  382. args.trt_opt_shape,
  383. args.trt_calib_mode,
  384. args.cpu_threads,
  385. args.enable_mkldnn,
  386. use_dark=False)
  387. self.kpt_buff = KeyPointBuff(skeleton_action_frames)
  388. if self.with_vehicleplate:
  389. vehicleplate_cfg = self.cfg['VEHICLE_PLATE']
  390. self.vehicleplate_detector = PlateRecognizer(args,
  391. vehicleplate_cfg)
  392. basemode = self.basemode['VEHICLE_PLATE']
  393. self.modebase[basemode] = True
  394. if self.with_mtmct:
  395. reid_cfg = self.cfg['REID']
  396. basemode = self.basemode['REID']
  397. self.modebase[basemode] = True
  398. self.reid_predictor = ReID.init_with_cfg(args, reid_cfg)
  399. if self.with_vehicle_retrograde:
  400. vehicleretrograde_cfg = self.cfg['VEHICLE_RETROGRADE']
  401. basemode = self.basemode['VEHICLE_RETROGRADE']
  402. self.modebase[basemode] = True
  403. self.vehicle_retrograde_predictor = VehicleRetrogradeRecognizer(
  404. vehicleretrograde_cfg)
  405. if self.with_mot or self.modebase["idbased"] or self.modebase[
  406. "skeletonbased"]:
  407. mot_cfg = self.cfg['MOT']
  408. model_dir = mot_cfg['model_dir']
  409. tracker_config = mot_cfg['tracker_config']
  410. batch_size = mot_cfg['batch_size']
  411. skip_frame_num = mot_cfg.get('skip_frame_num', -1)
  412. basemode = self.basemode['MOT']
  413. self.modebase[basemode] = True
  414. self.mot_predictor = SDE_Detector(
  415. model_dir,
  416. tracker_config,
  417. args.device,
  418. args.run_mode,
  419. batch_size,
  420. args.trt_min_shape,
  421. args.trt_max_shape,
  422. args.trt_opt_shape,
  423. args.trt_calib_mode,
  424. args.cpu_threads,
  425. args.enable_mkldnn,
  426. skip_frame_num=skip_frame_num,
  427. draw_center_traj=self.draw_center_traj,
  428. secs_interval=self.secs_interval,
  429. do_entrance_counting=self.do_entrance_counting,
  430. do_break_in_counting=self.do_break_in_counting,
  431. region_type=self.region_type,
  432. region_polygon=self.region_polygon)
  433. if self.with_video_action:
  434. video_action_cfg = self.cfg['VIDEO_ACTION']
  435. basemode = self.basemode['VIDEO_ACTION']
  436. self.modebase[basemode] = True
  437. self.video_action_predictor = VideoActionRecognizer.init_with_cfg(
  438. args, video_action_cfg)
  439. def set_file_name(self, path):
  440. if type(path) == int:
  441. self.file_name = path
  442. elif path is not None:
  443. self.file_name = os.path.split(path)[-1]
  444. if "." in self.file_name:
  445. self.file_name = self.file_name.split(".")[-2]
  446. else:
  447. # use camera id
  448. self.file_name = None
  449. def get_result(self):
  450. return self.collector.get_res()
  451. def run(self, input, thread_idx=0):
  452. if self.is_video:
  453. self.predict_video(input, thread_idx=thread_idx)
  454. else:
  455. self.predict_image(input)
  456. self.pipe_timer.info()
  457. self.mot_predictor.det_times.tracking_info(average=True)
  458. def predict_image(self, input):
  459. # det
  460. # det -> attr
  461. batch_loop_cnt = math.ceil(
  462. float(len(input)) / self.det_predictor.batch_size)
  463. self.warmup_frame = min(10, len(input) // 2) - 1
  464. for i in range(batch_loop_cnt):
  465. start_index = i * self.det_predictor.batch_size
  466. end_index = min((i + 1) * self.det_predictor.batch_size, len(input))
  467. batch_file = input[start_index:end_index]
  468. batch_input = [decode_image(f, {})[0] for f in batch_file]
  469. if i > self.warmup_frame:
  470. self.pipe_timer.total_time.start()
  471. self.pipe_timer.module_time['det'].start()
  472. # det output format: class, score, xmin, ymin, xmax, ymax
  473. det_res = self.det_predictor.predict_image(
  474. batch_input, visual=False)
  475. det_res = self.det_predictor.filter_box(det_res,
  476. self.cfg['crop_thresh'])
  477. if i > self.warmup_frame:
  478. self.pipe_timer.module_time['det'].end()
  479. self.pipe_timer.track_num += len(det_res['boxes'])
  480. self.pipeline_res.update(det_res, 'det')
  481. if self.with_human_attr:
  482. crop_inputs = crop_image_with_det(batch_input, det_res)
  483. attr_res_list = []
  484. if i > self.warmup_frame:
  485. self.pipe_timer.module_time['attr'].start()
  486. for crop_input in crop_inputs:
  487. attr_res = self.attr_predictor.predict_image(
  488. crop_input, visual=False)
  489. attr_res_list.extend(attr_res['output'])
  490. if i > self.warmup_frame:
  491. self.pipe_timer.module_time['attr'].end()
  492. attr_res = {'output': attr_res_list}
  493. self.pipeline_res.update(attr_res, 'attr')
  494. if self.with_vehicle_attr:
  495. crop_inputs = crop_image_with_det(batch_input, det_res)
  496. vehicle_attr_res_list = []
  497. if i > self.warmup_frame:
  498. self.pipe_timer.module_time['vehicle_attr'].start()
  499. for crop_input in crop_inputs:
  500. attr_res = self.vehicle_attr_predictor.predict_image(
  501. crop_input, visual=False)
  502. vehicle_attr_res_list.extend(attr_res['output'])
  503. if i > self.warmup_frame:
  504. self.pipe_timer.module_time['vehicle_attr'].end()
  505. attr_res = {'output': vehicle_attr_res_list}
  506. self.pipeline_res.update(attr_res, 'vehicle_attr')
  507. if self.with_vehicleplate:
  508. if i > self.warmup_frame:
  509. self.pipe_timer.module_time['vehicleplate'].start()
  510. crop_inputs = crop_image_with_det(batch_input, det_res)
  511. platelicenses = []
  512. for crop_input in crop_inputs:
  513. platelicense = self.vehicleplate_detector.get_platelicense(
  514. crop_input)
  515. platelicenses.extend(platelicense['plate'])
  516. if i > self.warmup_frame:
  517. self.pipe_timer.module_time['vehicleplate'].end()
  518. vehicleplate_res = {'vehicleplate': platelicenses}
  519. self.pipeline_res.update(vehicleplate_res, 'vehicleplate')
  520. if self.with_vehicle_press:
  521. vehicle_press_res_list = []
  522. if i > self.warmup_frame:
  523. self.pipe_timer.module_time['vehicle_press'].start()
  524. lanes, direction = self.laneseg_predictor.run(batch_input)
  525. if len(lanes) == 0:
  526. print(" no lanes!")
  527. continue
  528. lanes_res = {'output': lanes, 'direction': direction}
  529. self.pipeline_res.update(lanes_res, 'lanes')
  530. vehicle_press_res_list = self.vehicle_press_predictor.run(
  531. lanes, det_res)
  532. vehiclepress_res = {'output': vehicle_press_res_list}
  533. self.pipeline_res.update(vehiclepress_res, 'vehicle_press')
  534. self.pipe_timer.img_num += len(batch_input)
  535. if i > self.warmup_frame:
  536. self.pipe_timer.total_time.end()
  537. if self.cfg['visual']:
  538. self.visualize_image(batch_file, batch_input, self.pipeline_res)
  539. def capturevideo(self, capture, queue):
  540. frame_id = 0
  541. while (1):
  542. if queue.full():
  543. time.sleep(0.1)
  544. else:
  545. ret, frame = capture.read()
  546. if not ret:
  547. return
  548. frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
  549. queue.put(frame_rgb)
  550. def predict_video(self, video_file, thread_idx=0):
  551. # mot
  552. # mot -> attr
  553. # mot -> pose -> action
  554. capture = cv2.VideoCapture(video_file)
  555. # Get Video info : resolution, fps, frame count
  556. width = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH))
  557. height = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT))
  558. fps = int(capture.get(cv2.CAP_PROP_FPS))
  559. frame_count = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
  560. print("video fps: %d, frame_count: %d" % (fps, frame_count))
  561. if len(self.pushurl) > 0:
  562. video_out_name = 'output' if self.file_name is None else self.file_name
  563. pushurl = os.path.join(self.pushurl, video_out_name)
  564. print("the result will push stream to url:{}".format(pushurl))
  565. pushstream = PushStream(pushurl)
  566. pushstream.initcmd(fps, width, height)
  567. elif self.cfg['visual']:
  568. video_out_name = 'output' if (
  569. self.file_name is None or
  570. type(self.file_name) == int) else self.file_name
  571. if type(video_file) == str and "rtsp" in video_file:
  572. video_out_name = video_out_name + "_t" + str(thread_idx).zfill(
  573. 2) + "_rtsp"
  574. if not os.path.exists(self.output_dir):
  575. os.makedirs(self.output_dir)
  576. out_path = os.path.join(self.output_dir, video_out_name + ".mp4")
  577. fourcc = cv2.VideoWriter_fourcc(* 'mp4v')
  578. writer = cv2.VideoWriter(out_path, fourcc, fps, (width, height))
  579. frame_id = 0
  580. entrance, records, center_traj = None, None, None
  581. if self.draw_center_traj:
  582. center_traj = [{}]
  583. id_set = set()
  584. interval_id_set = set()
  585. in_id_list = list()
  586. out_id_list = list()
  587. prev_center = dict()
  588. records = list()
  589. if self.do_entrance_counting or self.do_break_in_counting or self.illegal_parking_time != -1:
  590. if self.region_type == 'horizontal':
  591. entrance = [0, height / 2., width, height / 2.]
  592. elif self.region_type == 'vertical':
  593. entrance = [width / 2, 0., width / 2, height]
  594. elif self.region_type == 'custom':
  595. entrance = []
  596. assert len(
  597. self.region_polygon
  598. ) % 2 == 0, "region_polygon should be pairs of coords points when do break_in counting."
  599. assert len(
  600. self.region_polygon
  601. ) > 6, 'region_type is custom, region_polygon should be at least 3 pairs of point coords.'
  602. for i in range(0, len(self.region_polygon), 2):
  603. entrance.append(
  604. [self.region_polygon[i], self.region_polygon[i + 1]])
  605. entrance.append([width, height])
  606. else:
  607. raise ValueError("region_type:{} unsupported.".format(
  608. self.region_type))
  609. video_fps = fps
  610. video_action_imgs = []
  611. if self.with_video_action:
  612. short_size = self.cfg["VIDEO_ACTION"]["short_size"]
  613. scale = ShortSizeScale(short_size)
  614. object_in_region_info = {
  615. } # store info for vehicle parking in region
  616. illegal_parking_dict = None
  617. cars_count = 0
  618. retrograde_traj_len = 0
  619. framequeue = queue.Queue(10)
  620. thread = threading.Thread(
  621. target=self.capturevideo, args=(capture, framequeue))
  622. thread.start()
  623. time.sleep(1)
  624. while (not framequeue.empty()):
  625. if frame_id % 10 == 0:
  626. print('Thread: {}; frame id: {}'.format(thread_idx, frame_id))
  627. frame_rgb = framequeue.get()
  628. if frame_id > self.warmup_frame:
  629. self.pipe_timer.total_time.start()
  630. if self.modebase["idbased"] or self.modebase["skeletonbased"]:
  631. if frame_id > self.warmup_frame:
  632. self.pipe_timer.module_time['mot'].start()
  633. mot_skip_frame_num = self.mot_predictor.skip_frame_num
  634. reuse_det_result = False
  635. if mot_skip_frame_num > 1 and frame_id > 0 and frame_id % mot_skip_frame_num > 0:
  636. reuse_det_result = True
  637. res = self.mot_predictor.predict_image(
  638. [copy.deepcopy(frame_rgb)],
  639. visual=False,
  640. reuse_det_result=reuse_det_result,
  641. frame_count=frame_id)
  642. # mot output format: id, class, score, xmin, ymin, xmax, ymax
  643. mot_res = parse_mot_res(res)
  644. if frame_id > self.warmup_frame:
  645. self.pipe_timer.module_time['mot'].end()
  646. self.pipe_timer.track_num += len(mot_res['boxes'])
  647. if frame_id % 10 == 0:
  648. print("Thread: {}; trackid number: {}".format(
  649. thread_idx, len(mot_res['boxes'])))
  650. # flow_statistic only support single class MOT
  651. boxes, scores, ids = res[0] # batch size = 1 in MOT
  652. mot_result = (frame_id + 1, boxes[0], scores[0],
  653. ids[0]) # single class
  654. statistic = flow_statistic(
  655. mot_result,
  656. self.secs_interval,
  657. self.do_entrance_counting,
  658. self.do_break_in_counting,
  659. self.region_type,
  660. video_fps,
  661. entrance,
  662. id_set,
  663. interval_id_set,
  664. in_id_list,
  665. out_id_list,
  666. prev_center,
  667. records,
  668. ids2names=self.mot_predictor.pred_config.labels)
  669. records = statistic['records']
  670. if self.illegal_parking_time != -1:
  671. object_in_region_info, illegal_parking_dict = update_object_info(
  672. object_in_region_info, mot_result, self.region_type,
  673. entrance, video_fps, self.illegal_parking_time)
  674. if len(illegal_parking_dict) != 0:
  675. # build relationship between id and plate
  676. for key, value in illegal_parking_dict.items():
  677. plate = self.collector.get_carlp(key)
  678. illegal_parking_dict[key]['plate'] = plate
  679. # nothing detected
  680. if len(mot_res['boxes']) == 0:
  681. frame_id += 1
  682. if frame_id > self.warmup_frame:
  683. self.pipe_timer.img_num += 1
  684. self.pipe_timer.total_time.end()
  685. if self.cfg['visual']:
  686. _, _, fps = self.pipe_timer.get_total_time()
  687. im = self.visualize_video(frame_rgb, mot_res, frame_id,
  688. fps, entrance, records,
  689. center_traj) # visualize
  690. if len(self.pushurl) > 0:
  691. pushstream.pipe.stdin.write(im.tobytes())
  692. else:
  693. writer.write(im)
  694. if self.file_name is None: # use camera_id
  695. cv2.imshow('Paddle-Pipeline', im)
  696. if cv2.waitKey(1) & 0xFF == ord('q'):
  697. break
  698. continue
  699. self.pipeline_res.update(mot_res, 'mot')
  700. crop_input, new_bboxes, ori_bboxes = crop_image_with_mot(
  701. frame_rgb, mot_res)
  702. if self.with_vehicleplate and frame_id % 10 == 0:
  703. if frame_id > self.warmup_frame:
  704. self.pipe_timer.module_time['vehicleplate'].start()
  705. plate_input, _, _ = crop_image_with_mot(
  706. frame_rgb, mot_res, expand=False)
  707. platelicense = self.vehicleplate_detector.get_platelicense(
  708. plate_input)
  709. if frame_id > self.warmup_frame:
  710. self.pipe_timer.module_time['vehicleplate'].end()
  711. self.pipeline_res.update(platelicense, 'vehicleplate')
  712. else:
  713. self.pipeline_res.clear('vehicleplate')
  714. if self.with_human_attr:
  715. if frame_id > self.warmup_frame:
  716. self.pipe_timer.module_time['attr'].start()
  717. attr_res = self.attr_predictor.predict_image(
  718. crop_input, visual=False)
  719. if frame_id > self.warmup_frame:
  720. self.pipe_timer.module_time['attr'].end()
  721. self.pipeline_res.update(attr_res, 'attr')
  722. if self.with_vehicle_attr:
  723. if frame_id > self.warmup_frame:
  724. self.pipe_timer.module_time['vehicle_attr'].start()
  725. attr_res = self.vehicle_attr_predictor.predict_image(
  726. crop_input, visual=False)
  727. if frame_id > self.warmup_frame:
  728. self.pipe_timer.module_time['vehicle_attr'].end()
  729. self.pipeline_res.update(attr_res, 'vehicle_attr')
  730. if self.with_vehicle_press or self.with_vehicle_retrograde:
  731. if frame_id == 0 or cars_count == 0 or cars_count > len(
  732. mot_res['boxes']):
  733. if frame_id > self.warmup_frame:
  734. self.pipe_timer.module_time['lanes'].start()
  735. lanes, directions = self.laneseg_predictor.run(
  736. [copy.deepcopy(frame_rgb)])
  737. lanes_res = {'output': lanes, 'directions': directions}
  738. if frame_id > self.warmup_frame:
  739. self.pipe_timer.module_time['lanes'].end()
  740. if frame_id == 0 or (len(lanes) > 0 and frame_id > 0):
  741. self.pipeline_res.update(lanes_res, 'lanes')
  742. cars_count = len(mot_res['boxes'])
  743. if self.with_vehicle_press:
  744. if frame_id > self.warmup_frame:
  745. self.pipe_timer.module_time['vehicle_press'].start()
  746. press_lane = copy.deepcopy(self.pipeline_res.get('lanes'))
  747. if press_lane is None:
  748. continue
  749. vehicle_press_res_list = self.vehicle_press_predictor.mot_run(
  750. press_lane, mot_res['boxes'])
  751. vehiclepress_res = {'output': vehicle_press_res_list}
  752. if frame_id > self.warmup_frame:
  753. self.pipe_timer.module_time['vehicle_press'].end()
  754. self.pipeline_res.update(vehiclepress_res, 'vehicle_press')
  755. if self.with_idbased_detaction:
  756. if frame_id > self.warmup_frame:
  757. self.pipe_timer.module_time['det_action'].start()
  758. det_action_res = self.det_action_predictor.predict(
  759. crop_input, mot_res)
  760. if frame_id > self.warmup_frame:
  761. self.pipe_timer.module_time['det_action'].end()
  762. self.pipeline_res.update(det_action_res, 'det_action')
  763. if self.cfg['visual']:
  764. self.det_action_visual_helper.update(det_action_res)
  765. if self.with_idbased_clsaction:
  766. if frame_id > self.warmup_frame:
  767. self.pipe_timer.module_time['cls_action'].start()
  768. cls_action_res = self.cls_action_predictor.predict_with_mot(
  769. crop_input, mot_res)
  770. if frame_id > self.warmup_frame:
  771. self.pipe_timer.module_time['cls_action'].end()
  772. self.pipeline_res.update(cls_action_res, 'cls_action')
  773. if self.cfg['visual']:
  774. self.cls_action_visual_helper.update(cls_action_res)
  775. if self.with_skeleton_action:
  776. if frame_id > self.warmup_frame:
  777. self.pipe_timer.module_time['kpt'].start()
  778. kpt_pred = self.kpt_predictor.predict_image(
  779. crop_input, visual=False)
  780. keypoint_vector, score_vector = translate_to_ori_images(
  781. kpt_pred, np.array(new_bboxes))
  782. kpt_res = {}
  783. kpt_res['keypoint'] = [
  784. keypoint_vector.tolist(), score_vector.tolist()
  785. ] if len(keypoint_vector) > 0 else [[], []]
  786. kpt_res['bbox'] = ori_bboxes
  787. if frame_id > self.warmup_frame:
  788. self.pipe_timer.module_time['kpt'].end()
  789. self.pipeline_res.update(kpt_res, 'kpt')
  790. self.kpt_buff.update(kpt_res, mot_res) # collect kpt output
  791. state = self.kpt_buff.get_state(
  792. ) # whether frame num is enough or lost tracker
  793. skeleton_action_res = {}
  794. if state:
  795. if frame_id > self.warmup_frame:
  796. self.pipe_timer.module_time[
  797. 'skeleton_action'].start()
  798. collected_keypoint = self.kpt_buff.get_collected_keypoint(
  799. ) # reoragnize kpt output with ID
  800. skeleton_action_input = parse_mot_keypoint(
  801. collected_keypoint, self.coord_size)
  802. skeleton_action_res = self.skeleton_action_predictor.predict_skeleton_with_mot(
  803. skeleton_action_input)
  804. if frame_id > self.warmup_frame:
  805. self.pipe_timer.module_time['skeleton_action'].end()
  806. self.pipeline_res.update(skeleton_action_res,
  807. 'skeleton_action')
  808. if self.cfg['visual']:
  809. self.skeleton_action_visual_helper.update(
  810. skeleton_action_res)
  811. if self.with_mtmct and frame_id % 10 == 0:
  812. crop_input, img_qualities, rects = self.reid_predictor.crop_image_with_mot(
  813. frame_rgb, mot_res)
  814. if frame_id > self.warmup_frame:
  815. self.pipe_timer.module_time['reid'].start()
  816. reid_res = self.reid_predictor.predict_batch(crop_input)
  817. if frame_id > self.warmup_frame:
  818. self.pipe_timer.module_time['reid'].end()
  819. reid_res_dict = {
  820. 'features': reid_res,
  821. "qualities": img_qualities,
  822. "rects": rects
  823. }
  824. self.pipeline_res.update(reid_res_dict, 'reid')
  825. else:
  826. self.pipeline_res.clear('reid')
  827. if self.with_video_action:
  828. # get the params
  829. frame_len = self.cfg["VIDEO_ACTION"]["frame_len"]
  830. sample_freq = self.cfg["VIDEO_ACTION"]["sample_freq"]
  831. if sample_freq * frame_len > frame_count: # video is too short
  832. sample_freq = int(frame_count / frame_len)
  833. # filter the warmup frames
  834. if frame_id > self.warmup_frame:
  835. self.pipe_timer.module_time['video_action'].start()
  836. # collect frames
  837. if frame_id % sample_freq == 0:
  838. # Scale image
  839. scaled_img = scale(frame_rgb)
  840. video_action_imgs.append(scaled_img)
  841. # the number of collected frames is enough to predict video action
  842. if len(video_action_imgs) == frame_len:
  843. classes, scores = self.video_action_predictor.predict(
  844. video_action_imgs)
  845. if frame_id > self.warmup_frame:
  846. self.pipe_timer.module_time['video_action'].end()
  847. video_action_res = {"class": classes[0], "score": scores[0]}
  848. self.pipeline_res.update(video_action_res, 'video_action')
  849. print("video_action_res:", video_action_res)
  850. video_action_imgs.clear() # next clip
  851. if self.with_vehicle_retrograde:
  852. # get the params
  853. frame_len = self.cfg["VEHICLE_RETROGRADE"]["frame_len"]
  854. sample_freq = self.cfg["VEHICLE_RETROGRADE"]["sample_freq"]
  855. if sample_freq * frame_len > frame_count: # video is too short
  856. sample_freq = int(frame_count / frame_len)
  857. # filter the warmup frames
  858. if frame_id > self.warmup_frame:
  859. self.pipe_timer.module_time['vehicle_retrograde'].start()
  860. if frame_id % sample_freq == 0:
  861. frame_mot_res = copy.deepcopy(self.pipeline_res.get('mot'))
  862. self.vehicle_retrograde_predictor.update_center_traj(
  863. frame_mot_res, max_len=frame_len)
  864. retrograde_traj_len = retrograde_traj_len + 1
  865. #the number of collected frames is enough to predict
  866. if retrograde_traj_len == frame_len:
  867. retrograde_mot_res = copy.deepcopy(
  868. self.pipeline_res.get('mot'))
  869. retrograde_lanes = copy.deepcopy(
  870. self.pipeline_res.get('lanes'))
  871. frame_shape = frame_rgb.shape
  872. if retrograde_lanes is None:
  873. continue
  874. retrograde_res, fence_line = self.vehicle_retrograde_predictor.mot_run(
  875. lanes_res=retrograde_lanes,
  876. det_res=retrograde_mot_res,
  877. frame_shape=frame_shape)
  878. retrograde_res_update = self.pipeline_res.get(
  879. 'vehicle_retrograde')
  880. if retrograde_res_update is not None:
  881. retrograde_res_update = retrograde_res_update['output']
  882. if retrograde_res is not None:
  883. for retrograde_res_id in retrograde_res:
  884. if retrograde_res_id not in retrograde_res_update:
  885. retrograde_res_update.append(
  886. retrograde_res_id)
  887. else:
  888. retrograde_res_update = []
  889. retrograde_res_dict = {
  890. 'output': retrograde_res_update,
  891. "fence_line": fence_line,
  892. }
  893. if retrograde_res is not None and len(retrograde_res) > 0:
  894. print("retrograde res:", retrograde_res)
  895. self.pipeline_res.update(retrograde_res_dict,
  896. 'vehicle_retrograde')
  897. if frame_id > self.warmup_frame:
  898. self.pipe_timer.module_time['vehicle_retrograde'].end()
  899. retrograde_traj_len = 0
  900. self.collector.append(frame_id, self.pipeline_res)
  901. if frame_id > self.warmup_frame:
  902. self.pipe_timer.img_num += 1
  903. self.pipe_timer.total_time.end()
  904. frame_id += 1
  905. if self.cfg['visual']:
  906. _, _, fps = self.pipe_timer.get_total_time()
  907. im = self.visualize_video(frame_rgb, self.pipeline_res,
  908. self.collector, frame_id, fps,
  909. entrance, records, center_traj,
  910. self.illegal_parking_time != -1,
  911. illegal_parking_dict) # visualize
  912. if len(self.pushurl) > 0:
  913. pushstream.pipe.stdin.write(im.tobytes())
  914. else:
  915. writer.write(im)
  916. if self.file_name is None: # use camera_id
  917. cv2.imshow('Paddle-Pipeline', im)
  918. if cv2.waitKey(1) & 0xFF == ord('q'):
  919. break
  920. if self.cfg['visual'] and len(self.pushurl) == 0:
  921. writer.release()
  922. print('save result to {}'.format(out_path))
  923. def visualize_video(self,
  924. image_rgb,
  925. result,
  926. collector,
  927. frame_id,
  928. fps,
  929. entrance=None,
  930. records=None,
  931. center_traj=None,
  932. do_illegal_parking_recognition=False,
  933. illegal_parking_dict=None):
  934. image = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR)
  935. mot_res = copy.deepcopy(result.get('mot'))
  936. if mot_res is not None:
  937. ids = mot_res['boxes'][:, 0]
  938. scores = mot_res['boxes'][:, 2]
  939. boxes = mot_res['boxes'][:, 3:]
  940. boxes[:, 2] = boxes[:, 2] - boxes[:, 0]
  941. boxes[:, 3] = boxes[:, 3] - boxes[:, 1]
  942. else:
  943. boxes = np.zeros([0, 4])
  944. ids = np.zeros([0])
  945. scores = np.zeros([0])
  946. # single class, still need to be defaultdict type for ploting
  947. num_classes = 1
  948. online_tlwhs = defaultdict(list)
  949. online_scores = defaultdict(list)
  950. online_ids = defaultdict(list)
  951. online_tlwhs[0] = boxes
  952. online_scores[0] = scores
  953. online_ids[0] = ids
  954. if mot_res is not None:
  955. image = plot_tracking_dict(
  956. image,
  957. num_classes,
  958. online_tlwhs,
  959. online_ids,
  960. online_scores,
  961. frame_id=frame_id,
  962. fps=fps,
  963. ids2names=self.mot_predictor.pred_config.labels,
  964. do_entrance_counting=self.do_entrance_counting,
  965. do_break_in_counting=self.do_break_in_counting,
  966. do_illegal_parking_recognition=do_illegal_parking_recognition,
  967. illegal_parking_dict=illegal_parking_dict,
  968. entrance=entrance,
  969. records=records,
  970. center_traj=center_traj)
  971. human_attr_res = result.get('attr')
  972. if human_attr_res is not None:
  973. boxes = mot_res['boxes'][:, 1:]
  974. human_attr_res = human_attr_res['output']
  975. image = visualize_attr(image, human_attr_res, boxes)
  976. image = np.array(image)
  977. vehicle_attr_res = result.get('vehicle_attr')
  978. if vehicle_attr_res is not None:
  979. boxes = mot_res['boxes'][:, 1:]
  980. vehicle_attr_res = vehicle_attr_res['output']
  981. image = visualize_attr(image, vehicle_attr_res, boxes)
  982. image = np.array(image)
  983. lanes_res = result.get('lanes')
  984. if lanes_res is not None:
  985. lanes = lanes_res['output'][0]
  986. image = visualize_lane(image, lanes)
  987. image = np.array(image)
  988. vehiclepress_res = result.get('vehicle_press')
  989. if vehiclepress_res is not None:
  990. press_vehicle = vehiclepress_res['output']
  991. if len(press_vehicle) > 0:
  992. image = visualize_vehiclepress(
  993. image, press_vehicle, threshold=self.cfg['crop_thresh'])
  994. image = np.array(image)
  995. if mot_res is not None:
  996. vehicleplate = False
  997. plates = []
  998. for trackid in mot_res['boxes'][:, 0]:
  999. plate = collector.get_carlp(trackid)
  1000. if plate != None:
  1001. vehicleplate = True
  1002. plates.append(plate)
  1003. else:
  1004. plates.append("")
  1005. if vehicleplate:
  1006. boxes = mot_res['boxes'][:, 1:]
  1007. image = visualize_vehicleplate(image, plates, boxes)
  1008. image = np.array(image)
  1009. kpt_res = result.get('kpt')
  1010. if kpt_res is not None:
  1011. image = visualize_pose(
  1012. image,
  1013. kpt_res,
  1014. visual_thresh=self.cfg['kpt_thresh'],
  1015. returnimg=True)
  1016. video_action_res = result.get('video_action')
  1017. if video_action_res is not None:
  1018. video_action_score = None
  1019. if video_action_res and video_action_res["class"] == 1:
  1020. video_action_score = video_action_res["score"]
  1021. mot_boxes = None
  1022. if mot_res:
  1023. mot_boxes = mot_res['boxes']
  1024. image = visualize_action(
  1025. image,
  1026. mot_boxes,
  1027. action_visual_collector=None,
  1028. action_text="SkeletonAction",
  1029. video_action_score=video_action_score,
  1030. video_action_text="Fight")
  1031. vehicle_retrograde_res = result.get('vehicle_retrograde')
  1032. if vehicle_retrograde_res is not None:
  1033. mot_retrograde_res = copy.deepcopy(result.get('mot'))
  1034. image = visualize_vehicle_retrograde(image, mot_retrograde_res,
  1035. vehicle_retrograde_res)
  1036. image = np.array(image)
  1037. visual_helper_for_display = []
  1038. action_to_display = []
  1039. skeleton_action_res = result.get('skeleton_action')
  1040. if skeleton_action_res is not None:
  1041. visual_helper_for_display.append(self.skeleton_action_visual_helper)
  1042. action_to_display.append("Falling")
  1043. det_action_res = result.get('det_action')
  1044. if det_action_res is not None:
  1045. visual_helper_for_display.append(self.det_action_visual_helper)
  1046. action_to_display.append("Smoking")
  1047. cls_action_res = result.get('cls_action')
  1048. if cls_action_res is not None:
  1049. visual_helper_for_display.append(self.cls_action_visual_helper)
  1050. action_to_display.append("Calling")
  1051. if len(visual_helper_for_display) > 0:
  1052. image = visualize_action(image, mot_res['boxes'],
  1053. visual_helper_for_display,
  1054. action_to_display)
  1055. return image
  1056. def visualize_image(self, im_files, images, result):
  1057. start_idx, boxes_num_i = 0, 0
  1058. det_res = result.get('det')
  1059. human_attr_res = result.get('attr')
  1060. vehicle_attr_res = result.get('vehicle_attr')
  1061. vehicleplate_res = result.get('vehicleplate')
  1062. lanes_res = result.get('lanes')
  1063. vehiclepress_res = result.get('vehicle_press')
  1064. for i, (im_file, im) in enumerate(zip(im_files, images)):
  1065. if det_res is not None:
  1066. det_res_i = {}
  1067. boxes_num_i = det_res['boxes_num'][i]
  1068. det_res_i['boxes'] = det_res['boxes'][start_idx:start_idx +
  1069. boxes_num_i, :]
  1070. im = visualize_box_mask(
  1071. im,
  1072. det_res_i,
  1073. labels=['target'],
  1074. threshold=self.cfg['crop_thresh'])
  1075. im = np.ascontiguousarray(np.copy(im))
  1076. im = cv2.cvtColor(im, cv2.COLOR_RGB2BGR)
  1077. if human_attr_res is not None:
  1078. human_attr_res_i = human_attr_res['output'][start_idx:start_idx
  1079. + boxes_num_i]
  1080. im = visualize_attr(im, human_attr_res_i, det_res_i['boxes'])
  1081. if vehicle_attr_res is not None:
  1082. vehicle_attr_res_i = vehicle_attr_res['output'][
  1083. start_idx:start_idx + boxes_num_i]
  1084. im = visualize_attr(im, vehicle_attr_res_i, det_res_i['boxes'])
  1085. if vehicleplate_res is not None:
  1086. plates = vehicleplate_res['vehicleplate']
  1087. det_res_i['boxes'][:, 4:6] = det_res_i[
  1088. 'boxes'][:, 4:6] - det_res_i['boxes'][:, 2:4]
  1089. im = visualize_vehicleplate(im, plates, det_res_i['boxes'])
  1090. if vehiclepress_res is not None:
  1091. press_vehicle = vehiclepress_res['output'][i]
  1092. if len(press_vehicle) > 0:
  1093. im = visualize_vehiclepress(
  1094. im, press_vehicle, threshold=self.cfg['crop_thresh'])
  1095. im = np.ascontiguousarray(np.copy(im))
  1096. if lanes_res is not None:
  1097. lanes = lanes_res['output'][i]
  1098. im = visualize_lane(im, lanes)
  1099. im = np.ascontiguousarray(np.copy(im))
  1100. img_name = os.path.split(im_file)[-1]
  1101. if not os.path.exists(self.output_dir):
  1102. os.makedirs(self.output_dir)
  1103. out_path = os.path.join(self.output_dir, img_name)
  1104. cv2.imwrite(out_path, im)
  1105. print("save result to: " + out_path)
  1106. start_idx += boxes_num_i
  1107. def main():
  1108. cfg = merge_cfg(FLAGS) # use command params to update config
  1109. print_arguments(cfg)
  1110. pipeline = Pipeline(FLAGS, cfg)
  1111. # pipeline.run()
  1112. pipeline.run_multithreads()
  1113. if __name__ == '__main__':
  1114. paddle.enable_static()
  1115. # parse params from command
  1116. parser = argsparser()
  1117. FLAGS = parser.parse_args()
  1118. FLAGS.device = FLAGS.device.upper()
  1119. assert FLAGS.device in ['CPU', 'GPU', 'XPU'
  1120. ], "device should be CPU, GPU or XPU"
  1121. main()