修复人脸画框

plugins/ai_face_det/ai_face_det_node.cpp

@@ -160,17 +160,18 @@ bool ExtractNc(const Tensor& t, int c, NcTensor& out) {
     if (elem_cnt == 0) return false;
 
     int n = 0;
-    bool transposed = false;
+    enum class Layout { FlatNc, CxN, NCHW, NHWC } layout = Layout::FlatNc;
+    int dN = 1, dH = 1, dW = 1;
     if (t.dims.size() == 3) {
         // Common: [1, C, N] or [1, N, C]
         const uint32_t d1 = t.dims[1];
         const uint32_t d2 = t.dims[2];
         if (static_cast<int>(d1) == c) {
             n = static_cast<int>(d2);
-            transposed = true;  // CxN
+            layout = Layout::CxN;  // [1, C, N]
         } else if (static_cast<int>(d2) == c) {
             n = static_cast<int>(d1);
-            transposed = false;  // NxC
+            layout = Layout::FlatNc;  // treat as contiguous [1, N, C]
         }
     } else if (t.dims.size() == 2) {
         // [N, C] or [C, N]
@@ -178,10 +179,31 @@ bool ExtractNc(const Tensor& t, int c, NcTensor& out) {
         const uint32_t d1 = t.dims[1];
         if (static_cast<int>(d1) == c) {
             n = static_cast<int>(d0);
-            transposed = false;
+            layout = Layout::FlatNc;
         } else if (static_cast<int>(d0) == c) {
             n = static_cast<int>(d1);
-            transposed = true;
+            layout = Layout::CxN;
+        }
+    } else if (t.dims.size() == 4) {
+        // Common: [N, C, H, W] (NCHW) or [N, H, W, C] (NHWC)
+        const int dn = static_cast<int>(t.dims[0]);
+        const int d1 = static_cast<int>(t.dims[1]);
+        const int d2 = static_cast<int>(t.dims[2]);
+        const int d3 = static_cast<int>(t.dims[3]);
+        if (d3 == c) {
+            // NHWC
+            layout = Layout::NHWC;
+            dN = std::max(1, dn);
+            dH = std::max(1, d1);
+            dW = std::max(1, d2);
+            n = dN * dH * dW;
+        } else if (d1 == c) {
+            // NCHW
+            layout = Layout::NCHW;
+            dN = std::max(1, dn);
+            dH = std::max(1, d2);
+            dW = std::max(1, d3);
+            n = dN * dH * dW;
         }
     }
 
@@ -217,11 +239,9 @@ bool ExtractNc(const Tensor& t, int c, NcTensor& out) {
         return Dequant(p[idx], t.zp, t.scale);
     };
 
-    if (!transposed) {
-        for (size_t i = 0; i < out.data.size(); ++i) {
-            out.data[i] = ReadElem(i);
-        }
-    } else {
+    if (layout == Layout::FlatNc) {
+        for (size_t i = 0; i < out.data.size(); ++i) out.data[i] = ReadElem(i);
+    } else if (layout == Layout::CxN) {
         // Input is [C, N] contiguous. Transpose to [N, C].
         for (int ci = 0; ci < c; ++ci) {
             for (int ni = 0; ni < n; ++ni) {
@@ -230,6 +250,37 @@ bool ExtractNc(const Tensor& t, int c, NcTensor& out) {
                 out.data[dst_idx] = ReadElem(src_idx);
             }
         }
+    } else if (layout == Layout::NHWC) {
+        // [N, H, W, C] contiguous; pack to [N*H*W, C]
+        size_t dst = 0;
+        for (int n0 = 0; n0 < dN; ++n0) {
+            for (int y = 0; y < dH; ++y) {
+                for (int x = 0; x < dW; ++x) {
+                    const size_t base = ((static_cast<size_t>(n0) * static_cast<size_t>(dH) + static_cast<size_t>(y)) *
+                                             static_cast<size_t>(dW) + static_cast<size_t>(x)) *
+                                        static_cast<size_t>(c);
+                    for (int ci = 0; ci < c; ++ci) {
+                        out.data[dst++] = ReadElem(base + static_cast<size_t>(ci));
+                    }
+                }
+            }
+        }
+    } else if (layout == Layout::NCHW) {
+        // [N, C, H, W] channel-first; interleave channels into [N*H*W, C]
+        size_t dst = 0;
+        for (int n0 = 0; n0 < dN; ++n0) {
+            for (int y = 0; y < dH; ++y) {
+                for (int x = 0; x < dW; ++x) {
+                    for (int ci = 0; ci < c; ++ci) {
+                        const size_t src_idx =
+                            ((static_cast<size_t>(n0) * static_cast<size_t>(c) + static_cast<size_t>(ci)) *
+                                 static_cast<size_t>(dH) + static_cast<size_t>(y)) *
+                                static_cast<size_t>(dW) + static_cast<size_t>(x);
+                        out.data[dst++] = ReadElem(src_idx);
+                    }
+                }
+            }
+        }
     }
 
     return true;
@@ -509,6 +560,24 @@ private:
             return;
         }
 
+        if (!logged_io_.exchange(true)) {
+            std::string s = "[ai_face_det] outputs:";
+            for (const auto& o : r.outputs) {
+                s += " idx=" + std::to_string(o.index) + " size=" + std::to_string(o.size);
+#if defined(RK3588_ENABLE_RKNN)
+                s += " type=" + std::to_string(static_cast<int>(o.type));
+                s += " zp=" + std::to_string(o.zp) + " scale=" + std::to_string(o.scale);
+                s += " dims=[";
+                for (size_t i = 0; i < o.dims.size(); ++i) {
+                    s += std::to_string(o.dims[i]);
+                    if (i + 1 < o.dims.size()) s += ",";
+                }
+                s += "]";
+#endif
+            }
+            LogInfo(s);
+        }
+
         std::vector<Tensor> tensors;
         tensors.reserve(r.outputs.size());
         for (const auto& o : r.outputs) {
@@ -711,6 +780,8 @@ private:
     int model_w_ = 320;
     int model_h_ = 320;
     uint32_t n_output_ = 0;
+
+    std::atomic<bool> logged_io_{false};
 };
 
 REGISTER_NODE(AiFaceDetNode, "ai_face_det");