Method Compute

tensorflow/core/kernels/encode_jpeg_op.cc:77–129

  }

  void Compute(OpKernelContext* context) override {
    const Tensor& image = context->input(0);
    OP_REQUIRES(context, image.dims() == 3,
                errors::InvalidArgument("image must be 3-dimensional",
                                        image.shape().DebugString()));

    OP_REQUIRES(
        context,
        FastBoundsCheck(image.NumElements(), std::numeric_limits<int32>::max()),
        errors::InvalidArgument(
            "Cannot encode images with >= max int32 elements"));

    const int32 dim_size0 = static_cast<int32>(image.dim_size(0));
    const int32 dim_size1 = static_cast<int32>(image.dim_size(1));
    const int32 dim_size2 = static_cast<int32>(image.dim_size(2));

    // Autodetect format if desired, otherwise make sure format and
    // image channels are consistent.
    int channels;
    jpeg::CompressFlags adjusted_flags = flags_;
    if (flags_.format == 0) {
      channels = dim_size2;
      if (channels == 1) {
        adjusted_flags.format = jpeg::FORMAT_GRAYSCALE;
      } else if (channels == 3) {
        adjusted_flags.format = jpeg::FORMAT_RGB;
      } else {
        OP_REQUIRES(
            context, false,
            errors::InvalidArgument("image must have 1 or 3 channels, got ",
                                    image.shape().DebugString()));
      }
    } else {
      if (flags_.format == jpeg::FORMAT_GRAYSCALE) {
        channels = 1;
      } else {  // RGB
        channels = 3;
      }
      OP_REQUIRES(context, channels == dim_size2,
                  errors::InvalidArgument("format ", format_, " expects ",
                                          channels, " channels, got ",
                                          image.shape().DebugString()));
    }

    // Encode image to jpeg string
    Tensor* output = nullptr;
    OP_REQUIRES_OK(context,
                   context->allocate_output(0, TensorShape({}), &output));
    OP_REQUIRES(context,
                jpeg::Compress(image.flat<uint8>().data(), dim_size1, dim_size0,
                               adjusted_flags, &output->scalar<string>()()),
                errors::Internal("JPEG encoding failed"));
  }

 private:
  string format_;