Support read video from stream (instead of file name).

Fixes 144

Signed-off-by: Yong Tang <[email protected]>

Author: yongtang

tensorflow_io/video/BUILD

@@ -18,8 +18,8 @@ cc_binary(
     includes = ["."],
     linkshared = 1,
     deps = [
+        "//tensorflow_io/core:dataset_ops",
         "@ffmpeg_3_4//:ffmpeg",
-         "//tensorflow_io/core:dataset_ops",
     ],
 )
 
@@ -39,8 +39,8 @@ cc_binary(
     includes = ["."],
     linkshared = 1,
     deps = [
+        "//tensorflow_io/core:dataset_ops",
         "@ffmpeg_2_8//:ffmpeg",
-         "//tensorflow_io/core:dataset_ops",
     ],
 )
 
@@ -60,7 +60,7 @@ cc_binary(
     includes = ["."],
     linkshared = 1,
     deps = [
+        "//tensorflow_io/core:dataset_ops",
         "@libav_9_20//:libav",
-         "//tensorflow_io/core:dataset_ops",
     ],
 )

tensorflow_io/video/kernels/ffmpeg.cc

@@ -34,8 +34,63 @@ namespace tensorflow {
 namespace data {
 namespace video {
 
+static int io_read_packet(void *opaque, uint8_t *buf, int buf_size) {
+  VideoReader *r = (VideoReader *)opaque;
+  StringPiece result;
+  Status status = r->stream_->Read(r->offset_, buf_size, &result, (char *)buf);
+  if (!(status.ok() || errors::IsOutOfRange(status))) {
+    return -1;
+  }
+  r->offset_ += result.size();
+  return result.size();
+}
+
+static int64_t io_seek(void *opaque, int64_t offset, int whence) {
+  VideoReader *r = (VideoReader *)opaque;
+  uint64 file_size = 0;
+  Status status = r->stream_->GetFileSize(&file_size);
+  if (!status.ok()) {
+    return -1;
+  }
+  switch (whence)
+  {
+  case SEEK_SET:
+    if (offset > file_size) {
+      return -1;
+    }
+    r->offset_ = offset;
+    return r->offset_;
+  case SEEK_CUR:
+    if (r->offset_ + offset > file_size) {
+      return -1;
+    }
+    r->offset_ += offset;
+    return r->offset_;
+  case SEEK_END:
+    if (offset > file_size) {
+      return -1;
+    }
+    r->offset_ = file_size - offset;
+    return r->offset_;
+  case AVSEEK_SIZE:
+    return file_size;
+  default:
+    break;
+  }
+  return -1;
+}
+
 Status VideoReader::ReadHeader()
 {
+    // Allocate format
+    if ((format_context_ = avformat_alloc_context()) == NULL) {
+      return errors::InvalidArgument("could not allocate format context");
+    }
+    // Allocate context
+    if ((io_context_ = avio_alloc_context(NULL, 0, 0, this, io_read_packet, NULL, io_seek)) == NULL) {
+      return errors::InvalidArgument("could not allocate io context");
+    }
+    format_context_->pb = io_context_;
     // Open input file, and allocate format context
     if (avformat_open_input(&format_context_, filename_.c_str(), NULL, NULL) < 0) {
       return errors::InvalidArgument("could not open video file: ", filename_);
@@ -206,6 +261,10 @@ VideoReader::~VideoReader() {
     avcodec_free_context(&codec_context_);
 #endif
     avformat_close_input(&format_context_);
+    av_free(format_context_);
+    if (io_context_ != NULL) {
+      av_free(io_context_);
+    }
 }
 
 }  // namespace

tensorflow_io/video/kernels/video_input.cc

@@ -13,7 +13,6 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include "kernels/dataset_ops.h"
 #include "kernels/video_reader.h"
 
 namespace tensorflow {
@@ -35,7 +34,7 @@ class VideoInput: public FileInput<video::VideoReader> {
   Status ReadRecord(io::InputStreamInterface* s, IteratorContext* ctx, std::unique_ptr<video::VideoReader>& state, int64 record_to_read, int64* record_read, std::vector<Tensor>* out_tensors) const override {
     if (state.get() == nullptr) {
       VideoReaderInit();
-      state.reset(new video::VideoReader(filename()));
+      state.reset(new video::VideoReader(dynamic_cast<SizedRandomAccessInputStreamInterface*>(s), filename()));
       TF_RETURN_IF_ERROR(state.get()->ReadHeader());
     }
     // Read the first frame to get height and width

tensorflow_io/video/kernels/video_reader.h

@@ -16,6 +16,7 @@ limitations under the License.
 #include "tensorflow/core/framework/dataset.h"
 #include "tensorflow/core/lib/io/buffered_inputstream.h"
 #include "tensorflow/core/platform/file_system.h"
+#include "kernels/dataset_ops.h"
 
 extern "C" {
 
@@ -33,7 +34,7 @@ namespace video {
 
 class VideoReader {
  public:
-  explicit VideoReader(const string &filename) : filename_(filename) {}
+  explicit VideoReader(SizedRandomAccessInputStreamInterface* s, const string& filename) : stream_(s), filename_(filename) {}
 
   Status ReadHeader();
 
@@ -43,6 +44,9 @@ class VideoReader {
 
   virtual ~VideoReader();
 
+ public:
+  SizedRandomAccessInputStreamInterface* stream_;
+  int64 offset_ = 0;
  private:
   std::string ahead_;
   std::string filename_;
@@ -58,6 +62,7 @@ class VideoReader {
   AVCodecContext *codec_context_ = 0;
   AVFrame *frame_ = 0;
   AVPacket packet_;
+  AVIOContext *io_context_ = NULL;
   TF_DISALLOW_COPY_AND_ASSIGN(VideoReader);
 };
 

tests/test_video.py

@@ -28,7 +28,7 @@
 
 video_path = os.path.join(
     os.path.dirname(os.path.abspath(__file__)), "test_video", "small.mp4")
-
+video_path = "file://" + video_path
 def test_video_predict():
   model = tf.keras.applications.resnet50.ResNet50(weights='imagenet')
   x = video_io.VideoDataset(video_path, batch=1).map(lambda x: tf.keras.applications.resnet50.preprocess_input(tf.image.resize(x, (224, 224))))

tests/test_video_eager.py

@@ -30,7 +30,7 @@
 
 video_path = os.path.join(
     os.path.dirname(os.path.abspath(__file__)), "test_video", "small.mp4")
-
+video_path = "file://" + video_path
 def test_video_dataset():
   """test_video_dataset"""
   num_repeats = 2