Add batch support for VideoDataset and support reading from streams (#295)

* Add batch support for VideoDataset

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

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

Fixes 144

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

* Address review feedback

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

Author: yongtang

tensorflow_io/video/BUILD

@@ -6,70 +6,61 @@ cc_binary(
     name = "python/ops/_video_ops_ffmpeg_3.4.so",
     srcs = [
         "kernels/ffmpeg.cc",
-        "kernels/video_dataset_ops.cc",
+        "kernels/video_input.cc",
         "kernels/video_reader.h",
-        "ops/dataset_ops.cc",
+        "ops/video_ops.cc",
     ],
     copts = [
         "-pthread",
         "-std=c++11",
         "-DNDEBUG",
     ],
-    data = [],
     includes = ["."],
-    linkopts = [],
     linkshared = 1,
     deps = [
+        "//tensorflow_io/core:dataset_ops",
         "@ffmpeg_3_4//:ffmpeg",
-        "@local_config_tf//:libtensorflow_framework",
-        "@local_config_tf//:tf_header_lib",
     ],
 )
 
 cc_binary(
     name = "python/ops/_video_ops_ffmpeg_2.8.so",
     srcs = [
         "kernels/ffmpeg.cc",
-        "kernels/video_dataset_ops.cc",
+        "kernels/video_input.cc",
         "kernels/video_reader.h",
-        "ops/dataset_ops.cc",
+        "ops/video_ops.cc",
     ],
     copts = [
         "-pthread",
         "-std=c++11",
         "-DNDEBUG",
     ],
-    data = [],
     includes = ["."],
-    linkopts = [],
     linkshared = 1,
     deps = [
+        "//tensorflow_io/core:dataset_ops",
         "@ffmpeg_2_8//:ffmpeg",
-        "@local_config_tf//:libtensorflow_framework",
-        "@local_config_tf//:tf_header_lib",
     ],
 )
 
 cc_binary(
     name = "python/ops/_video_ops_libav_9.20.so",
     srcs = [
         "kernels/ffmpeg.cc",
-        "kernels/video_dataset_ops.cc",
+        "kernels/video_input.cc",
         "kernels/video_reader.h",
-        "ops/dataset_ops.cc",
+        "ops/video_ops.cc",
     ],
     copts = [
         "-pthread",
         "-std=c++11",
         "-DNDEBUG",
     ],
-    data = [],
     includes = ["."],
-    linkopts = [],
     linkshared = 1,
     deps = [
+        "//tensorflow_io/core:dataset_ops",
         "@libav_9_20//:libav",
-        "@local_config_tf//:libtensorflow_framework",
-        "@local_config_tf//:tf_header_lib",
     ],
 )

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