bool Decoder::init()

in torchvision/csrc/io/decoder/decoder.cpp [227:406]

• 151 lines of code
• 31 McCabe index (conditional complexity)

bool Decoder::init(
    const DecoderParameters& params,
    DecoderInCallback&& in,
    std::vector<DecoderMetadata>* metadata) {
  cleanUp();

  if ((params.uri.empty() || in) && (!params.uri.empty() || !in)) {
    LOG(ERROR)
        << "uuid=" << params_.loggingUuid
        << " either external URI gets provided or explicit input callback";
    return false;
  }

  // set callback and params
  params_ = params;

  if (!(inputCtx_ = avformat_alloc_context())) {
    LOG(ERROR) << "uuid=" << params_.loggingUuid
               << " cannot allocate format context";
    return false;
  }

  AVInputFormat* fmt = nullptr;
  int result = 0;
  if (in) {
    ImageType type = ImageType::UNKNOWN;
    if ((result = seekableBuffer_.init(
             std::forward<DecoderInCallback>(in),
             params_.timeoutMs,
             params_.maxSeekableBytes,
             params_.isImage ? &type : nullptr)) < 0) {
      LOG(ERROR) << "uuid=" << params_.loggingUuid
                 << " can't initiate seekable buffer";
      cleanUp();
      return false;
    }

    if (params_.isImage) {
      const char* fmtName = "image2";
      switch (type) {
        case ImageType::JPEG:
          fmtName = "jpeg_pipe";
          break;
        case ImageType::PNG:
          fmtName = "png_pipe";
          break;
        case ImageType::TIFF:
          fmtName = "tiff_pipe";
          break;
        default:
          break;
      }

      fmt = av_find_input_format(fmtName);
    }

    const size_t avioCtxBufferSize = kIoBufferSize;
    uint8_t* avioCtxBuffer =
        (uint8_t*)av_malloc(avioCtxBufferSize + kIoPaddingSize);
    if (!avioCtxBuffer) {
      LOG(ERROR) << "uuid=" << params_.loggingUuid
                 << " av_malloc cannot allocate " << avioCtxBufferSize
                 << " bytes";
      cleanUp();
      return false;
    }

    if (!(avioCtx_ = avio_alloc_context(
              avioCtxBuffer,
              avioCtxBufferSize,
              0,
              reinterpret_cast<void*>(this),
              &Decoder::readFunction,
              nullptr,
              result == 1 ? &Decoder::seekFunction : nullptr))) {
      LOG(ERROR) << "uuid=" << params_.loggingUuid
                 << " avio_alloc_context failed";
      av_free(avioCtxBuffer);
      cleanUp();
      return false;
    }

    inputCtx_->pb = avioCtx_;
    inputCtx_->flags |= AVFMT_FLAG_CUSTOM_IO;
  }

  inputCtx_->opaque = reinterpret_cast<void*>(this);
  inputCtx_->interrupt_callback.callback = Decoder::shutdownFunction;
  inputCtx_->interrupt_callback.opaque = reinterpret_cast<void*>(this);

  // add network timeout
  inputCtx_->flags |= AVFMT_FLAG_NONBLOCK;

  AVDictionary* options = nullptr;
  if (params_.listen) {
    av_dict_set_int(&options, "listen", 1, 0);
  }
  if (params_.timeoutMs > 0) {
    av_dict_set_int(&options, "analyzeduration", params_.timeoutMs * 1000, 0);
    av_dict_set_int(&options, "stimeout", params_.timeoutMs * 1000, 0);
    av_dict_set_int(&options, "rw_timeout", params_.timeoutMs * 1000, 0);
    if (!params_.tlsCertFile.empty()) {
      av_dict_set(&options, "cert_file", params_.tlsCertFile.data(), 0);
    }
    if (!params_.tlsKeyFile.empty()) {
      av_dict_set(&options, "key_file", params_.tlsKeyFile.data(), 0);
    }
  }

  interrupted_ = false;

  // ffmpeg avformat_open_input call can hang if media source doesn't respond
  // set a guard for handle such situations, if requested
  std::promise<bool> p;
  std::future<bool> f = p.get_future();
  std::unique_ptr<std::thread> guard;
  if (params_.preventStaleness) {
    guard = std::make_unique<std::thread>([&f, this]() {
      auto timeout = std::chrono::milliseconds(params_.timeoutMs);
      if (std::future_status::timeout == f.wait_for(timeout)) {
        LOG(ERROR) << "uuid=" << params_.loggingUuid
                   << " cannot open stream within " << params_.timeoutMs
                   << " ms";
        interrupted_ = true;
      }
    });
  }

  if (fmt) {
    result = avformat_open_input(&inputCtx_, nullptr, fmt, &options);
  } else {
    result =
        avformat_open_input(&inputCtx_, params_.uri.c_str(), nullptr, &options);
  }

  av_dict_free(&options);

  if (guard) {
    p.set_value(true);
    guard->join();
    guard.reset();
  }

  if (result < 0 || interrupted_) {
    LOG(ERROR) << "uuid=" << params_.loggingUuid
               << " avformat_open_input failed, error="
               << Util::generateErrorDesc(result);
    cleanUp();
    return false;
  }

  result = avformat_find_stream_info(inputCtx_, nullptr);

  if (result < 0) {
    LOG(ERROR) << "uuid=" << params_.loggingUuid
               << " avformat_find_stream_info failed, error="
               << Util::generateErrorDesc(result);
    cleanUp();
    return false;
  }

  if (!openStreams(metadata)) {
    LOG(ERROR) << "uuid=" << params_.loggingUuid << " cannot activate streams";
    cleanUp();
    return false;
  }
  // SyncDecoder inherits Decoder which would override onInit.
  onInit();

  if (params.startOffset != 0) {
    auto offset = params.startOffset <= params.seekAccuracy
        ? 0
        : params.startOffset - params.seekAccuracy;

    av_seek_frame(inputCtx_, -1, offset, AVSEEK_FLAG_BACKWARD);
  }

  VLOG(1) << "Decoder initialized, log level: " << params_.logLevel;
  return true;
}