diff --git a/src/CLIPipeline.cpp b/src/CLIPipeline.cpp index 8b7623ff..e3117cb4 100644 --- a/src/CLIPipeline.cpp +++ b/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(verts.size() / 3); + int rigResolved = 0; + std::vector 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 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(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; diff --git a/src/EditModeController.cpp b/src/EditModeController.cpp index 4bbf036d..4d1e1ae0 100644 --- a/src/EditModeController.cpp +++ b/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 rigLabels; { int resolved = 0; - std::vector 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 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(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]() { diff --git a/src/MCPServer.cpp b/src/MCPServer.cpp index 4bac1e5c..e32453d6 100644 --- a/src/MCPServer.cpp +++ b/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(verts.size() / 3); + std::vector 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(indices.size()), modelPath, opts); + static_cast(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)); diff --git a/src/MeshSegmenter.cpp b/src/MeshSegmenter.cpp index e66d6b42..277a1a63 100644 --- a/src/MeshSegmenter.cpp +++ b/src/MeshSegmenter.cpp @@ -97,6 +97,71 @@ char boneSide(const QString& n) if (n.endsWith('r')) return 'r'; return 0; } + +void applyBoneHints(std::vector& labels, const int* boneProximity) +{ + if (!boneProximity) return; + const int count = static_cast(MeshSegmenter::Part::Count); + const int unknown = static_cast(MeshSegmenter::Part::Unknown); + for (int v = 0; v < static_cast(labels.size()); ++v) { + const int bp = boneProximity[v]; + if (bp > unknown && bp < count) + labels[v] = bp; + } +} + +void smoothLabelsByTopology(std::vector& labels, + const uint32_t* indices, + int indexCount, + int iterations) +{ + const int vertexCount = static_cast(labels.size()); + const int partCount = static_cast(MeshSegmenter::Part::Count); + if (vertexCount <= 0 || !indices || indexCount < 3 || iterations <= 0) + return; + + std::vector> neighbours(vertexCount); + for (int i = 0; i + 2 < indexCount; i += 3) { + const int a = static_cast(indices[i]); + const int b = static_cast(indices[i + 1]); + const int c = static_cast(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 next(labels.size()); + for (int it = 0; it < iterations; ++it) { + next = labels; + for (int v = 0; v < vertexCount; ++v) { + std::array(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 @@ MeshSegmenter::Result MeshSegmenter::segmentGeometric(const float* positions, in // 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(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 @@ MeshSegmenter::Result MeshSegmenter::predict(const float* positions, int vertexC } 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 pick(0, vertexCount - 1); std::vector pts(static_cast(N) * 3); std::vector 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(i) + 0.5) / static_cast(N); + v = std::clamp(static_cast(t * static_cast(vertexCount)), + 0, vertexCount - 1); + } srcVert[i] = v; for (int c = 0; c < 3; ++c) { const int a = axisFor[c]; @@ -545,13 +614,12 @@ MeshSegmenter::Result MeshSegmenter::predict(const float* positions, int vertexC // 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(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 @@ MeshSegmenter::Result MeshSegmenter::predict(const float* positions, int vertexC 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;