[codex] Improve mesh segmentation post-processing

src/CLIPipeline.cpp

@@ -8612,6 +8612,8 @@ int CLIPipeline::cmdSegment(int argc, char* argv[])
         err() << "Error: no readable geometry in " << inputPath << Qt::endl; return 1;
     }
     const int vertexCount = static_cast<int>(verts.size() / 3);
+    int rigResolved = 0;
+    std::vector<int> rigLabels = AutoRig::rigPriorPartLabels(entity, vertexCount, &rigResolved);
 
     // --- Training-data miner ------------------------------------------------
     // `--dump-training-data out.json` writes the mesh's point cloud + EXACT
@@ -8623,15 +8625,13 @@ int CLIPipeline::cmdSegment(int argc, char* argv[])
     // MODEL path (used on UNrigged meshes) keeps improving as more rigged
     // assets are mined. Requires a skinned mesh.
     if (!dumpPath.isEmpty()) {
-        int resolved = 0;
-        std::vector<int> labels = AutoRig::rigPriorPartLabels(entity, vertexCount, &resolved);
-        if (labels.empty()) {
+        if (rigLabels.empty()) {
             err() << "Error: --dump-training-data needs a SKINNED mesh (no skeleton in "
                   << inputPath << ")." << Qt::endl;
             return 1;
         }
-        if (resolved < (vertexCount + 1) / 2) {
-            err() << "Error: rig resolved only " << resolved << " / " << vertexCount
+        if (rigResolved < (vertexCount + 1) / 2) {
+            err() << "Error: rig resolved only " << rigResolved << " / " << vertexCount
                   << " vertices to body parts — too sparse to be reliable training data."
                   << Qt::endl;
             return 1;
@@ -8655,14 +8655,14 @@ int CLIPipeline::cmdSegment(int argc, char* argv[])
             pts.append((verts[3*v+0]-ctr[0])*invs);
             pts.append((verts[3*v+1]-ctr[1])*invs);
             pts.append((verts[3*v+2]-ctr[2])*invs);
-            labs.append(labels[v] < 0 ? 0 : labels[v]);   // bone-but-non-body → unknown(0)
+            labs.append(rigLabels[v] < 0 ? 0 : rigLabels[v]);   // bone-but-non-body → unknown(0)
         }
         QJsonObject root;
         root["schema"]      = "qtmesh-meshseg-training-v1";
         root["mesh"]        = fi.fileName();
         root["upAxis"]      = upAxis;     // miner records the source up axis
         root["vertexCount"] = vertexCount;
-        root["resolved"]    = resolved;
+        root["resolved"]    = rigResolved;
         root["partCount"]   = MeshSegmenter::partCount();
         root["points"]      = pts;        // flat [x,y,z, …] normalised to unit box
         root["labels"]      = labs;       // per-vertex part index (0=unknown)
@@ -8673,8 +8673,8 @@ int CLIPipeline::cmdSegment(int argc, char* argv[])
         out.write(QJsonDocument(root).toJson(QJsonDocument::Compact));
         out.close();
         cliWrite(QString("Wrote training sample %1 — %2 verts, %3 resolved (%4%)\n")
-                     .arg(dumpPath).arg(vertexCount).arg(resolved)
-                     .arg(100.0 * resolved / std::max(1, vertexCount), 0, 'f', 1));
+                     .arg(dumpPath).arg(vertexCount).arg(rigResolved)
+                     .arg(100.0 * rigResolved / std::max(1, vertexCount), 0, 'f', 1));
         return 0;
     }
 
@@ -8686,7 +8686,7 @@ int CLIPipeline::cmdSegment(int argc, char* argv[])
     opts.forceFallback = noModel;
     const MeshSegmenter::Result r = MeshSegmenter::predict(
         verts.data(), vertexCount, indices.data(), static_cast<int>(indices.size()),
-        modelPath, opts);
+        modelPath, opts, rigLabels.empty() ? nullptr : rigLabels.data());
     if (!r.ok) {
         err() << "Error: " << (r.error.isEmpty() ? QStringLiteral("segmentation failed") : r.error)
               << Qt::endl;

src/EditModeController.cpp

@@ -875,10 +875,10 @@ QString EditModeController::selectByPart(const QString& upAxis)
     // fall through to the worker (model / geometric fallback). The extraction is
     // shared with the training-data miner (AutoRig::rigPriorPartLabels) so the
     // GUI selection and the mined ground truth are bit-identical.
+    std::vector<int> rigLabels;
     {
         int resolved = 0;
-        std::vector<int> rigLabels =
-            AutoRig::rigPriorPartLabels(m_editEntity, vertexCount, &resolved);
+        rigLabels = AutoRig::rigPriorPartLabels(m_editEntity, vertexCount, &resolved);
         // Use the rig labels when they cover most of the mesh (a real rig
         // should). Below that, the rig is too sparse to trust → use the model.
         // (rigPriorPartLabels marks bone-but-non-body vertices as -1; map those
@@ -908,7 +908,7 @@ QString EditModeController::selectByPart(const QString& upAxis)
     auto cancel = m_segmentCancel;
     QPointer<EditModeController> self(this);
 
-    std::thread([self, verts, indices, vertexCount, cancel, upAxisIdx]() {
+    std::thread([self, verts, indices, rigLabels, vertexCount, cancel, upAxisIdx]() {
         // Progress callback (inference): marshal to the UI thread, honour cancel.
         auto progress = [self, cancel](int done, int total) -> bool {
             if (cancel->load()) return false;
@@ -938,7 +938,7 @@ QString EditModeController::selectByPart(const QString& upAxis)
         const MeshSegmenter::Result r = MeshSegmenter::predict(
             verts.data(), vertexCount, indices.data(),
             static_cast<int>(indices.size()), modelPath, opts,
-            /*boneProximity=*/nullptr, progress);
+            rigLabels.empty() ? nullptr : rigLabels.data(), progress);
 
         // Hand the result back to the main thread to touch Ogre selection.
         QMetaObject::invokeMethod(qApp, [self, r, cancel]() {

src/MCPServer.cpp

@@ -3554,6 +3554,8 @@ QJsonObject MCPServer::toolSegmentMesh(const QJsonObject &args)
         if (!AutoRig::gatherGeometry(entity, verts, indices) || verts.empty())
             return makeErrorResult("Error: no readable geometry");
         const int vertexCount = static_cast<int>(verts.size() / 3);
+        std::vector<int> rigLabels =
+            AutoRig::rigPriorPartLabels(entity, vertexCount);
 
         const bool noModel = args.value("no_model").toBool(false);
         QString modelPath;
@@ -3570,7 +3572,8 @@ QJsonObject MCPServer::toolSegmentMesh(const QJsonObject &args)
         }
         const MeshSegmenter::Result r = MeshSegmenter::predict(
             verts.data(), vertexCount, indices.data(),
-            static_cast<int>(indices.size()), modelPath, opts);
+            static_cast<int>(indices.size()), modelPath, opts,
+            rigLabels.empty() ? nullptr : rigLabels.data());
         if (!r.ok)
             return makeErrorResult(QString("Error: %1")
                 .arg(r.error.isEmpty() ? QStringLiteral("segmentation failed") : r.error));

src/MeshSegmenter.cpp

@@ -97,6 +97,71 @@
     if (n.endsWith('r')) return 'r';
     return 0;
 }
+
+void applyBoneHints(std::vector<int>& labels, const int* boneProximity)
+{
+    if (!boneProximity) return;
+    const int count = static_cast<int>(MeshSegmenter::Part::Count);
+    const int unknown = static_cast<int>(MeshSegmenter::Part::Unknown);
+    for (int v = 0; v < static_cast<int>(labels.size()); ++v) {
+        const int bp = boneProximity[v];
+        if (bp > unknown && bp < count)
+            labels[v] = bp;
+    }
+}
+
+void smoothLabelsByTopology(std::vector<int>& labels,
+                            const uint32_t* indices,
+                            int indexCount,
+                            int iterations)
+{
+    const int vertexCount = static_cast<int>(labels.size());
+    const int partCount = static_cast<int>(MeshSegmenter::Part::Count);
+    if (vertexCount <= 0 || !indices || indexCount < 3 || iterations <= 0)
+        return;
+
+    std::vector<std::vector<int>> neighbours(vertexCount);
+    for (int i = 0; i + 2 < indexCount; i += 3) {
+        const int a = static_cast<int>(indices[i]);
+        const int b = static_cast<int>(indices[i + 1]);
+        const int c = static_cast<int>(indices[i + 2]);
+        if (a < 0 || b < 0 || c < 0 || a >= vertexCount || b >= vertexCount || c >= vertexCount)
+            continue;
+        neighbours[a].push_back(b); neighbours[a].push_back(c);
+        neighbours[b].push_back(a); neighbours[b].push_back(c);
+        neighbours[c].push_back(a); neighbours[c].push_back(b);
+    }
+    for (auto& ns : neighbours) {
+        std::sort(ns.begin(), ns.end());
+        ns.erase(std::unique(ns.begin(), ns.end()), ns.end());
+    }
+
+    std::vector<int> next(labels.size());
+    for (int it = 0; it < iterations; ++it) {
+        next = labels;
+        for (int v = 0; v < vertexCount; ++v) {
+            std::array<int, static_cast<int>(MeshSegmenter::Part::Count)> votes{};
+            const int current = labels[v];
+            if (current >= 0 && current < partCount)
+                votes[current] += 3;
+            for (int n : neighbours[v]) {
+                const int label = labels[n];
+                if (label >= 0 && label < partCount)
+                    ++votes[label];
+            }
+            int best = current;
+            int bestVotes = (current >= 0 && current < partCount) ? votes[current] : -1;
+            for (int p = 0; p < partCount; ++p) {
+                if (votes[p] > bestVotes) {
+                    bestVotes = votes[p];
+                    best = p;
+                }
+            }
+            next[v] = best;
+        }
+        labels.swap(next);
+    }
+}
 } // namespace
 
 MeshSegmenter::Options::Options() = default;
@@ -328,13 +393,8 @@
     // 2) Override with VALID rig bone-proximity hints only. An unknown/invalid
     // hint (-1, or out of range) leaves the spatial label intact — a partially
     // hinted rig must not collapse its unhinted vertices to Torso.
-    if (haveBone) {
-        for (int v = 0; v < vertexCount; ++v) {
-            const int bp = boneProximity[v];
-            if (bp >= 0 && bp < static_cast<int>(Part::Count))
-                r.vertexLabels[v] = bp;
-        }
-    }
+    if (haveBone)
+        applyBoneHints(r.vertexLabels, boneProximity);
 
     r.faceLabels = facesFromVertexLabels(r.vertexLabels, indices, indexCount);
     r.ok = true;
@@ -467,13 +527,22 @@
         }
 
         const int N = std::max(256, opts.samplePoints);
-        // Deterministic point sample (with replacement if the mesh is small).
+        // Deterministic point sample. Small meshes include every vertex once so
+        // no vertex keeps the default label; large meshes are sampled across the
+        // whole vertex buffer instead of taking the first N vertices.
         std::mt19937 rng(0x5e6u);  // NOSONAR — non-crypto, fixed for reproducibility
         std::uniform_int_distribution<int> pick(0, vertexCount - 1);
         std::vector<float> pts(static_cast<size_t>(N) * 3);
         std::vector<int> srcVert(N);
         for (int i = 0; i < N; ++i) {
-            const int v = (vertexCount >= N) ? (i < vertexCount ? i : pick(rng)) : pick(rng);
+            int v = 0;
+            if (vertexCount <= N) {
+                v = (i < vertexCount) ? i : pick(rng);
+            } else {
+                const double t = (static_cast<double>(i) + 0.5) / static_cast<double>(N);
+                v = std::clamp(static_cast<int>(t * static_cast<double>(vertexCount)),
+                               0, vertexCount - 1);
+            }
             srcVert[i] = v;
             for (int c = 0; c < 3; ++c) {
                 const int a = axisFor[c];
@@ -545,13 +614,12 @@
 
         // Scatter sampled-point labels back to ALL vertices by nearest sampled
         // point (in normalised space). For samplePoints >= vertexCount the first
-        // N points already cover every vertex 1:1, so this is exact.
+        // vertexCount samples cover every vertex 1:1, so this is exact.
         Result r;
         r.vertexLabels.assign(vertexCount, static_cast<int>(Part::Torso));
         if (vertexCount <= N) {
-            for (int i = 0; i < N && i < vertexCount; ++i)
-                r.vertexLabels[srcVert[i]] = pointLabel[i];
-            // any vertex not directly sampled (shouldn't happen for vc<=N first-N) → nearest
+            for (int i = 0; i < vertexCount; ++i)
+                r.vertexLabels[i] = pointLabel[i];
         } else {
             // brute-force nearest sampled point per vertex (N is small, ~4k).
             for (int v = 0; v < vertexCount; ++v) {
@@ -569,6 +637,9 @@
                 r.vertexLabels[v] = pointLabel[best];
             }
         }
+        applyBoneHints(r.vertexLabels, boneProximity);
+        smoothLabelsByTopology(r.vertexLabels, indices, indexCount, 1);
+        applyBoneHints(r.vertexLabels, boneProximity);
         r.faceLabels = facesFromVertexLabels(r.vertexLabels, indices, indexCount);
         r.ok = true;
         r.usedModel = true;