From 1c3874c690b27d3aad85aabe6b6e8b074f8b2f08 Mon Sep 17 00:00:00 2001
From: Chen Liu <cliu7@meta.com>
Date: Thu, 3 Nov 2022 13:40:02 -0700
Subject: [PATCH] Set pytorch vision decoder probesize for getting stream info
 based on the value from decode setting (#6900)

Summary:
Pull Request resolved: https://github.com/pytorch/vision/pull/6900

1. Pass along the value of probesize from decode setting.
2. allow pytorch vision decoder to set the probesize for getting stream info

Reviewed By: jdsgomes

Differential Revision: D40967782

fbshipit-source-id: 642ae3b67c50b1fdafc20f9814a52fdf5db527d7
---
 torchvision/csrc/io/decoder/decoder.cpp | 2 ++
 torchvision/csrc/io/decoder/defs.h      | 6 ++++++
 2 files changed, 8 insertions(+)

diff --git a/torchvision/csrc/io/decoder/decoder.cpp b/torchvision/csrc/io/decoder/decoder.cpp
index f13e2c3ffcf..a895eed2d39 100644
--- a/torchvision/csrc/io/decoder/decoder.cpp
+++ b/torchvision/csrc/io/decoder/decoder.cpp
@@ -312,6 +312,8 @@ bool Decoder::init(
     }
   }
 
+  av_dict_set_int(&options, "probesize", params_.probeSize, 0);
+
   interrupted_ = false;
 
   // ffmpeg avformat_open_input call can hang if media source doesn't respond
diff --git a/torchvision/csrc/io/decoder/defs.h b/torchvision/csrc/io/decoder/defs.h
index 502e5762e46..ba82625c309 100644
--- a/torchvision/csrc/io/decoder/defs.h
+++ b/torchvision/csrc/io/decoder/defs.h
@@ -213,6 +213,12 @@ struct DecoderParameters {
 
   // Skip packets that fail with EPERM errors and continue decoding.
   bool skipOperationNotPermittedPackets{false};
+
+  // probing size in bytes, i.e. the size of the data to analyze to get stream
+  // information. A higher value will enable detecting more information in case
+  // it is dispersed into the stream, but will increase latency. Must be an
+  // integer not lesser than 32. It is 5000000 by default.
+  int64_t probeSize{5000000};
 };
 
 struct DecoderHeader {