提交 b9679975 编写于 作者: M mindspore-ci-bot 提交者: Gitee

!2858 JPEG decoder hangs if the image is sliced to 3 bytes

Merge pull request !2858 from h.farahat/slice_jpeg
...@@ -121,14 +121,14 @@ Status Resize(const std::shared_ptr<Tensor> &input, std::shared_ptr<Tensor> *out ...@@ -121,14 +121,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"; const unsigned char *kJpegMagic = (unsigned char *)"\xFF\xD8\xFF";
constexpr size_t kJpegMagicLen = 3; 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) { Status Decode(const std::shared_ptr<Tensor> &input, std::shared_ptr<Tensor> *output) {
if (HasJpegMagic(input)) { if (IsNonEmptyJPEG(input)) {
return JpegCropAndDecode(input, output); return JpegCropAndDecode(input, output);
} else { } else {
return DecodeCv(input, output); return DecodeCv(input, output);
......
...@@ -96,7 +96,7 @@ Status Decode(const std::shared_ptr<Tensor> &input, std::shared_ptr<Tensor> *out ...@@ -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); 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); void JpegSetSource(j_decompress_ptr c_info, const void *data, int64_t data_size);
......
...@@ -31,7 +31,7 @@ Status RandomCropDecodeResizeOp::Compute(const std::shared_ptr<Tensor> &input, s ...@@ -31,7 +31,7 @@ Status RandomCropDecodeResizeOp::Compute(const std::shared_ptr<Tensor> &input, s
if (input == nullptr) { if (input == nullptr) {
RETURN_STATUS_UNEXPECTED("input tensor is null"); RETURN_STATUS_UNEXPECTED("input tensor is null");
} }
if (!HasJpegMagic(input)) { if (!IsNonEmptyJPEG(input)) {
DecodeOp op(true); DecodeOp op(true);
std::shared_ptr<Tensor> decoded; std::shared_ptr<Tensor> decoded;
RETURN_IF_NOT_OK(op.Compute(input, &decoded)); RETURN_IF_NOT_OK(op.Compute(input, &decoded));
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册