Trouble with YUV420 to float tensor conversion

laganojunior: <p>I’m trying to use imageYUV420CenterCropToFloat32Tensor to convert an image in a YUV420 format, and I’m getting odd results. If I just try to convert the RGB tensor back directly to YUV, I get artifacts where the image seems interleaved with the image duplicated 4 times with some color changes.</p>
<p>I’m probably misunderstanding something about YUV420 encoding. Looking at the code here (<a href="https://github.com/pytorch/pytorch/blob/master/android/pytorch_android_torchvision/src/main/cpp/pytorch_vision_jni.cpp#L104" class="inline-onebox" rel="noopener nofollow ugc">pytorch/pytorch_vision_jni.cpp at master · pytorch/pytorch · GitHub</a>). The y dimension seems to be halved twice, once in the (yBeforeRtn &gt;&gt; 1) and another time in a previous line (int uvRowStride = uRowStride &gt;&gt; 1). However, the x dimension isn’t halved at all.</p>
<p>My understanding is that both x and y dimensions are halved in 420 subsampling, so this struck me as odd. For fun, if I just take my YUV image and read the UV indices as this function does and then write back to the UV buffers as I expected the format to be, I get similar artifacts.</p>
<p>E.x:</p>
<pre><code class="lang-auto">for (i in 0 until (imageHeight shr 1)) {
  for (j in 0 until imageWidth) {
    uBuffer.put((j / 2) * uPixelStride + i * uRowStride, uBufferOrig!!.get(j * uPixelStride + i * uRowStride / 2))
    vBuffer.put((j / 2) * vPixelStride + i * vRowStride, vBufferOrig!!.get(j * vPixelStride + i * vRowStride / 2))
  }
}
</code></pre>

ptrblck: <p>The “interleaved” artifacts could point towards a layout mismatch (channels-first vs. channels-last) or a <code>view</code> operations, while a <code>permute</code> op would be necessary. Could you check, if the memory layout matches the expected one in the code?</p>

laganojunior: <p>Yeah, it took me a while to reason about what expects channels-first/channels-last. I eventually got this working by copying the c++ code from the pytorch torchvision utils file and modifying it to be the way I expected (by reducing the halving of the vertical dimension to  only once, and then adding a halving of the horizontal dimension)</p>
<p>I can isolate the problem away from any sort of ordering of RGB channel order or even reading the RGB output. I think that the read access pattern for the U/V buffer seems off.</p>
<p>Here’s the code from the pytorch repo: <a href="https://github.com/pytorch/pytorch/blob/104b2c6/android/pytorch_android_torchvision/src/main/cpp/pytorch_vision_jni.cpp" rel="noopener nofollow ugc">https://github.com/pytorch/pytorch/blob/104b2c6/android/pytorch_android_torchvision/src/main/cpp/pytorch_vision_jni.cpp</a>. I just added two lines to set 0s in the U/V buffer to indices as they’re being read: <a href="https://gist.github.com/laganojunior/75f39be0c6843e20b235c49f4aa8ee22#file-mod_pytorch_vision_jni-cpp-L105-L106" class="inline-onebox" rel="noopener nofollow ugc">activity.kt · GitHub</a></p>
<p>On the android side, I’m just calling the function on an image. I don’t even try to use its output: <a href="https://gist.github.com/laganojunior/75f39be0c6843e20b235c49f4aa8ee22#file-activity-kt" class="inline-onebox" rel="noopener nofollow ugc">activity.kt · GitHub</a></p>
<p>This image is what I get on my phone (Pixel 3XL)<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/6/66631fc06b43a46e35a5c5052f9e504109e8a6a4.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/66631fc06b43a46e35a5c5052f9e504109e8a6a4" title="DSC01592"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/6/6/66631fc06b43a46e35a5c5052f9e504109e8a6a4_2_690x388.jpeg" alt="DSC01592" data-base62-sha1="eBL4Bnx7167y0IuG9ySYsQmZfqA" width="690" height="388" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/6/6/66631fc06b43a46e35a5c5052f9e504109e8a6a4_2_690x388.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/6/6/66631fc06b43a46e35a5c5052f9e504109e8a6a4_2_1035x582.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/6/6/66631fc06b43a46e35a5c5052f9e504109e8a6a4_2_1380x776.jpeg 2x" data-dominant-color="3A5E2C"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">DSC01592</span><span class="informations">6000×3376 2.11 MB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div> . Only half the screen turns green, which is what you expect if u and v are 0. The rest of the image seems untouched, so I don’t even think the function is reading all the U/V buffer. This behavior matches what I expect if indeed the vertical dimension is being halved twice as I think the current code in the repo is doing.</p>

IvanKobzarev: <p>Hello <a class="mention" href="/u/laganojunior">@laganojunior</a></p>
<p>Thanks for reporting it.</p>
<p>There was a bug in yuv to rgb conversion, the fix was landed in master recently <a href="https://github.com/pytorch/pytorch/pull/50871" class="inline-onebox" rel="noopener nofollow ugc">[android] Fix YUV camera image to tensor by IvanKobzarev · Pull Request #50871 · pytorch/pytorch · GitHub</a></p>
<p>Next nightly build of</p>
<pre><code class="lang-auto">    implementation 'org.pytorch:pytorch_android_torchvision:1.8.0-SNAPSHOT'
</code></pre>
<p>will contain that fix.</p>

laganojunior: <p><a class="mention" href="/u/ivankobzarev">@IvanKobzarev</a> Thanks for the turnaround and fix! That does seem to be the change I implemented on my end to get it working. I’ll have to test this change out from nightly snapshot some time.</p>
<p>I have another issue that I think I had fixed on my end as well. When the image gets particularly dark or bright, I would see errant splotches of color. I think I see the artifacts on your image on the pull request as well, around the shadows in the right leg and the chair rest.</p>
<p>On my end, I had this issue fixed by wrapping the clamp0255 macro in parenthesis. E.g:</p>
<pre><code class="lang-auto">#define clamp0255(x) (x &gt; 255 ? 255 : x &lt; 0 ? 0 : x)
</code></pre>
<p>I think in this context:</p>
<pre><code class="lang-auto">outData[wr++] = (clamp0255(ri) - normMeanRm255) / normStdRm255;
</code></pre>
<p>The way the macro expands and the precedence of the ternary operator, I think that if ri clips below 0 or over 255, it doesn’t actually get subtracted by normMeanRm255.</p>

IvanKobzarev: <p>Thanks a lot for pointing it. The fix was landed to master <a href="https://github.com/pytorch/pytorch/pull/50951" class="inline-onebox" rel="noopener nofollow ugc">[android] fix yuv conversion - remove define by IvanKobzarev · Pull Request #50951 · pytorch/pytorch · GitHub</a></p>