| 25 | } |
| 26 | |
| 27 | def decoder(addr, vipc_server, vst, nvidia, W, H, debug=False): |
| 28 | sock_name = ENCODE_SOCKETS[vst] |
| 29 | if debug: |
| 30 | print(f"start decoder for {sock_name}, {W}x{H}") |
| 31 | |
| 32 | if nvidia: |
| 33 | os.environ["NV_LOW_LATENCY"] = "3" # both bLowLatency and CUVID_PKT_ENDOFPICTURE |
| 34 | sys.path += os.environ["LD_LIBRARY_PATH"].split(":") |
| 35 | import PyNvCodec as nvc |
| 36 | |
| 37 | nvDec = nvc.PyNvDecoder(W, H, nvc.PixelFormat.NV12, nvc.CudaVideoCodec.HEVC, 0) |
| 38 | cc1 = nvc.ColorspaceConversionContext(nvc.ColorSpace.BT_709, nvc.ColorRange.JPEG) |
| 39 | conv_yuv = nvc.PySurfaceConverter(W, H, nvc.PixelFormat.NV12, nvc.PixelFormat.YUV420, 0) |
| 40 | nvDwn_yuv = nvc.PySurfaceDownloader(W, H, nvc.PixelFormat.YUV420, 0) |
| 41 | img_yuv = np.ndarray((H*W//2*3), dtype=np.uint8) |
| 42 | else: |
| 43 | codec = av.CodecContext.create("hevc", "r") |
| 44 | |
| 45 | os.environ["ZMQ"] = "1" |
| 46 | messaging.reset_context() |
| 47 | sock = messaging.sub_sock(sock_name, None, addr=addr, conflate=False) |
| 48 | cnt = 0 |
| 49 | last_idx = -1 |
| 50 | seen_iframe = False |
| 51 | |
| 52 | time_q = [] |
| 53 | while 1: |
| 54 | msgs = messaging.drain_sock(sock, wait_for_one=True) |
| 55 | for evt in msgs: |
| 56 | evta = getattr(evt, evt.which()) |
| 57 | if debug and evta.idx.encodeId != 0 and evta.idx.encodeId != (last_idx+1): |
| 58 | print("DROP PACKET!") |
| 59 | last_idx = evta.idx.encodeId |
| 60 | if not seen_iframe and not (evta.idx.flags & V4L2_BUF_FLAG_KEYFRAME): |
| 61 | if debug: |
| 62 | print("waiting for iframe") |
| 63 | continue |
| 64 | time_q.append(time.monotonic()) |
| 65 | network_latency = (int(time.time()*1e9) - evta.unixTimestampNanos)/1e6 # noqa: TID251 |
| 66 | frame_latency = ((evta.idx.timestampEof/1e9) - (evta.idx.timestampSof/1e9))*1000 |
| 67 | process_latency = ((evt.logMonoTime/1e9) - (evta.idx.timestampEof/1e9))*1000 |
| 68 | |
| 69 | # put in header (first) |
| 70 | if not seen_iframe: |
| 71 | if nvidia: |
| 72 | nvDec.DecodeSurfaceFromPacket(np.frombuffer(evta.header, dtype=np.uint8)) |
| 73 | else: |
| 74 | codec.decode(av.packet.Packet(evta.header)) |
| 75 | seen_iframe = True |
| 76 | |
| 77 | if nvidia: |
| 78 | rawSurface = nvDec.DecodeSurfaceFromPacket(np.frombuffer(evta.data, dtype=np.uint8)) |
| 79 | if rawSurface.Empty(): |
| 80 | if debug: |
| 81 | print("DROP SURFACE") |
| 82 | continue |
| 83 | convSurface = conv_yuv.Execute(rawSurface, cc1) |
| 84 | nvDwn_yuv.DownloadSingleSurface(convSurface, img_yuv) |