deepstream-no-osd-queue-streammux-nvinfer-nvtracker.py 15 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387
  1. # Import required libraries
  2. import argparse
  3. import sys
  4. import os
  5. sys.path.append('../')
  6. sys.path.append('../source_code')
  7. import gi
  8. import configparser
  9. gi.require_version('Gst', '1.0')
  10. from gi.repository import GObject, Gst
  11. from gi.repository import GLib
  12. from ctypes import *
  13. import time
  14. import sys
  15. import math
  16. import platform
  17. from common.bus_call import bus_call
  18. from common.FPS import GETFPS
  19. import pyds
  20. PGIE_CLASS_ID_VEHICLE = 0
  21. PGIE_CLASS_ID_BICYCLE = 1
  22. PGIE_CLASS_ID_PERSON = 2
  23. PGIE_CLASS_ID_ROADSIGN = 3
  24. g_args=None
  25. # Define variables to be used later
  26. fps_streams_new={}
  27. pgie_classes_str= ["Vehicle", "TwoWheeler", "Person","RoadSign"]
  28. ################ Three Stream Pipeline ###########
  29. # Define Input and output Stream information
  30. INPUT_VIDEO = 'file:///opt/nvidia/deepstream/deepstream-5.0/samples/streams/sample_720p.h264'
  31. def parse_args():
  32. parser = argparse.ArgumentParser()
  33. parser.add_argument("--num-sources", type=int, default=1, help="Number of sources, it replicates inputs if its is greater than length of inputs")
  34. parser.add_argument("--sgie-batch-size", type=int, default=1, help="Batch size for Seconday inference")
  35. parser.add_argument("--prof", type=bool, default=False, help="Profiling Mode , Profiles with a shorter video clip.")
  36. args = parser.parse_args()
  37. return args
  38. def cb_newpad(decodebin, decoder_src_pad,data):
  39. print("In cb_newpad\n")
  40. caps=decoder_src_pad.get_current_caps()
  41. gststruct=caps.get_structure(0)
  42. gstname=gststruct.get_name()
  43. source_bin=data
  44. features=caps.get_features(0)
  45. # Need to check if the pad created by the decodebin is for video and not
  46. # audio.
  47. print("gstname=",gstname)
  48. if(gstname.find("video")!=-1):
  49. # Link the decodebin pad only if decodebin has picked nvidia
  50. # decoder plugin nvdec_*. We do this by checking if the pad caps contain
  51. # NVMM memory features.
  52. print("features=",features)
  53. if features.contains("memory:NVMM"):
  54. # Get the source bin ghost pad
  55. bin_ghost_pad=source_bin.get_static_pad("src")
  56. if not bin_ghost_pad.set_target(decoder_src_pad):
  57. sys.stderr.write("Failed to link decoder src pad to source bin ghost pad\n")
  58. else:
  59. sys.stderr.write(" Error: Decodebin did not pick nvidia decoder plugin.\n")
  60. def decodebin_child_added(child_proxy,Object,name,user_data):
  61. print("Decodebin child added:", name, "\n")
  62. if(name.find("decodebin") != -1):
  63. Object.connect("child-added",decodebin_child_added,user_data)
  64. def create_source_bin(args, index,uri):
  65. print("Creating source bin")
  66. # Create a source GstBin to abstract this bin's content from the rest of the
  67. # pipeline
  68. bin_name="source-bin-%02d" %index
  69. print(bin_name)
  70. nbin=Gst.Bin.new(bin_name)
  71. if not nbin:
  72. sys.stderr.write(" Unable to create source bin \n")
  73. # Source element for reading from the uri.
  74. # We will use decodebin and let it figure out the container format of the
  75. # stream and the codec and plug the appropriate demux and decode plugins.
  76. uri_decode_bin=Gst.ElementFactory.make("uridecodebin", "uri-decode-bin")
  77. if not uri_decode_bin:
  78. sys.stderr.write(" Unable to create uri decode bin \n")
  79. # We set the input uri to the source element
  80. uri_decode_bin.set_property("uri",uri)
  81. # Connect to the "pad-added" signal of the decodebin which generates a
  82. # callback once a new pad for raw data has beed created by the decodebin
  83. uri_decode_bin.connect("pad-added",cb_newpad,nbin)
  84. uri_decode_bin.connect("child-added",decodebin_child_added,nbin)
  85. # We need to create a ghost pad for the source bin which will act as a proxy
  86. # for the video decoder src pad. The ghost pad will not have a target right
  87. # now. Once the decode bin creates the video decoder and generates the
  88. # cb_newpad callback, we will set the ghost pad target to the video decoder
  89. # src pad.
  90. Gst.Bin.add(nbin,uri_decode_bin)
  91. bin_pad=nbin.add_pad(Gst.GhostPad.new_no_target("src",Gst.PadDirection.SRC))
  92. if not bin_pad:
  93. sys.stderr.write(" Failed to add ghost pad in source bin \n")
  94. return None
  95. return nbin
  96. ## Make Element or Print Error and any other detail
  97. def make_elm_or_print_err(factoryname, name, printedname, detail=""):
  98. print("Creating", printedname)
  99. elm = Gst.ElementFactory.make(factoryname, name)
  100. if not elm:
  101. sys.stderr.write("Unable to create " + printedname + " \n")
  102. if detail:
  103. sys.stderr.write(detail)
  104. return elm
  105. # src_pad_buffer_probe
  106. def src_pad_buffer_probe(pad,info,u_data):
  107. #Intiallizing object counter with 0.
  108. obj_counter = {
  109. PGIE_CLASS_ID_VEHICLE:0,
  110. PGIE_CLASS_ID_PERSON:0,
  111. PGIE_CLASS_ID_BICYCLE:0,
  112. PGIE_CLASS_ID_ROADSIGN:0
  113. }
  114. # Set frame_number & rectangles to draw as 0
  115. frame_number=0
  116. num_rects=0
  117. gst_buffer = info.get_buffer()
  118. if not gst_buffer:
  119. print("Unable to get GstBuffer ")
  120. return
  121. # Retrieve batch metadata from the gst_buffer
  122. # Note that pyds.gst_buffer_get_nvds_batch_meta() expects the
  123. # C address of gst_buffer as input, which is obtained with hash(gst_buffer)
  124. batch_meta = pyds.gst_buffer_get_nvds_batch_meta(hash(gst_buffer))
  125. l_frame = batch_meta.frame_meta_list
  126. while l_frame is not None:
  127. try:
  128. # Note that l_frame.data needs a cast to pyds.NvDsFrameMeta
  129. frame_meta = pyds.NvDsFrameMeta.cast(l_frame.data)
  130. except StopIteration:
  131. break
  132. # Get frame number , number of rectables to draw and object metadata
  133. frame_number=frame_meta.frame_num
  134. num_rects = frame_meta.num_obj_meta
  135. l_obj=frame_meta.obj_meta_list
  136. while l_obj is not None:
  137. try:
  138. # Casting l_obj.data to pyds.NvDsObjectMeta
  139. obj_meta=pyds.NvDsObjectMeta.cast(l_obj.data)
  140. except StopIteration:
  141. break
  142. # Increment Object class by 1
  143. obj_counter[obj_meta.class_id] += 1
  144. try:
  145. l_obj=l_obj.next
  146. except StopIteration:
  147. break
  148. print("Frame Number={} Number of Objects={} Vehicle_count={} Person_count={}".format(frame_number, num_rects, obj_counter[PGIE_CLASS_ID_VEHICLE], obj_counter[PGIE_CLASS_ID_PERSON]))
  149. # FPS Probe
  150. fps_streams_new["stream{0}".format(frame_meta.pad_index)].get_fps()
  151. try:
  152. l_frame=l_frame.next
  153. except StopIteration:
  154. break
  155. return Gst.PadProbeReturn.OK
  156. def main():
  157. args = parse_args()
  158. global g_args
  159. g_args = args
  160. num_sources = args.num_sources
  161. sgie_batch_size = args.sgie_batch_size
  162. path = os.path.abspath(os.getcwd())
  163. if (args.prof):
  164. INPUT_VIDEO = 'file://' + path +'/../source_code/dataset/sample_720p_prof.mp4'
  165. else :
  166. INPUT_VIDEO = 'file:///opt/nvidia/deepstream/deepstream-5.0/samples/streams/sample_720p.h264'
  167. print("Creating pipeline with "+str(num_sources)+" streams")
  168. # Initialise FPS
  169. for i in range(0,num_sources):
  170. fps_streams_new["stream{0}".format(i)]=GETFPS(i)
  171. # Standard GStreamer initialization
  172. Gst.init(None)
  173. # Create gstreamer elements */
  174. # Create Pipeline element that will form a connection of other elements
  175. print("Creating Pipeline \n ")
  176. pipeline = Gst.Pipeline()
  177. if not pipeline:
  178. sys.stderr.write(" Unable to create Pipeline \n")
  179. ########### Create Elements required for the Pipeline ###########
  180. # Create nvstreammux instance to form batches from one or more sources.
  181. streammux = make_elm_or_print_err("nvstreammux", "Stream-muxer","Stream-muxer")
  182. pipeline.add(streammux)
  183. for i in range(num_sources):
  184. print("Creating source_bin ",i," \n ")
  185. uri_name=INPUT_VIDEO
  186. if uri_name.find("rtsp://") == 0 :
  187. is_live = True
  188. source_bin=create_source_bin(args, i, uri_name)
  189. if not source_bin:
  190. sys.stderr.write("Unable to create source bin \n")
  191. pipeline.add(source_bin)
  192. padname="sink_%u" %i
  193. sinkpad = streammux.get_request_pad(padname)
  194. if not sinkpad:
  195. sys.stderr.write("Unable to create sink pad bin \n")
  196. srcpad = source_bin.get_static_pad("src")
  197. if not srcpad:
  198. sys.stderr.write("Unable to create src pad bin \n")
  199. srcpad.link(sinkpad)
  200. # Use nvinfer to run inferencing on decoder's output, behaviour of inferencing is set through config file
  201. pgie = make_elm_or_print_err("nvinfer", "primary-inference" ,"pgie")
  202. # Use nvtracker to give objects unique-ids
  203. tracker = make_elm_or_print_err("nvtracker", "tracker",'tracker')
  204. # Seconday inference for Finding Car Color
  205. sgie1 = make_elm_or_print_err("nvinfer", "secondary1-nvinference-engine",'sgie1')
  206. # Seconday inference for Finding Car Make
  207. sgie2 = make_elm_or_print_err("nvinfer", "secondary2-nvinference-engine",'sgie2')
  208. # # Seconday inference for Finding Car Type
  209. sgie3 = make_elm_or_print_err("nvinfer", "secondary3-nvinference-engine",'sgie3')
  210. # Create Sink for storing the output
  211. fakesink = make_elm_or_print_err("fakesink", "fakesink", "Sink")
  212. # Queues to enable buffering
  213. queue1=make_elm_or_print_err("queue","queue1","queue1")
  214. queue2=make_elm_or_print_err("queue","queue2","queue2")
  215. queue3=make_elm_or_print_err("queue","queue3","queue3")
  216. queue4=make_elm_or_print_err("queue","queue4","queue4")
  217. queue5=make_elm_or_print_err("queue","queue5","queue5")
  218. queue6=make_elm_or_print_err("queue","queue6","queue6")
  219. ############ Set properties for the Elements ############
  220. # Set Input Width , Height and Batch Size
  221. streammux.set_property('width', 1280)
  222. streammux.set_property('height', 720)
  223. streammux.set_property('batch-size', num_sources)
  224. # Timeout in microseconds to wait after the first buffer is available
  225. # to push the batch even if a complete batch is not formed.
  226. streammux.set_property('batched-push-timeout', 4000000)
  227. # Set configuration file for nvinfer
  228. pgie.set_property('config-file-path', "../source_code/N2/dstest4_pgie_config.txt")
  229. sgie1.set_property('config-file-path', "../source_code/N2/dstest4_sgie1_config.txt")
  230. sgie2.set_property('config-file-path', "../source_code/N2/dstest4_sgie2_config.txt")
  231. sgie3.set_property('config-file-path', "../source_code/N2/dstest4_sgie3_config.txt")
  232. # Setting batch_size for pgie
  233. pgie_batch_size=pgie.get_property("batch-size")
  234. if(pgie_batch_size != num_sources):
  235. print("WARNING: Overriding infer-config batch-size",pgie_batch_size," with number of sources ", num_sources," \n")
  236. pgie.set_property("batch-size",num_sources)
  237. #################### Secondary Batch size ######################
  238. # Setting batch_size for sgie1
  239. sgie1_batch_size=sgie1.get_property("batch-size")
  240. if(sgie1_batch_size != sgie_batch_size):
  241. print("WARNING: Overriding infer-config batch-size",sgie1_batch_size," with number of sources ", sgie_batch_size," \n")
  242. sgie1.set_property("batch-size",sgie_batch_size)
  243. # Setting batch_size for sgie2
  244. sgie2_batch_size=sgie2.get_property("batch-size")
  245. if(sgie2_batch_size != sgie_batch_size):
  246. print("WARNING: Overriding infer-config batch-size",sgie2_batch_size," with number of sources ", sgie_batch_size," \n")
  247. sgie2.set_property("batch-size",sgie_batch_size)
  248. # Setting batch_size for sgie2
  249. sgie3_batch_size=sgie3.get_property("batch-size")
  250. if(sgie3_batch_size != sgie_batch_size):
  251. print("WARNING: Overriding infer-config batch-size",sgie3_batch_size," with number of sources ", sgie_batch_size," \n")
  252. sgie3.set_property("batch-size",sgie_batch_size)
  253. #Set properties of tracker from tracker_config
  254. config = configparser.ConfigParser()
  255. config.read('../source_code/N1/dstest4_tracker_config.txt')
  256. config.sections()
  257. for key in config['tracker']:
  258. if key == 'tracker-width' :
  259. tracker_width = config.getint('tracker', key)
  260. tracker.set_property('tracker-width', tracker_width)
  261. if key == 'tracker-height' :
  262. tracker_height = config.getint('tracker', key)
  263. tracker.set_property('tracker-height', tracker_height)
  264. if key == 'gpu-id' :
  265. tracker_gpu_id = config.getint('tracker', key)
  266. tracker.set_property('gpu_id', tracker_gpu_id)
  267. if key == 'll-lib-file' :
  268. tracker_ll_lib_file = config.get('tracker', key)
  269. tracker.set_property('ll-lib-file', tracker_ll_lib_file)
  270. if key == 'll-config-file' :
  271. tracker_ll_config_file = config.get('tracker', key)
  272. tracker.set_property('ll-config-file', tracker_ll_config_file)
  273. if key == 'enable-batch-process' :
  274. tracker_enable_batch_process = config.getint('tracker', key)
  275. tracker.set_property('enable_batch_process', tracker_enable_batch_process)
  276. # Fake sink properties
  277. fakesink.set_property("sync", 0)
  278. fakesink.set_property("async", 0)
  279. ########## Add and Link ELements in the Pipeline ##########
  280. print("Adding elements to Pipeline \n")
  281. pipeline.add(queue1)
  282. pipeline.add(pgie)
  283. pipeline.add(queue2)
  284. pipeline.add(tracker)
  285. pipeline.add(queue3)
  286. pipeline.add(sgie1)
  287. pipeline.add(queue4)
  288. pipeline.add(sgie2)
  289. pipeline.add(queue5)
  290. pipeline.add(sgie3)
  291. pipeline.add(queue6)
  292. pipeline.add(fakesink)
  293. print("Linking elements in the Pipeline \n")
  294. streammux.link(queue1)
  295. queue1.link(pgie)
  296. pgie.link(queue2)
  297. queue2.link(tracker)
  298. tracker.link(queue3)
  299. queue3.link(sgie1)
  300. sgie1.link(queue4)
  301. queue4.link(sgie2)
  302. sgie2.link(queue5)
  303. queue5.link(sgie3)
  304. sgie3.link(queue6)
  305. queue6.link(fakesink)
  306. # create an event loop and feed gstreamer bus mesages to it
  307. loop = GLib.MainLoop()
  308. bus = pipeline.get_bus()
  309. bus.add_signal_watch()
  310. bus.connect ("message", bus_call, loop)
  311. print("Added and Linked elements to pipeline")
  312. src_pad=queue6.get_static_pad("src")
  313. if not src_pad:
  314. sys.stderr.write(" Unable to get src pad \n")
  315. else:
  316. src_pad.add_probe(Gst.PadProbeType.BUFFER, src_pad_buffer_probe, 0)
  317. # List the sources
  318. print("Now playing...")
  319. print("Starting pipeline \n")
  320. # start play back and listed to events
  321. pipeline.set_state(Gst.State.PLAYING)
  322. start_time = time.time()
  323. try:
  324. loop.run()
  325. except:
  326. pass
  327. # cleanup
  328. print("Exiting app\n")
  329. pipeline.set_state(Gst.State.NULL)
  330. print("--- %s seconds ---" % (time.time() - start_time))
  331. if __name__ == '__main__':
  332. sys.exit(main())