From c45bce4e8f0d7e82dfda88c90596512649d2899b Mon Sep 17 00:00:00 2001
From: hesham <h.farahat@huawei.com>
Date: Fri, 3 Jul 2020 23:26:36 -0400
Subject: [PATCH] JPEG decoder hangs if the image is sliced to 3 bytes only.

---
 mindspore/ccsrc/dataset/kernels/image/image_utils.cc        | 6 +++---
 mindspore/ccsrc/dataset/kernels/image/image_utils.h         | 2 +-
 .../dataset/kernels/image/random_crop_decode_resize_op.cc   | 2 +-
 3 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/mindspore/ccsrc/dataset/kernels/image/image_utils.cc b/mindspore/ccsrc/dataset/kernels/image/image_utils.cc
index 9a20743e6c..28c4240d21 100644
--- a/mindspore/ccsrc/dataset/kernels/image/image_utils.cc
+++ b/mindspore/ccsrc/dataset/kernels/image/image_utils.cc
@@ -120,14 +120,14 @@ Status Resize(const std::shared_ptr<Tensor> &input, std::shared_ptr<Tensor> *out
   }
 }
 
-bool HasJpegMagic(const std::shared_ptr<Tensor> &input) {
+bool IsNonEmptyJPEG(const std::shared_ptr<Tensor> &input) {
   const unsigned char *kJpegMagic = (unsigned char *)"\xFF\xD8\xFF";
   constexpr size_t kJpegMagicLen = 3;
-  return input->SizeInBytes() >= kJpegMagicLen && memcmp(input->GetBuffer(), kJpegMagic, kJpegMagicLen) == 0;
+  return input->SizeInBytes() > kJpegMagicLen && memcmp(input->GetBuffer(), kJpegMagic, kJpegMagicLen) == 0;
 }
 
 Status Decode(const std::shared_ptr<Tensor> &input, std::shared_ptr<Tensor> *output) {
-  if (HasJpegMagic(input)) {
+  if (IsNonEmptyJPEG(input)) {
     return JpegCropAndDecode(input, output);
   } else {
     return DecodeCv(input, output);
diff --git a/mindspore/ccsrc/dataset/kernels/image/image_utils.h b/mindspore/ccsrc/dataset/kernels/image/image_utils.h
index 57ffce6a12..6ebd13ad55 100644
--- a/mindspore/ccsrc/dataset/kernels/image/image_utils.h
+++ b/mindspore/ccsrc/dataset/kernels/image/image_utils.h
@@ -96,7 +96,7 @@ Status Decode(const std::shared_ptr<Tensor> &input, std::shared_ptr<Tensor> *out
 
 Status DecodeCv(const std::shared_ptr<Tensor> &input, std::shared_ptr<Tensor> *output);
 
-bool HasJpegMagic(const std::shared_ptr<Tensor> &input);
+bool IsNonEmptyJPEG(const std::shared_ptr<Tensor> &input);
 
 void JpegSetSource(j_decompress_ptr c_info, const void *data, int64_t data_size);
 
diff --git a/mindspore/ccsrc/dataset/kernels/image/random_crop_decode_resize_op.cc b/mindspore/ccsrc/dataset/kernels/image/random_crop_decode_resize_op.cc
index 74aa91ea7e..36d80aea98 100644
--- a/mindspore/ccsrc/dataset/kernels/image/random_crop_decode_resize_op.cc
+++ b/mindspore/ccsrc/dataset/kernels/image/random_crop_decode_resize_op.cc
@@ -31,7 +31,7 @@ Status RandomCropDecodeResizeOp::Compute(const std::shared_ptr<Tensor> &input, s
   if (input == nullptr) {
     RETURN_STATUS_UNEXPECTED("input tensor is null");
   }
-  if (!HasJpegMagic(input)) {
+  if (!IsNonEmptyJPEG(input)) {
     DecodeOp op(true);
     std::shared_ptr<Tensor> decoded;
     RETURN_IF_NOT_OK(op.Compute(input, &decoded));