Initial snapshot: hand motion pipeline (Dyn-HaMR + dex-retargeting + SPIDER)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
+18
@@ -0,0 +1,18 @@
|
||||
# Source snapshot: runtime exclusions anchored to repository root.
|
||||
/output/
|
||||
/results/
|
||||
/logs/
|
||||
/weights/
|
||||
/dist/
|
||||
/.venv/
|
||||
/venv/
|
||||
/.dex/
|
||||
/.spider/
|
||||
/.cuda/
|
||||
/.hf_cache/
|
||||
/.torch_cache/
|
||||
/.spider_cache/
|
||||
/.uv_cache_spider/
|
||||
/.env
|
||||
__pycache__/
|
||||
*.pyc
|
||||
+5763
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,250 @@
|
||||
[
|
||||
{
|
||||
"path": "third_party/l20_assets",
|
||||
"bundled": false,
|
||||
"files": [
|
||||
{
|
||||
"path": "third_party/l20_assets/L20/RIGHT/g20_right_G20.json",
|
||||
"bytes": 111584,
|
||||
"sha256": "314865d5bc359996bb182942fefc597f989ebfa6c9b3c97a9086acf38b0712ee"
|
||||
},
|
||||
{
|
||||
"path": "third_party/l20_assets/L20/RIGHT/linkerhand_g20_right.urdf",
|
||||
"bytes": 28687,
|
||||
"sha256": "aa198155651eb7e815f1c27bc8bd9c902e969e81c1165aed756779db378e715d"
|
||||
},
|
||||
{
|
||||
"path": "third_party/l20_assets/L20/RIGHT/linkerhand_g20_right_mj.urdf",
|
||||
"bytes": 28801,
|
||||
"sha256": "5a79601debec6013bda85ce65ea020878b645ca01e20440a2c56a78c3ecbee2d"
|
||||
},
|
||||
{
|
||||
"path": "third_party/l20_assets/L20/RIGHT/meshes/hand_base_link.STL",
|
||||
"bytes": 3433184,
|
||||
"sha256": "79e36d4c7e9282c79601976d41144557d1928810202589e31bed335f3d308cd1"
|
||||
},
|
||||
{
|
||||
"path": "third_party/l20_assets/L20/RIGHT/meshes/index_distal.STL",
|
||||
"bytes": 984884,
|
||||
"sha256": "324aa4b852e0c469bfc0c0bb9b19e0ccffc7dcb401d6a678481340588465fa26"
|
||||
},
|
||||
{
|
||||
"path": "third_party/l20_assets/L20/RIGHT/meshes/index_metacarpals.STL",
|
||||
"bytes": 151784,
|
||||
"sha256": "106ad58d52420c3216ccb87865f55062d4b4b558078dda95210a085a9bbe277c"
|
||||
},
|
||||
{
|
||||
"path": "third_party/l20_assets/L20/RIGHT/meshes/index_middle.STL",
|
||||
"bytes": 805584,
|
||||
"sha256": "27f9ccd91c5bdb1f6a1c42aeee4ab79e7e0aaf1e124a36999c75715b70d0aee6"
|
||||
},
|
||||
{
|
||||
"path": "third_party/l20_assets/L20/RIGHT/meshes/index_proximal.STL",
|
||||
"bytes": 883184,
|
||||
"sha256": "d4f85e69dd94d56bdc4646c945ca8a237c652fd5dffba86c7b9ce1d8067ae8e1"
|
||||
},
|
||||
{
|
||||
"path": "third_party/l20_assets/L20/RIGHT/meshes/middle_distal.STL",
|
||||
"bytes": 984884,
|
||||
"sha256": "688be3d9dd7fda49a8f9dbba824fedef5cf427c94e2e9e1ea96834eafbf6945a"
|
||||
},
|
||||
{
|
||||
"path": "third_party/l20_assets/L20/RIGHT/meshes/middle_metacarpals.STL",
|
||||
"bytes": 151784,
|
||||
"sha256": "65bc65257a01e0ccd284dfdf3b2a66f2a865bf80faeb4da746cd1fd05114fe40"
|
||||
},
|
||||
{
|
||||
"path": "third_party/l20_assets/L20/RIGHT/meshes/middle_middle.STL",
|
||||
"bytes": 805584,
|
||||
"sha256": "600b8e7486abc8a2707a5ecc66d4b5bea1dc0895c7173e911ba994f5cac30c10"
|
||||
},
|
||||
{
|
||||
"path": "third_party/l20_assets/L20/RIGHT/meshes/middle_proximal.STL",
|
||||
"bytes": 883184,
|
||||
"sha256": "edd3af02319cd4a3b4e07a6f5d33d501887d5ab78625489254f289785b6f57e1"
|
||||
},
|
||||
{
|
||||
"path": "third_party/l20_assets/L20/RIGHT/meshes/pinky_distal.STL",
|
||||
"bytes": 984884,
|
||||
"sha256": "3b22871d948ef80447f93866af3643ffa87cb22c1c172750fca9357f3aaa46de"
|
||||
},
|
||||
{
|
||||
"path": "third_party/l20_assets/L20/RIGHT/meshes/pinky_metacarpals.STL",
|
||||
"bytes": 151784,
|
||||
"sha256": "b63d10a0e18af0169c75b85eb740c1f25123f6c4c5be0c5368afc65e8c73771f"
|
||||
},
|
||||
{
|
||||
"path": "third_party/l20_assets/L20/RIGHT/meshes/pinky_middle.STL",
|
||||
"bytes": 805584,
|
||||
"sha256": "c6696bfcd2dbef8f248f86a6f45a90d5ba7b40e11cdb13a86b96347240a7b751"
|
||||
},
|
||||
{
|
||||
"path": "third_party/l20_assets/L20/RIGHT/meshes/pinky_proximal.STL",
|
||||
"bytes": 883184,
|
||||
"sha256": "311ce029c4972530d3b234307f4ea6ca1789577c2786d56f969a5dfac2ef6a57"
|
||||
},
|
||||
{
|
||||
"path": "third_party/l20_assets/L20/RIGHT/meshes/ring_distal.STL",
|
||||
"bytes": 984884,
|
||||
"sha256": "fca7bf31fa3f377110d31f7623e27f8c7845faedbdca7c09f98c94a28867bd69"
|
||||
},
|
||||
{
|
||||
"path": "third_party/l20_assets/L20/RIGHT/meshes/ring_metacarpals.STL",
|
||||
"bytes": 151784,
|
||||
"sha256": "11797279afee7f856e472755d3a96593baea659d2ec0014edb02006ae48b5d71"
|
||||
},
|
||||
{
|
||||
"path": "third_party/l20_assets/L20/RIGHT/meshes/ring_middle.STL",
|
||||
"bytes": 805584,
|
||||
"sha256": "2b3a0815e8daf95e76c9f54f84c13709a3951aa44ec6bb863facf60ba32c3c16"
|
||||
},
|
||||
{
|
||||
"path": "third_party/l20_assets/L20/RIGHT/meshes/ring_proximal.STL",
|
||||
"bytes": 883184,
|
||||
"sha256": "168df91e5a5643a899556f96ca3f3b7fe8a3349436be42dfb5632c0d524f9a47"
|
||||
},
|
||||
{
|
||||
"path": "third_party/l20_assets/L20/RIGHT/meshes/thumb_distal.STL",
|
||||
"bytes": 1108484,
|
||||
"sha256": "026dd2fac43c4461c5a90561ef29039f4a08802c0e47e5c422ecac2f34babe1e"
|
||||
},
|
||||
{
|
||||
"path": "third_party/l20_assets/L20/RIGHT/meshes/thumb_metacarpals.STL",
|
||||
"bytes": 1158784,
|
||||
"sha256": "0a82998ece3c784bba7861b850d00ed483d93434f94f72c6298fea0917e328ca"
|
||||
},
|
||||
{
|
||||
"path": "third_party/l20_assets/L20/RIGHT/meshes/thumb_metacarpals_base1.STL",
|
||||
"bytes": 852984,
|
||||
"sha256": "8416e5f826f12c4597754fbe3052983635345a76e3d98a4cdbc38ed363a9da78"
|
||||
},
|
||||
{
|
||||
"path": "third_party/l20_assets/L20/RIGHT/meshes/thumb_metacarpals_base2.STL",
|
||||
"bytes": 227884,
|
||||
"sha256": "a5679a28a612ab6aa62a4f7fa1c67be42f807bba7c2ca00ba062ca2e4ef0070d"
|
||||
},
|
||||
{
|
||||
"path": "third_party/l20_assets/L20/RIGHT/meshes/thumb_proximal.STL",
|
||||
"bytes": 716984,
|
||||
"sha256": "4cd7d6a283d2e5c3f6dfbd88aa1e18f9f481375d1e6b47f2c373a285184b123d"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "third_party/bottle_model_parametric",
|
||||
"bundled": false,
|
||||
"files": [
|
||||
{
|
||||
"path": "third_party/bottle_model_parametric/README.md",
|
||||
"bytes": 1309,
|
||||
"sha256": "f0a7942736018fb3b8cb6758665cabbc09a0a232bdb86d65b6812c9949f9da77"
|
||||
},
|
||||
{
|
||||
"path": "third_party/bottle_model_parametric/assets/bottle.obj",
|
||||
"bytes": 32331,
|
||||
"sha256": "e4d48245adfed5ea0a72ee3dc247b68c6d419cf1ccf5651cc8f7576fe2f440f2"
|
||||
},
|
||||
{
|
||||
"path": "third_party/bottle_model_parametric/assets/shoulder.obj",
|
||||
"bytes": 15949,
|
||||
"sha256": "d27b809f8510e37fe20c14c50f43a6977e4b6c1c62edbbaf258f8df76b7a3c92"
|
||||
},
|
||||
{
|
||||
"path": "third_party/bottle_model_parametric/bottle.glb",
|
||||
"bytes": 19252,
|
||||
"sha256": "a18c28ca8afa6bf00da54aa4dcc81396cb8b52dc14f23fcbc0cfc638cc570689"
|
||||
},
|
||||
{
|
||||
"path": "third_party/bottle_model_parametric/bottle.stl",
|
||||
"bytes": 51284,
|
||||
"sha256": "3a6c64e0acd9b9cbc11704cab2ecbc572e4972f04764a082137dd61d8a8d3efa"
|
||||
},
|
||||
{
|
||||
"path": "third_party/bottle_model_parametric/bottle.xml",
|
||||
"bytes": 962,
|
||||
"sha256": "a151a8c4b54bc29e466a77dd831778708f014447887ad6ee281af10c5c02a15e"
|
||||
},
|
||||
{
|
||||
"path": "third_party/bottle_model_parametric/preview.png",
|
||||
"bytes": 11177,
|
||||
"sha256": "a8387a2d72dbd88eb558919194585fe71311988485c13c282818404fc61adcc0"
|
||||
},
|
||||
{
|
||||
"path": "third_party/bottle_model_parametric/profile.npz",
|
||||
"bytes": 1062,
|
||||
"sha256": "d9b28a9f0acc06c7df8221b618bc44fff7e7eb54eac2bd70b46101e867a22dc3"
|
||||
},
|
||||
{
|
||||
"path": "third_party/bottle_model_parametric/validation.json",
|
||||
"bytes": 512,
|
||||
"sha256": "bedf450277eba0e195ebc82ebb0ce9adf6ed29950c655db9c264406f3c503f14"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "third_party/dex-assets",
|
||||
"bundled": false,
|
||||
"files": [
|
||||
{
|
||||
"path": "third_party/dex-assets/robots/hands/leap_hand/LICENSE.txt",
|
||||
"bytes": 1262,
|
||||
"sha256": "fad9015ea2b25dcc14d7188c01de23aadf1180e49e7afae4beb566e21a39afbd"
|
||||
},
|
||||
{
|
||||
"path": "third_party/dex-assets/robots/hands/leap_hand/leap_hand_right.urdf",
|
||||
"bytes": 24927,
|
||||
"sha256": "7170056b44553779af566a33a3f80fc44c431f739d85fff809e9a16ed4887f76"
|
||||
},
|
||||
{
|
||||
"path": "third_party/dex-assets/robots/hands/leap_hand/meshes/collision/white_tip.obj",
|
||||
"bytes": 640739,
|
||||
"sha256": "bbc1d4b64f4f4115fb0fd776a14abd9496e339804bf42cb645d30aee13e61f88"
|
||||
},
|
||||
{
|
||||
"path": "third_party/dex-assets/robots/hands/leap_hand/meshes/collision/white_tip_thumb.obj",
|
||||
"bytes": 402712,
|
||||
"sha256": "80e7c3de59c0b1b141e2c91ffa36374bef11636e362173d7a68a79022829020f"
|
||||
},
|
||||
{
|
||||
"path": "third_party/dex-assets/robots/hands/leap_hand/meshes/visual/dip.obj",
|
||||
"bytes": 1013886,
|
||||
"sha256": "472e054348c5a0e987829e721d5ba8486ddd47d867e416bb45d8d2a2ff917acf"
|
||||
},
|
||||
{
|
||||
"path": "third_party/dex-assets/robots/hands/leap_hand/meshes/visual/fingertip.obj",
|
||||
"bytes": 359831,
|
||||
"sha256": "6bfa2afb2a06557ce0c54b03baec13522bfe561aabb576bd4701d4704f93ec0f"
|
||||
},
|
||||
{
|
||||
"path": "third_party/dex-assets/robots/hands/leap_hand/meshes/visual/mcp_joint.obj",
|
||||
"bytes": 1065471,
|
||||
"sha256": "3d7194a6a33626ffcdfce78a0ca854e67b6b49f2e09a3bd62656f94515fdfbec"
|
||||
},
|
||||
{
|
||||
"path": "third_party/dex-assets/robots/hands/leap_hand/meshes/visual/palm_lower.obj",
|
||||
"bytes": 2978219,
|
||||
"sha256": "90d1d9b2a1c9eea15fbeb2232019c6191a55a3c4afa7f7427ce7afa0e3108b07"
|
||||
},
|
||||
{
|
||||
"path": "third_party/dex-assets/robots/hands/leap_hand/meshes/visual/pip.obj",
|
||||
"bytes": 673869,
|
||||
"sha256": "3fff19dc1f0100bb5f34eaf8a40fcf3db40f96fd7179bb7336e538b21ff8e067"
|
||||
},
|
||||
{
|
||||
"path": "third_party/dex-assets/robots/hands/leap_hand/meshes/visual/thumb_dip.obj",
|
||||
"bytes": 968312,
|
||||
"sha256": "5b402876601f2b0b36631e6a54ad20588f88e74f52320b6c902a8bc694578689"
|
||||
},
|
||||
{
|
||||
"path": "third_party/dex-assets/robots/hands/leap_hand/meshes/visual/thumb_fingertip.obj",
|
||||
"bytes": 734609,
|
||||
"sha256": "4695d72a70285e8a0132ed4de49876cf99c813182e3a24b172b53631224890de"
|
||||
},
|
||||
{
|
||||
"path": "third_party/dex-assets/robots/hands/leap_hand/meshes/visual/thumb_pip.obj",
|
||||
"bytes": 294575,
|
||||
"sha256": "d28876cc3b674323aa7d8a628a4240cda2757c61263b05e86a66280035a2689d"
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
@@ -0,0 +1,21 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2026 Mingxi Xu <mxu176@connect.hkust-gz.edu.cn>, HKUST(GZ) Di² Lab
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
@@ -0,0 +1,437 @@
|
||||
[
|
||||
{
|
||||
"path": "scripts/audit_l20_jitter.py",
|
||||
"line": 10,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "scripts/audit_l20_jitter.py",
|
||||
"line": 11,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "scripts/build_l20_full_replay.py",
|
||||
"line": 47,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "scripts/export_dynhamr_dex.py",
|
||||
"line": 15,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "scripts/export_dynhamr_dex.py",
|
||||
"line": 16,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "scripts/export_l20_right_hdf5.py",
|
||||
"line": 23,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "scripts/export_l20_right_hdf5.py",
|
||||
"line": 51,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "scripts/export_l20_right_hdf5.py",
|
||||
"line": 110,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "scripts/l20_bottle_physics.py",
|
||||
"line": 13,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "scripts/l20_bottle_sequence.py",
|
||||
"line": 9,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "scripts/l20_bottle_sequence.py",
|
||||
"line": 67,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "scripts/l20_contact.py",
|
||||
"line": 18,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "scripts/l20_kinematic_grasp_check.py",
|
||||
"line": 8,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "scripts/package_source_release.py",
|
||||
"line": 60,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "scripts/prepare_spider_l20.py",
|
||||
"line": 61,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "scripts/render_l20.py",
|
||||
"line": 13,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "scripts/render_l20_bottle.py",
|
||||
"line": 11,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "scripts/render_l20_jitter_comparison.py",
|
||||
"line": 9,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "scripts/render_l20_stabilization.py",
|
||||
"line": 13,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "scripts/render_l20_stabilization.py",
|
||||
"line": 15,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "scripts/resume_dynhamr_2047635068.py",
|
||||
"line": 41,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "scripts/resume_dynhamr_2047635068.py",
|
||||
"line": 46,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "scripts/resume_dynhamr_2047635068.py",
|
||||
"line": 47,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "scripts/resume_dynhamr_2047635068.py",
|
||||
"line": 49,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "scripts/run_dex_l20.py",
|
||||
"line": 15,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "scripts/run_dex_l20.py",
|
||||
"line": 46,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "scripts/run_dex_leap.py",
|
||||
"line": 19,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "scripts/run_handflow_15886123.py",
|
||||
"line": 9,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "scripts/run_handflow_15886123.py",
|
||||
"line": 23,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "scripts/smooth_l20_offline.py",
|
||||
"line": 8,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "scripts/stabilize_l20_temporal.py",
|
||||
"line": 19,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "scripts/stabilize_l20_temporal.py",
|
||||
"line": 20,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "scripts/verify_dex_video.py",
|
||||
"line": 6,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "scripts/verify_l20_bottle.py",
|
||||
"line": 6,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/hamer/hamer/configs/datasets_eval.yaml",
|
||||
"line": 4,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/hamer/hamer/configs/datasets_eval.yaml",
|
||||
"line": 10,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/hamer/hamer/configs/datasets_eval.yaml",
|
||||
"line": 16,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/hamer/hamer/configs/datasets_eval.yaml",
|
||||
"line": 22,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/hamer/hamer/configs/datasets_eval.yaml",
|
||||
"line": 28,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/hamer/hamer/configs/datasets_eval.yaml",
|
||||
"line": 34,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/hamer/hamer/configs/datasets_eval.yaml",
|
||||
"line": 40,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/hamer/hamer/configs/datasets_eval.yaml",
|
||||
"line": 46,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/hamer/hamer/configs/datasets_eval.yaml",
|
||||
"line": 52,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/hamer/hamer/configs/datasets_eval.yaml",
|
||||
"line": 58,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/hamer/hamer/configs/datasets_eval.yaml",
|
||||
"line": 64,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/hamer/third-party/ViTPose/docker/serve/entrypoint.sh",
|
||||
"line": 6,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/Dyn-HaMR/dyn-hamr/HMP/fitting.py",
|
||||
"line": 846,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/Dyn-HaMR/dyn-hamr/HMP/fitting.py",
|
||||
"line": 847,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/Dyn-HaMR/dyn-hamr/HMP/fitting.py",
|
||||
"line": 857,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/Dyn-HaMR/dyn-hamr/HMP/fitting_utils.py",
|
||||
"line": 317,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/Dyn-HaMR/dyn-hamr/confs/data/video_driod.yaml",
|
||||
"line": 3,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/Dyn-HaMR/dyn-hamr/confs/data/video_vipe.yaml",
|
||||
"line": 3,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/Dyn-HaMR/dyn-hamr/confs/data/video_vipe.yaml",
|
||||
"line": 28,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/Dyn-HaMR/dyn-hamr/test.sh",
|
||||
"line": 4,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/Dyn-HaMR/dyn-hamr/test.sh",
|
||||
"line": 5,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/Dyn-HaMR/third-party/DROID-SLAM/thirdparty/eigen/scripts/eigen_monitor_perf.sh",
|
||||
"line": 25,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/Dyn-HaMR/third-party/DROID-SLAM/thirdparty/lietorch/eigen/scripts/eigen_monitor_perf.sh",
|
||||
"line": 25,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/Dyn-HaMR/third-party/hamer/demo.py",
|
||||
"line": 179,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/Dyn-HaMR/third-party/hamer/hamer/configs/datasets_eval.yaml",
|
||||
"line": 4,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/Dyn-HaMR/third-party/hamer/hamer/configs/datasets_eval.yaml",
|
||||
"line": 10,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/Dyn-HaMR/third-party/hamer/hamer/configs/datasets_eval.yaml",
|
||||
"line": 16,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/Dyn-HaMR/third-party/hamer/hamer/configs/datasets_eval.yaml",
|
||||
"line": 22,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/Dyn-HaMR/third-party/hamer/hamer/configs/datasets_eval.yaml",
|
||||
"line": 28,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/Dyn-HaMR/third-party/hamer/hamer/configs/datasets_eval.yaml",
|
||||
"line": 34,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/Dyn-HaMR/third-party/hamer/hamer/configs/datasets_eval.yaml",
|
||||
"line": 40,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/Dyn-HaMR/third-party/hamer/hamer/configs/datasets_eval.yaml",
|
||||
"line": 46,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/Dyn-HaMR/third-party/hamer/hamer/configs/datasets_eval.yaml",
|
||||
"line": 52,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/Dyn-HaMR/third-party/hamer/hamer/configs/datasets_eval.yaml",
|
||||
"line": 58,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/Dyn-HaMR/third-party/hamer/hamer/configs/datasets_eval.yaml",
|
||||
"line": 64,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/Dyn-HaMR/third-party/hamer/third-party/ViTPose/docker/serve/entrypoint.sh",
|
||||
"line": 6,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/spider/examples/run_hdmi.sh",
|
||||
"line": 3,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/spider/spider/postprocess/evaluate_maniptrans.py",
|
||||
"line": 47,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/spider/spider/postprocess/read_to_hdmi.py",
|
||||
"line": 206,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/spider/spider/postprocess/render_maniptrans_trajectory.py",
|
||||
"line": 48,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/spider/spider/postprocess/render_maniptrans_trajectory.py",
|
||||
"line": 69,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/spider/spider/postprocess/render_maniptrans_trajectory.py",
|
||||
"line": 72,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/spider/spider/postprocess/render_maniptrans_trajectory.py",
|
||||
"line": 80,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/spider/spider/postprocess/render_maniptrans_trajectory.py",
|
||||
"line": 165,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/spider/spider/postprocess/render_maniptrans_trajectory.py",
|
||||
"line": 211,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/spider/spider/preprocess/locomujoco.py",
|
||||
"line": 299,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/spider/spider/process_datasets/gmr.py",
|
||||
"line": 13,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/spider/spider/process_datasets/gmr.py",
|
||||
"line": 18,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/spider/spider/process_datasets/load_and_viz_data.py",
|
||||
"line": 173,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/spider/spider/simulators/isaac.py",
|
||||
"line": 39,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/spider/spider/simulators/maniptrans.py",
|
||||
"line": 240,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/spider/spider/simulators/maniptrans.py",
|
||||
"line": 1282,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
},
|
||||
{
|
||||
"path": "third_party/spider/spider/simulators/maniptrans.py",
|
||||
"line": 1317,
|
||||
"kind": "machine path or fixed example identifier"
|
||||
}
|
||||
]
|
||||
@@ -0,0 +1,48 @@
|
||||
# Hand Motion Pipeline:Dyn-HaMR · dex-retargeting · SPIDER
|
||||
|
||||
在 HandFlow 工作区中集成单目人手重建、L20 右手重定向、接触优化实验和 HDF5 轨迹导出。
|
||||
|
||||
```text
|
||||
视频 → Dyn-HaMR / HandFlow + ViPE → 人手世界轨迹
|
||||
→ dex-retargeting + L20 模型约束 → 机器人腕部与关节轨迹
|
||||
→ SPIDER 接触引导实验 / 完整运动学回放 → HDF5
|
||||
```
|
||||
|
||||
## 从这里开始
|
||||
|
||||
- [集成结构、环境、运行顺序与上传步骤](docs/INTEGRATED_PROJECT.md)
|
||||
- [重定向方法、验证与限制](docs/retargeting_workflow_and_lessons_zh.md)
|
||||
- [HDF5 原始接口要求](docs/HDF5_REQUIREMENTS.md):原文为左手;当前导出为明确标记的右手变体。
|
||||
- [HandFlow 原始 README](docs/HANDFLOW_UPSTREAM_README.md)
|
||||
|
||||
## 当前能力
|
||||
|
||||
- 已有 L20 右手重定向、腕部运动、平滑、模型约束检查、可视化及 HDF5 导出代码。
|
||||
- 1333 帧完整回放包含运动学接触配准,不能等同于完整 SPIDER 动力学抓取成功。
|
||||
- SPIDER 实验和完整运动学回放是不同产物;只根据各自的实际评估结果作结论。
|
||||
- 部分脚本仍使用 `2047635068` / `15886123` 示例路径,需要按数据修改;尚未整理为任意视频的一键通用工具。
|
||||
|
||||
## 目录
|
||||
|
||||
| 目录 | 内容 |
|
||||
|---|---|
|
||||
| `scripts/` | 重建导出、L20 适配、SPIDER 场景/运行、回放验证、HDF5 导出 |
|
||||
| `model/`、`utils/`、`preprocessing/`、`visualization/` | HandFlow 推理代码 |
|
||||
| `third_party/` | 上游源码;源码交付包使用带来源记录的普通目录快照 |
|
||||
| `configs/`、`requirements/` | 配置、重定向依赖和本机观测版本 |
|
||||
| `docs/` | 集成、数据与方法说明 |
|
||||
| `output/`、权重和本地环境 | 本地运行资源,不进入源码交付 |
|
||||
|
||||
## 生成源码交付包
|
||||
|
||||
在当前完整开发工作区运行:
|
||||
|
||||
```bash
|
||||
python3 scripts/package_source_release.py
|
||||
```
|
||||
|
||||
输出至 `dist/hand-motion-pipeline/` 和同名 `.tar.gz`。包中保存第三方版本、源码文件哈希、外部资源清单、排除文件及可移植性检查报告。此命令只整理文件,不运行推理,不上传远程仓库。
|
||||
|
||||
## 来源和授权
|
||||
|
||||
保留 [HandFlow 原 LICENSE](LICENSE);Dyn-HaMR、HaMeR、ViPE、SPIDER 等使用各自目录的上游 LICENSE。dex-retargeting 作为版本依赖使用。本地代码是组件适配与实验集成,上游算法和资源保持原归属。
|
||||
@@ -0,0 +1,10 @@
|
||||
{
|
||||
"local_python_syntax_checked": 67,
|
||||
"shell_syntax": "PASS",
|
||||
"source_files": 7000,
|
||||
"bytes": 59113672,
|
||||
"portability_findings": 87,
|
||||
"excluded_files": 1440,
|
||||
"end_to_end_reconstruction_rerun": false,
|
||||
"remote_upload_performed": false
|
||||
}
|
||||
+7001
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,71 @@
|
||||
[
|
||||
{
|
||||
"path": ".",
|
||||
"commit": "67fa7df536db233408fe6270ca5d2de28d5959c3",
|
||||
"origin": "https://github.com/mxxu00/HandFlow.git",
|
||||
"snapshot": "selected integration source and documentation from current working tree"
|
||||
},
|
||||
{
|
||||
"path": "third_party/hamer",
|
||||
"commit": "3a01849f4148352e9260b69bf28b65d1671a4905",
|
||||
"origin": "https://github.com/geopavlakos/hamer.git",
|
||||
"local_changes": "",
|
||||
"snapshot": "working tree content; tracked files plus source additions"
|
||||
},
|
||||
{
|
||||
"path": "third_party/hamer/third-party/ViTPose",
|
||||
"commit": "d5216452796c90c6bc29f5c5ec0bdba94366768a",
|
||||
"origin": "https://github.com/ViTAE-Transformer/ViTPose.git",
|
||||
"local_changes": "",
|
||||
"snapshot": "working tree content; tracked files plus source additions"
|
||||
},
|
||||
{
|
||||
"path": "third_party/vipe",
|
||||
"commit": "b7cac64616763bc755009db79a1815c7cdc9b130",
|
||||
"origin": "https://github.com/nv-tlabs/vipe.git",
|
||||
"local_changes": "M vipe/pipeline/default.py\n M vipe/slam/networks/droid_net.py\n M vipe/streams/base.py\n?? configs/pipeline/dynhamr_cameras.yaml",
|
||||
"snapshot": "working tree content; tracked files plus source additions"
|
||||
},
|
||||
{
|
||||
"path": "third_party/Dyn-HaMR",
|
||||
"commit": "fa9cd7412c205fd15ee4139c8caacf79bf6167e6",
|
||||
"origin": "https://github.com/ZhengdiYu/Dyn-HaMR.git",
|
||||
"local_changes": "M dyn-hamr/HMP/fitting.py\n M dyn-hamr/data/vidproc.py\n M dyn-hamr/optim/base_scene.py\n M dyn-hamr/run_opt.py\n?? dyn-hamr/HMP/windowed.py\n?? third-party/Hand-BMC-pytorch-main/",
|
||||
"snapshot": "working tree content; tracked files plus source additions"
|
||||
},
|
||||
{
|
||||
"path": "third_party/Dyn-HaMR/third-party/DROID-SLAM",
|
||||
"commit": "385bffda51190fad66e135ff337ff95fd9c049dc",
|
||||
"origin": "https://github.com/ZhengdiYu/DROID-SLAM.git",
|
||||
"local_changes": "",
|
||||
"snapshot": "working tree content; tracked files plus source additions"
|
||||
},
|
||||
{
|
||||
"path": "third_party/Dyn-HaMR/third-party/hamer",
|
||||
"commit": "2a5b8ee852d87fdb56d65df7fb0e473c0d4f7648",
|
||||
"origin": "https://github.com/ZhengdiYu/hamer.git",
|
||||
"local_changes": "",
|
||||
"snapshot": "working tree content; tracked files plus source additions"
|
||||
},
|
||||
{
|
||||
"path": "third_party/Dyn-HaMR/third-party/hamer/third-party/ViTPose",
|
||||
"commit": "d5216452796c90c6bc29f5c5ec0bdba94366768a",
|
||||
"origin": "https://github.com/ViTAE-Transformer/ViTPose.git",
|
||||
"local_changes": "",
|
||||
"snapshot": "working tree content; tracked files plus source additions"
|
||||
},
|
||||
{
|
||||
"path": "third_party/Dyn-HaMR/third-party/vipe",
|
||||
"commit": "fdc79359ec6d063c86b9d203ca3d7054a9636823",
|
||||
"origin": "https://github.com/nv-tlabs/vipe",
|
||||
"local_changes": "",
|
||||
"snapshot": "working tree content; tracked files plus source additions"
|
||||
},
|
||||
{
|
||||
"path": "third_party/spider",
|
||||
"commit": "71238456bf97a7eeb3d0471aa31974e2d404d4ae",
|
||||
"origin": "https://github.com/facebookresearch/spider.git",
|
||||
"local_changes": "",
|
||||
"snapshot": "working tree content; tracked files plus source additions"
|
||||
}
|
||||
]
|
||||
@@ -0,0 +1,21 @@
|
||||
# HandFlow inference configuration
|
||||
# Shared by scripts/demo.py
|
||||
|
||||
model_yaml: configs/model.yaml
|
||||
|
||||
# The denoiser requires window_size (only this field; other dataset fields are eval-only and omitted)
|
||||
dataset:
|
||||
window_size: 16
|
||||
|
||||
inference:
|
||||
ode_steps: 3
|
||||
overlap_size: 8
|
||||
overlap_method: vblend # avg | vblend
|
||||
window_size: 16
|
||||
conf_threshold: 0.1 # HaMeR confidence threshold
|
||||
|
||||
# MANO model directory (environment variable, see README; overridable by run-script CLI args)
|
||||
eval:
|
||||
mano_root: ${oc.env:MANO_ROOT}
|
||||
|
||||
# fm_ckpt is passed via the run script's --fm_ckpt CLI arg (specified per inference)
|
||||
@@ -0,0 +1,29 @@
|
||||
model:
|
||||
fm:
|
||||
beta_dim: 10
|
||||
pose_dim: 48 # global_aa(3) + joints_aa(45), directly regressed per frame
|
||||
trans_dim: 3 # translation, directly regressed
|
||||
normalization_stats: ${oc.env:HANDFLOW_NORMALIZATION_STATS}
|
||||
|
||||
condition:
|
||||
hamer_embed_dim: 1280
|
||||
hamer_n_tokens: 192 # 16×12 patches
|
||||
n_joints: 21
|
||||
skeleton_pe_freqs: 6
|
||||
cmask_ratio: 0.2
|
||||
|
||||
frame_compressor:
|
||||
num_layers: 8 # cross-attention pooling layers (replaces Refiner + Compressor)
|
||||
|
||||
transformer:
|
||||
hidden_size: 512
|
||||
num_heads: 16
|
||||
depth: 8 # double-stream blocks
|
||||
depth_single_blocks: 16 # single-stream blocks
|
||||
mlp_ratio: 4.0
|
||||
qkv_bias: true
|
||||
dropout: 0.1
|
||||
|
||||
hamer:
|
||||
checkpoint: ${oc.env:HAMER_CKPT}
|
||||
frozen: true
|
||||
@@ -0,0 +1,185 @@
|
||||
# HandFlow
|
||||
|
||||
**Fully Generative 4D Hand Recovery with Flow Matching.**
|
||||
|
||||
HandFlow reconstructs 4D (temporal 3D) hand poses from monocular RGB video. At its core is a Rectified Flow Denoiser: HaMeR ViT-H extracts image features, combined with 2D skeletons as multimodal conditioning, to directly regress a sequence of MANO parameters. ViPE SLAM can optionally supply camera geometry to reconstruct the hand trajectory in the world coordinate frame.
|
||||
|
||||
> This repository is the **V1 open-source release**: it includes inference + visualization demos only, without training code.
|
||||
|
||||
---
|
||||
|
||||
## Pipeline
|
||||
|
||||
A single entry point, `scripts/demo.py`, produces two mp4s per run: a camera-view mesh **overlay** and an **orthographic** view. The `--fix_camera` flag selects how camera geometry is obtained:
|
||||
|
||||
```
|
||||
demo.py
|
||||
┌──────────────────────────────────────────────────────────────────┐
|
||||
│ Video (--input) │
|
||||
│ ├─ Intrinsics: --intrinsics fx,fy,cx,cy | auto | │
|
||||
│ │ unset → generic default 600,600,W/2,H/2 │
|
||||
│ ├─ Camera geometry (c2w): │
|
||||
│ │ --fix_camera → c2w = identity, ViPE skipped │
|
||||
│ │ (default) → ViPE SLAM estimates c2w │
|
||||
│ ├─ Online HaMeR → FM Denoiser → MANO FK → verts_cam (m) │
|
||||
│ └─ verts_world = c2w · verts_cam │
|
||||
│ │
|
||||
│ → overlay.mp4 clean camera-space MANO mesh over the RGB frames │
|
||||
│ → ortho.mp4 --fix_camera : level side view, fixed camera │
|
||||
│ (default) : world topdown — camera trajectory │
|
||||
│ line + per-frame hand mesh │
|
||||
└──────────────────────────────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
- **`--fix_camera`** (fixed camera, e.g. a tripod-mounted rig like DexYCB): ViPE is skipped, `c2w` is identity, and the orthographic view assumes a **level, fixed camera** — a side view of the camera-space hand motion (the camera stays at the origin; only the hand moves).
|
||||
- **default** (moving camera, e.g. a head-mounted rig like HOT3D): ViPE SLAM estimates `c2w`; the orthographic view is a world-space trajectory view (camera trajectory line + per-frame hand mesh). Because SLAM trajectory error accumulates, the world result looks slightly more "skewed" than the overlay; the orthographic view is exactly where this camera motion becomes visible.
|
||||
|
||||
---
|
||||
|
||||
## Installation
|
||||
|
||||
### 1. Clone (with submodules)
|
||||
|
||||
```bash
|
||||
git clone --recurse-submodules https://github.com/mxxu00/HandFlow.git
|
||||
cd HandFlow
|
||||
# If you already cloned without submodules:
|
||||
git submodule update --init --recursive
|
||||
```
|
||||
|
||||
Submodules:
|
||||
- `third_party/hamer` — HaMeR ViT-H backbone ([geopavlakos/hamer](https://github.com/geopavlakos/hamer))
|
||||
- `third_party/vipe` — ViPE SLAM ([nv-tlabs/vipe](https://github.com/nv-tlabs/vipe), only needed in moving-camera mode)
|
||||
|
||||
### 2. Inference + rendering environment
|
||||
|
||||
```bash
|
||||
bash setup_env.sh # creates conda env handflow (Python 3.10 + torch 2.7.0+cu128)
|
||||
conda activate handflow
|
||||
pip install -e third_party/hamer # HaMeR and its dependencies
|
||||
```
|
||||
|
||||
> `setup_env.sh` also builds **pytorch3d** (the Phong renderer) from source, which requires `nvcc`; if the build fails, see the [pytorch3d install guide](https://github.com/facebookresearch/pytorch3d/blob/main/INSTALL.md). `--fix_camera` mode then works without the ViPE environment.
|
||||
|
||||
### 3. ViPE environment (only needed in moving-camera mode, optional)
|
||||
|
||||
ViPE ships a CUDA extension that must be compiled in a **separate conda environment**; `demo.py` invokes it automatically via `conda run -n vipe` whenever `--fix_camera` is **unset**.
|
||||
|
||||
```bash
|
||||
bash setup_vipe_env.sh # creates conda env vipe (with nvcc + Eigen, compiles vipe_ext)
|
||||
```
|
||||
|
||||
> The first build requires internet access to download Eigen 3.4. ViPE's torch version (2.7.0+cu128) matches handflow, so the extension compiles correctly.
|
||||
|
||||
---
|
||||
|
||||
## Weights & Data
|
||||
|
||||
Download the following weights yourself and specify their paths via environment variables:
|
||||
|
||||
| Environment variable | Description | Source |
|
||||
|---|---|---|
|
||||
| `HAMER_CKPT` | `hamer.ckpt` (HaMeR ViT-H weights) | [HaMeR releases](https://github.com/geopavlakos/hamer) |
|
||||
| `DETECTOR_CKPT` | `detector.pt` (YOLO hand detector, **from WiLoR**) | [WiLoR](https://github.com/rolpotamias/WiLoR) `pretrained_models/detector.pt` |
|
||||
| `MANO_ROOT` | MANO model directory (`.pkl`) | [mano.is.tue.mpg.de](https://mano.is.tue.mpg.de) (register to download) |
|
||||
| `HANDFLOW_NORMALIZATION_STATS` | `normalization_stats.npz` | [HuggingFace](https://huggingface.co/mxxu00/HandFlow) |
|
||||
| `--fm_ckpt` | HandFlow FM checkpoint (`handflow_denoiser.pt`) | [HuggingFace](https://huggingface.co/mxxu00/HandFlow) (passed as a CLI argument) |
|
||||
|
||||
Example setup (add to `~/.bashrc`):
|
||||
```bash
|
||||
export HAMER_CKPT=/path/to/hamer.ckpt
|
||||
export DETECTOR_CKPT=/path/to/detector.pt
|
||||
export MANO_ROOT=/path/to/mano
|
||||
export HANDFLOW_NORMALIZATION_STATS=/path/to/normalization_stats.npz
|
||||
```
|
||||
|
||||
> **HandFlow weights** (`normalization_stats.npz` + `handflow_denoiser.pt`) are hosted on [HuggingFace](https://huggingface.co/mxxu00/HandFlow). Download both with:
|
||||
>
|
||||
> ```bash
|
||||
> hf download mxxu00/HandFlow --local-dir ./weights
|
||||
> ```
|
||||
>
|
||||
> Then point `HANDFLOW_NORMALIZATION_STATS` at `normalization_stats.npz` and pass `--fm_ckpt ./weights/handflow_denoiser.pt` to `demo.py`.
|
||||
|
||||
---
|
||||
|
||||
## Demo
|
||||
|
||||
Two short sample clips are bundled in `demo/` — one from [DexYCB](https://dex-ycb.github.io/) (fixed camera) and one from [HOT3D](https://hot3d.github.io/) (moving camera). The commands below run the full pipeline on each.
|
||||
|
||||
**DexYCB — fixed camera (no ViPE needed):**
|
||||
```bash
|
||||
conda activate handflow
|
||||
python scripts/demo.py --input demo/dexycb_sample.mp4 \
|
||||
--fm_ckpt ./weights/handflow_denoiser.pt \
|
||||
--intrinsics 616.2495,615.87665,321.6139,244.8281 \
|
||||
--fix_camera --output_dir output/dexycb
|
||||
# → output/dexycb/overlay.mp4
|
||||
# → output/dexycb/ortho_side.mp4
|
||||
```
|
||||
|
||||
**HOT3D — moving camera (requires the vipe env):**
|
||||
```bash
|
||||
python scripts/demo.py --input demo/hot3d_sample.mp4 \
|
||||
--fm_ckpt ./weights/handflow_denoiser.pt \
|
||||
--intrinsics 609.7035,609.7035,707.6459,704.7342 \
|
||||
--output_dir output/hot3d
|
||||
# → output/hot3d/overlay.mp4
|
||||
# → output/hot3d/ortho_topdown.mp4
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Usage
|
||||
|
||||
```bash
|
||||
# Your own video, known intrinsics, fixed camera (e.g. webcam on a stand)
|
||||
python scripts/demo.py --input video.mp4 --fm_ckpt <ckpt> \
|
||||
--intrinsics fx,fy,cx,cy --fix_camera --output_dir output/demo
|
||||
|
||||
# Your own video, known intrinsics, moving camera (ViPE estimates c2w only)
|
||||
python scripts/demo.py --input video.mp4 --fm_ckpt <ckpt> \
|
||||
--intrinsics fx,fy,cx,cy --output_dir output/demo
|
||||
|
||||
# Your own video, unknown intrinsics, moving camera (ViPE estimates intrinsics + c2w)
|
||||
python scripts/demo.py --input video.mp4 --fm_ckpt <ckpt> \
|
||||
--intrinsics auto --output_dir output/demo
|
||||
```
|
||||
|
||||
Key options:
|
||||
- `--intrinsics` `fx,fy,cx,cy` (recommended) | `auto` (ViPE estimates; moving-camera only) | unset → generic default `600,600,W/2,H/2`.
|
||||
- `--fix_camera` : fixed-camera mode (skip ViPE, `c2w=identity`, level side ortho view). Incompatible with `--intrinsics auto`.
|
||||
- `--view` `topdown|side` : ortho view direction. Defaults to `side` under `--fix_camera`, else `topdown`.
|
||||
|
||||
> Left-hand videos must be mirrored to right-hand first (the FM model is trained on right-hand data).
|
||||
|
||||
---
|
||||
|
||||
## Project Layout
|
||||
|
||||
```
|
||||
HandFlow/
|
||||
├── scripts/
|
||||
│ └── demo.py # unified demo: cam overlay + orthographic view
|
||||
├── demo/ # bundled demo clips (dexycb_sample.mp4, hot3d_sample.mp4)
|
||||
├── model/ # Rectified Flow Denoiser + feature extractors
|
||||
├── utils/ # inference / MANO FK / online HaMeR / ViPE worker
|
||||
├── visualization/ # pytorch3d Phong renderer (overlay + orthographic trajectory) + video I/O
|
||||
├── preprocessing/ # crop utilities (used by online_hamer)
|
||||
├── configs/ # model.yaml + inference.yaml (paths via ${oc.env:})
|
||||
├── third_party/
|
||||
│ ├── hamer/ # submodule
|
||||
│ └── vipe/ # submodule (optional, moving-camera mode only)
|
||||
├── setup_env.sh # inference environment
|
||||
└── setup_vipe_env.sh # ViPE environment (optional, moving-camera mode only)
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## TODO List
|
||||
|
||||
The current release (V1) ships **inference + visualization only**. The following components are planned for future release:
|
||||
|
||||
- [ ] Training code
|
||||
- [ ] Data preprocessing code
|
||||
- [ ] Evaluation code
|
||||
@@ -0,0 +1,217 @@
|
||||
# L20 左手专家轨迹 HDF5 交付要求
|
||||
|
||||
**接口版本:`l20_tracking_v1`**
|
||||
|
||||
**交付对象:视频重定向/示教数据同事**
|
||||
|
||||
**当前用途:浮动腕部与全部手指参考关节的轨迹回放,后续用于学习策略。**
|
||||
|
||||
本文是可独立分享的数据交付说明,与当前仓库校验器一致。它不是通用 HDF5 标准,也不是完整 DexSchema。
|
||||
请先交付 **1 条短样例**,确认坐标、尺度、关节映射后再批量导出。
|
||||
|
||||
## 1. 要交什么
|
||||
|
||||
交付重定向后的 **L20 左手机器人参考状态轨迹**:
|
||||
|
||||
- 腕部/手根的三维位置与姿态。
|
||||
- 约定模型的全部手指关节角,包括从动关节。
|
||||
- 同步时间戳与逐帧有效标记。
|
||||
- 模型身份、坐标变换、尺度和数据来源说明。
|
||||
|
||||
不是原始人手关键点、MANO 参数或电机实测指令。关节状态不能直接等同于执行器控制命令。
|
||||
本阶段不要求图像、物体位姿、奖励、接触标签、速度、加速度、电流或力矩;缺少这些内容不需要补零伪造。
|
||||
|
||||
## 2. 文件结构(必填)
|
||||
|
||||
文件扩展名建议使用 `.hdf5`;HDF5 根属性通过 `file.attrs` 保存,**不要创建名为 `attrs` 的 group**。
|
||||
|
||||
```text
|
||||
demonstrations.hdf5
|
||||
├── 根属性:schema_version、embodiment、hand_side、asset_sha256、root_link、
|
||||
│ provenance、source_description、metric_scale_provenance、scale_to_meters
|
||||
├── metadata/
|
||||
│ ├── joint_names UTF-8 [J]
|
||||
│ └── world_from_source float64 [4, 4]
|
||||
└── episodes/
|
||||
├── demo_000000/
|
||||
│ ├── time float64 [T]
|
||||
│ ├── wrist_position float32 [T, 3]
|
||||
│ ├── wrist_quaternion float32 [T, 4]
|
||||
│ ├── joint_position float32 [T, J]
|
||||
│ └── valid bool [T]
|
||||
└── demo_000001/
|
||||
└── 同样结构,可以有不同帧数 T
|
||||
```
|
||||
|
||||
- 至少一个 episode;名称必须为 `demo_` 加六位数字,如 `demo_000012`。
|
||||
- 每条 episode 至少两帧;各字段的 T 必须相同。
|
||||
- 一个文件共享模型、关节顺序及源世界标定;标定、模型或关节布局改变时另建文件。
|
||||
- 当前约定资产的 **J=21**,不是从“L20”名称推断出的自由度。21 个状态关节不代表 21 个独立执行器。
|
||||
|
||||
## 3. 根属性
|
||||
|
||||
除 `scale_to_meters` 外,下列属性均为非空 UTF-8 字符串。
|
||||
|
||||
| 属性 | 内容与要求 |
|
||||
| --- | --- |
|
||||
| `schema_version` | 固定为 `l20_tracking_v1` |
|
||||
| `embodiment` | 固定为 `L20` |
|
||||
| `hand_side` | 固定为 `left` |
|
||||
| `asset_sha256` | 双方约定的原始模型依赖包哈希,64 位小写十六进制;见第 6 节 |
|
||||
| `root_link` | 固定为当前模型的 `hand_base_link` |
|
||||
| `provenance` | 实际专家重定向数据填 `expert_retargeted`;程序生成的测试数据填 `synthetic` |
|
||||
| `source_description` | 视频/轨迹标识、重定向工具及版本、世界原点与 X/Y 朝向、源坐标约定、腕点到机器人根部的映射依据、质量限制 |
|
||||
| `metric_scale_provenance` | 米制尺度如何获得,标定方法/版本或估算方法、可信程度与已知误差;不得把估算写成实测 |
|
||||
| `scale_to_meters` | **有限正浮点标量**,记录已应用的源位置尺度系数;不是要求读取方再缩放 |
|
||||
|
||||
例:源位置以毫米表示,则 `scale_to_meters=0.001`;输出位置仍须已经是米。
|
||||
源数据已经是米时可填 `1.0`,但必须有实际依据,不能仅为通过校验填写。
|
||||
|
||||
## 4. 字段与坐标语义
|
||||
|
||||
| 路径 | 类型 / shape | 单位及语义 |
|
||||
| --- | --- | --- |
|
||||
| `metadata/joint_names` | HDF5 UTF-8 string `[J]` | `joint_position` 每列对应的模型关节名称,严格使用第 5 节顺序 |
|
||||
| `metadata/world_from_source` | `float64 [4,4]` | 已应用的源世界系到交付世界系刚体变换;矩阵平移部分单位 m |
|
||||
| `episodes/.../time` | `float64 [T]` | 秒;每条轨迹从 0 开始,严格递增 |
|
||||
| `episodes/.../wrist_position` | `float32 [T,3]` | `hand_base_link` 的 **link 原点**在交付世界系中的 `[x,y,z]`,单位 m |
|
||||
| `episodes/.../wrist_quaternion` | `float32 [T,4]` | **`[w,x,y,z]`**;将根 link 局部向量主动旋转到交付世界系 |
|
||||
| `episodes/.../joint_position` | `float32 [T,J]` | 参考关节角,单位 rad;零位和正方向与约定模型一致 |
|
||||
| `episodes/.../valid` | `bool [T]` | 整帧是否有效;不是每关节标记,不使用 0/1 整数数组代替 bool |
|
||||
|
||||
### 坐标与姿态
|
||||
|
||||
1. 交付世界系采用 **右手系、Z 向上**,在一个 episode 内固定,不随腕部移动。
|
||||
2. 根部位姿不是人手腕点、机器人质心或 USD default prim 位姿。重定向端必须先完成到
|
||||
`hand_base_link` 坐标系的映射,并说明映射依据。
|
||||
3. 记录的尺度和刚体变换满足:
|
||||
|
||||
```text
|
||||
p_world = R_world_from_source @ (scale_to_meters * p_source) + t_world_from_source
|
||||
```
|
||||
|
||||
此式用于源位置点;人手腕点到机器人根部的额外映射仍须由重定向端完成。
|
||||
导出的 `wrist_position` 和 `wrist_quaternion` **已经在交付世界系**。
|
||||
接收端不会再次应用该矩阵或尺度。
|
||||
4. `world_from_source` 最后一行为 `[0,0,0,1]`,旋转部分正交且行列式为 +1,不能混入尺度或镜像。
|
||||
已统一世界系时可以使用单位矩阵,但不能用单位矩阵掩盖未知标定。
|
||||
5. 有效四元数必须单位化,范数容差 `1e-4`;相邻两帧都有效时,四元数点积必须非负。
|
||||
`q` 与 `-q` 表示同一旋转,导出时应统一相邻符号;不能混用 `xyzw` 或欧拉角。
|
||||
6. 不要为了适配当前合成测试而把真实轨迹自动平移到 `z=0.4`,或把起始姿态强制设为单位四元数。
|
||||
场景对齐必须有明确、可追溯的坐标约定。
|
||||
|
||||
## 5. 关节顺序与联动
|
||||
|
||||
当前 `joint_names` 必须严格等于下列列表;这是数据清单顺序,**不是仿真运行时 DOF 顺序**。
|
||||
|
||||
```python
|
||||
joint_names = [
|
||||
"index_dip", "index_mcp_pitch", "index_mcp_roll", "index_pip",
|
||||
"middle_dip", "middle_mcp_pitch", "middle_mcp_roll", "middle_pip",
|
||||
"pinky_dip", "pinky_mcp_pitch", "pinky_mcp_roll", "pinky_pip",
|
||||
"ring_dip", "ring_mcp_pitch", "ring_mcp_roll", "ring_pip",
|
||||
"thumb_cmc_pitch", "thumb_cmc_roll", "thumb_cmc_yaw", "thumb_ip", "thumb_mcp",
|
||||
]
|
||||
```
|
||||
|
||||
- 不允许重名、缺列、多列、静默重排或用其他手型的关节代替。
|
||||
- 每个有效帧必须满足清单 `joints[].lower_rad/upper_rad`,校验容差为 `1e-6 rad`。
|
||||
精确限位以随交付约定的 JSON 清单为准,避免使用文档四舍五入值作为限位。
|
||||
- 当前源模型定义以下联动;数据中仍须包含这些从动关节列:
|
||||
|
||||
```text
|
||||
thumb_ip = 1.02 * thumb_mcp
|
||||
index_dip = 0.89 * index_pip
|
||||
middle_dip = 0.89 * middle_pip
|
||||
ring_dip = 0.89 * ring_pip
|
||||
pinky_dip = 0.89 * pinky_pip
|
||||
```
|
||||
|
||||
偏置均为 0 rad,有效帧等式残差绝对值不得超过 `1e-3 rad`。
|
||||
- 联动关系和各关节限位必须同时满足;不能单独裁剪从动关节而破坏联动。
|
||||
- 这些是约定模型的约束,不是已经核实的硬件电机协议。重定向映射有疑问时先确认,不要擅自独立控制从动关节。
|
||||
|
||||
## 6. 模型身份与随附清单
|
||||
|
||||
本仓库当前清单:
|
||||
|
||||
[`assets/robots/dex_hand/linkerhand_g20_left/tracking_manifest.json`](assets/robots/dex_hand/linkerhand_g20_left/tracking_manifest.json)
|
||||
|
||||
当前原始模型依赖包 `asset_sha256`:
|
||||
|
||||
```text
|
||||
6c8f35358f481cf604ee802588017f30c1661c0f38e1187c49e5340db3830538
|
||||
```
|
||||
|
||||
这是**当前版本快照**。请由仿真方提供同版 `tracking_manifest.json` 与约定模型,数据方确认使用该模型后填写哈希。
|
||||
不得把上述哈希复制到实际使用了其他模型的数据中。
|
||||
|
||||
- 哈希不是 HDF5 文件自身哈希,不是仅入口 USD 的哈希,也不是浮动控制覆盖层的哈希。
|
||||
- 它由完整解析资源包的相对路径与各文件 SHA-256 汇总生成,使用仓库工具获取,不自行猜测算法。
|
||||
- 模型依赖内容或名称变化时必须重新对齐。分享本文给同事时,**请同时附上同版清单**,不要只发送失效的仓库相对链接。
|
||||
|
||||
## 7. 时间同步、无效帧与缺失数据
|
||||
|
||||
- `time[0]` 必须精确等于 `0.0`,之后严格递增;不得有重复、倒序时间戳。
|
||||
- 允许非均匀采样,不强制视频帧率等于仿真控制频率。保留实际时间间隔,不通过伪造时间戳“拉齐”数据。
|
||||
- 同一帧腕部、姿态和关节角必须已经同步;异步数据须由交付方先明确对齐方式。
|
||||
- 任一必需部分缺失或不可信,整帧 `valid=false`。至少一帧有效;全部无效的 episode 会被拒绝。
|
||||
- **所有数值都必须有限,包括无效帧,不能出现 NaN/Inf。** 无效帧可以用明确的有限占位值,
|
||||
但必须保留 `valid=false`,并在来源说明中解释。无效帧四元数不要求单位化。
|
||||
- 必填 dataset 不能省略;禁止把缺失帧补零后标成有效,或把插值/估算值冒充测量。
|
||||
- 当前重采样 CLI 只接受全有效 episode,不跨无效片段插值。建议交付方将连续有效片段显式切成
|
||||
独立 episode,并为每段重新设置从 0 开始的时间;原片段来源需可追溯。
|
||||
|
||||
## 8. 导出注意事项
|
||||
|
||||
使用 `h5py` 时,字符串必须显式声明 UTF-8,数值类型按契约转换:
|
||||
|
||||
```python
|
||||
# 仅展示 dtype 写法;不是完整数据生成器,不包含真实轨迹。
|
||||
import h5py
|
||||
import numpy as np
|
||||
|
||||
utf8 = h5py.string_dtype(encoding="utf-8")
|
||||
# metadata.create_dataset("joint_names", data=joint_names, dtype=utf8)
|
||||
# group.create_dataset("time", data=np.asarray(time, dtype=np.float64))
|
||||
# group.create_dataset("wrist_position", data=np.asarray(position, dtype=np.float32))
|
||||
# group.create_dataset("wrist_quaternion", data=np.asarray(quaternion_wxyz, dtype=np.float32))
|
||||
# group.create_dataset("joint_position", data=np.asarray(q, dtype=np.float32))
|
||||
# group.create_dataset("valid", data=np.asarray(valid, dtype=np.bool_))
|
||||
```
|
||||
|
||||
不要用 `dtype="S"` 或普通 ASCII bytes 数组写 `joint_names`。不要用 float64 替代契约中的 float32,
|
||||
也不要用整数时间戳或整数 `valid` 代替约定类型。
|
||||
|
||||
## 9. 交付及验收流程
|
||||
|
||||
### 数据方交付
|
||||
|
||||
1. 一份包含一条短轨迹的 `.hdf5` 样例,优先选择连续、全有效、缓慢运动的片段。
|
||||
2. 对应模型清单/版本确认。
|
||||
3. 明确的标定与重定向说明;写入根属性,也可附独立说明文件。
|
||||
4. 可选:对应原视频片段或可视化,便于人工核对方向、尺度与动作,不是本格式必填字段。
|
||||
|
||||
### 仿真方校验
|
||||
|
||||
从仓库根目录运行:
|
||||
|
||||
```bash
|
||||
export PYTHONPATH="$PWD/source/dex_workbench${PYTHONPATH:+:$PYTHONPATH}"
|
||||
~/isaacsim/python.sh -m dex_workbench_tracking.cli validate /path/to/demonstrations.hdf5 \
|
||||
--manifest assets/robots/dex_hand/linkerhand_g20_left/tracking_manifest.json
|
||||
```
|
||||
|
||||
该校验不启动仿真。仿真方本地的 `python.sh` 路径不是数据文件的依赖;具备项目模块、NumPy、h5py
|
||||
的普通 Python 环境也可执行同一 `-m` 命令。
|
||||
|
||||
- **PASS**:schema、模型身份、关节顺序、限位和已记录联动一致。
|
||||
- **FAIL**:修正具体报错后重新交付,不绕过校验器或放宽断言。
|
||||
- 不加 `--manifest` 只检查格式,不能证明模型/关节兼容。
|
||||
- 校验通过后仍需人工核对坐标和标定,再做受限动力学回放。格式通过不等于重定向正确、动力学可执行、
|
||||
策略已经可训练或真机可运行。
|
||||
|
||||
**注意:** 当前诊断脚本要求腕部及全部五组联动都有足够运动,用于发现失效约束;这是诊断用例要求,
|
||||
不是 HDF5 格式要求。合法的静止或局部手指轨迹不应为了满足诊断而添加伪造动作。
|
||||
|
||||
仿真实现与已有验证结果另见 [`L20_TRACKING.md`](L20_TRACKING.md)。本文只约定数据交付,不要求数据方运行 Isaac Sim。
|
||||
@@ -0,0 +1,99 @@
|
||||
# Dyn-HaMR → dex-retargeting → SPIDER 集成说明
|
||||
|
||||
## 组件和职责
|
||||
|
||||
| 层 | 上游/实现 | 输入与输出 |
|
||||
|---|---|---|
|
||||
| 人手重建 | Dyn-HaMR、HaMeR、ViPE | 视频 → MANO 参数与估计世界轨迹 |
|
||||
| 另一条重建入口 | HandFlow | 视频 → 时序 MANO 与相机/世界坐标结果 |
|
||||
| 机器人重定向 | dex-retargeting 0.5.0 + scripts/l20_* | 21 点人手几何 → L20 右手关节及浮动腕部 |
|
||||
| 接触优化实验 | SPIDER 上游 runner + L20 场景适配 | 机器人/物体参考轨迹 → 接触引导优化与回放检查 |
|
||||
| 数据交付 | export_l20_right_hdf5.py | 机器人参考轨迹 → HDF5 |
|
||||
|
||||
集成脚本与上游项目分别保留;不能把上游算法或整个仓库归为本地个人实现。
|
||||
|
||||
## 当前入口和顺序
|
||||
|
||||
以下命令从项目根目录执行,输入文件、权重、模型和环境须先准备。这里只列出已存在入口,不代表干净机器已端到端复现。
|
||||
|
||||
### 1. 重建结果
|
||||
|
||||
Dyn-HaMR 原生入口为 `third_party/Dyn-HaMR/dyn-hamr/run_opt.py`,参数见其 README 和配置。`scripts/resume_dynhamr_2047635068.py` 是已有缓存的恢复脚本,依赖 `output/vipe_weights_downloads.txt` 及先前生成的缓存,不能当作新视频通用入口。
|
||||
|
||||
本地 ViPE 调用已采用 Python CLI;相关修改随源码快照保留。环境与数据路径仍需按部署机器配置。
|
||||
|
||||
本例期望重建产物位于 `output/results_2047635068/world_results.npz`。
|
||||
|
||||
```bash
|
||||
third_party/Dyn-HaMR/.dynhamr/bin/python scripts/export_dynhamr_dex.py
|
||||
```
|
||||
|
||||
产物:`output/dex_2047635068/human_joints.npz`。导出器要求恰好一条完整右手轨迹;MANO convention 按 Dyn-HaMR 原模型计算。
|
||||
|
||||
### 2. 重定向与平滑
|
||||
|
||||
```bash
|
||||
MUJOCO_GL=osmesa .dex/bin/python scripts/run_dex_l20.py
|
||||
MUJOCO_GL=osmesa .dex/bin/python scripts/run_dex_l20.py \
|
||||
--output-dir output/l20_2047635068/jitter_audit/corrected
|
||||
MUJOCO_GL=osmesa .dex/bin/python scripts/smooth_l20_offline.py
|
||||
```
|
||||
|
||||
该分支使用 JSON 非线性标定联动。后续 SPIDER/完整回放构建器按原 URDF 重新展开线性联动;两种约束不能混称为同一个模型。
|
||||
|
||||
HandFlow 分支的入口为 `scripts/demo.py` → `export_handflow_dex.py` → `retarget_l20_video.py`;后者支持 `--input` 和 `--output-dir`,直接使用 URDF 线性联动。
|
||||
|
||||
### 3. SPIDER 实验
|
||||
|
||||
先准备 `third_party/l20_assets/L20/RIGHT/` 和 `third_party/bottle_model_parametric/bottle.xml` 及依赖资源。
|
||||
|
||||
```bash
|
||||
MUJOCO_GL=osmesa .spider/bin/python scripts/prepare_spider_l20.py
|
||||
MUJOCO_GL=osmesa .spider/bin/python scripts/run_spider_l20_contact.py
|
||||
MUJOCO_GL=osmesa .spider/bin/python scripts/evaluate_spider_l20.py
|
||||
```
|
||||
|
||||
准备脚本选择第 600–659 帧的 2 秒片段,并估算瓶子相对位姿。必须查看实际运行日志和评估结果;程序开始运行不等于优化完成或稳定抓取。
|
||||
|
||||
### 4. 1333 帧完整运动学回放和导出
|
||||
|
||||
```bash
|
||||
MUJOCO_GL=osmesa .spider/bin/python scripts/build_l20_full_replay.py
|
||||
MUJOCO_GL=osmesa .spider/bin/python scripts/verify_l20_full_replay.py
|
||||
MUJOCO_GL=osmesa .spider/bin/python scripts/play_l20_full.py --render
|
||||
MUJOCO_GL=osmesa .dex/bin/python scripts/export_l20_right_hdf5.py \
|
||||
--source output/l20_full_replay/motion.npz \
|
||||
--scene output/l20_full_replay/scene.xml \
|
||||
--output output/hdf5_delivery/right_1333
|
||||
```
|
||||
|
||||
`.dex` 环境还需可导入 h5py。最终完整回放为运动学指定状态及接触配准,不是 44 秒 SPIDER 动力学成功轨迹。瓶子位姿是估计注册;HDF5 只导出手部状态,右手约定与原左手验收要求不同。详见 HDF5 根属性和 `docs/HDF5_REQUIREMENTS.md`。
|
||||
|
||||
## 环境
|
||||
|
||||
本机使用独立环境:`.venv`(HandFlow/ViPE)、`third_party/Dyn-HaMR/.dynhamr`、`.dex`、`.spider`。保留兼容的 Torch/CUDA,避免将多个上游依赖盲目装入一个环境。
|
||||
|
||||
- 原 HandFlow 的 `requirements.txt` 和 setup 脚本保留历史默认值,不代表当前集成环境锁。
|
||||
- `requirements/retarget.txt` 为重定向与导出依赖列表;`retarget_observed.json` 为本机观测版本,不是完整锁文件。
|
||||
- Dyn-HaMR/SPIDER/ViPE 的安装步骤分别见所附上游 README、requirements 或 pyproject 文件;GPU 扩展需匹配实际 Python、Torch、CUDA。
|
||||
- 示例解释器路径表示角色约定,可替换为部署环境的对应解释器。不要复制本机 venv 或其中的绝对路径 `.pth` 文件。
|
||||
|
||||
## 资源与代码托管
|
||||
|
||||
权重、MANO 参数、用户视频、缓存和实验输出不进源码包。L20 模型及标定文件也作为独立资源管理,接收方按 `EXTERNAL_ASSETS.json` 恢复原相对路径。资源使用条件分别由提供方约定;根 LICENSE 不覆盖所有第三方资源。
|
||||
|
||||
源码导出脚本将本机实际第三方源码以普通目录快照保存,包含本地改动,保留 LICENSE。它不依赖平台能递归拉取 Git 子模块。第三方版本和文件 SHA-256 记录在交付清单中;dex-retargeting 通过版本依赖安装。被排除的上游文件记录在 `EXCLUDED_FILES.json`,不可把源码包误认为含全部运行资源的离线安装包。
|
||||
|
||||
## 上传到本地代码托管平台
|
||||
|
||||
先在平台创建空仓库,解压源码包,在解压后的项目目录执行:
|
||||
|
||||
```bash
|
||||
git init -b main
|
||||
git add .
|
||||
git commit -m "Import integrated hand reconstruction and L20 retargeting pipeline"
|
||||
git remote add origin <本地平台仓库地址>
|
||||
git push -u origin main
|
||||
```
|
||||
|
||||
源码包没有原 Git 历史或 remote;当前工作目录的上游 origin 未修改。提交前可查看 `git status --short`;大数据放平台附件或独立存储。
|
||||
@@ -0,0 +1,478 @@
|
||||
# HandFlow 项目:手部动作重定向流程与踩坑记录
|
||||
|
||||
整理日期:2026-09-15。依据当前本地脚本、输出 README 与 JSON 验证报告整理;本次仅整理文档,没有重新执行重建、重定向或动力学优化。文中数值是保存的运行结果,不代表新一轮实验。
|
||||
|
||||
## 1. 当前实现与结论边界
|
||||
|
||||
本项目的机器人重定向流程,将视频重建的人手动作转换为机器人手的关节轨迹,再保留世界坐标下的腕部运动,并增加物体接触处理。
|
||||
|
||||
当前机器人分支实际消费的是 **Dyn-HaMR 的 MANO 重建结果**。仓库主 README 中 HandFlow FM 推理入口与这里新增的机器人脚本应区分;不能将这次全部输出描述为 HandFlow FM 模型直接生成的机器人动作。
|
||||
|
||||
| 阶段 | 已有结果 | 能支持的结论 | 尚不能支持的结论 |
|
||||
|---|---|---|---|
|
||||
| Dyn-HaMR 重建 | 1333 帧、两条手轨迹、相机参数与世界坐标参数 | 全长重建及渲染流程完成 | 已达到实测三维精度 |
|
||||
| LEAP 基线 | 四指、固定腕部的 dex 重定向 | 人手关键点到机器人关节链路可运行 | 移动腕部或抓瓶成功 |
|
||||
| L20 基础重定向 | 16 主动 + 5 被动关节,移动腕部 | 限位、联动与前向运动学一致性通过 | 真实硬件可直接执行 |
|
||||
| 离线平滑 | 完整 1333 帧手部轨迹 | 保存指标显示连续性改善 | 实时滤波或无偏恢复真实运动 |
|
||||
| 全长接触回放 | 44.43 秒完整视频 | 指定姿态下接触几何与连续性检查完成 | 44 秒动力学稳定抓取 |
|
||||
| SPIDER | 2 秒片段优化及独立 CPU/GPU 回放 | GPU 上保持接触,明显优于原控制基线 | 通过保持门限、跨后端可靠抓取或硬件验证 |
|
||||
|
||||
**当前完整视频是运动学接触回放。SPIDER 只验证了两秒片段,保持门限未通过,独立 CPU 回放出现滑脱。**
|
||||
|
||||
## 2. 数据流、版本与关键文件
|
||||
|
||||
```text
|
||||
2047635068.mp4
|
||||
→ Dyn-HaMR / VIPE / 运动先验等上游处理
|
||||
→ output/results_2047635068/world_results.npz
|
||||
→ export_dynhamr_dex.py
|
||||
→ output/dex_2047635068/human_joints.npz
|
||||
├─ run_dex_leap.py → LEAP 固定腕部基线
|
||||
└─ run_dex_l20.py → L20 JSON 非线性联动、移动腕部
|
||||
→ jitter_audit/corrected/motion.npz
|
||||
→ smooth_l20_offline.py
|
||||
→ jitter_audit/temporal/motion.npz
|
||||
├─ prepare_spider_l20.py → 线性联动、2 秒参考与场景
|
||||
│ → run_spider_l20_contact.py → SPIDER 动力学结果
|
||||
│ → 独立 GPU / CPU 控制回放
|
||||
└─ build_l20_full_replay.py
|
||||
→ 复用已准备的 SPIDER 场景及初始配准
|
||||
→ 线性联动、全长接触配准
|
||||
→ output/l20_full_replay/
|
||||
```
|
||||
|
||||
这里存在真实的文件依赖:全长回放构建脚本读取 SPIDER **准备阶段**生成的场景和初始参考,不要求先跑完 SPIDER 优化。它没有把 SPIDER 的两秒优化结果扩展成整段物理结果。
|
||||
|
||||
| 用途 | 文件 |
|
||||
|---|---|
|
||||
| 导出人手关键点 | [export_dynhamr_dex.py](../scripts/export_dynhamr_dex.py) |
|
||||
| L20 重定向入口 | [run_dex_l20.py](../scripts/run_dex_l20.py) |
|
||||
| JSON 标定与非线性联动 | [l20_calibrated.py](../scripts/l20_calibrated.py) |
|
||||
| URDF/MJCF 与关键点构建 | [l20_model_source.py](../scripts/l20_model_source.py) |
|
||||
| dex 目标函数修复 | [l20_consistent_optimizer.py](../scripts/l20_consistent_optimizer.py) |
|
||||
| 离线平滑 | [smooth_l20_offline.py](../scripts/smooth_l20_offline.py) |
|
||||
| 抖动诊断 | [audit_l20_jitter.py](../scripts/audit_l20_jitter.py) |
|
||||
| 全长接触回放构建 | [build_l20_full_replay.py](../scripts/build_l20_full_replay.py) |
|
||||
| 全长独立验证 | [verify_l20_full_replay.py](../scripts/verify_l20_full_replay.py) |
|
||||
| 交互回放与视频导出 | [play_l20_full.py](../scripts/play_l20_full.py) |
|
||||
| SPIDER 任务准备 | [prepare_spider_l20.py](../scripts/prepare_spider_l20.py) |
|
||||
| SPIDER 官方入口封装 | [run_spider_l20_contact.py](../scripts/run_spider_l20_contact.py) |
|
||||
|
||||
## 3. 第一步:从 MANO 参数得到 21 点与腕部轨迹
|
||||
|
||||
输入 `world_results.npz` 包含 `pose_body`、`betas`、`root_orient`、`trans`、`is_right` 等字段。导出脚本要求恰好有一条整段一致的右手轨迹,否则断言失败。
|
||||
|
||||
对选定轨迹,按批次使用 Dyn-HaMR 自带 MANO wrapper 计算:
|
||||
|
||||
1. 将手指姿态展开为每帧 45 维 axis-angle 参数,使用该轨迹的固定手形 `betas`。
|
||||
2. 根旋转和整体平移设为零,运行 MANO 前向运动学。
|
||||
3. 所有点减去腕点,得到局部 21 点 `joints`,形状为 `(T,21,3)`。
|
||||
4. 使用重建的根旋转、平移再次运行 MANO,取第 0 点作为 `wrist_world`。
|
||||
5. 保存根旋转、原始平移、源轨迹编号和源文件路径,便于追溯。
|
||||
|
||||
世界腕点需要前向计算,因为模型的腕点定义包含手形等影响;不能直接用 MANO `trans` 代替。沿用原 wrapper 也避免额外叠加手部均值造成姿态约定变化。
|
||||
|
||||
当前脚本把帧率写为 30,输入视频恰好也是 30 fps。换视频时必须修改这一约定以及所有下游时间轴。
|
||||
|
||||
来源:[上游结果说明](../output/results_2047635068/README.md)、[导出脚本](../scripts/export_dynhamr_dex.py)。
|
||||
|
||||
## 4. 第二步:对齐掌面,重新构造机器人目标
|
||||
|
||||
### 4.1 为什么需要重新构造目标
|
||||
|
||||
人手与 L20 的骨长、指根位置、自由度不同。直接复制人手位置会要求机器人伸到无法到达的位置;直接复制 MANO 角度也没有一一对应的机械含义。
|
||||
|
||||
当前做法是保留手指骨段方向,使用机器人自身骨长和指根位置。
|
||||
|
||||
### 4.2 掌面坐标系
|
||||
|
||||
令腕点为 `p[0]`,食指、中指、小指掌指关节点为 `p[5]`、`p[9]`、`p[17]`:
|
||||
|
||||
- `z = normalize(p[9] - p[0])`;
|
||||
- 从 `p[5] - p[17]` 中去掉沿 z 的分量,归一化得到 y;
|
||||
- `x = cross(y,z)`;
|
||||
- 组成基底矩阵 `B = [x,y,z]`。
|
||||
|
||||
机器人中立姿态得到 `B_robot`,每帧人手得到 `B_human`,代码计算:
|
||||
|
||||
```python
|
||||
align = B_human @ B_robot.T
|
||||
aligned = (human_points - human_points[0]) @ align
|
||||
```
|
||||
|
||||
上式关键点采用行向量表达;解释旋转时不要将其与列向量乘法顺序混用。
|
||||
|
||||
### 4.3 骨段方向映射
|
||||
|
||||
目标先复制机器人中立姿态。五条链是 `(1,2,3,4)`、`(5,6,7,8)`、`(9,10,11,12)`、`(13,14,15,16)`、`(17,18,19,20)`。链起点保持机器人中立位置,随后逐节构造:
|
||||
|
||||
\[
|
||||
p_b^{target}=p_a^{target}+L_{ab}^{robot}
|
||||
\frac{p_b^{aligned}-p_a^{aligned}}{\|p_b^{aligned}-p_a^{aligned}\|}.
|
||||
\]
|
||||
|
||||
所以代码使用的是逐骨段长度适配,不是一个整体缩放系数。腕点与指根位置基本固定在机器人掌面局部空间,手指方向随视频改变。
|
||||
|
||||
机器人侧建立 `landmark_00` 至 `landmark_20`:关节标记主要来自关节原点,指尖来自末端 STL 端部表面的近似点。报告中的指尖拟合误差依赖这些标记定义。
|
||||
|
||||
## 5. 第三步:有限位逆运动学与联动
|
||||
|
||||
### 5.1 PositionOptimizer 的求解对象
|
||||
|
||||
基础 L20 使用 dex-retargeting 0.5.0 的 position 模式,SLSQP 优化 16 个主动角,通过前向运动学匹配 21 个目标点。核心目标是:
|
||||
|
||||
\[
|
||||
\min_{q_t}\ \mathrm{SmoothL1}(FK(q_t)-P_t^{target})
|
||||
+\lambda\|q_t-q_{t-1}\|^2,
|
||||
\quad q_{min}\le q_t\le q_{max}.
|
||||
\]
|
||||
|
||||
这里 SmoothL1 按位置分量聚合;不是只匹配五个指尖的欧氏距离。`normal_delta=0.002`,dex 内置低通通过 `low_pass_alpha=-1` 关闭,后续由脚本显式平滑。以上目标表达的是修复后的实现。
|
||||
|
||||
初始姿态来自 JSON 基准命令,主动角先裁剪到有效范围,再通过联动适配器展开。后续顺序逐帧求解,利用前一帧状态。
|
||||
|
||||
### 5.2 JSON 非线性联动版本
|
||||
|
||||
提供资产位于 `third_party/l20_assets/L20/RIGHT/`,使用 `linkerhand_g20_right.urdf` 与 `g20_right_G20.json`。
|
||||
|
||||
标定解释为:
|
||||
|
||||
\[
|
||||
q_j(u)=angle\_rad_j[u]+urdf\_zero\_offset\_rad_j.
|
||||
\]
|
||||
|
||||
具体实现:
|
||||
|
||||
1. 按 `motor_index` 排序 16 个主动关节。
|
||||
2. 每个被动关节找到共享电机编号的主动关节。
|
||||
3. 遍历 0–255 命令,排除会使同电机任一关节超过 URDF 限位的值。
|
||||
4. 检查有效命令区间连续,并形成主动角上下限。
|
||||
5. 从同命令下的主动角、被动角建立分段线性插值曲线。
|
||||
6. 前向计算时按曲线展开被动角;反向 Jacobian 使用当前插值段斜率。
|
||||
|
||||
生成的 URDF 去除原线性 mimic,生成的 MJCF 去除对应 equality,避免原线性关系和 JSON 关系重复生效。原始资产保留不变。
|
||||
|
||||
### 5.3 量化命令
|
||||
|
||||
求解后按主动角寻找最近的合法命令值,生成 20 槽位 `command_u8`。11–14 槽位保持 baseline;20 槽位不等于 20 个独立电机。命令再解码成全部关节角,检查量化误差。
|
||||
|
||||
修正版本报告中的最大量化角误差约 0.00604 rad。该值只是当前 LUT 下的往返误差,不是硬件执行误差。JSON 中还存在复用小指曲线的字段,不能描述成每根手指均独立实测。
|
||||
|
||||
### 5.4 后续线性联动版本
|
||||
|
||||
完整接触回放与 SPIDER 按既有需求使用 URDF 线性关系:
|
||||
|
||||
\[
|
||||
q_{passive}=a q_{active}+b.
|
||||
\]
|
||||
|
||||
它们从平滑轨迹读取主动角,再按线性关系重算被动角,同时收紧主动角范围以满足被动角限位。**它们并非重新从人手目标完整执行一次线性版 dex 拟合,也没有继续使用 JSON 被动角。**
|
||||
|
||||
因此基础 JSON 轨迹的拟合指标和命令文件,不能直接作为最终线性回放的拟合精度与硬件命令。
|
||||
|
||||
## 6. 第四步:保留腕部运动
|
||||
|
||||
每帧机器人腕部世界旋转由 `root_orient` 对应旋转矩阵乘以 `align` 得到。腕部世界位置使用 MANO 计算的 `wrist_world`。
|
||||
|
||||
展示时以首帧腕部旋转的逆作为固定场景旋转,将首帧腕部移至 `(0,0,0.2)`。这一变换只固定一次,整段运动仍然保留。保存原始世界腕点、世界旋转、场景旋转和平移,便于追溯。
|
||||
|
||||
四元数保存顺序是 WXYZ;SciPy 的默认顺序是 XYZW,读写需要显式转换。相邻四元数若点积小于零,则翻转符号,避免等价四元数符号跳变。
|
||||
|
||||
基础回放使用 freejoint;SPIDER 和最终全长场景采用三个平移、三个旋转标量坐标。其 CSV/NPZ 布局与基础手部文件不同,不可按列号直接混读。
|
||||
|
||||
这里没有机械臂模型、机械臂 IK 或机械臂动力学,腕部运动不能直接转换成真实机械臂执行轨迹。
|
||||
|
||||
## 7. 第五步:抖动诊断与离线平滑
|
||||
|
||||
基础脚本先对主动角做 `sigma=1` 帧高斯平滑,再重算被动角;腕部最初直接沿用重建结果。
|
||||
|
||||
后续从修复版基础输出出发,再对主动角做 `sigma=2.5` 帧离线高斯平滑。被动角仍由主动角重新计算,不能独立滤波破坏联动。腕部位置同步平滑;旋转在相对首帧的旋转向量空间处理,脚本断言该片段旋转范数小于 π/2。
|
||||
|
||||
| 版本 | 最大关节帧间变化 | 关节二阶差分 RMS | 平均指尖拟合误差 |
|
||||
|---|---:|---:|---:|
|
||||
| 原基础 L20 | 4.139° | 0.12148° | 10.675 mm |
|
||||
| 仅修复目标函数 | 2.418° | 0.11907° | 10.680 mm |
|
||||
| 修复 + 离线平滑 | 1.529° | 0.03146° | 10.701 mm |
|
||||
| 旧接触修正 | 54.375° | 1.67298° | 目标改变,不直接比较 |
|
||||
|
||||
平滑版本相对原基础关节二阶差分下降约 74.1%;腕部位置二阶差分分量 RMS 降至 0.114 mm,但腕点相对原轨迹最大改变 11.92 mm。
|
||||
|
||||
二阶差分包含真实动作加速度,不能把它直接叫作测得的噪声。高斯平滑使用未来帧,当前实现属于离线处理;输出不是零延迟实时控制器。
|
||||
|
||||
证据:[抖动审查](../output/l20_2047635068/jitter_audit/README.md)、[平滑验证](../output/l20_2047635068/jitter_audit/temporal/validation.json)。
|
||||
|
||||
## 8. 第六步:完整视频接触配准
|
||||
|
||||
### 8.1 物体与位姿假设
|
||||
|
||||
当前使用参数化瓶子,高 157.856 mm、最大直径 54.110 mm;瓶肩与瓶颈是规则设计形状。它替换了早期轮廓异常的瓶子模型。
|
||||
|
||||
没有当前视频逐帧物体位姿,也缺少原三目相机到 VIPE 世界坐标的外参。当前构造固定的瓶子到手掌相对变换:平移 `[0.12,-0.065,0.15]` 米,绕 x 轴旋转 `-π/2`。再使用第 600 帧做场景整体对齐。
|
||||
|
||||
第 130–160 帧被设为估计抓取过渡,使用平滑权重启用修正。160 帧前瓶子目标固定,后续随估计手掌关系运动。这是人为配准假设,不是视频物体跟踪。
|
||||
|
||||
### 8.2 求解与连续性
|
||||
|
||||
每个关键帧优化 22 个修正量:16 个主动角、3 个腕部平移、3 个腕部旋转。
|
||||
|
||||
残差包含指末端表面距离、全手对物体穿透惩罚、姿态偏移正则和相邻关键帧修正量差异。当前全长脚本的主要尺度为:
|
||||
|
||||
| 项 | 代码含义 |
|
||||
|---|---|
|
||||
| 指末端距离 | `(distance + 0.0002) / 0.004`,目标轻微压入约 0.2 mm |
|
||||
| 穿透惩罚 | 超过约 0.7 mm 的穿透进入惩罚,以 1 mm 归一化 |
|
||||
| 手指连续性 | 与上一求解修正量的差异进入残差 |
|
||||
| 腕部修正边界 | 平移各轴 ±65 mm,旋转向量各分量 ±0.55 rad |
|
||||
| 求解频率 | 从第 160 帧开始,每两帧一个关键帧,补入末帧 |
|
||||
| 求解器 | SciPy `least_squares`,每关键帧最多 65 次函数评估 |
|
||||
| 后处理 | 插值修正量,`sigma=1.5` 帧高斯平滑,再乘过渡权重 |
|
||||
|
||||
前一关键帧解作为下一次初值。插值和平滑后再次展开关节、检查限位与接触。该流程有连续性代价,但不是带硬速度/加速度约束的全局轨迹最优化。
|
||||
|
||||
### 8.3 回放与检查
|
||||
|
||||
输出场景移除驱动器,由播放器每帧赋值 `qpos` 后调用 `mj_forward`。瓶子位姿同样由文件指定;场景地板只提供参照,起始瓶子没有物理支撑。
|
||||
|
||||
保存报告显示:1333 帧有限,限位违规为零,线性联动误差为零,主动角最大帧间变化 1.465°;持瓶阶段 1173 帧中 1172 帧存在手瓶接触。
|
||||
|
||||
独立 `mj_forward` 接触流形检查最大穿透约 1.008 mm。距离查询保存了约 32.018 mm 的孤立尖峰,与接触流形不一致,应保留两者并说明差异,不能只挑有利数字,也不能据此单独确认同等深度的真实几何穿透。
|
||||
|
||||
来源:[完整回放说明](../output/l20_full_replay/README.md)、[验证报告](../output/l20_full_replay/validation.json)。
|
||||
|
||||
## 9. 第七步:SPIDER 两秒动力学实验
|
||||
|
||||
准备阶段使用第 600–659 帧平滑轨迹、URDF 线性联动和参数化瓶子。初始帧通过多初值 least-squares 配准,并将同一初始修正施加到整个片段,避免逐帧独立接触修补。以五个末端最近表面点建立接触标记和物体参考点。
|
||||
|
||||
参考插值到 2.5 ms,生成 801 个参考状态,执行 800 步。腕部和手指有位置驱动器,瓶子有搜索用虚拟驱动器。
|
||||
|
||||
| 配置 | 保存值 |
|
||||
|---|---:|
|
||||
| 后端 | MuJoCo Warp GPU |
|
||||
| 仿真步长 | 0.0025 s |
|
||||
| 控制更新周期 | 0.1 s |
|
||||
| 预测时长 | 0.6 s |
|
||||
| 样本数 | 512 |
|
||||
| 每次最大迭代数 | 6 |
|
||||
| 指关节力矩限制 | ±1 Nm |
|
||||
| 腕平移力限制 | ±200 N |
|
||||
| 腕旋转力矩限制 | ±20 Nm |
|
||||
| 瓶子质量 / 摩擦系数 | 0.2308 kg / 0.6,未实测 |
|
||||
|
||||
运行封装调用官方 `examples/run_mjwp.py`。快速入口不支持 `contact_guidance`,所以使用标准入口。搜索阶段物体虚拟辅助增益随迭代衰减,最后清零;独立执行回放的物体增益与偏置为零,没有焊接或每步重设瓶子位姿。
|
||||
|
||||
保持判据:全程手—物体相对位置相对参考偏差小于 20 mm,且末帧仍有接触。
|
||||
|
||||
| 测试 | 平均瓶子位置误差 | 接触持续时间 | 结果 |
|
||||
|---|---:|---:|---|
|
||||
| 原控制 GPU 基线 | 266.90 mm | 0.035 s | 落地 |
|
||||
| SPIDER GPU 优化 | 17.76 mm | 2.000 s | 保持接触,门限未通过 |
|
||||
| 独立 GPU 控制回放 | 20.14 mm | 2.000 s | 保持接触,门限未通过 |
|
||||
| 独立 CPU 控制回放 | 75.97 mm | 1.3675 s | 滑脱落地 |
|
||||
|
||||
SPIDER GPU 最大手—物体相对位置偏差为 76.64 mm,最大穿透为 3.31 mm。日志有 64 次非有限候选回报警告;最终执行轨迹有限不等于全部候选采样数值稳定。
|
||||
|
||||
独立 GPU 回放从新模型与初始状态出发,复用控制输入而不调用优化器,仍保持接触。这说明保持现象不只存在于优化器内部状态,但 CPU 差异说明接触与控制仍然敏感。
|
||||
|
||||
来源:[SPIDER 说明](../output/spider_l20_contact/README.md)、[配置](../output/spider_l20_contact/config.json)、[验证](../output/spider_l20_contact/validation.json)。
|
||||
|
||||
## 10. 踩过的坑:现象、证据与处理
|
||||
|
||||
### 10.1 环境冲突与误用启动脚本
|
||||
|
||||
现有重定向运行约定清空 `PYTHONPATH`、`LD_LIBRARY_PATH`,避免外部 ROS 路径与 Pinocchio 动态库造成冲突。重建导出使用 `.dynhamr`,dex 使用 `.dex`,SPIDER 使用 `.spider`;不要默认它们可以互换。
|
||||
|
||||
`run_dex_2047635068.sh` 实际运行的是 **LEAP** 分支。重新跑 L20 应显式调用 `run_dex_l20.py`。该差异能直接在脚本中确认。
|
||||
|
||||
### 10.2 将模型名字、关节数量和命令槽位混为一谈
|
||||
|
||||
目录名 L20、JSON `model=G20`、URDF 内部历史名称并不完全一致。实现按文件关节名工作,尚未硬件确认具体型号。21 个模型关节、16 个主动控制量、20 个命令槽位是不同概念。
|
||||
|
||||
### 10.3 重复施加联动约束
|
||||
|
||||
JSON 非线性 LUT 与 URDF 线性 mimic 不相同。基础版移除生成文件中的线性约束,由 adaptor 接管。后续切换线性版时重算被动角,不能同时保留非线性被动角并期待线性约束自动一致。
|
||||
|
||||
### 10.4 坐标正确不等于精度可靠
|
||||
|
||||
MANO wrapper、腕点定义、行列向量乘法、WXYZ/XYZW 都需要显式保持一致。已有 Pinocchio/MuJoCo FK 检查约达 1e-13 m,一致性很好,但它只支持实现约定相符,不能证明单目重建正确。
|
||||
|
||||
已有坐标诊断未发现该输入掌面轴逐帧翻转,不能未经检查就把所有抖动归因于坐标翻转或角度单位错误。
|
||||
|
||||
### 10.5 优化器目标值与梯度不一致
|
||||
|
||||
安装的 PositionOptimizer 返回值缺少时间正则,梯度却包含它。修复补回 `norm_delta * sum((x-last_qpos)**2)`,有限差分最大误差从约 1.20e-4 降至 1.36e-13。
|
||||
|
||||
但同目标对照中全段关节二阶差分只改善约 1.99%,排除接触前阶段后还略升。它是确定的实现缺陷,不能当作大幅抖动的充分解释。
|
||||
|
||||
### 10.6 单帧接触修补引入剧烈跳解
|
||||
|
||||
旧版每八帧独立优化,插值后再修穿透帧,缺少邻帧连续性代价。第 460→461 帧源 L20 最大变化约 0.270°,旧接触修正后无名指 PIP 跳变 54.375°,下一帧反跳 51.614°。
|
||||
|
||||
旧方案将全段穿透压到约 1.499 mm,但轨迹仍不连续。后续全长方案使用连续热启动、每两帧求解、邻帧修正正则及平滑插值。小穿透与小关节步长必须分开验证。
|
||||
|
||||
### 10.7 原瓶子网格与碰撞轮廓问题
|
||||
|
||||
早期提供的 OBJ 封口存在越界、退化三角形,仅在项目副本修复索引和朝向,保留 6146 个顶点,修复后 12288 个三角形。原始目录未修改。
|
||||
|
||||
早期碰撞用十段圆柱,半径存在突然缩小又扩大的台阶;视觉网格和碰撞近似也有差异。最近碰撞片段切换可能加重优化不连续,但未通过完整物理 A/B 实验量化其独立贡献。后续使用参数化瓶子,不将理想化瓶肩瓶颈宣称为实测形状。
|
||||
|
||||
### 10.8 几何贴合误当成物理抓稳
|
||||
|
||||
早期球体接口与旧瓶子实验均出现几何接触收敛但自由物体滑脱。旧瓶子三秒测试末尾相对位移约 17.56 m,与滑落后持续自由落体有关。即使初始静态接触力分配残差很小,也未保证动态保持。
|
||||
|
||||
处理原则是分别报告目标点误差、表面间距、穿透、自由物体保持和控制回放结果。`mj_forward` 回放视频无法替代 `mj_step` 动力学检查。
|
||||
|
||||
### 10.9 平均误差下降误当成门限通过
|
||||
|
||||
SPIDER 的平均位置误差从约 267 mm 降到约 18 mm,但保持判据使用全程最大相对偏差和末帧接触。平均值改善不等于全程达到 20 mm 门限。
|
||||
|
||||
### 10.10 GPU 与 CPU 回放不一致
|
||||
|
||||
初始配准改善、仿真步长从 5 ms 降到 2.5 ms 后,CPU/GPU 复核更接近,但最终 CPU 仍滑脱。只能确认后端或接触数值敏感,尚未独立定位唯一原因。下一步应先改进接触与执行器稳定性,而非只延长视频。
|
||||
|
||||
### 10.11 距离查询尖峰与接触流形不一致
|
||||
|
||||
全长结果中 `mj_geomDistance` 出现约 32 mm 的孤立穿透量,而独立 `mj_forward` 接触流形最大约 1 mm。保存原始查询和独立验证,明确报告口径;不要静默删尖峰或直接把查询异常等同真实穿模。
|
||||
|
||||
### 10.12 有限输出不能抹掉采样警告
|
||||
|
||||
SPIDER 保存轨迹有限,日志仍有非有限候选。验证需同时记录完成标志、最终数组有限性、候选警告和独立回放,不能只看进程完成。
|
||||
|
||||
### 10.13 离线平滑和硬编码限制
|
||||
|
||||
当前平滑使用未来帧并改变轨迹;帧率 30、源片段 600–659、抓取区间 130–160、右手选择和瓶子相对位姿均有特定输入假设。换视频不能仅替换一个文件名。
|
||||
|
||||
### 10.14 旧报告与新脚本可能属于不同版本
|
||||
|
||||
原 `l20_2047635068/retarget_validation.json` 保存的是历史运行,当前脚本默认已开启目标函数修复。比较修复效果应使用 `jitter_audit/corrected/` 的报告。旧阶段 README 中“尚未执行”的语句可能只描述当时阶段;当前状态应结合下游独立目录及其验证报告读取。
|
||||
|
||||
## 11. 输出格式与读取注意事项
|
||||
|
||||
| 文件 | 内容与约定 |
|
||||
|---|---|
|
||||
| `human_joints.npz` | `(T,21,3)` 局部人手点,米;另含世界腕点与根姿态 |
|
||||
| 基础 `motion.npz` | 21 关节角、16 主动角、命令、目标点、实际点、腕位姿与场景变换 |
|
||||
| 基础 `trajectory.csv` | 秒、腕部 XYZ、WXYZ 四元数、按关节名排列的角度,rad |
|
||||
| 基础 `commands_u8.csv` | 帧号及 20 命令槽位,仅对应生成该文件的 JSON 版本 |
|
||||
| 平滑 `temporal/motion.npz` | 平滑后的主字段与部分继承字段;`qpos_raw` 等仍可能是上游历史数据 |
|
||||
| SPIDER `optimized_trajectory.csv` | 时间、33 个标量位置坐标、28 个控制输入;不等于硬件命令 |
|
||||
| 全长 `motion.npz` | 对应 `scene.xml` 的完整 qpos,含手腕、手指和瓶子坐标 |
|
||||
| 全长 `scene.xml` | 场景与资产引用;单独打开只显示模型,需要 NPZ 才能播放 |
|
||||
|
||||
读取时以各文件字段、`joint_names` 或模型关节顺序为准。平滑脚本会更新 NPZ 内命令数组,但没有另外写出平滑版 `commands_u8.csv`,不要误用基础目录旧 CSV。
|
||||
|
||||
## 12. 复现步骤
|
||||
|
||||
以下命令依据现有脚本和保存的运行说明整理,本次未重跑。它们会重新生成对应输出。工作目录固定为 `/home/timessage/HandFlow`,并假定模型、环境与上游重建已存在。
|
||||
|
||||
### 12.1 基础导出、重定向、平滑
|
||||
|
||||
```bash
|
||||
cd /home/timessage/HandFlow
|
||||
|
||||
env PYTHONPATH= LD_LIBRARY_PATH= \
|
||||
third_party/Dyn-HaMR/.dynhamr/bin/python scripts/export_dynhamr_dex.py
|
||||
|
||||
env PYTHONPATH= LD_LIBRARY_PATH= .dex/bin/python \
|
||||
scripts/run_dex_l20.py --consistent-objective \
|
||||
--output-dir output/l20_2047635068/jitter_audit/corrected
|
||||
|
||||
env PYTHONPATH= LD_LIBRARY_PATH= \
|
||||
.dex/bin/python scripts/smooth_l20_offline.py
|
||||
```
|
||||
|
||||
### 12.2 准备瓶子场景与初始配准
|
||||
|
||||
参数化瓶子已存在于 `third_party/bottle_model_parametric/`。若需要重新生成,先核对并运行 `scripts/build_parametric_bottle.py`。随后:
|
||||
|
||||
```bash
|
||||
env PYTHONPATH= LD_LIBRARY_PATH= \
|
||||
.spider/bin/python scripts/prepare_spider_l20.py
|
||||
```
|
||||
|
||||
该步骤生成全长回放依赖的 `scene_act.xml` 和 `trajectory_kinematic_act.npz`,同时准备 SPIDER 两秒任务。
|
||||
|
||||
### 12.3 全长运动学回放
|
||||
|
||||
```bash
|
||||
env PYTHONPATH= LD_LIBRARY_PATH= MUJOCO_GL=osmesa \
|
||||
.spider/bin/python scripts/build_l20_full_replay.py
|
||||
|
||||
env PYTHONPATH= LD_LIBRARY_PATH= MUJOCO_GL=osmesa \
|
||||
.spider/bin/python scripts/verify_l20_full_replay.py
|
||||
|
||||
env PYTHONPATH= LD_LIBRARY_PATH= MUJOCO_GL=osmesa \
|
||||
.spider/bin/python scripts/play_l20_full.py --render
|
||||
```
|
||||
|
||||
构建脚本的初始报告使用距离查询;独立验证脚本会补充接触流形检查。所以生成后应执行验证,不应把构建结束的初始 JSON 当成最终验证报告。
|
||||
|
||||
### 12.4 SPIDER 优化与独立复核
|
||||
|
||||
在可访问 GPU 的环境中运行:
|
||||
|
||||
```bash
|
||||
env PYTHONPATH= LD_LIBRARY_PATH= \
|
||||
.spider/bin/python -u scripts/run_spider_l20_contact.py
|
||||
|
||||
env PYTHONPATH= LD_LIBRARY_PATH= \
|
||||
.spider/bin/python scripts/replay_spider_l20_gpu.py
|
||||
|
||||
env PYTHONPATH= LD_LIBRARY_PATH= \
|
||||
.spider/bin/python scripts/evaluate_spider_l20.py
|
||||
```
|
||||
|
||||
优化完成应结合 `SPIDER_RUN_COMPLETE`、输出数组、验证报告与独立控制回放确认,不能只看日志正在运行。
|
||||
|
||||
### 12.5 直接查看已有成品
|
||||
|
||||
```bash
|
||||
/home/timessage/HandFlow/output/l20_full_replay/replay.sh
|
||||
```
|
||||
|
||||
空格暂停/继续,左右方向键逐帧,R 重播,鼠标调整视角。
|
||||
|
||||
- [完整视频](../output/l20_full_replay/full_replay.mp4)
|
||||
- [完整回放说明](../output/l20_full_replay/README.md)
|
||||
- [平滑基础 CSV](../output/l20_2047635068/jitter_audit/temporal/trajectory.csv)
|
||||
- [SPIDER 对照视频](../output/spider_l20_contact/spider_contact_comparison.mp4)
|
||||
|
||||
## 13. 换视频或换机器人时需要改什么
|
||||
|
||||
| 修改对象 | 必须检查的内容 |
|
||||
|---|---|
|
||||
| 输入视频 | 帧率、长度、左右手选择、重建文件路径、单位与坐标系 |
|
||||
| 手模型 | URDF 关节名、限位、指尖标记、掌面轴、主动/被动关系 |
|
||||
| 标定表 | 电机槽位、零位偏置、有效命令范围、曲线是否单调及是否实测 |
|
||||
| 联动模式 | 明确采用 JSON 还是 URDF,生成与验证使用同一模式 |
|
||||
| 物体 | 几何尺寸、视觉/碰撞一致性、物体位姿来源、质量与摩擦 |
|
||||
| 接触时段 | 当前 130–160 与 600–659 帧假设不能直接沿用 |
|
||||
| 动态执行 | 执行器力矩/速度能力、控制频率、接触稳定性及独立回放 |
|
||||
|
||||
当前代码没有覆盖任意输入:退化骨段可能导致方向归一化失败;右手轨迹选择有严格断言;短视频无法使用第 600 帧初始化;较大旋转可能触发平滑断言。新增输入应先做小段数据与模型检查,再运行全段。
|
||||
|
||||
## 14. 后续工作的验证顺序
|
||||
|
||||
以下为建议,尚未完成:
|
||||
|
||||
1. 对齐真实型号、联动与标定解释,确定最终采用哪一种控制模型。
|
||||
2. 获取或核验物体逐帧位姿与坐标外参,降低手掌关系假设的影响。
|
||||
3. 在短片段中定位接触布局、碰撞几何、执行器与 CPU/GPU 差异,先通过无辅助物体驱动的独立回放保持门限。
|
||||
4. 增加明确的关节速度、加速度与控制平滑约束,单独报告接触代价和动作偏离。
|
||||
5. 短片段通过后,再扩展接近、闭合、持握、抬升以及全长动力学实验。
|
||||
|
||||
不能用延长运动学视频、增大采样预算或降低平均误差,替代上述稳定性验证。
|
||||
|
||||
## 15. 主要证据索引
|
||||
|
||||
- [Dyn-HaMR 上游结果](../output/results_2047635068/README.md)
|
||||
- [LEAP 基线与环境约定](../output/dex_2047635068/README.md)
|
||||
- [L20 基础实现说明](../output/l20_2047635068/README.md)
|
||||
- [修正后的基础验证](../output/l20_2047635068/jitter_audit/corrected/retarget_validation.json)
|
||||
- [抖动诊断与消融](../output/l20_2047635068/jitter_audit/README.md)
|
||||
- [离线平滑验证](../output/l20_2047635068/jitter_audit/temporal/validation.json)
|
||||
- [旧瓶子网格修复、接触跳变与物理失败](../output/l20_2047635068/bottle_grasp/README.md)
|
||||
- [全长回放最终验证](../output/l20_full_replay/validation.json)
|
||||
- [SPIDER 配准与运行边界](../output/spider_l20_contact/README.md)
|
||||
- [SPIDER 最终验证](../output/spider_l20_contact/validation.json)
|
||||
|
||||
本文件汇总的是仓库中第三方组件与本地集成脚本的行为和已有证据,不用于推断个人贡献归属。
|
||||
@@ -0,0 +1,30 @@
|
||||
"""HandFlow model package — attention function with flash-attn fallback."""
|
||||
|
||||
import torch
|
||||
from einops import rearrange
|
||||
from torch import Tensor
|
||||
|
||||
FA3_FOUND = True
|
||||
try:
|
||||
from flash_attn_interface import flash_attn_func # noqa: F401
|
||||
except ImportError:
|
||||
FA3_FOUND = False
|
||||
|
||||
if FA3_FOUND:
|
||||
|
||||
def attention(q: Tensor, k: Tensor, v: Tensor) -> Tensor:
|
||||
q = rearrange(q, "B H L D -> B L H D")
|
||||
k = rearrange(k, "B H L D -> B L H D")
|
||||
v = rearrange(v, "B H L D -> B L H D")
|
||||
x = flash_attn_func(q, k, v)
|
||||
if not isinstance(x, Tensor):
|
||||
x = x[0]
|
||||
x = rearrange(x, "B L H D -> B L (H D)")
|
||||
return x
|
||||
|
||||
else:
|
||||
|
||||
def attention(q: Tensor, k: Tensor, v: Tensor) -> Tensor:
|
||||
x = torch.nn.functional.scaled_dot_product_attention(q, k, v)
|
||||
x = rearrange(x, "B H L D -> B L (H D)")
|
||||
return x
|
||||
@@ -0,0 +1 @@
|
||||
"""HandFlow feature extractors package."""
|
||||
@@ -0,0 +1,67 @@
|
||||
"""
|
||||
Condition Builder (with cmask mechanism).
|
||||
|
||||
Concatenates image token + skeleton token into the condition stream.
|
||||
cmask: learnable mask token; both training and inference use HaMeR confidence as a continuous mask.
|
||||
|
||||
Inputs:
|
||||
- image_tokens: (B, T, d_model) — from Image Refiner + Compression
|
||||
- skeleton_tokens: (B, T, d_model) — from Skeleton Encoder
|
||||
- hamer_confidence: (B, T) — HaMeR detection confidence
|
||||
|
||||
Output: condition (B, 2*T, d_model)
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class ConditionBuilder(nn.Module):
|
||||
def __init__(self, d_model: int, mask_ratio: float = 0.2):
|
||||
super().__init__()
|
||||
self.d_model = d_model
|
||||
self.mask_ratio = mask_ratio
|
||||
|
||||
# Learnable mask token — initialized with truncated normal
|
||||
self.c_mask = nn.Parameter(torch.empty(1, 1, d_model))
|
||||
nn.init.trunc_normal_(self.c_mask, std=0.02)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
image_tokens: torch.Tensor,
|
||||
skeleton_tokens: torch.Tensor,
|
||||
hamer_confidence: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Args:
|
||||
image_tokens: (B, T, d_model)
|
||||
skeleton_tokens: (B, T, d_model)
|
||||
hamer_confidence: (B, T) — used as m_t during inference
|
||||
|
||||
Returns:
|
||||
condition: (B, 2*T, d_model)
|
||||
m: (B, T) — actual mask values used (for logging)
|
||||
"""
|
||||
B, T, D = image_tokens.shape
|
||||
|
||||
# concat: (B, T, 2, d_model) -> (B, 2*T, d_model)
|
||||
tokens = torch.stack([image_tokens, skeleton_tokens], dim=2) # (B, T, 2, D)
|
||||
tokens = tokens.reshape(B, 2 * T, D)
|
||||
|
||||
# Use HaMeR confidence as a continuous mask (consistent between training and inference)
|
||||
m = hamer_confidence.float() # (B, T)
|
||||
|
||||
# Optional: randomly dropout some frames during training as regularization
|
||||
if self.training and self.mask_ratio > 0:
|
||||
keep = (torch.rand(B, T, device=tokens.device) >= self.mask_ratio).float()
|
||||
m = m * keep
|
||||
|
||||
# Expand m to match 2*T tokens (the two tokens per frame share the same m)
|
||||
# m: (B, T) -> (B, 2*T)
|
||||
m_expanded = m.unsqueeze(-1).expand(-1, -1, 2).reshape(B, 2 * T) # (B, 2*T)
|
||||
m_expanded = m_expanded.unsqueeze(-1) # (B, 2*T, 1)
|
||||
|
||||
# cmask: c'_t = m_t * c_t + (1 - m_t) * c_mask
|
||||
condition = m_expanded * tokens + (1 - m_expanded) * self.c_mask
|
||||
|
||||
return condition, m
|
||||
@@ -0,0 +1,110 @@
|
||||
"""
|
||||
HaMeR ViT-Huge backbone wrapper (frozen).
|
||||
|
||||
Provides image token extraction for the 2D condition stream.
|
||||
Input: (B, 3, 256, 256) float32 [0,1]
|
||||
Output: (B, 192, 1280) — 16x12 patches x 1280D
|
||||
"""
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
# Add HaMeR repository to the import path
|
||||
_HAMER_ROOT = Path(__file__).resolve().parents[2] / "third_party" / "hamer"
|
||||
if _HAMER_ROOT.exists() and str(_HAMER_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(_HAMER_ROOT))
|
||||
|
||||
|
||||
class HaMeRBackbone(nn.Module):
|
||||
"""
|
||||
Frozen HaMeR ViT-Huge backbone.
|
||||
|
||||
Internal pipeline:
|
||||
1. ImageNet normalization (mean/std)
|
||||
2. Crop 32px from each horizontal side: (B,3,256,256) -> (B,3,256,192)
|
||||
3. Patch embedding: 192/16 x 256/16 = 12x16 = 192 tokens x 1280D
|
||||
4. 32-layer ViT-H Transformer
|
||||
5. Reshape output -> (B, 192, 1280)
|
||||
"""
|
||||
|
||||
# ImageNet normalization parameters
|
||||
_IMAGENET_MEAN = torch.tensor([0.485, 0.456, 0.406])
|
||||
_IMAGENET_STD = torch.tensor([0.229, 0.224, 0.225])
|
||||
|
||||
def __init__(self, checkpoint_path: Optional[str] = None):
|
||||
super().__init__()
|
||||
# Import the ViT backbone directly, bypassing hamer/models/__init__.py (avoids smplx dependency)
|
||||
import importlib.util
|
||||
vit_path = _HAMER_ROOT / "hamer" / "models" / "backbones" / "vit.py"
|
||||
spec = importlib.util.spec_from_file_location("hamer_vit", str(vit_path))
|
||||
vit_mod = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(vit_mod)
|
||||
ViT = vit_mod.ViT
|
||||
|
||||
self.backbone = ViT(
|
||||
img_size=(256, 192),
|
||||
patch_size=16,
|
||||
embed_dim=1280,
|
||||
depth=32,
|
||||
num_heads=16,
|
||||
ratio=1,
|
||||
use_checkpoint=False,
|
||||
mlp_ratio=4,
|
||||
qkv_bias=True,
|
||||
drop_path_rate=0.55,
|
||||
)
|
||||
self.backbone.eval()
|
||||
|
||||
# Load pretrained weights
|
||||
if checkpoint_path is not None:
|
||||
ckpt = torch.load(checkpoint_path, map_location="cpu")
|
||||
if "state_dict" in ckpt:
|
||||
state_dict = {
|
||||
k.replace("backbone.", ""): v
|
||||
for k, v in ckpt["state_dict"].items()
|
||||
if k.startswith("backbone.")
|
||||
}
|
||||
else:
|
||||
state_dict = ckpt
|
||||
self.backbone.load_state_dict(state_dict, strict=False)
|
||||
|
||||
# Freeze all parameters
|
||||
self.backbone.requires_grad_(False)
|
||||
for p in self.backbone.parameters():
|
||||
p.requires_grad = False
|
||||
|
||||
# Register normalization buffers
|
||||
self.register_buffer("_mean", self._IMAGENET_MEAN.view(1, 3, 1, 1))
|
||||
self.register_buffer("_std", self._IMAGENET_STD.view(1, 3, 1, 1))
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Args:
|
||||
x: (B, 3, 256, 256) float32 in range [0, 1]
|
||||
Returns:
|
||||
tokens: (B, 192, 1280) patch tokens
|
||||
"""
|
||||
# ImageNet normalization
|
||||
x = (x - self._mean) / self._std
|
||||
|
||||
# HaMeR internal crop: 32px from each side (256 -> 192)
|
||||
x = x[:, :, :, 32:-32]
|
||||
|
||||
# Backbone forward
|
||||
feat_map = self.backbone.forward_features(x) # (B, 1280, Hp, Wp)
|
||||
B, C, Hp, Wp = feat_map.shape
|
||||
|
||||
# Flatten into tokens
|
||||
tokens = feat_map.flatten(2).transpose(1, 2) # (B, Hp*Wp, 1280)
|
||||
|
||||
return tokens
|
||||
|
||||
def train(self, mode: bool = True):
|
||||
"""Always stay in eval mode."""
|
||||
super().train(False)
|
||||
self.backbone.eval()
|
||||
return self
|
||||
@@ -0,0 +1,282 @@
|
||||
"""
|
||||
Image Spatio-Temporal Refiner + Frame Compression.
|
||||
|
||||
Step 1: 8+8 alternating spatio-temporal attention over patch tokens of all frames
|
||||
Step 2: 4 layers of cross-attention pooling compress 192 tokens per frame -> 1 token
|
||||
|
||||
Input: HaMeR backbone output (B, T, 192, 1280)
|
||||
Output: (B, T, d_model) — one condensed image condition token per frame
|
||||
"""
|
||||
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from einops import rearrange
|
||||
from model import attention
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
# 2D RoPE (consistent with the legacy TemporalRefiner)
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class RotaryEmbedding1D(nn.Module):
|
||||
"""1D RoPE: used for temporal positional encoding."""
|
||||
|
||||
def __init__(self, dim: int, base: int = 10000):
|
||||
super().__init__()
|
||||
assert dim % 2 == 0
|
||||
self.dim = dim
|
||||
inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
|
||||
self.register_buffer("inv_freq", inv_freq, persistent=False)
|
||||
|
||||
def forward(self, x: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
x: (B, H, L, D), positions: (L,)
|
||||
"""
|
||||
freqs = torch.einsum("i,j->ij", positions.float(), self.inv_freq.to(positions))
|
||||
sin = freqs.sin().unsqueeze(0).unsqueeze(0)
|
||||
cos = freqs.cos().unsqueeze(0).unsqueeze(0)
|
||||
x1, x2 = x[..., 0::2], x[..., 1::2]
|
||||
return torch.stack([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1).flatten(-2)
|
||||
|
||||
|
||||
class RotaryEmbedding2D(nn.Module):
|
||||
"""2D RoPE: split head_dim in half and apply 1D RoPE with y/x coordinates respectively."""
|
||||
|
||||
def __init__(self, dim: int, base: int = 10000):
|
||||
super().__init__()
|
||||
assert dim % 2 == 0
|
||||
self.dim = dim
|
||||
self.dim_per_axis = dim // 2
|
||||
inv_freq = 1.0 / (base ** (torch.arange(0, self.dim_per_axis, 2).float() / self.dim_per_axis))
|
||||
self.register_buffer("inv_freq", inv_freq, persistent=False)
|
||||
|
||||
def forward(self, x: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
x: (B, H, L, D), positions: (L, 2) (y, x)
|
||||
"""
|
||||
x_y, x_x = x.chunk(2, dim=-1)
|
||||
x_y = self._apply_1d(x_y, positions[:, 0])
|
||||
x_x = self._apply_1d(x_x, positions[:, 1])
|
||||
return torch.cat([x_y, x_x], dim=-1)
|
||||
|
||||
def _apply_1d(self, x: torch.Tensor, pos: torch.Tensor) -> torch.Tensor:
|
||||
freqs = torch.einsum("i,j->ij", pos.float(), self.inv_freq.to(pos))
|
||||
sin = freqs.sin().unsqueeze(0).unsqueeze(0)
|
||||
cos = freqs.cos().unsqueeze(0).unsqueeze(0)
|
||||
x1, x2 = x[..., 0::2], x[..., 1::2]
|
||||
return torch.stack([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1).flatten(-2)
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
# Building blocks
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class RMSNorm(nn.Module):
|
||||
def __init__(self, dim: int):
|
||||
super().__init__()
|
||||
self.scale = nn.Parameter(torch.ones(dim))
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
x_dtype = x.dtype
|
||||
x = x.float()
|
||||
rrms = torch.rsqrt(torch.mean(x ** 2, dim=-1, keepdim=True) + 1e-6)
|
||||
return (x * rrms).to(dtype=x_dtype) * self.scale
|
||||
|
||||
|
||||
class RefinerBlock(nn.Module):
|
||||
"""Pre-LN Transformer block with optional 2D RoPE."""
|
||||
|
||||
def __init__(self, d_model: int, num_heads: int, mlp_ratio: float = 4.0, dropout: float = 0.1):
|
||||
super().__init__()
|
||||
self.num_heads = num_heads
|
||||
self.head_dim = d_model // num_heads
|
||||
self.norm1 = nn.LayerNorm(d_model)
|
||||
self.qkv = nn.Linear(d_model, d_model * 3)
|
||||
self.q_norm = RMSNorm(self.head_dim)
|
||||
self.k_norm = RMSNorm(self.head_dim)
|
||||
self.proj = nn.Linear(d_model, d_model)
|
||||
self.norm2 = nn.LayerNorm(d_model)
|
||||
mlp_hidden = int(d_model * mlp_ratio)
|
||||
self.mlp = nn.Sequential(nn.Linear(d_model, mlp_hidden), nn.GELU(), nn.Linear(mlp_hidden, d_model))
|
||||
self.drop = nn.Dropout(dropout)
|
||||
|
||||
def forward(self, x: torch.Tensor, rope: Optional[RotaryEmbedding2D] = None, positions: Optional[torch.Tensor] = None) -> torch.Tensor:
|
||||
residual = x
|
||||
x = self.norm1(x)
|
||||
qkv = self.qkv(x)
|
||||
q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
|
||||
q = self.q_norm(q)
|
||||
k = self.k_norm(k)
|
||||
if rope is not None and positions is not None:
|
||||
q = rope(q, positions)
|
||||
k = rope(k, positions)
|
||||
x = attention(q, k, v)
|
||||
x = self.proj(x)
|
||||
x = residual + self.drop(x)
|
||||
x = x + self.drop(self.mlp(self.norm2(x)))
|
||||
return x
|
||||
|
||||
|
||||
class CrossAttentionPoolBlock(nn.Module):
|
||||
"""Cross-attention pooling: 1 learnable query token cross-attends K patches."""
|
||||
|
||||
def __init__(self, d_model: int, num_heads: int, mlp_ratio: float = 4.0, dropout: float = 0.1):
|
||||
super().__init__()
|
||||
self.num_heads = num_heads
|
||||
self.head_dim = d_model // num_heads
|
||||
self.norm_q = nn.LayerNorm(d_model)
|
||||
self.norm_kv = nn.LayerNorm(d_model)
|
||||
self.q_proj = nn.Linear(d_model, d_model)
|
||||
self.kv_proj = nn.Linear(d_model, d_model * 2)
|
||||
self.q_norm = RMSNorm(self.head_dim)
|
||||
self.k_norm = RMSNorm(self.head_dim)
|
||||
self.out_proj = nn.Linear(d_model, d_model)
|
||||
self.norm_ff = nn.LayerNorm(d_model)
|
||||
mlp_hidden = int(d_model * mlp_ratio)
|
||||
self.ff = nn.Sequential(nn.Linear(d_model, mlp_hidden), nn.GELU(), nn.Linear(mlp_hidden, d_model))
|
||||
self.drop = nn.Dropout(dropout)
|
||||
|
||||
def forward(self, query: torch.Tensor, patches: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
query: (B, 1, D), patches: (B, N, D)
|
||||
Returns: (B, 1, D)
|
||||
"""
|
||||
residual = query
|
||||
q = self.q_proj(self.norm_q(query))
|
||||
kv = self.kv_proj(self.norm_kv(patches))
|
||||
k, v = rearrange(kv, "B N (K H D) -> K B H N D", K=2, H=self.num_heads)
|
||||
q = rearrange(q, "B L (H D) -> B H L D", H=self.num_heads)
|
||||
q = self.q_norm(q)
|
||||
k = self.k_norm(k)
|
||||
out = attention(q, k, v) # (B, 1, H*D) = (B, 1, D)
|
||||
out = self.out_proj(out)
|
||||
query = residual + self.drop(out)
|
||||
query = query + self.drop(self.ff(self.norm_ff(query)))
|
||||
return query
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
# Spatio-Temporal Refiner
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class SpatioTemporalRefiner(nn.Module):
|
||||
"""
|
||||
8+8 alternating spatio-temporal attention.
|
||||
|
||||
Even layers: intra-frame spatial attention (B*T, 192, dim)
|
||||
Odd layers: cross-frame same-position attention (B, T, dim) per spatial position
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
d_model: int,
|
||||
hamer_dim: int = 1280,
|
||||
spatial_layers: int = 8,
|
||||
temporal_layers: int = 8,
|
||||
num_heads: int = 8,
|
||||
mlp_ratio: float = 4.0,
|
||||
dropout: float = 0.1,
|
||||
n_patches_h: int = 16,
|
||||
n_patches_w: int = 12,
|
||||
):
|
||||
super().__init__()
|
||||
self.d_model = d_model
|
||||
self.n_patches = n_patches_h * n_patches_w
|
||||
self.depth = spatial_layers + temporal_layers
|
||||
self.spatial_layers = spatial_layers
|
||||
self.temporal_layers = temporal_layers
|
||||
|
||||
self.input_proj = nn.Linear(hamer_dim, d_model)
|
||||
|
||||
# 2D spatial positions
|
||||
y = torch.arange(n_patches_h)
|
||||
x = torch.arange(n_patches_w)
|
||||
self.register_buffer("spatial_positions", torch.cartesian_prod(y, x))
|
||||
|
||||
self.rope = RotaryEmbedding2D(dim=d_model // num_heads)
|
||||
self.temporal_rope = RotaryEmbedding1D(dim=d_model // num_heads)
|
||||
self.blocks = nn.ModuleList([
|
||||
RefinerBlock(d_model, num_heads, mlp_ratio, dropout)
|
||||
for _ in range(self.depth)
|
||||
])
|
||||
self.use_checkpoint = False
|
||||
|
||||
def gradient_checkpointing_enable(self):
|
||||
self.use_checkpoint = True
|
||||
|
||||
def forward(self, hamer_features: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Args:
|
||||
hamer_features: (B, T, 192, 1280)
|
||||
Returns:
|
||||
(B, T, 192, d_model)
|
||||
"""
|
||||
B, T, N, _ = hamer_features.shape
|
||||
x = self.input_proj(hamer_features)
|
||||
|
||||
spatial_pos = self.spatial_positions # (N, 2)
|
||||
temporal_pos = torch.arange(T, device=x.device) # (T,)
|
||||
temporal_pos_expanded = temporal_pos.unsqueeze(1).expand(T, N).reshape(T * N) # (T*N,)
|
||||
|
||||
for i, blk in enumerate(self.blocks):
|
||||
if i % 2 == 0:
|
||||
# Intra-frame spatial attention (2D RoPE: y, x)
|
||||
x = x.reshape(B * T, N, self.d_model)
|
||||
x = blk(x, rope=self.rope, positions=spatial_pos)
|
||||
x = x.reshape(B, T, N, self.d_model)
|
||||
else:
|
||||
# Cross-frame global attention (1D temporal RoPE: frame index t)
|
||||
x = x.reshape(B, T * N, self.d_model)
|
||||
x = blk(x, rope=self.temporal_rope, positions=temporal_pos_expanded)
|
||||
x = x.reshape(B, T, N, self.d_model)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
# Frame Compressor (Cross-Attention Pooling)
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class FrameCompressor(nn.Module):
|
||||
"""
|
||||
Compress N patch tokens of each frame into 1 condensed token.
|
||||
Uses a learnable query token with multi-layer cross-attention.
|
||||
Includes input_proj: supports arbitrary input dims (e.g. HaMeR 1280D -> d_model).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
d_model: int,
|
||||
input_dim: int | None = None,
|
||||
num_heads: int = 8,
|
||||
num_layers: int = 8,
|
||||
mlp_ratio: float = 4.0,
|
||||
dropout: float = 0.1,
|
||||
):
|
||||
super().__init__()
|
||||
self.input_proj = nn.Linear(input_dim, d_model) if input_dim else nn.Identity()
|
||||
self.query = nn.Parameter(torch.randn(1, 1, d_model))
|
||||
nn.init.trunc_normal_(self.query, std=0.02)
|
||||
self.layers = nn.ModuleList([
|
||||
CrossAttentionPoolBlock(d_model, num_heads, mlp_ratio, dropout)
|
||||
for _ in range(num_layers)
|
||||
])
|
||||
|
||||
def forward(self, patches: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
patches: (B, T, N, input_dim or d_model)
|
||||
Returns: (B, T, d_model)
|
||||
"""
|
||||
B, T, N, _ = patches.shape
|
||||
patches = self.input_proj(patches) # (B, T, N, d_model)
|
||||
patches = patches.reshape(B * T, N, -1) # (B*T, N, D)
|
||||
q = self.query.expand(B * T, -1, -1) # (B*T, 1, D)
|
||||
for layer in self.layers:
|
||||
q = layer(q, patches)
|
||||
return q.reshape(B, T, -1)
|
||||
@@ -0,0 +1,96 @@
|
||||
"""
|
||||
Ray Direction Skeleton Encoder.
|
||||
|
||||
2D skeleton -> inverse intrinsics -> ray direction -> high-frequency sinusoidal PE -> flatten -> MLP -> embed_dim.
|
||||
Outputs one skeleton token per frame.
|
||||
|
||||
Input: landmarks (B, T, 21, 2) + crop_intrinsics (B, T, 4)
|
||||
Output: (B, T, embed_dim)
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class RayDirectionSkeletonEncoder(nn.Module):
|
||||
"""
|
||||
Encodes a 2D skeleton into one token per frame via ray direction.
|
||||
|
||||
Pipeline:
|
||||
1. un-normalize landmarks -> pixel coordinates (in crop coordinate system)
|
||||
2. ray = [(u-cx)/fx, (v-cy)/fy, 1] -> normalize -> (21, 3)
|
||||
3. High-frequency sinusoidal PE: [sin(2^L * ray), cos(2^L * ray)] L=0..num_freqs-1
|
||||
4. flatten: 21 joints x (3 x 2 x num_freqs) per joint
|
||||
5. MLP -> embed_dim
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
n_joints: int = 21,
|
||||
num_freqs: int = 6,
|
||||
embed_dim: int = 512,
|
||||
crop_h: int = 256,
|
||||
crop_w: int = 256,
|
||||
mlp_hidden: int = 512,
|
||||
):
|
||||
super().__init__()
|
||||
self.n_joints = n_joints
|
||||
self.num_freqs = num_freqs
|
||||
self.crop_h = crop_h
|
||||
self.crop_w = crop_w
|
||||
|
||||
# Per-joint dim after PE: 3 (ray) x 2 (sin+cos) x num_freqs
|
||||
pe_dim = 3 * 2 * num_freqs # = 36
|
||||
flat_dim = n_joints * pe_dim # = 756
|
||||
|
||||
self.mlp = nn.Sequential(
|
||||
nn.Linear(flat_dim, mlp_hidden),
|
||||
nn.SiLU(),
|
||||
nn.Linear(mlp_hidden, embed_dim),
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
landmarks: torch.Tensor,
|
||||
crop_intrinsics: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Args:
|
||||
landmarks: (B, T, 21, 2) normalized [0, 1]
|
||||
crop_intrinsics: (B, T, 4) = [fx, fy, cx, cy] in crop coords
|
||||
Returns:
|
||||
(B, T, embed_dim)
|
||||
"""
|
||||
B, T, J, _ = landmarks.shape
|
||||
fx = crop_intrinsics[..., 0:1] # (B, T, 1)
|
||||
fy = crop_intrinsics[..., 1:2]
|
||||
cx = crop_intrinsics[..., 2:3]
|
||||
cy = crop_intrinsics[..., 3:4]
|
||||
|
||||
# un-normalize -> crop pixel coordinates
|
||||
u = landmarks[..., 0:1] * self.crop_w # (B, T, 21, 1)
|
||||
v = landmarks[..., 1:2] * self.crop_h
|
||||
|
||||
# ray direction: [(u-cx)/fx, (v-cy)/fy, 1]
|
||||
ray_x = (u - cx.unsqueeze(-2)) / (fx.unsqueeze(-2) + 1e-8)
|
||||
ray_y = (v - cy.unsqueeze(-2)) / (fy.unsqueeze(-2) + 1e-8)
|
||||
ray_z = torch.ones_like(ray_x)
|
||||
ray = torch.cat([ray_x, ray_y, ray_z], dim=-1) # (B, T, 21, 3)
|
||||
|
||||
# normalize ray direction
|
||||
ray = ray / (ray.norm(dim=-1, keepdim=True) + 1e-8)
|
||||
|
||||
# High-frequency sinusoidal PE
|
||||
pe_parts = []
|
||||
for L in range(self.num_freqs):
|
||||
freq = 2.0 ** L
|
||||
pe_parts.append(torch.sin(freq * ray))
|
||||
pe_parts.append(torch.cos(freq * ray))
|
||||
# pe: (B, T, 21, 3 * 2 * num_freqs) = (B, T, 21, 36)
|
||||
pe = torch.cat(pe_parts, dim=-1)
|
||||
|
||||
# flatten all joints: (B, T, 21 * 36) = (B, T, 756)
|
||||
flat = pe.reshape(B, T, -1)
|
||||
|
||||
# MLP -> embed_dim
|
||||
return self.mlp(flat)
|
||||
@@ -0,0 +1 @@
|
||||
"""HandFlow flow matching package."""
|
||||
@@ -0,0 +1,345 @@
|
||||
"""
|
||||
HandFlow Rectified Flow Denoiser — all-frame direct regression.
|
||||
|
||||
Pose (48D axis-angle) and Trans (3D) are both regressed per-frame directly,
|
||||
without velocity/cumsum. Beta (10D) serves as a global token.
|
||||
|
||||
Flow Target layout (flat packing per window, managed by mano_slice):
|
||||
x_1 = [beta(10) | pose_0(48) | ... | pose_{T-1}(48) | trans_0(3) | ... | trans_{T-1}(3)]
|
||||
flat_dim = 10 + T × 51
|
||||
|
||||
z-stream: 1 beta token + T pose tokens + T trans tokens = 1+2T tokens
|
||||
c-stream: 2 × T tokens (1 image + 1 skeleton per frame, with cmask)
|
||||
|
||||
Condition pipeline:
|
||||
HaMeR backbone (frozen) → FrameCompressor (8-layer cross-attention pooling) → image tokens
|
||||
2D skeleton → ray direction → PE → MLP → skeleton tokens
|
||||
image + skeleton → cmask → condition stream
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from omegaconf import DictConfig
|
||||
|
||||
from model.flow_matching.dualstream_transformer import (
|
||||
FlowMatchingTransformer,
|
||||
MLPEmbedder,
|
||||
)
|
||||
from model.flow_matching.helpers.model_wrapper import ModelWrapper
|
||||
from model.flow_matching.helpers.scheduler import FluxTimeSampler
|
||||
from model.feature_extractors.hamer_extractor import HaMeRBackbone
|
||||
from model.feature_extractors.image_refiner import FrameCompressor
|
||||
from model.feature_extractors.skeleton_encoder import RayDirectionSkeletonEncoder
|
||||
from model.feature_extractors.condition_builder import ConditionBuilder
|
||||
from model.pose_embedding import RotaryEmbedding
|
||||
|
||||
|
||||
class ManoSlice(nn.Module):
|
||||
"""Manage slice positions and normalization of each component in the flow target flat vector.
|
||||
|
||||
Flat layout: [beta(10) | poses(T,48) | trans(T,3)]
|
||||
flat_dim = 10 + T * 51
|
||||
"""
|
||||
|
||||
def __init__(self, beta_dim: int, pose_dim: int,
|
||||
trans_dim: int, window_size: int, stats_path: str | None = None):
|
||||
super().__init__()
|
||||
self.beta_dim = beta_dim # 10
|
||||
self.pose_dim = pose_dim # 48
|
||||
self.trans_dim = trans_dim # 3
|
||||
self.window_size = window_size
|
||||
|
||||
self.flat_dim = beta_dim + window_size * (pose_dim + trans_dim)
|
||||
|
||||
# Build slices
|
||||
s = 0
|
||||
self.beta = slice(s, s + beta_dim) # [0:10]
|
||||
s += beta_dim
|
||||
self.poses = slice(s, s + window_size * pose_dim) # [10:10+48T]
|
||||
s += window_size * pose_dim
|
||||
self.trans = slice(s, s + window_size * trans_dim) # [10+48T:10+51T]
|
||||
|
||||
# Normalization statistics
|
||||
flat_mean = torch.zeros(self.flat_dim)
|
||||
flat_std = torch.ones(self.flat_dim)
|
||||
if stats_path is not None:
|
||||
import numpy as np
|
||||
stats = np.load(stats_path)
|
||||
beta_mean = torch.from_numpy(stats["beta_mean"]) # (10,)
|
||||
beta_std = torch.from_numpy(stats["beta_std"]) # (10,)
|
||||
pose_mean = torch.from_numpy(stats["pose_mean"]) # (48,)
|
||||
pose_std = torch.from_numpy(stats["pose_std"]) # (48,)
|
||||
trans_mean = torch.from_numpy(stats["trans_mean"]) # (3,)
|
||||
trans_std = torch.from_numpy(stats["trans_std"]) # (3,)
|
||||
flat_mean = torch.cat([
|
||||
beta_mean,
|
||||
pose_mean.repeat(window_size),
|
||||
trans_mean.repeat(window_size),
|
||||
])
|
||||
flat_std = torch.cat([
|
||||
beta_std,
|
||||
pose_std.repeat(window_size),
|
||||
trans_std.repeat(window_size),
|
||||
])
|
||||
self.register_buffer("flat_mean", flat_mean)
|
||||
self.register_buffer("flat_std", flat_std)
|
||||
|
||||
# per-component statistics for auxiliary loss denormalization
|
||||
self.register_buffer("beta_mean_", flat_mean[self.beta])
|
||||
self.register_buffer("beta_std_", flat_std[self.beta])
|
||||
|
||||
if stats_path is not None:
|
||||
self.register_buffer("pose_mean_", pose_mean)
|
||||
self.register_buffer("pose_std_", pose_std)
|
||||
self.register_buffer("trans_mean_", trans_mean)
|
||||
self.register_buffer("trans_std_", trans_std)
|
||||
else:
|
||||
self.register_buffer("pose_mean_", torch.zeros(pose_dim))
|
||||
self.register_buffer("pose_std_", torch.ones(pose_dim))
|
||||
self.register_buffer("trans_mean_", torch.zeros(trans_dim))
|
||||
self.register_buffer("trans_std_", torch.ones(trans_dim))
|
||||
|
||||
def normalize(self, x_flat: torch.Tensor) -> torch.Tensor:
|
||||
"""Normalize the flat vector to ~N(0,1)."""
|
||||
return (x_flat - self.flat_mean) / self.flat_std
|
||||
|
||||
def denormalize(self, x_flat: torch.Tensor) -> torch.Tensor:
|
||||
"""Denormalize the flat vector back to original scale."""
|
||||
return x_flat * self.flat_std + self.flat_mean
|
||||
|
||||
def denormalize_components(
|
||||
self, beta: torch.Tensor, poses: torch.Tensor, trans: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""Denormalize each component back to original scale."""
|
||||
beta_orig = beta * self.beta_std_ + self.beta_mean_
|
||||
poses_orig = poses * self.pose_std_ + self.pose_mean_ # (48,) broadcast
|
||||
trans_orig = trans * self.trans_std_ + self.trans_mean_ # (3,) broadcast
|
||||
return beta_orig, poses_orig, trans_orig
|
||||
|
||||
def unpack(self, x_flat: torch.Tensor):
|
||||
"""Unpack flat → (beta, poses, trans)"""
|
||||
B = x_flat.shape[0]
|
||||
beta = x_flat[:, self.beta] # (B, 10)
|
||||
poses = x_flat[:, self.poses].reshape(B, self.window_size, self.pose_dim) # (B, T, 48)
|
||||
trans = x_flat[:, self.trans].reshape(B, self.window_size, self.trans_dim) # (B, T, 3)
|
||||
return beta, poses, trans
|
||||
|
||||
def pack(self, beta: torch.Tensor, poses: torch.Tensor,
|
||||
trans: torch.Tensor) -> torch.Tensor:
|
||||
"""Pack → flat"""
|
||||
B = beta.shape[0]
|
||||
return torch.cat([
|
||||
beta.reshape(B, -1),
|
||||
poses.reshape(B, -1),
|
||||
trans.reshape(B, -1),
|
||||
], dim=-1)
|
||||
|
||||
|
||||
class HandPoseDenoiser(nn.Module):
|
||||
"""
|
||||
Rectified Flow Denoiser — all-frame direct regression.
|
||||
"""
|
||||
|
||||
def __init__(self, config: DictConfig):
|
||||
super().__init__()
|
||||
fm_cfg = config.model.fm
|
||||
cond_cfg = fm_cfg.condition
|
||||
comp_cfg = fm_cfg.frame_compressor
|
||||
tf_cfg = fm_cfg.transformer
|
||||
|
||||
self.d_model = tf_cfg.hidden_size
|
||||
self.beta_dim = fm_cfg.beta_dim # 10
|
||||
self.pose_dim = fm_cfg.pose_dim # 48
|
||||
self.trans_dim = fm_cfg.trans_dim # 3
|
||||
self.window_size = config.dataset.window_size
|
||||
|
||||
# mano_slice (includes normalization statistics)
|
||||
stats_path = fm_cfg.get("normalization_stats", None)
|
||||
self.mano_slice = ManoSlice(
|
||||
beta_dim=self.beta_dim,
|
||||
pose_dim=self.pose_dim,
|
||||
trans_dim=self.trans_dim,
|
||||
window_size=self.window_size,
|
||||
stats_path=stats_path,
|
||||
)
|
||||
|
||||
# z-stream input projections: each component → d_model
|
||||
self.beta_proj_in = nn.Linear(self.beta_dim, self.d_model)
|
||||
self.pose_proj_in = nn.Linear(self.pose_dim, self.d_model)
|
||||
self.trans_proj_in = nn.Linear(self.trans_dim, self.d_model)
|
||||
|
||||
# z-stream output projections: d_model → native dims
|
||||
self.beta_proj_out = nn.Linear(self.d_model, self.beta_dim)
|
||||
self.pose_proj_out = nn.Linear(self.d_model, self.pose_dim)
|
||||
self.trans_proj_out = nn.Linear(self.d_model, self.trans_dim)
|
||||
|
||||
# HaMeR backbone (frozen)
|
||||
hamer_ckpt = config.model.get("hamer", {}).get("checkpoint", None)
|
||||
self.hamer_backbone = HaMeRBackbone(checkpoint_path=hamer_ckpt)
|
||||
|
||||
# Frame Compressor (directly compresses HaMeR patch tokens)
|
||||
self.frame_compressor = FrameCompressor(
|
||||
d_model=self.d_model,
|
||||
input_dim=cond_cfg.hamer_embed_dim,
|
||||
num_heads=tf_cfg.num_heads,
|
||||
num_layers=comp_cfg.num_layers,
|
||||
)
|
||||
|
||||
# Skeleton Encoder
|
||||
self.skeleton_encoder = RayDirectionSkeletonEncoder(
|
||||
n_joints=cond_cfg.n_joints,
|
||||
num_freqs=cond_cfg.skeleton_pe_freqs,
|
||||
embed_dim=self.d_model,
|
||||
)
|
||||
|
||||
# Condition Builder (cmask; mask_ratio only used for optional regularization)
|
||||
self.condition_builder = ConditionBuilder(
|
||||
d_model=self.d_model,
|
||||
mask_ratio=cond_cfg.get("cmask_ratio", 0.0),
|
||||
)
|
||||
|
||||
# RoPE for Denoiser
|
||||
self.rope = RotaryEmbedding(dim=self.d_model // tf_cfg.num_heads)
|
||||
|
||||
# Flow Matching Transformer
|
||||
self.transformer = FlowMatchingTransformer(
|
||||
in_channels=self.d_model,
|
||||
out_channels=self.d_model,
|
||||
use_context_in=True,
|
||||
context_in_dim=self.d_model,
|
||||
use_txt_in=False,
|
||||
vec_in_dim=0,
|
||||
use_pre_text_attn=False,
|
||||
config=tf_cfg,
|
||||
)
|
||||
|
||||
def _extract_image_features(self, images: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Image feature extraction: HaMeR backbone → Frame Compressor.
|
||||
images: (B, T, 3, 256, 256) uint8
|
||||
Returns: (B, T, d_model) — 1 image token per frame
|
||||
"""
|
||||
B, T = images.shape[:2]
|
||||
imgs = images.reshape(B * T, 3, images.shape[-2], images.shape[-1]).float() / 255.0
|
||||
with torch.no_grad():
|
||||
patches = self.hamer_backbone(imgs)
|
||||
patches = patches.reshape(B, T, 192, -1)
|
||||
compressed = self.frame_compressor(patches)
|
||||
return compressed
|
||||
|
||||
def _build_condition(self, batch: dict) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Build the condition stream and return (condition, m)."""
|
||||
if "image_tokens" in batch:
|
||||
image_tokens = batch["image_tokens"]
|
||||
else:
|
||||
image_tokens = self._extract_image_features(batch["images"])
|
||||
skeleton_tokens = self.skeleton_encoder(
|
||||
batch["hamer_landmarks"],
|
||||
batch["crop_intrinsics"],
|
||||
)
|
||||
condition, m = self.condition_builder(
|
||||
image_tokens,
|
||||
skeleton_tokens,
|
||||
batch["hamer_confidence"],
|
||||
)
|
||||
return condition, m
|
||||
|
||||
def _build_z_positions(self, B: int, T: int, device: torch.device) -> torch.Tensor:
|
||||
"""z-stream RoPE positions: [0, 1, 1, 2, 2, ..., T, T]"""
|
||||
pos = torch.zeros(1 + 2 * T, dtype=torch.long, device=device)
|
||||
for t in range(T):
|
||||
pos[1 + t] = t + 1 # pose token
|
||||
pos[1 + T + t] = t + 1 # trans token
|
||||
return pos
|
||||
|
||||
def _build_cond_positions(self, B: int, T: int, device: torch.device) -> torch.Tensor:
|
||||
"""c-stream RoPE positions: the 2 tokens per frame share position t"""
|
||||
pos = torch.zeros(2 * T, dtype=torch.long, device=device)
|
||||
for t in range(T):
|
||||
pos[2 * t] = t
|
||||
pos[2 * t + 1] = t
|
||||
return pos
|
||||
|
||||
def forward(self, x_t_flat: torch.Tensor, t: torch.Tensor, batch: dict):
|
||||
"""
|
||||
Args:
|
||||
x_t_flat: (B, flat_dim) — noisy flow target
|
||||
t: (B,) — flow timestep
|
||||
batch: dict with all condition data
|
||||
Returns:
|
||||
v_pred_flat: (B, flat_dim) — predicted flow velocity
|
||||
"""
|
||||
B, T = x_t_flat.shape[0], self.window_size
|
||||
device = x_t_flat.device
|
||||
|
||||
# 1. Unpack flat → components
|
||||
beta_t, poses_t, trans_t = self.mano_slice.unpack(x_t_flat)
|
||||
|
||||
# 2. Project to d_model → z-stream tokens
|
||||
beta_tok = self.beta_proj_in(beta_t).unsqueeze(1) # (B, 1, D)
|
||||
pose_tok = self.pose_proj_in(poses_t) # (B, T, D)
|
||||
trans_tok = self.trans_proj_in(trans_t) # (B, T, D)
|
||||
z_stream = torch.cat([beta_tok, pose_tok, trans_tok], dim=1) # (B, 1+2T, D)
|
||||
|
||||
# 3. Build condition → c-stream
|
||||
condition, _ = self._build_condition(batch) # (B, 2*T, D)
|
||||
|
||||
# 4. RoPE positions
|
||||
z_positions = self._build_z_positions(B, T, device)
|
||||
cond_positions = self._build_cond_positions(B, T, device)
|
||||
|
||||
# 5. Transformer forward
|
||||
pc_out, _ = self.transformer(
|
||||
pc=z_stream,
|
||||
dino=condition,
|
||||
timesteps=t,
|
||||
y=None,
|
||||
txt_tokens=None,
|
||||
z_positions=z_positions,
|
||||
cond_positions=cond_positions,
|
||||
rope_module=self.rope,
|
||||
)
|
||||
|
||||
# 6. Project back to native dims
|
||||
v_beta = self.beta_proj_out(pc_out[:, 0:1]) # (B, 1, 10)
|
||||
v_pose = self.pose_proj_out(pc_out[:, 1:1 + T]) # (B, T, 48)
|
||||
v_trans = self.trans_proj_out(pc_out[:, 1 + T:]) # (B, T, 3)
|
||||
|
||||
# 7. Pack flat
|
||||
v_pred_flat = self.mano_slice.pack(
|
||||
v_beta.squeeze(1), v_pose, v_trans
|
||||
)
|
||||
return v_pred_flat
|
||||
|
||||
@torch.no_grad()
|
||||
def build_flow_target(self, batch: dict) -> torch.Tensor:
|
||||
"""Build the GT flow target x_1 (normalized space) from the batch."""
|
||||
beta = batch["mano_betas"] # (B, 10)
|
||||
poses = batch["mano_params"] # (B, T, 48)
|
||||
trans = batch["mano_trans"] # (B, T, 3)
|
||||
flat = self.mano_slice.pack(beta, poses, trans)
|
||||
return self.mano_slice.normalize(flat)
|
||||
|
||||
@torch.no_grad()
|
||||
def reconstruct_params(self, x_1_flat: torch.Tensor):
|
||||
"""
|
||||
Reconstruct MANO parameters (original scale) from the predicted x_1 (normalized space).
|
||||
|
||||
All values are used directly, without cumsum.
|
||||
|
||||
Returns:
|
||||
params: (B, T, 51) = cat[pose(48), trans(3)]
|
||||
beta: (B, 10)
|
||||
"""
|
||||
x_1_orig = self.mano_slice.denormalize(x_1_flat)
|
||||
beta, poses, trans = self.mano_slice.unpack(x_1_orig)
|
||||
params = torch.cat([poses, trans], dim=-1) # (B, T, 51)
|
||||
return params, beta
|
||||
|
||||
def get_model_wrapper(self):
|
||||
"""Return a ModelWrapper for the ODE solver."""
|
||||
return ModelWrapper(self)
|
||||
|
||||
def get_time_sampler(self):
|
||||
"""Return the time sampler."""
|
||||
return FluxTimeSampler()
|
||||
@@ -0,0 +1,464 @@
|
||||
"""
|
||||
Dual-stream transformer for flow matching on hand pose sequences.
|
||||
|
||||
Derived from ShapeR/model/flow_matching/dualstream_transformer.py with changes:
|
||||
- Import path adjusted from `from model import attention` (same, works as-is)
|
||||
- Renamed internal variables: img→z (latent stream), txt→cond (condition stream)
|
||||
- Added RoPE support: position IDs + rope_module passed through block calls
|
||||
|
||||
Architecture (Flux-style):
|
||||
- DoubleStreamBlock: joint attention over (latent z_t, condition context)
|
||||
- SingleStreamBlock: self-attention on concatenated tokens
|
||||
- LastLayer: final projection to output dimension
|
||||
"""
|
||||
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# This source code is licensed under the CC BY-NC 4.0 license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
#
|
||||
# This file is inspired by the Flux dual stream transformer architecture.
|
||||
# Original source: https://github.com/black-forest-labs/flux
|
||||
# Original license: Apache License 2.0
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import torch
|
||||
from einops import rearrange
|
||||
from model import attention
|
||||
from model.pose_embedding import RotaryEmbedding
|
||||
from torch import nn, Tensor
|
||||
|
||||
|
||||
class FlowMatchingTransformer(nn.Module):
|
||||
"""Transformer model for flow matching on sequences."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
use_context_in,
|
||||
context_in_dim,
|
||||
use_txt_in,
|
||||
vec_in_dim,
|
||||
use_pre_text_attn,
|
||||
config,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.config = config
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = out_channels
|
||||
self.use_txt_in = use_txt_in
|
||||
self.use_context_in = use_context_in
|
||||
self.use_pre_text_attn = use_pre_text_attn
|
||||
if config.hidden_size % config.num_heads != 0:
|
||||
raise ValueError(
|
||||
f"Hidden size {config.hidden_size} must be divisible by num_heads {config.num_heads}"
|
||||
)
|
||||
self.hidden_size = config.hidden_size
|
||||
self.num_heads = config.num_heads
|
||||
self.pc_in = nn.Linear(self.in_channels, self.hidden_size, bias=True)
|
||||
self.time_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size)
|
||||
if self.use_txt_in:
|
||||
self.vector_in = MLPEmbedder(vec_in_dim, self.hidden_size)
|
||||
|
||||
self.dropout = getattr(config, 'dropout', 0.1)
|
||||
|
||||
if self.use_pre_text_attn:
|
||||
self.text_attn_in = nn.Linear(context_in_dim, self.hidden_size)
|
||||
self.text_blocks = nn.ModuleList(
|
||||
[
|
||||
DoubleStreamBlock(
|
||||
self.hidden_size,
|
||||
self.num_heads,
|
||||
mlp_ratio=config.mlp_ratio,
|
||||
qkv_bias=config.qkv_bias,
|
||||
dropout=self.dropout,
|
||||
)
|
||||
for _ in range(4)
|
||||
]
|
||||
)
|
||||
|
||||
if self.use_context_in:
|
||||
self.dino_in = nn.Linear(context_in_dim, self.hidden_size)
|
||||
self.double_blocks = nn.ModuleList(
|
||||
[
|
||||
DoubleStreamBlock(
|
||||
self.hidden_size,
|
||||
self.num_heads,
|
||||
mlp_ratio=config.mlp_ratio,
|
||||
qkv_bias=config.qkv_bias,
|
||||
dropout=self.dropout,
|
||||
)
|
||||
for _ in range(config.depth)
|
||||
]
|
||||
)
|
||||
|
||||
self.single_blocks = nn.ModuleList(
|
||||
[
|
||||
SingleStreamBlock(
|
||||
self.hidden_size, self.num_heads, mlp_ratio=config.mlp_ratio,
|
||||
dropout=self.dropout,
|
||||
)
|
||||
for _ in range(
|
||||
config.depth_single_blocks
|
||||
+ (config.depth if not self.use_context_in else 0)
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
self.final_layer = LastLayer(self.hidden_size, 1, self.out_channels)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
pc: Tensor,
|
||||
dino: Tensor,
|
||||
timesteps: Tensor,
|
||||
y: Tensor,
|
||||
txt_tokens: Tensor,
|
||||
z_positions: Tensor = None,
|
||||
cond_positions: Tensor = None,
|
||||
rope_module=None,
|
||||
) -> Tensor:
|
||||
if pc.ndim != 3 or (dino is not None and dino.ndim != 3):
|
||||
raise ValueError("Input pc and dino tensors must have 3 dimensions.")
|
||||
|
||||
pc = self.pc_in(pc)
|
||||
vec = self.time_in(timestep_embedding(timesteps, 256))
|
||||
intermediates = None
|
||||
if self.use_txt_in and y is not None:
|
||||
vec = vec + self.vector_in(y)
|
||||
|
||||
if self.use_pre_text_attn:
|
||||
txt_tokens = self.text_attn_in(txt_tokens)
|
||||
for block in self.text_blocks:
|
||||
pc, txt_tokens = block(z=pc, cond=txt_tokens, vec=vec)
|
||||
|
||||
cond_len = 0
|
||||
if self.use_context_in:
|
||||
dino = self.dino_in(dino)
|
||||
cond_len = dino.shape[1]
|
||||
for block in self.double_blocks:
|
||||
pc, dino = block(
|
||||
z=pc, cond=dino, vec=vec,
|
||||
z_positions=z_positions,
|
||||
cond_positions=cond_positions,
|
||||
rope_module=rope_module,
|
||||
)
|
||||
# Build combined positions (matching concatenation order: [cond, z])
|
||||
if z_positions is not None and cond_positions is not None:
|
||||
combined_positions = torch.cat([cond_positions, z_positions], dim=0)
|
||||
else:
|
||||
combined_positions = None
|
||||
pc = torch.cat((dino, pc), 1)
|
||||
|
||||
for block_idx, block in enumerate(self.single_blocks):
|
||||
pc = block(pc, vec=vec, positions=combined_positions, rope_module=rope_module)
|
||||
if block_idx == 3:
|
||||
intermediates = pc
|
||||
|
||||
if self.use_context_in:
|
||||
pc = pc[:, cond_len:, ...]
|
||||
|
||||
pc = self.final_layer(pc, vec)
|
||||
return pc, intermediates
|
||||
|
||||
|
||||
def timestep_embedding(t: Tensor, dim, max_period=10000, time_factor: float = 1000.0):
|
||||
"""Create sinusoidal timestep embeddings."""
|
||||
t = time_factor * t
|
||||
half = dim // 2
|
||||
freqs = torch.exp(
|
||||
-math.log(max_period)
|
||||
* torch.arange(start=0, end=half, dtype=torch.float32)
|
||||
/ half
|
||||
).to(t.device)
|
||||
|
||||
args = t[:, None].float() * freqs[None]
|
||||
embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
|
||||
if dim % 2:
|
||||
embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
|
||||
if torch.is_floating_point(t):
|
||||
embedding = embedding.to(t)
|
||||
return embedding
|
||||
|
||||
|
||||
class MLPEmbedder(nn.Module):
|
||||
"""Two-layer MLP for embedding."""
|
||||
|
||||
def __init__(self, in_dim: int, hidden_dim: int):
|
||||
super().__init__()
|
||||
self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True)
|
||||
self.silu = nn.SiLU()
|
||||
self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True)
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
return self.out_layer(self.silu(self.in_layer(x)))
|
||||
|
||||
|
||||
class RMSNorm(torch.nn.Module):
|
||||
"""Root Mean Square Layer Normalization."""
|
||||
|
||||
def __init__(self, dim: int):
|
||||
super().__init__()
|
||||
self.scale = nn.Parameter(torch.ones(dim))
|
||||
|
||||
def forward(self, x: Tensor):
|
||||
x_dtype = x.dtype
|
||||
x = x.float()
|
||||
rrms = torch.rsqrt(torch.mean(x**2, dim=-1, keepdim=True) + 1e-6)
|
||||
return (x * rrms).to(dtype=x_dtype) * self.scale
|
||||
|
||||
|
||||
class QKNorm(torch.nn.Module):
|
||||
"""Query-Key normalization for stable attention."""
|
||||
|
||||
def __init__(self, dim: int):
|
||||
super().__init__()
|
||||
self.query_norm = RMSNorm(dim)
|
||||
self.key_norm = RMSNorm(dim)
|
||||
|
||||
def forward(self, q: Tensor, k: Tensor, v: Tensor) -> Tuple[Tensor, Tensor]:
|
||||
q = self.query_norm(q)
|
||||
k = self.key_norm(k)
|
||||
return q.to(v), k.to(v)
|
||||
|
||||
|
||||
class SelfAttention(nn.Module):
|
||||
"""Standard self-attention with QK normalization."""
|
||||
|
||||
def __init__(self, dim: int, num_heads: int = 8, qkv_bias: bool = False):
|
||||
super().__init__()
|
||||
self.num_heads = num_heads
|
||||
head_dim = dim // num_heads
|
||||
|
||||
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
|
||||
self.norm = QKNorm(head_dim)
|
||||
self.proj = nn.Linear(dim, dim)
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
qkv = self.qkv(x)
|
||||
q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
|
||||
q, k = self.norm(q, k, v)
|
||||
x = attention(q, k, v)
|
||||
x = self.proj(x)
|
||||
return x
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModulationOut:
|
||||
shift: Tensor
|
||||
scale: Tensor
|
||||
gate: Tensor
|
||||
|
||||
|
||||
class Modulation(nn.Module):
|
||||
"""Adaptive layer norm modulation from timestep embedding."""
|
||||
|
||||
def __init__(self, dim: int, double: bool):
|
||||
super().__init__()
|
||||
self.is_double = double
|
||||
self.multiplier = 6 if double else 3
|
||||
self.lin = nn.Linear(dim, self.multiplier * dim, bias=True)
|
||||
|
||||
def forward(self, vec: Tensor) -> Tuple[ModulationOut, Optional[ModulationOut]]:
|
||||
out = self.lin(nn.functional.silu(vec))[:, None, :].chunk(
|
||||
self.multiplier, dim=-1
|
||||
)
|
||||
return (
|
||||
ModulationOut(*out[:3]),
|
||||
ModulationOut(*out[3:]) if self.is_double else None,
|
||||
)
|
||||
|
||||
|
||||
class DoubleStreamBlock(nn.Module):
|
||||
"""Double-stream attention block with separate z (latent) and cond (condition) streams."""
|
||||
|
||||
def __init__(
|
||||
self, hidden_size: int, num_heads: int, mlp_ratio: float,
|
||||
qkv_bias: bool = False, dropout: float = 0.1,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
mlp_hidden_dim = int(hidden_size * mlp_ratio)
|
||||
self.num_heads = num_heads
|
||||
self.hidden_size = hidden_size
|
||||
# z stream (latent) modules
|
||||
self.z_mod = Modulation(hidden_size, double=True)
|
||||
self.z_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
|
||||
self.z_attn = SelfAttention(
|
||||
dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias
|
||||
)
|
||||
self.z_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
|
||||
self.z_mlp = nn.Sequential(
|
||||
nn.Linear(hidden_size, mlp_hidden_dim, bias=True),
|
||||
nn.GELU(approximate="tanh"),
|
||||
nn.Linear(mlp_hidden_dim, hidden_size, bias=True),
|
||||
)
|
||||
self.z_attn_drop = nn.Dropout(dropout)
|
||||
self.z_mlp_drop = nn.Dropout(dropout)
|
||||
|
||||
# cond stream (condition) modules
|
||||
self.cond_mod = Modulation(hidden_size, double=True)
|
||||
self.cond_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
|
||||
self.cond_attn = SelfAttention(
|
||||
dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias
|
||||
)
|
||||
self.cond_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
|
||||
self.cond_mlp = nn.Sequential(
|
||||
nn.Linear(hidden_size, mlp_hidden_dim, bias=True),
|
||||
nn.GELU(approximate="tanh"),
|
||||
nn.Linear(mlp_hidden_dim, hidden_size, bias=True),
|
||||
)
|
||||
self.cond_attn_drop = nn.Dropout(dropout)
|
||||
self.cond_mlp_drop = nn.Dropout(dropout)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
z: Tensor,
|
||||
cond: Tensor,
|
||||
vec: Tensor,
|
||||
z_positions: Tensor = None,
|
||||
cond_positions: Tensor = None,
|
||||
rope_module=None,
|
||||
) -> Tuple[Tensor, Tensor]:
|
||||
z_mod1, z_mod2 = self.z_mod(vec)
|
||||
cond_mod1, cond_mod2 = self.cond_mod(vec)
|
||||
|
||||
# prepare z stream for attention
|
||||
z_modulated = self.z_norm1(z)
|
||||
z_modulated = (1 + z_mod1.scale) * z_modulated + z_mod1.shift
|
||||
z_qkv = self.z_attn.qkv(z_modulated)
|
||||
z_q, z_k, z_v = rearrange(
|
||||
z_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads
|
||||
)
|
||||
z_q, z_k = self.z_attn.norm(z_q, z_k, z_v)
|
||||
|
||||
# prepare cond stream for attention
|
||||
cond_modulated = self.cond_norm1(cond)
|
||||
cond_modulated = (1 + cond_mod1.scale) * cond_modulated + cond_mod1.shift
|
||||
cond_qkv = self.cond_attn.qkv(cond_modulated)
|
||||
cond_q, cond_k, cond_v = rearrange(
|
||||
cond_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads
|
||||
)
|
||||
cond_q, cond_k = self.cond_attn.norm(cond_q, cond_k, cond_v)
|
||||
|
||||
# Apply RoPE (after QK-norm, before concatenation)
|
||||
if rope_module is not None and z_positions is not None:
|
||||
sin, cos = RotaryEmbedding.get_sin_cos_from_positions(
|
||||
z_positions, rope_module.inv_freq
|
||||
)
|
||||
z_q = RotaryEmbedding.apply_rope(z_q, sin, cos)
|
||||
z_k = RotaryEmbedding.apply_rope(z_k, sin, cos)
|
||||
if rope_module is not None and cond_positions is not None:
|
||||
sin, cos = RotaryEmbedding.get_sin_cos_from_positions(
|
||||
cond_positions, rope_module.inv_freq
|
||||
)
|
||||
cond_q = RotaryEmbedding.apply_rope(cond_q, sin, cos)
|
||||
cond_k = RotaryEmbedding.apply_rope(cond_k, sin, cos)
|
||||
|
||||
# run actual attention (concat [cond, z] streams)
|
||||
q = torch.cat((cond_q, z_q), dim=2)
|
||||
k = torch.cat((cond_k, z_k), dim=2)
|
||||
v = torch.cat((cond_v, z_v), dim=2)
|
||||
|
||||
attn = attention(q, k, v)
|
||||
cond_attn, z_attn = attn[:, : cond.shape[1]], attn[:, cond.shape[1] :]
|
||||
|
||||
# calculate the z stream blocks
|
||||
z = z + z_mod1.gate * self.z_attn_drop(self.z_attn.proj(z_attn))
|
||||
z = z + z_mod2.gate * self.z_mlp_drop(self.z_mlp(
|
||||
(1 + z_mod2.scale) * self.z_norm2(z) + z_mod2.shift
|
||||
))
|
||||
|
||||
# calculate the cond stream blocks
|
||||
cond = cond + cond_mod1.gate * self.cond_attn_drop(self.cond_attn.proj(cond_attn))
|
||||
cond = cond + cond_mod2.gate * self.cond_mlp_drop(self.cond_mlp(
|
||||
(1 + cond_mod2.scale) * self.cond_norm2(cond) + cond_mod2.shift
|
||||
))
|
||||
return z, cond
|
||||
|
||||
|
||||
class SingleStreamBlock(nn.Module):
|
||||
"""
|
||||
A DiT block with parallel linear layers as described in
|
||||
https://arxiv.org/abs/2302.05442 and adapted modulation interface.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size: int,
|
||||
num_heads: int,
|
||||
mlp_ratio: float = 4.0,
|
||||
qk_scale: Optional[float] = None,
|
||||
dropout: float = 0.1,
|
||||
):
|
||||
super().__init__()
|
||||
self.hidden_dim = hidden_size
|
||||
self.num_heads = num_heads
|
||||
head_dim = hidden_size // num_heads
|
||||
self.scale = qk_scale or head_dim**-0.5
|
||||
|
||||
self.mlp_hidden_dim = int(hidden_size * mlp_ratio)
|
||||
self.linear1 = nn.Linear(hidden_size, hidden_size * 3 + self.mlp_hidden_dim)
|
||||
self.linear2 = nn.Linear(hidden_size + self.mlp_hidden_dim, hidden_size)
|
||||
|
||||
self.norm = QKNorm(head_dim)
|
||||
|
||||
self.hidden_size = hidden_size
|
||||
self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
|
||||
|
||||
self.mlp_act = nn.GELU(approximate="tanh")
|
||||
self.modulation = Modulation(hidden_size, double=False)
|
||||
self.attn_drop = nn.Dropout(dropout)
|
||||
self.mlp_drop = nn.Dropout(dropout)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: Tensor,
|
||||
vec: Tensor,
|
||||
positions: Tensor = None,
|
||||
rope_module=None,
|
||||
) -> Tensor:
|
||||
mod, _ = self.modulation(vec)
|
||||
x_mod = (1 + mod.scale) * self.pre_norm(x) + mod.shift
|
||||
qkv, mlp = torch.split(
|
||||
self.linear1(x_mod), [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1
|
||||
)
|
||||
|
||||
q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
|
||||
q, k = self.norm(q, k, v)
|
||||
|
||||
# Apply RoPE (after QK-norm, before attention)
|
||||
if rope_module is not None and positions is not None:
|
||||
sin, cos = RotaryEmbedding.get_sin_cos_from_positions(
|
||||
positions, rope_module.inv_freq
|
||||
)
|
||||
q = RotaryEmbedding.apply_rope(q, sin, cos)
|
||||
k = RotaryEmbedding.apply_rope(k, sin, cos)
|
||||
|
||||
attn = self.attn_drop(attention(q, k, v))
|
||||
output = self.mlp_drop(self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2)))
|
||||
return x + mod.gate * output
|
||||
|
||||
|
||||
class LastLayer(nn.Module):
|
||||
"""Final projection layer with adaptive layer norm."""
|
||||
|
||||
def __init__(self, hidden_size: int, patch_size: int, out_channels: int):
|
||||
super().__init__()
|
||||
self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
|
||||
self.linear = nn.Linear(
|
||||
hidden_size, patch_size * patch_size * out_channels, bias=True
|
||||
)
|
||||
self.adaLN_modulation = nn.Sequential(
|
||||
nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True)
|
||||
)
|
||||
|
||||
def forward(self, x: Tensor, vec: Tensor) -> Tensor:
|
||||
shift, scale = self.adaLN_modulation(vec).chunk(2, dim=1)
|
||||
x = (1 + scale[:, None, :]) * self.norm_final(x) + shift[:, None, :]
|
||||
x = self.linear(x)
|
||||
return x
|
||||
@@ -0,0 +1,20 @@
|
||||
"""Flow matching helpers package — direct reuse from ShapeR."""
|
||||
|
||||
from model.flow_matching.helpers.path import AffineProbPath, PathSample
|
||||
from model.flow_matching.helpers.scheduler import (
|
||||
CondOTScheduler,
|
||||
FluxTimeSampler,
|
||||
TimeSampler,
|
||||
)
|
||||
from model.flow_matching.helpers.solver import ODESolver
|
||||
from model.flow_matching.helpers.model_wrapper import ModelWrapper
|
||||
|
||||
__all__ = [
|
||||
"AffineProbPath",
|
||||
"PathSample",
|
||||
"CondOTScheduler",
|
||||
"FluxTimeSampler",
|
||||
"TimeSampler",
|
||||
"ODESolver",
|
||||
"ModelWrapper",
|
||||
]
|
||||
@@ -0,0 +1,46 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# This source code is licensed under the CC BY-NC 4.0 license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
#
|
||||
# This file is derived from https://github.com/facebookresearch/flow_matching
|
||||
# Licensed under CC BY-NC 4.0: https://creativecommons.org/licenses/by-nc/4.0/
|
||||
|
||||
# pyre-unsafe
|
||||
|
||||
from abc import ABC
|
||||
|
||||
from torch import nn, Tensor
|
||||
|
||||
|
||||
class ModelWrapper(ABC, nn.Module):
|
||||
"""
|
||||
This class is used to wrap around another model, adding custom forward pass logic.
|
||||
"""
|
||||
|
||||
def __init__(self, model: nn.Module):
|
||||
super().__init__()
|
||||
self.model = model
|
||||
|
||||
def forward(self, x: Tensor, t: Tensor, **extras) -> Tensor:
|
||||
r"""
|
||||
This method defines how inputs should be passed through the wrapped model.
|
||||
Here, we're assuming that the wrapped model takes both :math:`x` and :math:`t` as input,
|
||||
along with any additional keyword arguments.
|
||||
|
||||
Optional things to do here:
|
||||
- check that t is in the dimensions that the model is expecting.
|
||||
- add a custom forward pass logic.
|
||||
- call the wrapped model.
|
||||
|
||||
| given x, t
|
||||
| returns the model output for input x at time t, with extra information `extra`.
|
||||
|
||||
Args:
|
||||
x (Tensor): input data to the model (batch_size, ...).
|
||||
t (Tensor): time (batch_size).
|
||||
**extras: additional information forwarded to the model, e.g., text condition.
|
||||
|
||||
Returns:
|
||||
Tensor: model output.
|
||||
"""
|
||||
return self.model(x=x, t=t, **extras)
|
||||
@@ -0,0 +1,134 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# This source code is licensed under the CC BY-NC 4.0 license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
#
|
||||
# This file is derived from https://github.com/facebookresearch/flow_matching
|
||||
# Licensed under CC BY-NC 4.0: https://creativecommons.org/licenses/by-nc/4.0/
|
||||
|
||||
# pyre-unsafe
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from torch import Tensor
|
||||
|
||||
|
||||
def expand_tensor_like(input_tensor: Tensor, expand_to: Tensor) -> Tensor:
|
||||
"""`input_tensor` is a 1d vector of length equal to the batch size of `expand_to`,
|
||||
expand `input_tensor` to have the same shape as `expand_to` along all remaining dimensions.
|
||||
|
||||
Args:
|
||||
input_tensor (Tensor): (batch_size,).
|
||||
expand_to (Tensor): (batch_size, ...).
|
||||
|
||||
Returns:
|
||||
Tensor: (batch_size, ...).
|
||||
"""
|
||||
assert input_tensor.ndim == 1, "Input tensor must be a 1d vector."
|
||||
assert (
|
||||
input_tensor.shape[0] == expand_to.shape[0]
|
||||
), f"The first (batch_size) dimension must match. Got shape {input_tensor.shape} and {expand_to.shape}."
|
||||
|
||||
dim_diff = expand_to.ndim - input_tensor.ndim
|
||||
|
||||
t_expanded = input_tensor.clone()
|
||||
t_expanded = t_expanded.reshape(-1, *([1] * dim_diff))
|
||||
|
||||
return t_expanded.expand_as(expand_to)
|
||||
|
||||
|
||||
@dataclass
|
||||
class PathSample:
|
||||
r"""Represents a sample of a conditional-flow generated probability path.
|
||||
|
||||
Attributes:
|
||||
x_1 (Tensor): the target sample :math:`X_1`.
|
||||
x_0 (Tensor): the source sample :math:`X_0`.
|
||||
t (Tensor): the time sample :math:`t`.
|
||||
x_t (Tensor): samples :math:`X_t \sim p_t(X_t)`, shape (batch_size, ...).
|
||||
dx_t (Tensor): conditional target :math:`\frac{\partial X}{\partial t}`, shape: (batch_size, ...).
|
||||
|
||||
"""
|
||||
|
||||
x_1: Tensor = field(metadata={"help": "target samples X_1 (batch_size, ...)."})
|
||||
x_0: Tensor = field(metadata={"help": "source samples X_0 (batch_size, ...)."})
|
||||
t: Tensor = field(metadata={"help": "time samples t (batch_size, ...)."})
|
||||
x_t: Tensor = field(
|
||||
metadata={"help": "samples x_t ~ p_t(X_t), shape (batch_size, ...)."}
|
||||
)
|
||||
dx_t: Tensor = field(
|
||||
metadata={"help": "conditional target dX_t, shape: (batch_size, ...)."}
|
||||
)
|
||||
|
||||
|
||||
class AffineProbPath:
|
||||
r"""The ``AffineProbPath`` class represents a specific type of probability path where the transformation between distributions is affine.
|
||||
An affine transformation can be represented as:
|
||||
|
||||
.. math::
|
||||
|
||||
X_t = \alpha_t X_1 + \sigma_t X_0,
|
||||
|
||||
where :math:`X_t` is the transformed data point at time `t`. :math:`X_0` and :math:`X_1` are the source and target data points, respectively. :math:`\alpha_t` and :math:`\sigma_t` are the parameters of the affine transformation at time `t`.
|
||||
|
||||
The scheduler is responsible for providing the time-dependent parameters :math:`\alpha_t` and :math:`\sigma_t`, as well as their derivatives, which define the affine transformation at any given time `t`.
|
||||
|
||||
Using ``AffineProbPath`` in the flow matching framework:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# Instantiates a probability path
|
||||
my_path = AffineProbPath(...)
|
||||
mse_loss = torch.nn.MSELoss()
|
||||
|
||||
for x_1 in dataset:
|
||||
# Sets x_0 to random noise
|
||||
x_0 = torch.randn()
|
||||
|
||||
# Sets t to a random value in [0,1]
|
||||
t = torch.rand()
|
||||
|
||||
# Samples the conditional path X_t ~ p_t(X_t|X_0,X_1)
|
||||
path_sample = my_path.sample(x_0=x_0, x_1=x_1, t=t)
|
||||
|
||||
# Computes the MSE loss w.r.t. the velocity
|
||||
loss = mse_loss(path_sample.dx_t, my_model(x_t, t))
|
||||
loss.backward()
|
||||
|
||||
Args:
|
||||
scheduler (Scheduler): An instance of a scheduler that provides the parameters :math:`\alpha_t`, :math:`\sigma_t`, and their derivatives over time.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, scheduler):
|
||||
self.scheduler = scheduler
|
||||
|
||||
def sample(self, x_0: Tensor, x_1: Tensor, t: Tensor):
|
||||
r"""Sample from the affine probability path:
|
||||
|
||||
| given :math:`(X_0,X_1) \sim \pi(X_0,X_1)` and a scheduler :math:`(\alpha_t,\sigma_t)`.
|
||||
| return :math:`X_0, X_1, X_t = \alpha_t X_1 + \sigma_t X_0`, and the conditional velocity at :math:`X_t, \dot{X}_t = \dot{\alpha}_t X_1 + \dot{\sigma}_t X_0`.
|
||||
|
||||
Args:
|
||||
x_0 (Tensor): source data point, shape (batch_size, ...).
|
||||
x_1 (Tensor): target data point, shape (batch_size, ...).
|
||||
t (Tensor): times in [0,1], shape (batch_size).
|
||||
|
||||
Returns:
|
||||
PathSample: a conditional sample at :math:`X_t \sim p_t`.
|
||||
"""
|
||||
scheduler_output = self.scheduler(t)
|
||||
alpha_t = expand_tensor_like(
|
||||
input_tensor=scheduler_output.alpha_t, expand_to=x_1
|
||||
)
|
||||
sigma_t = expand_tensor_like(
|
||||
input_tensor=scheduler_output.sigma_t, expand_to=x_1
|
||||
)
|
||||
d_alpha_t = expand_tensor_like(
|
||||
input_tensor=scheduler_output.d_alpha_t, expand_to=x_1
|
||||
)
|
||||
d_sigma_t = expand_tensor_like(
|
||||
input_tensor=scheduler_output.d_sigma_t, expand_to=x_1
|
||||
)
|
||||
x_t = sigma_t * x_0 + alpha_t * x_1
|
||||
dx_t = d_sigma_t * x_0 + d_alpha_t * x_1
|
||||
return PathSample(x_t=x_t, dx_t=dx_t, x_1=x_1, x_0=x_0, t=t)
|
||||
@@ -0,0 +1,146 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# This source code is licensed under the CC BY-NC 4.0 license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
#
|
||||
# This file is derived from https://github.com/facebookresearch/flow_matching
|
||||
# Licensed under CC BY-NC 4.0: https://creativecommons.org/licenses/by-nc/4.0/
|
||||
|
||||
# pyre-unsafe
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from torch.distributions import LogisticNormal, Uniform
|
||||
|
||||
|
||||
def get_lin_function(
|
||||
x1: float = 256, y1: float = 0.5, x2: float = 4096, y2: float = 1.15
|
||||
):
|
||||
m = (y2 - y1) / (x2 - x1)
|
||||
b = y1 - m * x1
|
||||
return lambda x: m * x + b
|
||||
|
||||
|
||||
def time_shift(mu: float, sigma: float, t: Tensor):
|
||||
return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
|
||||
|
||||
|
||||
@dataclass
|
||||
class SchedulerOutput:
|
||||
r"""Represents a sample of a conditional-flow generated probability path.
|
||||
|
||||
Attributes:
|
||||
alpha_t (Tensor): :math:`\alpha_t`, shape (...).
|
||||
sigma_t (Tensor): :math:`\sigma_t`, shape (...).
|
||||
d_alpha_t (Tensor): :math:`\frac{\partial}{\partial t}\alpha_t`, shape (...).
|
||||
d_sigma_t (Tensor): :math:`\frac{\partial}{\partial t}\sigma_t`, shape (...).
|
||||
|
||||
"""
|
||||
|
||||
alpha_t: Tensor = field(metadata={"help": "alpha_t"})
|
||||
sigma_t: Tensor = field(metadata={"help": "sigma_t"})
|
||||
d_alpha_t: Tensor = field(metadata={"help": "Derivative of alpha_t."})
|
||||
d_sigma_t: Tensor = field(metadata={"help": "Derivative of sigma_t."})
|
||||
|
||||
|
||||
class CondOTScheduler:
|
||||
"""CondOT Scheduler."""
|
||||
|
||||
def __call__(self, t: Tensor) -> SchedulerOutput:
|
||||
r"""
|
||||
Args:
|
||||
t (Tensor): times in [0,1], shape (...).
|
||||
|
||||
Returns:
|
||||
SchedulerOutput: :math:`\alpha_t,\sigma_t,\frac{\partial}{\partial t}\alpha_t,\frac{\partial}{\partial t}\sigma_t`
|
||||
"""
|
||||
return SchedulerOutput(
|
||||
alpha_t=t,
|
||||
sigma_t=1 - t,
|
||||
d_alpha_t=torch.ones_like(t),
|
||||
d_sigma_t=-torch.ones_like(t),
|
||||
)
|
||||
|
||||
def snr_inverse(self, snr: Tensor) -> Tensor:
|
||||
r"""
|
||||
Computes :math:`t` from the signal-to-noise ratio :math:`\frac{\alpha_t}{\sigma_t}`.
|
||||
|
||||
Args:
|
||||
snr (Tensor): The signal-to-noise, shape (...)
|
||||
|
||||
Returns:
|
||||
Tensor: t, shape (...)
|
||||
"""
|
||||
kappa_t = snr / (1.0 + snr)
|
||||
|
||||
return self.kappa_inverse(kappa=kappa_t)
|
||||
|
||||
def kappa_inverse(self, kappa: Tensor) -> Tensor:
|
||||
"""
|
||||
Computes :math:`t` from :math:`kappa_t`.
|
||||
|
||||
Args:
|
||||
kappa (Tensor): :math:`kappa`, shape (...)
|
||||
|
||||
Returns:
|
||||
Tensor: t, shape (...)
|
||||
"""
|
||||
return kappa
|
||||
|
||||
|
||||
class TimeSampler:
|
||||
def __init__(self, dist):
|
||||
self.dist_name = dist
|
||||
if dist == "uniform":
|
||||
self.distribution = Uniform(0.0, 1.0)
|
||||
elif dist == "lognormal":
|
||||
self.distribution = LogisticNormal(0.0, 1.0)
|
||||
elif dist == "flux":
|
||||
self.distribution = FluxTimeSampler(mode="train")
|
||||
else:
|
||||
raise ValueError(f"Unknown distribution {dist}")
|
||||
|
||||
def __call__(self, shape, num_tokens, device):
|
||||
if self.dist_name == "uniform":
|
||||
return self.distribution.sample(sample_shape=(shape,)).to(device).squeeze(0)
|
||||
elif self.dist_name == "lognormal":
|
||||
return self.distribution.sample(sample_shape=(shape,))[:, 0].to(device)
|
||||
elif self.dist_name == "flux":
|
||||
return self.distribution(shape, num_tokens, device)
|
||||
|
||||
|
||||
class FluxTimeSampler:
|
||||
def __init__(self, **kwargs):
|
||||
self.mode = kwargs.get("mode", "train")
|
||||
|
||||
def __call__(self, bs, num_tokens, device):
|
||||
if self.mode == "train":
|
||||
logits_norm = torch.randn(bs, device=device)
|
||||
timesteps = logits_norm.sigmoid()
|
||||
mu = get_lin_function(y1=0.5, y2=1.15)(num_tokens)
|
||||
timesteps = time_shift(mu, 1.0, timesteps)
|
||||
timesteps = 1 - timesteps
|
||||
elif self.mode == "inference":
|
||||
timesteps = torch.linspace(1, 0, bs + 1).to(device)
|
||||
mu = get_lin_function(y1=0.5, y2=1.15)(num_tokens)
|
||||
timesteps = time_shift(mu, 1.0, timesteps)
|
||||
timesteps = 1 - timesteps
|
||||
return timesteps
|
||||
|
||||
|
||||
def visualize_sampling():
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
flux_sampler = FluxTimeSampler(mode="inference")
|
||||
timesteps = flux_sampler(25, 2048, torch.device("cuda"))
|
||||
x = torch.linspace(0, 1, 25 + 1)
|
||||
plt.scatter(x, y=timesteps.cpu().numpy())
|
||||
plt.savefig("scatter_plot.png")
|
||||
|
||||
flux_sampler = FluxTimeSampler(mode="train")
|
||||
timesteps = flux_sampler(5000, 4096, torch.device("cuda"))
|
||||
plt.hist(timesteps.cpu().numpy(), bins=50, edgecolor="black")
|
||||
plt.savefig("scatter_plot_4096.png")
|
||||
print(timesteps)
|
||||
@@ -0,0 +1,78 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# This source code is licensed under the CC BY-NC 4.0 license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
#
|
||||
# This file is derived from https://github.com/facebookresearch/flow_matching
|
||||
# Licensed under CC BY-NC 4.0: https://creativecommons.org/licenses/by-nc/4.0/
|
||||
|
||||
# pyre-unsafe
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from torchdiffeq import odeint
|
||||
|
||||
|
||||
class ODESolver:
|
||||
"""A class to solve ordinary differential equations (ODEs) using a specified velocity model.
|
||||
|
||||
This class utilizes a velocity field model to solve ODEs over a given time grid using numerical ode solvers.
|
||||
|
||||
Args:
|
||||
velocity_model (Union[ModelWrapper, Callable]): a velocity field model receiving :math:`(x,t)` and returning :math:`u_t(x)`
|
||||
"""
|
||||
|
||||
def __init__(self, velocity_model):
|
||||
super().__init__()
|
||||
self.velocity_model = velocity_model
|
||||
|
||||
def sample(
|
||||
self,
|
||||
x_init: Tensor,
|
||||
# pyre-fixme[9]: step_size has type `float`; used as `None`.
|
||||
step_size: float = None,
|
||||
method: str = "euler",
|
||||
atol: float = 1e-5,
|
||||
rtol: float = 1e-5,
|
||||
time_grid: Tensor = torch.tensor([0.0, 1.0]),
|
||||
return_intermediates: bool = False,
|
||||
enable_grad: bool = False,
|
||||
**model_extras,
|
||||
):
|
||||
r"""Solve the ODE with the velocity field.
|
||||
|
||||
Args:
|
||||
x_init (Tensor): initial conditions (e.g., source samples :math:`X_0 \sim p`). Shape: [batch_size, ...].
|
||||
step_size (Optional[float]): The step size. Must be None for adaptive step solvers.
|
||||
method (str): A method supported by torchdiffeq. Defaults to "euler". Other commonly used solvers are "dopri5", "midpoint" and "heun3". For a complete list, see torchdiffeq.
|
||||
atol (float): Absolute tolerance, used for adaptive step solvers.
|
||||
rtol (float): Relative tolerance, used for adaptive step solvers.
|
||||
time_grid (Tensor): The process is solved in the interval [min(time_grid, max(time_grid)] and if step_size is None then time discretization is set by the time grid. May specify a descending time_grid to solve in the reverse direction. Defaults to torch.tensor([0.0, 1.0]).
|
||||
return_intermediates (bool, optional): If True then return intermediate time steps according to time_grid. Defaults to False.
|
||||
enable_grad (bool, optional): Whether to compute gradients during sampling. Defaults to False.
|
||||
**model_extras: Additional input for the model.
|
||||
|
||||
Returns:
|
||||
The last timestep when return_intermediates=False, otherwise all values specified in time_grid.
|
||||
"""
|
||||
|
||||
def ode_func(t, x):
|
||||
return self.velocity_model(x=x, t=t, **model_extras)
|
||||
|
||||
ode_opts = {"step_size": step_size} if step_size is not None else {}
|
||||
|
||||
with torch.set_grad_enabled(enable_grad):
|
||||
# Approximate ODE solution with numerical ODE solver
|
||||
sol = odeint(
|
||||
ode_func,
|
||||
x_init,
|
||||
time_grid,
|
||||
method=method,
|
||||
options=ode_opts,
|
||||
atol=atol,
|
||||
rtol=rtol,
|
||||
)
|
||||
|
||||
if return_intermediates:
|
||||
return sol
|
||||
else:
|
||||
return sol[-1]
|
||||
@@ -0,0 +1,58 @@
|
||||
"""
|
||||
RoPE rotary positional encoding, reused by the Denoiser and Transformer.
|
||||
Extracted from model/vae/autoencoder.py.
|
||||
"""
|
||||
|
||||
from typing import Tuple
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class RotaryEmbedding(nn.Module):
|
||||
"""RoPE rotary positional encoding (applied to Q/K)."""
|
||||
|
||||
def __init__(self, dim: int, base: int = 10000):
|
||||
super().__init__()
|
||||
if dim % 2 != 0:
|
||||
raise ValueError("RoPE dim must be even")
|
||||
inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
|
||||
self.register_buffer("inv_freq", inv_freq, persistent=False)
|
||||
|
||||
def get_sin_cos(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
t = torch.arange(seq_len, device=device, dtype=dtype)
|
||||
freqs = torch.einsum("i,j->ij", t, self.inv_freq.to(device=device, dtype=dtype))
|
||||
sin = freqs.sin()
|
||||
cos = freqs.cos()
|
||||
return sin, cos
|
||||
|
||||
@staticmethod
|
||||
def get_sin_cos_from_positions(
|
||||
positions: torch.Tensor,
|
||||
inv_freq: torch.Tensor,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Compute sin/cos from arbitrary position indices (supports repeated indices).
|
||||
|
||||
Args:
|
||||
positions: (L,) int — arbitrary position indices, duplicates allowed
|
||||
inv_freq: (D/2,) — from RotaryEmbedding.inv_freq
|
||||
|
||||
Returns:
|
||||
sin, cos: each (L, D/2)
|
||||
"""
|
||||
freqs = torch.einsum("i,j->ij", positions.float(), inv_freq.to(positions))
|
||||
return freqs.sin(), freqs.cos()
|
||||
|
||||
@staticmethod
|
||||
def apply_rope(x: torch.Tensor, sin: torch.Tensor, cos: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
x: (B, H, L, D)
|
||||
sin/cos: (L, D/2)
|
||||
"""
|
||||
x1 = x[..., 0::2]
|
||||
x2 = x[..., 1::2]
|
||||
sin = sin.unsqueeze(0).unsqueeze(0)
|
||||
cos = cos.unsqueeze(0).unsqueeze(0)
|
||||
x_rot = torch.stack([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
|
||||
return x_rot.flatten(-2)
|
||||
@@ -0,0 +1,592 @@
|
||||
"""
|
||||
HaMeR full preprocessing (LMDB mode) — unified crop LMDB + packed npz
|
||||
|
||||
Reads images from the DexYCB + HOT3D RGB LMDBs, runs HaMeR inference, and saves:
|
||||
- 256x256 crop images -> unified crop LMDB (all_hamer_crop_lmdb/)
|
||||
- landmarks (21,2) normalized [0,1] + confidence -> all_hamer_packed.npz
|
||||
- crop_intrinsics (4,) -> all_crop_params.npz
|
||||
|
||||
Both npz formats match the legacy version, so dataset.py can switch seamlessly.
|
||||
|
||||
Usage (GPU node; training data preprocessing, not required for V1 demo):
|
||||
export HANDFLOW_DATA_ROOT=/path/to/dataset_v2
|
||||
python preprocessing/hamer_crop_and_skeleton.py \
|
||||
--mano_npz $HANDFLOW_DATA_ROOT/all_mano_aa_packed.npz \
|
||||
--lmdb_hot3d $HANDFLOW_DATA_ROOT/hot3d_rgb_lmdb \
|
||||
--lmdb_dexycb $HANDFLOW_DATA_ROOT/dexycb_rgb_lmdb \
|
||||
--intrinsics_hot3d $HANDFLOW_DATA_ROOT/hot3d_intrinsics.npz \
|
||||
--intrinsics_dexycb $HANDFLOW_DATA_ROOT/dexycb_intrinsics.npz \
|
||||
--output_dir $HANDFLOW_DATA_ROOT \
|
||||
--crop_lmdb $HANDFLOW_DATA_ROOT/all_hamer_crop_lmdb \
|
||||
--batch_size 32
|
||||
|
||||
Note: the V1 demo only uses the crop utility functions in this file (imported by utils/online_hamer.py);
|
||||
main() / run_hamer_batch_full() / setup_hamer() are for training data preprocessing and are optional.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
from concurrent.futures import Future, ThreadPoolExecutor
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import lmdb
|
||||
import numpy as np
|
||||
import torch
|
||||
from scipy.spatial.transform import Rotation as _R
|
||||
from tqdm import tqdm
|
||||
|
||||
_PROJECT_ROOT = Path(__file__).parent.parent.resolve()
|
||||
sys.path.insert(0, str(_PROJECT_ROOT))
|
||||
|
||||
from utils.dual_conf import select_target_box_dual_conf
|
||||
|
||||
# Paths are specified by environment variables (see README); HaMeR defaults to third_party/hamer
|
||||
HAMER_ROOT = Path(os.environ.get("HAMER_ROOT", str(_PROJECT_ROOT / "third_party" / "hamer")))
|
||||
HAMER_DATA = Path(os.environ.get("HAMER_DATA", str(HAMER_ROOT / "_DATA")))
|
||||
HAMER_CKPT = Path(os.environ.get("HAMER_CKPT", str(HAMER_DATA / "hamer_ckpts" / "checkpoints" / "hamer.ckpt")))
|
||||
PYRENDER_STUB = HAMER_ROOT / "pyrender_stub"
|
||||
DETECTOR_PT = Path(os.environ.get("DETECTOR_CKPT", str(HAMER_ROOT / "detector.pt")))
|
||||
WILOR_ROOT = DETECTOR_PT.parent
|
||||
|
||||
|
||||
# ── Crop utility functions ─────────────────────────────────────────────────────
|
||||
|
||||
def expand_to_aspect_ratio(input_shape, target_aspect_ratio=None):
|
||||
if target_aspect_ratio is None:
|
||||
return input_shape
|
||||
w, h = input_shape
|
||||
w_t, h_t = target_aspect_ratio
|
||||
if h / w < h_t / w_t:
|
||||
h_new = max(w * h_t / w_t, h)
|
||||
w_new = w
|
||||
else:
|
||||
h_new = h
|
||||
w_new = max(h * w_t / h_t, w)
|
||||
return np.array([w_new, h_new])
|
||||
|
||||
|
||||
def gen_trans_from_patch_cv(c_x, c_y, src_width, src_height,
|
||||
dst_width, dst_height, scale=1.0, rot=0.0):
|
||||
src_w = src_width * scale
|
||||
src_h = src_height * scale
|
||||
rot_rad = np.pi * rot / 180.0
|
||||
sn, cs = np.sin(rot_rad), np.cos(rot_rad)
|
||||
src_center = np.array([c_x, c_y], dtype=np.float32)
|
||||
src_downdir = np.array([0, src_h * 0.5], dtype=np.float32)
|
||||
src_downdir = np.array([src_downdir[0]*cs - src_downdir[1]*sn,
|
||||
src_downdir[0]*sn + src_downdir[1]*cs], dtype=np.float32)
|
||||
src_rightdir = np.array([src_w * 0.5, 0], dtype=np.float32)
|
||||
src_rightdir = np.array([src_rightdir[0]*cs - src_rightdir[1]*sn,
|
||||
src_rightdir[0]*sn + src_rightdir[1]*cs], dtype=np.float32)
|
||||
|
||||
dst_center = np.array([dst_width * 0.5, dst_height * 0.5], dtype=np.float32)
|
||||
dst_downdir = np.array([0, dst_height * 0.5], dtype=np.float32)
|
||||
dst_rightdir = np.array([dst_width * 0.5, 0], dtype=np.float32)
|
||||
|
||||
src = np.zeros((3, 2), dtype=np.float32)
|
||||
src[0, :] = src_center
|
||||
src[1, :] = src_center + src_downdir
|
||||
src[2, :] = src_center + src_rightdir
|
||||
|
||||
dst = np.zeros((3, 2), dtype=np.float32)
|
||||
dst[0, :] = dst_center
|
||||
dst[1, :] = dst_center + dst_downdir
|
||||
dst[2, :] = dst_center + dst_rightdir
|
||||
|
||||
trans = cv2.getAffineTransform(src, dst)
|
||||
return trans
|
||||
|
||||
|
||||
def compute_crop_intrinsics(trans, K_orig_4):
|
||||
"""K_crop = S @ K_orig"""
|
||||
fx, fy, cx, cy = K_orig_4
|
||||
K_orig = np.array([[fx, 0, cx],
|
||||
[0, fy, cy],
|
||||
[0, 0, 1]], dtype=np.float64)
|
||||
S = np.eye(3, dtype=np.float64)
|
||||
S[:2, :] = trans.astype(np.float64)
|
||||
K_crop = S @ K_orig
|
||||
return np.array([K_crop[0, 0], K_crop[1, 1],
|
||||
K_crop[0, 2], K_crop[1, 2]], dtype=np.float32)
|
||||
|
||||
|
||||
# ══════════════════════════════════════════════════════════════════════════════
|
||||
# Model initialization
|
||||
# ══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
def setup_hamer(device: str):
|
||||
sys.path.insert(0, str(HAMER_ROOT))
|
||||
os.environ["HAMER_DATA"] = str(HAMER_DATA)
|
||||
|
||||
import torch as _t
|
||||
_orig = _t.load
|
||||
def _p(*a, **k):
|
||||
k.setdefault("weights_only", False)
|
||||
return _orig(*a, **k)
|
||||
_t.load = _p
|
||||
from ultralytics import YOLO
|
||||
detector = YOLO(str(DETECTOR_PT))
|
||||
_t.load = _orig
|
||||
|
||||
import hamer.configs as _hcfg
|
||||
_hcfg.CACHE_DIR_HAMER = str(HAMER_DATA)
|
||||
|
||||
from pathlib import Path as _P
|
||||
from hamer.configs import get_config
|
||||
from hamer.models.hamer import HAMER
|
||||
model_cfg_path = str(_P(HAMER_CKPT).parent.parent / "model_config.yaml")
|
||||
model_cfg = get_config(model_cfg_path, update_cachedir=True)
|
||||
if (model_cfg.MODEL.BACKBONE.TYPE == "vit") and ("BBOX_SHAPE" not in model_cfg.MODEL):
|
||||
model_cfg.defrost()
|
||||
model_cfg.MODEL.BBOX_SHAPE = [192, 256]
|
||||
model_cfg.freeze()
|
||||
if "PRETRAINED_WEIGHTS" in model_cfg.MODEL.BACKBONE:
|
||||
model_cfg.defrost()
|
||||
model_cfg.MODEL.BACKBONE.pop("PRETRAINED_WEIGHTS")
|
||||
model_cfg.freeze()
|
||||
model = HAMER.load_from_checkpoint(str(HAMER_CKPT), strict=False, cfg=model_cfg,
|
||||
init_renderer=False)
|
||||
|
||||
model = model.to(device).eval()
|
||||
detector.to(device)
|
||||
print(f"HaMeR loaded on {device}")
|
||||
return detector, model, model_cfg
|
||||
|
||||
|
||||
# ══════════════════════════════════════════════════════════════════════════════
|
||||
# Coordinate conversion helpers
|
||||
# ══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
def _cam_crop_to_full(cam_bbox, box_center, box_size, img_size, focal_length):
|
||||
img_w, img_h = img_size[:, 0], img_size[:, 1]
|
||||
cx, cy, b = box_center[:, 0], box_center[:, 1], box_size
|
||||
w_2, h_2 = img_w / 2.0, img_h / 2.0
|
||||
bs = b * cam_bbox[:, 0] + 1e-9
|
||||
tz = 2.0 * focal_length / bs
|
||||
tx = (2.0 * (cx - w_2) / bs) + cam_bbox[:, 1]
|
||||
ty = (2.0 * (cy - h_2) / bs) + cam_bbox[:, 2]
|
||||
return torch.stack([tx, ty, tz], dim=-1)
|
||||
|
||||
|
||||
def _project_full_img(points, cam_trans, focal_length, img_w, img_h):
|
||||
cx, cy = img_w / 2.0, img_h / 2.0
|
||||
K = np.array([[focal_length, 0.0, cx],
|
||||
[0.0, focal_length, cy],
|
||||
[0.0, 0.0, 1.0]], dtype=np.float64)
|
||||
pts = points.astype(np.float64) + cam_trans
|
||||
pts = pts / pts[:, 2:3]
|
||||
return (K @ pts.T).T[:, :2]
|
||||
|
||||
|
||||
# ══════════════════════════════════════════════════════════════════════════════
|
||||
# HaMeR batch inference + crop saving
|
||||
# ══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
_EMPTY = (None, 0.0, None, None, None, None, None)
|
||||
|
||||
|
||||
def run_hamer_batch_full(
|
||||
imgs_bgr: list,
|
||||
detector,
|
||||
model,
|
||||
model_cfg,
|
||||
device: str,
|
||||
det_conf: float = 0.3,
|
||||
batch_size: int = 32,
|
||||
crop_txn=None,
|
||||
crop_lmdb_seq_key: str | None = None,
|
||||
frame_indices: list | None = None,
|
||||
intrinsics_4d_list: list | None = None,
|
||||
) -> list[tuple]:
|
||||
"""Run HaMeR on a batch of BGR images, returning predictions + saving crops"""
|
||||
results_per_frame = [None] * len(imgs_bgr)
|
||||
|
||||
# ── 1. YOLO batch detection ──
|
||||
yolo_out = detector(imgs_bgr, conf=det_conf, verbose=False)
|
||||
|
||||
valid_imgs, valid_bboxes, valid_confs, valid_hw, valid_fi = [], [], [], [], []
|
||||
for fi, (img_bgr, yres) in enumerate(zip(imgs_bgr, yolo_out)):
|
||||
if len(yres.boxes) == 0:
|
||||
results_per_frame[fi] = _EMPTY
|
||||
continue
|
||||
|
||||
clses = yres.boxes.cls.cpu().numpy()
|
||||
confs = yres.boxes.conf.cpu().numpy()
|
||||
bboxes = yres.boxes.xyxy.cpu().numpy()
|
||||
|
||||
# Pick highest-conf right-hand box; but if it heavily overlaps (IoU>thr) a left-hand box
|
||||
# with higher conf, the right box is defeated (actually the left hand) -> no detection.
|
||||
sel = select_target_box_dual_conf(clses, confs, bboxes, target_cls=1)
|
||||
if sel is None:
|
||||
results_per_frame[fi] = _EMPTY
|
||||
continue
|
||||
|
||||
valid_imgs.append(img_bgr)
|
||||
valid_bboxes.append(bboxes[sel:sel+1].astype(np.float32))
|
||||
valid_confs.append(float(confs[sel]))
|
||||
valid_hw.append(img_bgr.shape[:2])
|
||||
valid_fi.append(fi)
|
||||
|
||||
if not valid_fi:
|
||||
for i in range(len(results_per_frame)):
|
||||
if results_per_frame[i] is None:
|
||||
results_per_frame[i] = _EMPTY
|
||||
return results_per_frame
|
||||
|
||||
# ── 1.5 Generate raw crops + affine transform + save to LMDB ──
|
||||
img_size = model_cfg.MODEL.IMAGE_SIZE
|
||||
BBOX_SHAPE = model_cfg.MODEL.get('BBOX_SHAPE', None)
|
||||
rescale_factor = 2.0
|
||||
|
||||
crop_bgr_list = []
|
||||
crop_intr_list = []
|
||||
box_center_list = []
|
||||
bbox_size_list = []
|
||||
|
||||
for vi, (img, bbox) in enumerate(zip(valid_imgs, valid_bboxes)):
|
||||
box = bbox[0].astype(np.float32)
|
||||
center = (box[2:4] + box[0:2]) / 2.0
|
||||
scale = rescale_factor * (box[2:4] - box[0:2]) / 200.0
|
||||
bbox_size = expand_to_aspect_ratio(scale * 200, target_aspect_ratio=BBOX_SHAPE).max()
|
||||
center_x, center_y = center
|
||||
box_center_list.append(center.astype(np.float32))
|
||||
bbox_size_list.append(float(bbox_size))
|
||||
|
||||
trans = gen_trans_from_patch_cv(center_x, center_y, bbox_size, bbox_size,
|
||||
img_size, img_size, 1.0, 0)
|
||||
crop_bgr = cv2.warpAffine(img, trans, (img_size, img_size),
|
||||
flags=cv2.INTER_LINEAR,
|
||||
borderMode=cv2.BORDER_CONSTANT)
|
||||
crop_bgr_list.append(crop_bgr)
|
||||
|
||||
fi_global = valid_fi[vi]
|
||||
if intrinsics_4d_list is not None and fi_global < len(intrinsics_4d_list):
|
||||
crop_intr = compute_crop_intrinsics(trans, intrinsics_4d_list[fi_global])
|
||||
else:
|
||||
crop_intr = np.zeros(4, dtype=np.float32)
|
||||
crop_intr_list.append(crop_intr)
|
||||
|
||||
if crop_txn is not None and crop_lmdb_seq_key is not None and frame_indices is not None:
|
||||
frame_idx = frame_indices[fi_global]
|
||||
_, img_enc = cv2.imencode(".jpg", crop_bgr, [cv2.IMWRITE_JPEG_QUALITY, 95])
|
||||
key = f"img/{crop_lmdb_seq_key}/{frame_idx:06d}".encode("utf-8")
|
||||
crop_txn.put(key, img_enc.tobytes())
|
||||
|
||||
# ── 2. HaMeR ViT batch inference (uses crop_bgr directly, skipping ViTDetDataset) ──
|
||||
MEAN = torch.tensor([123.675, 116.28, 103.53], device=device).view(1, 3, 1, 1)
|
||||
STD = torch.tensor([58.395, 57.12, 57.375], device=device).view(1, 3, 1, 1)
|
||||
|
||||
# (N, 256, 256, 3) BGR uint8 -> (N, 3, 256, 256) RGB normalized
|
||||
crops_np = np.stack(crop_bgr_list)[:, :, :, ::-1].astype(np.float32)
|
||||
crops_t = torch.from_numpy(crops_np).permute(0, 3, 1, 2).to(device)
|
||||
crops_t = (crops_t - MEAN) / STD
|
||||
|
||||
# metadata for camera unprojection
|
||||
box_center_t = torch.tensor(np.stack(box_center_list), dtype=torch.float32, device=device)
|
||||
bbox_size_t = torch.tensor(bbox_size_list, dtype=torch.float32, device=device)
|
||||
img_size_t = torch.tensor(
|
||||
[[float(w), float(h)] for h, w in valid_hw],
|
||||
dtype=torch.float32, device=device,
|
||||
)
|
||||
|
||||
scaled_fl = float(
|
||||
model_cfg.EXTRA.FOCAL_LENGTH
|
||||
/ model_cfg.MODEL.IMAGE_SIZE
|
||||
* img_size_t[0].max().item()
|
||||
)
|
||||
|
||||
all_kp_norms, all_global_aa, all_hand_aa, all_betas_l, all_cam_t = [], [], [], [], []
|
||||
|
||||
for s in range(0, len(crops_t), batch_size):
|
||||
e = min(s + batch_size, len(crops_t))
|
||||
with torch.no_grad():
|
||||
out = model({'img': crops_t[s:e]})
|
||||
|
||||
pred_cam = out["pred_cam"].clone()
|
||||
joints_3d = out["pred_keypoints_3d"].detach().cpu().numpy()
|
||||
|
||||
cam_t_full_np = _cam_crop_to_full(
|
||||
pred_cam, box_center_t[s:e], bbox_size_t[s:e],
|
||||
img_size_t[s:e], scaled_fl,
|
||||
).detach().cpu().numpy()
|
||||
|
||||
g_np = out["pred_mano_params"]["global_orient"].detach().cpu().numpy()
|
||||
h_np = out["pred_mano_params"]["hand_pose"].detach().cpu().numpy()
|
||||
b_np = out["pred_mano_params"]["betas"].detach().cpu().numpy()
|
||||
B = g_np.shape[0]
|
||||
|
||||
g_aa = _R.from_matrix(g_np.reshape(-1, 3, 3)).as_rotvec().reshape(B, 3).astype(np.float32)
|
||||
h_aa = _R.from_matrix(h_np.reshape(-1, 3, 3)).as_rotvec().reshape(B, 45).astype(np.float32)
|
||||
|
||||
for i in range(B):
|
||||
pos = s + i
|
||||
H, W = valid_hw[pos]
|
||||
kp2d = _project_full_img(joints_3d[i], cam_t_full_np[i], scaled_fl, W, H)
|
||||
kp_n = np.clip(kp2d / np.array([W, H], dtype=np.float32), 0.0, 1.0).astype(np.float32)
|
||||
all_kp_norms.append(kp_n)
|
||||
all_global_aa.append(g_aa[i])
|
||||
all_hand_aa.append(h_aa[i])
|
||||
all_betas_l.append(b_np[i].astype(np.float32))
|
||||
all_cam_t.append(cam_t_full_np[i].astype(np.float32))
|
||||
|
||||
for vi, (fi, kp_n, conf, gaa, haa, bet, camt) in enumerate(zip(
|
||||
valid_fi, all_kp_norms, valid_confs,
|
||||
all_global_aa, all_hand_aa, all_betas_l, all_cam_t)):
|
||||
results_per_frame[fi] = (kp_n, conf, gaa, haa, bet, camt, crop_intr_list[vi])
|
||||
|
||||
for i in range(len(results_per_frame)):
|
||||
if results_per_frame[i] is None:
|
||||
results_per_frame[i] = _EMPTY
|
||||
return results_per_frame
|
||||
|
||||
|
||||
# ══════════════════════════════════════════════════════════════════════════════
|
||||
# Main function
|
||||
# ══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="HaMeR full preprocessing — unified crop LMDB + packed npz"
|
||||
)
|
||||
parser.add_argument("--mano_npz", type=str, required=True,
|
||||
help="Unified MANO npz (to obtain seq_keys and num_frames)")
|
||||
parser.add_argument("--lmdb_hot3d", type=str, default=None)
|
||||
parser.add_argument("--lmdb_dexycb", type=str, default=None)
|
||||
parser.add_argument("--intrinsics_hot3d", type=str, default=None)
|
||||
parser.add_argument("--intrinsics_dexycb", type=str, default=None)
|
||||
parser.add_argument("--output_dir", type=str,
|
||||
default=os.environ.get("HANDFLOW_DATA_ROOT", "data"))
|
||||
parser.add_argument("--crop_lmdb", type=str, default=None,
|
||||
help="Unified Crop LMDB output path")
|
||||
parser.add_argument("--device", type=str, default="cuda")
|
||||
parser.add_argument("--batch_size", type=int, default=32)
|
||||
parser.add_argument("--det_conf", type=float, default=0.3)
|
||||
parser.add_argument("--seq_start", type=int, default=0)
|
||||
parser.add_argument("--seq_end", type=int, default=999999)
|
||||
parser.add_argument("--prefetch", type=int, default=8,
|
||||
help="Number of LMDB prefetch threads (higher keeps the GPU busier but uses more memory)")
|
||||
args = parser.parse_args()
|
||||
|
||||
out_dir = Path(args.output_dir)
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
device = args.device if torch.cuda.is_available() else "cpu"
|
||||
|
||||
# ── Read sequence list ──
|
||||
print(f"Loading sequence list: {args.mano_npz}")
|
||||
ref = np.load(args.mano_npz, allow_pickle=True)
|
||||
all_seq_keys = [str(k) for k in ref["seq_keys"]]
|
||||
all_num_frames = ref["num_frames"].tolist()
|
||||
|
||||
seq_keys = all_seq_keys[args.seq_start : args.seq_end]
|
||||
num_frames = all_num_frames[args.seq_start : args.seq_end]
|
||||
print(f"Total sequences: {len(all_seq_keys)}, this shard: [{args.seq_start}, {args.seq_end}), {len(seq_keys)} sequences")
|
||||
|
||||
# ── Determine which LMDB each seq belongs to ──
|
||||
def _lkey(sk):
|
||||
return "hot3d" if sk.startswith("clip-") else "dexycb"
|
||||
|
||||
# ── Open LMDBs ──
|
||||
lmdb_envs, txns = {}, {}
|
||||
if args.lmdb_hot3d:
|
||||
print(f"Opening HOT3D LMDB: {args.lmdb_hot3d}")
|
||||
lmdb_envs["hot3d"] = lmdb.open(args.lmdb_hot3d, readonly=True,
|
||||
lock=False, readahead=False, meminit=False)
|
||||
txns["hot3d"] = lmdb_envs["hot3d"].begin(write=False)
|
||||
if args.lmdb_dexycb:
|
||||
print(f"Opening DexYCB LMDB: {args.lmdb_dexycb}")
|
||||
lmdb_envs["dexycb"] = lmdb.open(args.lmdb_dexycb, readonly=True,
|
||||
lock=False, readahead=False, meminit=False)
|
||||
txns["dexycb"] = lmdb_envs["dexycb"].begin(write=False)
|
||||
|
||||
# ── Load intrinsics ──
|
||||
intrinsics_per_seq: dict[str, dict[int, np.ndarray]] = {}
|
||||
|
||||
def _load_intr(npz_path, label):
|
||||
print(f"Loading {label} intrinsics: {npz_path}")
|
||||
d = np.load(npz_path, allow_pickle=True)
|
||||
keys = [str(k) for k in d["seq_keys"]]
|
||||
offsets = d["offsets"]
|
||||
all_intr = d["all_intrinsics"]
|
||||
for i, sk in enumerate(keys):
|
||||
s, e = int(offsets[i]), int(offsets[i + 1])
|
||||
intrinsics_per_seq[sk] = {j - s: all_intr[j] for j in range(s, e)}
|
||||
print(f" Loaded intrinsics for {len(keys)} sequences")
|
||||
|
||||
if args.intrinsics_hot3d:
|
||||
_load_intr(args.intrinsics_hot3d, "HOT3D")
|
||||
if args.intrinsics_dexycb:
|
||||
_load_intr(args.intrinsics_dexycb, "DexYCB")
|
||||
|
||||
# ── Load HaMeR ──
|
||||
print("Loading HaMeR model...")
|
||||
detector, model, model_cfg = setup_hamer(device)
|
||||
|
||||
# ── Open Crop LMDB ──
|
||||
crop_env = None
|
||||
if args.crop_lmdb:
|
||||
total_est = sum(num_frames)
|
||||
map_size = max(total_est * 80 * 1024, 1 << 30) # >=80KB/crop, at least 1GB
|
||||
print(f"Crop LMDB: {args.crop_lmdb} (map_size={map_size / 1024 / 1024:.0f} MB)")
|
||||
Path(args.crop_lmdb).parent.mkdir(parents=True, exist_ok=True)
|
||||
crop_env = lmdb.open(args.crop_lmdb, map_size=map_size)
|
||||
|
||||
# ── Accumulate results (unified packed format) ───────────────────────────────
|
||||
result_seq_keys: list[str] = []
|
||||
result_seq_offsets = [0]
|
||||
result_num_frames: list[int] = []
|
||||
|
||||
all_inds_list: list[np.ndarray] = []
|
||||
all_landmarks_list: list[np.ndarray] = []
|
||||
all_confidence_list: list[np.ndarray] = []
|
||||
all_crop_intrinsics_list: list[np.ndarray] = []
|
||||
|
||||
# ── Prefetch pipeline ──
|
||||
def _read_seq(seq_key, nf, lkey):
|
||||
"""Background thread: read all frames of a sequence from the corresponding LMDB"""
|
||||
t = txns[lkey]
|
||||
imgs, fids, intrs = [], [], []
|
||||
il = intrinsics_per_seq.get(seq_key)
|
||||
for fi in range(nf):
|
||||
key = f"img/{seq_key}/{fi:06d}".encode()
|
||||
val = t.get(key)
|
||||
if val is None:
|
||||
continue
|
||||
arr = np.frombuffer(val, dtype=np.uint8)
|
||||
img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
|
||||
if img is not None:
|
||||
imgs.append(img)
|
||||
fids.append(fi)
|
||||
intrs.append(il.get(fi, np.zeros(4, np.float32)) if il else np.zeros(4, np.float32))
|
||||
return imgs, fids, intrs
|
||||
|
||||
pending = list(zip(seq_keys, num_frames))
|
||||
total_det = total_loaded = 0
|
||||
prefetch_n = args.prefetch
|
||||
|
||||
with ThreadPoolExecutor(max_workers=prefetch_n) as pool:
|
||||
futures: dict[str, Future] = {}
|
||||
|
||||
def _submit_batch(start_idx):
|
||||
for idx in range(start_idx, min(start_idx + prefetch_n, len(pending))):
|
||||
sk, nf = pending[idx]
|
||||
if sk not in futures:
|
||||
futures[sk] = pool.submit(_read_seq, sk, nf, _lkey(sk))
|
||||
|
||||
_submit_batch(0)
|
||||
|
||||
for pi, (seq_key, nf) in enumerate(tqdm(pending, desc="seqs")):
|
||||
_submit_batch(pi + prefetch_n)
|
||||
|
||||
try:
|
||||
imgs_bgr, frame_idx_list, intr_4d_list = futures.pop(seq_key).result()
|
||||
|
||||
if not imgs_bgr:
|
||||
result_seq_keys.append(seq_key)
|
||||
result_num_frames.append(nf)
|
||||
result_seq_offsets.append(result_seq_offsets[-1])
|
||||
continue
|
||||
|
||||
crop_txn = crop_env.begin(write=True) if crop_env else None
|
||||
|
||||
# Chunked inference
|
||||
all_results: list[tuple] = []
|
||||
for i in range(0, len(imgs_bgr), args.batch_size):
|
||||
chunk = imgs_bgr[i:i + args.batch_size]
|
||||
chunk_intr = intr_4d_list[i:i + args.batch_size] if intr_4d_list else None
|
||||
all_results.extend(
|
||||
run_hamer_batch_full(
|
||||
chunk, detector, model, model_cfg, device,
|
||||
args.det_conf, args.batch_size,
|
||||
crop_txn=crop_txn,
|
||||
crop_lmdb_seq_key=seq_key,
|
||||
frame_indices=frame_idx_list[i:i + args.batch_size],
|
||||
intrinsics_4d_list=chunk_intr,
|
||||
)
|
||||
)
|
||||
|
||||
if crop_txn:
|
||||
crop_txn.commit()
|
||||
|
||||
# Aggregate this sequence
|
||||
inds_out, lm_out, conf_out, cintr_out = [], [], [], []
|
||||
for fi, (kp_n, conf, gaa, haa, bet, camt, cintr) in zip(frame_idx_list, all_results):
|
||||
if kp_n is not None:
|
||||
inds_out.append(fi)
|
||||
lm_out.append(kp_n)
|
||||
conf_out.append(conf)
|
||||
cintr_out.append(cintr)
|
||||
|
||||
n_det = len(inds_out)
|
||||
total_det += n_det
|
||||
total_loaded += len(frame_idx_list)
|
||||
|
||||
# Append to unified results
|
||||
result_seq_keys.append(seq_key)
|
||||
result_num_frames.append(nf)
|
||||
result_seq_offsets.append(result_seq_offsets[-1] + n_det)
|
||||
|
||||
if inds_out:
|
||||
all_inds_list.append(np.array(inds_out, np.int32))
|
||||
all_landmarks_list.append(np.stack(lm_out).astype(np.float32))
|
||||
all_confidence_list.append(np.array(conf_out, np.float32))
|
||||
all_crop_intrinsics_list.append(np.stack(cintr_out).astype(np.float32))
|
||||
|
||||
if n_det == 0:
|
||||
tqdm.write(f" [WARN] no detection: {seq_key}")
|
||||
|
||||
except Exception as e:
|
||||
tqdm.write(f" [ERROR] {seq_key}: {e}")
|
||||
import traceback; traceback.print_exc()
|
||||
result_seq_keys.append(seq_key)
|
||||
result_num_frames.append(nf)
|
||||
result_seq_offsets.append(result_seq_offsets[-1])
|
||||
|
||||
# Close LMDBs
|
||||
for env in lmdb_envs.values():
|
||||
env.close()
|
||||
if crop_env is not None:
|
||||
crop_env.close()
|
||||
|
||||
# ── Write unified packed npz ──
|
||||
print("\nWriting unified packed npz...")
|
||||
|
||||
all_inds = np.concatenate(all_inds_list) if all_inds_list else np.zeros(0, np.int32)
|
||||
all_landmarks = np.concatenate(all_landmarks_list) if all_landmarks_list else np.zeros((0, 21, 2), np.float32)
|
||||
all_confidence = np.concatenate(all_confidence_list) if all_confidence_list else np.zeros(0, np.float32)
|
||||
all_crop_intrinsics = np.concatenate(all_crop_intrinsics_list) if all_crop_intrinsics_list else np.zeros((0, 4), np.float32)
|
||||
|
||||
hamer_path = out_dir / "all_hamer_packed.npz"
|
||||
np.savez_compressed(
|
||||
str(hamer_path),
|
||||
seq_keys=np.array(result_seq_keys, dtype=object),
|
||||
seq_offsets=np.array(result_seq_offsets, dtype=np.int64),
|
||||
all_inds=all_inds,
|
||||
all_landmarks=all_landmarks,
|
||||
all_confidence=all_confidence,
|
||||
)
|
||||
|
||||
crop_params_path = out_dir / "all_crop_params.npz"
|
||||
np.savez_compressed(
|
||||
str(crop_params_path),
|
||||
seq_keys=np.array(result_seq_keys, dtype=object),
|
||||
seq_offsets=np.array(result_seq_offsets, dtype=np.int64),
|
||||
all_inds=all_inds,
|
||||
all_crop_intrinsics=all_crop_intrinsics,
|
||||
)
|
||||
|
||||
hamer_mb = hamer_path.stat().st_size / 1024 / 1024
|
||||
crop_mb = crop_params_path.stat().st_size / 1024 / 1024
|
||||
print(f"\n[saved] {hamer_path} ({hamer_mb:.1f} MB)")
|
||||
print(f"[saved] {crop_params_path} ({crop_mb:.1f} MB)")
|
||||
|
||||
rate = 100.0 * total_det / max(total_loaded, 1)
|
||||
print(f"Detection rate: {total_det}/{total_loaded} ({rate:.1f}%)")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,31 @@
|
||||
# HandFlow V1 — inference + render deps (no training deps)
|
||||
# One-line install: pip install -r requirements.txt
|
||||
# Or use setup_env.sh to create a conda environment
|
||||
|
||||
--extra-index-url https://download.pytorch.org/whl/cu128
|
||||
|
||||
# -- PyTorch (CUDA 12.8; matches the ViPE environment, see setup_vipe_env.sh) --
|
||||
torch==2.7.0+cu128
|
||||
torchvision==0.22.0+cu128
|
||||
|
||||
# -- Inference core --
|
||||
numpy
|
||||
scipy
|
||||
omegaconf
|
||||
tqdm
|
||||
|
||||
# -- Hand detection + MANO FK --
|
||||
ultralytics # YOLO detection (detector.pt weights from WiLoR, see README)
|
||||
manopth # MANO forward kinematics
|
||||
|
||||
# -- Rendering (demo.py -> mp4, pytorch3d Phong shading) --
|
||||
# pytorch3d must match the torch/CUDA above; no prebuilt wheel for torch 2.7+cu128, so build from source:
|
||||
# pip install "git+https://github.com/facebookresearch/pytorch3d.git@v0.7.8"
|
||||
pytorch3d
|
||||
opencv-python
|
||||
imageio[ffmpeg]
|
||||
matplotlib
|
||||
|
||||
# The following two are installed separately as submodules:
|
||||
# HaMeR backbone: pip install -e third_party/hamer
|
||||
# ViPE SLAM (optional): see setup_vipe_env.sh (standalone environment)
|
||||
@@ -0,0 +1,10 @@
|
||||
# Install into a compatible environment after selecting its Torch/CUDA runtime.
|
||||
# Observed versions, not a complete cross-platform dependency lock.
|
||||
dex-retargeting==0.5.0
|
||||
mujoco==3.5.0
|
||||
scipy==1.17.1
|
||||
numpy==2.4.6
|
||||
h5py==3.16.0
|
||||
imageio
|
||||
imageio-ffmpeg
|
||||
Pillow
|
||||
@@ -0,0 +1,12 @@
|
||||
{
|
||||
"dex-retargeting": "0.5.0",
|
||||
"mujoco": "3.5.0",
|
||||
"numpy": "2.4.6",
|
||||
"scipy": "1.17.1",
|
||||
"torch": "2.7.1+cu128",
|
||||
"torchvision": "0.22.1+cu128",
|
||||
"nlopt": "2.11.0",
|
||||
"pin": "4.1.0",
|
||||
"imageio": "2.37.3",
|
||||
"Pillow": "12.3.0"
|
||||
}
|
||||
@@ -0,0 +1,37 @@
|
||||
"""Measure jitter by pipeline stage and verify objective/gradient consistency."""
|
||||
import json
|
||||
from pathlib import Path
|
||||
import numpy as np
|
||||
from scipy.spatial.transform import Rotation
|
||||
from l20_calibrated import ROOT,Calibration,CalibratedAdaptor,palm_basis
|
||||
from l20_consistent_optimizer import ConsistentPositionOptimizer
|
||||
from dex_retargeting.retargeting_config import RetargetingConfig
|
||||
|
||||
base=ROOT/'output/l20_2047635068';out=base/'jitter_audit'
|
||||
old=np.load(base/'motion.npz');new=np.load(out/'corrected/motion.npz');contact=np.load(base/'bottle_grasp/registered_motion.npz');human=np.load(ROOT/'output/dex_2047635068/human_joints.npz')
|
||||
def stats(a,scale=1):
|
||||
a=a*scale;v=np.diff(a,axis=0);acc=np.diff(a,n=2,axis=0)
|
||||
return dict(max_component_step=float(np.abs(v).max()),rms_component_step=float(np.sqrt(np.mean(v*v))),rms_second_difference=float(np.sqrt(np.mean(acc*acc))),p99_component_step=float(np.percentile(np.abs(v),99)))
|
||||
r={'units':'Position in mm; joint and rotation angles in degrees; differences are per frame at 30 fps', 'stages':{}}
|
||||
for name,a,scale in [('human_root_local_joints',human['joints'],1000),('scaled_robot_targets',old['targets'],1000),('old_robot_fk',old['actual'],1000),('corrected_robot_fk',new['actual'],1000),('old_raw_joint',old['qpos_raw'],180/np.pi),('old_filtered_joint',old['qpos'],180/np.pi),('corrected_raw_joint',new['qpos_raw'],180/np.pi),('corrected_filtered_joint',new['qpos'],180/np.pi),('old_contact_joint',contact['qpos'],180/np.pi),('source_wrist',old['wrist_pos'],1000),('contact_wrist',contact['wrist_pos'],1000)]:r['stages'][name]=stats(a,scale)
|
||||
r['target_arrays_identical']=bool(np.array_equal(old['targets'],new['targets']))
|
||||
r['wrist_arrays_identical']=bool(np.array_equal(old['wrist_pos'],new['wrist_pos']) and np.array_equal(old['wrist_quat_wxyz'],new['wrist_quat_wxyz']))
|
||||
r['tip_errors_mm']={key:float(np.linalg.norm(m['actual'][:,[4,8,12,16,20]]-m['targets'][:,[4,8,12,16,20]],axis=-1).mean()*1000) for key,m in [('old',old),('corrected',new)]}
|
||||
steps=np.max(np.abs(np.diff(contact['qpos'],axis=0)),axis=1);ix=np.argsort(steps)[-10:][::-1]
|
||||
r['largest_contact_jumps']=[dict(from_frame=int(t),to_frame=int(t+1),time_s=float((t+1)/30),joint=str(contact['joint_names'][np.argmax(np.abs(contact['qpos'][t+1]-contact['qpos'][t]))]),joint_step_deg=float(steps[t]*180/np.pi),source_target_max_step_mm=float(np.linalg.norm(np.diff(old['targets'][t:t+2],axis=0)[0],axis=-1).max()*1000),source_joint_max_step_deg=float(np.max(np.abs(old['qpos'][t+1]-old['qpos'][t]))*180/np.pi)) for t in ix]
|
||||
basis=Rotation.from_matrix(np.array([palm_basis(p) for p in human['joints']]))
|
||||
root=Rotation.from_rotvec(human['root_orient'])
|
||||
r['palm_basis_max_step_deg']=float(np.rad2deg((basis[:-1].inv()*basis[1:]).magnitude().max()))
|
||||
r['source_root_max_step_deg']=float(np.rad2deg((root[:-1].inv()*root[1:]).magnitude().max()))
|
||||
# Check gradient away from the previous iterate, where the omitted penalty is observable.
|
||||
cfg=json.loads((out/'corrected/dex_config.json').read_text());rt=RetargetingConfig.from_dict(cfg).build();cal=Calibration();ad=CalibratedAdaptor(rt.optimizer.robot,cal);rt.optimizer.set_kinematic_adaptor(ad)
|
||||
names=list(new['joint_names']);last=np.array([new['qpos'][666,names.index(n)] for n in cal.active]);bounds=np.array([cal.bounds[n] for n in cal.active]);x=np.clip(last+.03,bounds[:,0]+.001,bounds[:,1]-.001)
|
||||
def check(op):
|
||||
fun=op.get_objective_function(new['targets'][666],np.array([]),last);g=np.zeros(16);fun(x,g);fd=[];eps=1e-5
|
||||
for j in range(16):
|
||||
d=np.eye(16)[j]*eps;fd.append((fun(x+d,np.array([]))-fun(x-d,np.array([])))/(2*eps))
|
||||
return float(np.max(np.abs(g-fd)))
|
||||
r['objective_gradient_max_error_before']=check(rt.optimizer);rt.optimizer.__class__=ConsistentPositionOptimizer;r['objective_gradient_max_error_after']=check(rt.optimizer)
|
||||
r['corrected_joint_second_difference_reduction_pct']=100*(1-r['stages']['corrected_filtered_joint']['rms_second_difference']/r['stages']['old_filtered_joint']['rms_second_difference'])
|
||||
r['holding_phase_comparison']={key:dict(holding_max_joint_step_deg=float(np.abs(np.diff(m['qpos'][160:],axis=0)).max()*180/np.pi),holding_second_diff_rms_deg=float(np.sqrt(np.mean(np.diff(m['qpos'][160:],n=2,axis=0)**2))*180/np.pi)) for key,m in [('original',old),('corrected',new),('contact',contact)]}
|
||||
(out/'audit.json').write_text(json.dumps(r,indent=2));print(json.dumps(r,indent=2))
|
||||
@@ -0,0 +1,95 @@
|
||||
"""Build full video kinematic contact playback with linear URDF coupling."""
|
||||
import os
|
||||
os.environ.setdefault('MUJOCO_GL', 'osmesa')
|
||||
import json, shutil, xml.etree.ElementTree as E
|
||||
from pathlib import Path
|
||||
import numpy as np
|
||||
import mujoco
|
||||
from scipy.spatial.transform import Rotation
|
||||
from scipy.ndimage import gaussian_filter1d
|
||||
from scipy.optimize import least_squares
|
||||
ROOT=Path(__file__).resolve().parents[1]
|
||||
OUT=ROOT/'output/l20_full_replay';OUT.mkdir(exist_ok=True)
|
||||
TASK=ROOT/'output/spider_l20_contact/datasets/processed/current/l20/right/bottle_hold'
|
||||
tree=E.parse(TASK/'scene_act.xml');root=tree.getroot();world=root.find('worldbody');asset=root.find('asset')
|
||||
# A playback scene has no object assistance or hidden controller.
|
||||
root.remove(root.find('actuator'))
|
||||
for item in list(world.findall('light')):world.remove(item)
|
||||
for item in list(asset.findall("texture[@type='skybox']")):asset.remove(item)
|
||||
E.SubElement(asset,'texture',name='daylight',type='skybox',builtin='gradient',rgb1='.55 .67 .8',rgb2='.88 .91 .94',width='512',height='3072')
|
||||
E.SubElement(asset,'texture',name='floor_grid',type='2d',builtin='checker',rgb1='.54 .59 .62',rgb2='.64 .69 .72',width='512',height='512')
|
||||
E.SubElement(asset,'material',name='floor_mat',texture='floor_grid',texrepeat='8 8',texuniform='true',reflectance='.08',rgba='1 1 1 1')
|
||||
world.find("geom[@name='floor']").attrib.update(material='floor_mat',rgba='1 1 1 1',size='2 2 .05',pos='0 0 0')
|
||||
E.SubElement(world,'light',name='key',pos='-.5 -.5 2',dir='.3 .3 -1',directional='true',diffuse='.45 .42 .38',ambient='.04 .05 .06',castshadow='false')
|
||||
E.SubElement(world,'light',name='fill',pos='1 1 1.5',dir='-.5 -.5 -1',directional='true',diffuse='.15 .18 .21',castshadow='false')
|
||||
root.find('visual/headlight').attrib.update(ambient='.28 .28 .28',diffuse='.25 .25 .25',specular='.15 .15 .15')
|
||||
root.find('visual/rgba').set('haze','.8 .86 .91 1')
|
||||
root.find('visual/quality').attrib.update(shadowsize='2048',offsamples='2')
|
||||
E.SubElement(root.find('visual'),'map',shadowclip='1',shadowscale='.8',zfar='10',znear='.01')
|
||||
(OUT/'assets').mkdir(exist_ok=True)
|
||||
for item in asset:
|
||||
if item.get('file'):
|
||||
src=Path(item.get('file'));dest=OUT/'assets'/(item.get('name')+src.suffix);shutil.copy2(src,dest);item.set('file',str(dest.relative_to(OUT)))
|
||||
scene=OUT/'scene.xml';tree.write(scene)
|
||||
model=mujoco.MjModel.from_xml_path(str(scene));data=mujoco.MjData(model)
|
||||
addr={model.joint(j).name:int(model.jnt_qposadr[j]) for j in range(model.njnt)}
|
||||
wn=['right_hand_'+s+'_'+a for s in ['pos','rot'] for a in 'xyz'];on=['right_object_'+s+'_'+a for s in ['pos','rot'] for a in 'xyz']
|
||||
wa=[addr[n] for n in wn];oa=[addr[n] for n in on]
|
||||
urdf=E.parse(ROOT/'third_party/l20_assets/L20/RIGHT/linkerhand_g20_right.urdf').getroot();js=[j for j in urdf.findall('joint') if j.get('type')!='fixed']
|
||||
active=[j.get('name') for j in js if j.find('mimic') is None];mimic={}
|
||||
limits={j.get('name'):np.array([float(j.find('limit').get(k)) for k in ['lower','upper']]) for j in js}
|
||||
for j in js:
|
||||
mi=j.find('mimic')
|
||||
if mi is not None:
|
||||
n=j.get('name');parent=mi.get('joint');s=float(mi.get('multiplier','1'));off=float(mi.get('offset','0'));mimic[n]=(parent,s,off)
|
||||
lo,hi=np.sort((limits[n]-off)/s);limits[parent]=np.array([max(limits[parent][0],lo),min(limits[parent][1],hi)])
|
||||
low=np.array([limits[n][0] for n in active]);high=np.array([limits[n][1] for n in active]);aa=[addr[n] for n in active]
|
||||
src=np.load(ROOT/'output/l20_2047635068/jitter_audit/temporal/motion.npz');N=len(src['qpos']);sn=list(src['joint_names']);base=np.clip(np.stack([src['qpos'][:,sn.index(n)] for n in active],1),low,high)
|
||||
wr0=Rotation.from_quat(src['wrist_quat_wxyz'][:,[1,2,3,0]]);wp0=src['wrist_pos'];rel=Rotation.from_euler('x',-np.pi/2);orr0=wr0*rel;op0=wp0+wr0.apply(np.tile([.12,-.065,.15],(N,1)))
|
||||
align=orr0[600].inv();shift=np.array([0,0,.3])-align.apply(op0[600]);wp=align.apply(wp0)+shift;wr=align*wr0;op=align.apply(op0)+shift;orr=align*orr0
|
||||
ref=np.load(TASK/'0/trajectory_kinematic_act.npz')['qpos'][0]
|
||||
offset=ref[aa]-base[600];dp=ref[wa[:3]]-wp[600];dr=(wr[600].inv()*Rotation.from_euler('XYZ',ref[wa[3:]])).as_rotvec()
|
||||
onset=160;weight=np.clip((np.arange(N)-130)/30,0,1);weight=weight*weight*(3-2*weight)
|
||||
# Estimated acquisition bracket, not a measured bottle track.
|
||||
op[:onset]=op[onset];oq=orr.as_quat();oq[:onset]=oq[onset];orr=Rotation.from_quat(oq)
|
||||
handgeoms=np.flatnonzero(model.geom_contype==1);objgeoms=np.flatnonzero(model.geom_contype==2)
|
||||
fingers=['thumb','index','middle','ring','pinky'];tips=[next(i for i,g in enumerate(handgeoms) if model.geom_bodyid[g]==model.body(f+'_distal').id) for f in fingers]
|
||||
def setpose(t,x):
|
||||
vals=dict(zip(active,np.clip(base[t]+x[:16],low,high)))
|
||||
vals.update({n:s*vals[p]+o for n,(p,s,o) in mimic.items()})
|
||||
for n,v in vals.items():data.qpos[addr[n]]=v
|
||||
data.qpos[wa]=np.r_[wp[t]+x[16:19],(wr[t]*Rotation.from_rotvec(x[19:])).as_euler('XYZ')]
|
||||
data.qpos[oa]=np.r_[op[t],orr[t].as_euler('XYZ')]
|
||||
mujoco.mj_kinematics(model,data)
|
||||
def distances():return np.array([min(mujoco.mj_geomDistance(model,data,int(g),int(o),.3,None) for o in objgeoms) for g in handgeoms])
|
||||
prev=np.r_[offset,dp,dr];fits=[];frames=np.unique(np.r_[np.arange(onset,N,2),N-1])
|
||||
for k,t in enumerate(frames):
|
||||
previous=prev.copy()
|
||||
def fun(x):
|
||||
setpose(t,x);d=distances()
|
||||
return np.r_[(d[tips]+.0002)/.004,np.maximum(-d-.0007,0)/.001,.18*x[:16],.8*(x[:16]-previous[:16]),.2*x[16:19]/.03,.15*x[19:]/.3,.5*(x[16:19]-previous[16:19])/.01,.5*(x[19:]-previous[19:])/.1]
|
||||
lb=np.r_[low-base[t],[-.065]*3,[-.55]*3];ub=np.r_[high-base[t],[.065]*3,[.55]*3]
|
||||
fit=least_squares(fun,np.clip(prev,lb+1e-9,ub-1e-9),bounds=(lb,ub),max_nfev=65,diff_step=1e-4)
|
||||
prev=fit.x;fits.append(prev.copy())
|
||||
if k%20==0:print('Contact keyframes',k,len(frames),flush=True)
|
||||
fit=np.array(fits);correction=np.stack([np.interp(np.arange(N),frames,fit[:,j]) for j in range(22)],1)
|
||||
correction=gaussian_filter1d(correction,1.5,axis=0)*weight[:,None]
|
||||
rows=[];penetration=[];gaps=[];minz=1e9
|
||||
for t in range(N):
|
||||
setpose(t,correction[t]);rows.append(data.qpos.copy());d=distances();penetration.append(max(0,-d.min()));gaps.append(d[tips])
|
||||
for g in np.r_[handgeoms,objgeoms]:
|
||||
mesh=model.geom_dataid[g]
|
||||
if model.geom_type[g]==mujoco.mjtGeom.mjGEOM_MESH:
|
||||
vertices=model.mesh_vert[model.mesh_vertadr[mesh]:model.mesh_vertadr[mesh]+model.mesh_vertnum[mesh]]
|
||||
minz=min(minz,float((vertices@data.geom_xmat[g].reshape(3,3).T+data.geom_xpos[g])[:,2].min()))
|
||||
minz=min(minz,float(op[t,2]-.16))
|
||||
q=np.array(rows);zshift=max(0,.025-minz);q[:,wa[2]]+=zshift;q[:,oa[2]]+=zshift
|
||||
q[:,wa[3:]]=np.unwrap(q[:,wa[3:]],axis=0);q[:,oa[3:]]=np.unwrap(q[:,oa[3:]],axis=0)
|
||||
center=np.mean(np.r_[q[:,wa[:3]],q[:,oa[:3]]],axis=0);center[2]+=.04
|
||||
root.find('statistic').attrib.update(center=' '.join(map(str,center)),extent='1')
|
||||
root.find('visual/global').attrib.update(azimuth='135',elevation='-22')
|
||||
E.SubElement(E.SubElement(root,'keyframe'),'key',name='video_start',qpos=' '.join(map(str,q[0])))
|
||||
tree.write(scene)
|
||||
np.savez_compressed(OUT/'motion.npz',qpos=q,fps=30,time=np.arange(N)/30,wrist_pos=q[:,wa[:3]],bottle_pos=q[:,oa[:3]],camera_center=center,correction=correction,contact_gaps_m=gaps,penetration_m=penetration)
|
||||
report=dict(frames=N,fps=30,duration_seconds=N/30,mode='kinematic prescribed-pose playback; not a full SPIDER dynamic rollout',bottle_pose='estimated registration; acquisition bracket 130-160, not video object tracking',nonlinear_coupling=False,linear_mimic=True,wrist_range_m=np.ptp(q[:,wa[:3]],axis=0).tolist(),finite=bool(np.isfinite(q).all()),max_finger_step_deg=float(np.rad2deg(np.abs(np.diff(q[:,aa],axis=0))).max()),mean_holding_fingertip_gap_mm=float(np.abs(gaps[onset:]).mean()*1000),max_holding_penetration_mm=float(np.max(penetration[onset:])*1000),max_all_penetration_mm=float(np.max(penetration)*1000),bottle_dimensions_m=[.05411,.05411,.157856],floor_z=0,ambient_lighting=True)
|
||||
(OUT/'validation.json').write_text(json.dumps(report,indent=2));print(json.dumps(report,indent=2),flush=True)
|
||||
@@ -0,0 +1,49 @@
|
||||
"""Size-constrained idealized bottle; no use of the corrupted radial profile."""
|
||||
import os
|
||||
os.environ.setdefault('MUJOCO_GL','osmesa')
|
||||
from pathlib import Path
|
||||
import json,xml.etree.ElementTree as E
|
||||
import numpy as np
|
||||
import trimesh,mujoco
|
||||
from PIL import Image
|
||||
|
||||
OUT=Path(__file__).resolve().parents[1]/'third_party/bottle_model_parametric'
|
||||
OUT.mkdir(parents=True,exist_ok=True);(OUT/'assets').mkdir(exist_ok=True)
|
||||
H=.157856;R=.05411/2;neck=.014
|
||||
# Only H and 2R are inherited dimensions. Shoulder/neck proportions are design choices.
|
||||
profile=np.array([[0,R],[.103,R],[.139,neck],[H,neck]])
|
||||
def lathe(points):
|
||||
n=128;a=np.arange(n)*2*np.pi/n
|
||||
vertices=np.concatenate([np.c_[r*np.cos(a),r*np.sin(a),np.full(n,z)] for z,r in points])
|
||||
faces=[]
|
||||
for k in range(len(points)-1):
|
||||
for i in range(n):
|
||||
j=(i+1)%n;faces.extend([[k*n+i,k*n+j,(k+1)*n+j],[k*n+i,(k+1)*n+j,(k+1)*n+i]])
|
||||
bottom=len(vertices);top=bottom+1
|
||||
vertices=np.r_[vertices,[[0,0,points[0,0]],[0,0,points[-1,0]]]]
|
||||
for i in range(n):
|
||||
j=(i+1)%n;faces.extend([[bottom,j,i],[top,(len(points)-1)*n+i,(len(points)-1)*n+j]])
|
||||
return trimesh.Trimesh(vertices=vertices,faces=faces,process=True)
|
||||
mesh=lathe(profile);mesh.export(OUT/'assets/bottle.obj');mesh.export(OUT/'bottle.stl');mesh.export(OUT/'bottle.glb')
|
||||
shoulder=lathe(profile[1:3]);shoulder.export(OUT/'assets/shoulder.obj')
|
||||
root=E.Element('mujoco',model='dimension_based_bottle');E.SubElement(root,'option',gravity='0 0 -9.81')
|
||||
asset=E.SubElement(root,'asset')
|
||||
E.SubElement(asset,'mesh',name='bottle_visual',file='assets/bottle.obj')
|
||||
E.SubElement(asset,'mesh',name='bottle_shoulder',file='assets/shoulder.obj')
|
||||
world=E.SubElement(root,'worldbody');body=E.SubElement(world,'body',name='bottle')
|
||||
E.SubElement(body,'freejoint',name='bottle_freejoint')
|
||||
mesh.density=.2308/mesh.volume
|
||||
E.SubElement(body,'inertial',mass='.2308',pos=' '.join(map(str,mesh.center_mass)),diaginertia=' '.join(map(str,np.diag(mesh.moment_inertia))))
|
||||
E.SubElement(body,'geom',type='mesh',mesh='bottle_visual',group='1',contype='0',conaffinity='0',rgba='.68 .74 .8 1')
|
||||
common=dict(group='3',friction='.6 .005 .001',condim='6')
|
||||
E.SubElement(body,'geom',type='cylinder',size=f'{R} {.103/2}',pos=f'0 0 {.103/2}',**common)
|
||||
E.SubElement(body,'geom',type='mesh',mesh='bottle_shoulder',**common)
|
||||
E.SubElement(body,'geom',type='cylinder',size=f'{neck} {(H-.139)/2}',pos=f'0 0 {(H+.139)/2}',**common)
|
||||
E.indent(root);E.ElementTree(root).write(OUT/'bottle.xml')
|
||||
np.savez(OUT/'profile.npz',h=profile[:,0],r=profile[:,1],height=H,diameter=2*R)
|
||||
model=mujoco.MjModel.from_xml_path(str(OUT/'bottle.xml'));data=mujoco.MjData(model);mujoco.mj_forward(model,data)
|
||||
renderer=mujoco.Renderer(model,height=480,width=640);cam=mujoco.MjvCamera();cam.lookat[:]=[0,0,H/2];cam.distance=.36;cam.azimuth=60;cam.elevation=-8
|
||||
opt=mujoco.MjvOption();opt.geomgroup[3]=0;model.vis.headlight.ambient[:]=.7
|
||||
renderer.update_scene(data,camera=cam,scene_option=opt);Image.fromarray(renderer.render()).save(OUT/'preview.png');renderer.close()
|
||||
report=dict(height_mm=H*1000,max_diameter_mm=2*R*1000,neck_diameter_mm=neck*2000,watertight=bool(mesh.is_watertight),winding_consistent=bool(mesh.is_winding_consistent),volume_m3=mesh.volume,mujoco_load_passed=True,collision_parts=3,extent_m=mesh.extents.tolist(),shape_source='Idealized cylinder body, tapered shoulder, cylinder neck; only height and maximum diameter inherited',mass_source='Previous XML mass 0.2308 kg retained; uniform solid inertia recomputed, not measured')
|
||||
(OUT/'validation.json').write_text(json.dumps(report,indent=2));print(json.dumps(report,indent=2))
|
||||
+359
@@ -0,0 +1,359 @@
|
||||
#!/usr/bin/env python3
|
||||
"""demo.py — HandFlow unified demo: cam overlay + orthographic view in one run.
|
||||
|
||||
Two camera modes (selected by --fix_camera):
|
||||
--fix_camera Fixed camera (e.g. DexYCB on a tripod). ViPE is skipped, c2w is
|
||||
identity, and the orthographic view assumes a level, fixed camera:
|
||||
a side orthographic view of the camera-space hand motion (the
|
||||
camera marker stays at the origin; only the hand moves).
|
||||
(default) Moving camera (e.g. HOT3D head-mounted). ViPE SLAM estimates c2w
|
||||
and the orthographic view is a world-space trajectory view: camera
|
||||
trajectory line + per-frame hand mesh.
|
||||
|
||||
Both modes write two mp4s into --output_dir:
|
||||
overlay.mp4 Clean camera-space MANO mesh overlaid onto the RGB frames.
|
||||
ortho.mp4 Orthographic (side / topdown) view of the hand (+ camera motion).
|
||||
|
||||
Intrinsics (--intrinsics):
|
||||
fx,fy,cx,cy Explicit pinhole intrinsics (recommended; e.g. the demo GT values).
|
||||
auto Let ViPE estimate intrinsics (only valid without --fix_camera).
|
||||
(unset) Generic default [600, 600, W/2, H/2].
|
||||
|
||||
Usage:
|
||||
# Fixed camera (no ViPE needed)
|
||||
python scripts/demo.py --input video.mp4 --fm_ckpt <ckpt> \
|
||||
--intrinsics fx,fy,cx,cy --fix_camera --output_dir output/demo
|
||||
|
||||
# Moving camera (ViPE estimates c2w; requires the vipe env)
|
||||
python scripts/demo.py --input video.mp4 --fm_ckpt <ckpt> \
|
||||
--intrinsics fx,fy,cx,cy --output_dir output/demo
|
||||
|
||||
# Moving camera, unknown intrinsics -> ViPE estimates both intr + c2w
|
||||
python scripts/demo.py --input video.mp4 --fm_ckpt <ckpt> \
|
||||
--intrinsics auto --output_dir output/demo
|
||||
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import gc
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
from omegaconf import OmegaConf
|
||||
from pytorch3d.transforms import axis_angle_to_quaternion, quaternion_to_axis_angle
|
||||
|
||||
_PROJECT_ROOT = Path(__file__).resolve().parent.parent
|
||||
sys.path.insert(0, str(_PROJECT_ROOT))
|
||||
|
||||
from utils.checkpoint_utils import load_denoiser_from_ckpt
|
||||
from utils.inference_utils import run_fm_inference_with_overlap, split_sequence_into_windows
|
||||
from utils.mano_utils import MANOForwardKinematics
|
||||
from utils.online_hamer import OnlineHaMeRPipeline
|
||||
from utils.vipe_worker import run_vipe_slam
|
||||
from visualization.renderer_p3d import PhongRenderer
|
||||
from visualization.video_io import read_video_frames, write_video_ffmpeg
|
||||
|
||||
|
||||
def resolve_intrinsics(arg_str, W: int, H: int):
|
||||
"""Resolve --intrinsics into (intr(4,) | None, vipe_estimate: bool).
|
||||
|
||||
None / "" / "default" -> generic [600, 600, W/2, H/2], vipe_estimate=False
|
||||
"auto" -> (None, True) (ViPE estimates; only without --fix_camera)
|
||||
"fx,fy,cx,cy" -> explicit (4,), vipe_estimate=False
|
||||
"""
|
||||
if arg_str in (None, "", "default"):
|
||||
return np.array([600.0, 600.0, W / 2.0, H / 2.0], dtype=np.float32), False
|
||||
if arg_str == "auto":
|
||||
return None, True
|
||||
parts = [float(x) for x in arg_str.split(",")]
|
||||
if len(parts) != 4:
|
||||
raise ValueError(f"--intrinsics expects 4 values fx,fy,cx,cy or 'auto', got {arg_str!r}")
|
||||
return np.array(parts, dtype=np.float32), False
|
||||
|
||||
|
||||
def extract_frames_to_dir(frames_bgr, img_dir: str) -> None:
|
||||
"""Write frames as jpg into a directory (input for ViPE SLAM)."""
|
||||
os.makedirs(img_dir, exist_ok=True)
|
||||
for i, f in enumerate(frames_bgr):
|
||||
cv2.imwrite(str(Path(img_dir) / f"{i:06d}.jpg"), f)
|
||||
|
||||
|
||||
def transform_verts_to_world(verts_cam_m: np.ndarray, c2w: np.ndarray) -> np.ndarray:
|
||||
"""verts_cam (T,N,3) meters + c2w (T,4,4) -> verts_world (T,N,3) meters."""
|
||||
T, N, _ = verts_cam_m.shape
|
||||
homog = np.concatenate([verts_cam_m, np.ones((T, N, 1), dtype=np.float32)], axis=-1)
|
||||
return np.einsum("tij,tnj->tni", c2w.astype(np.float32), homog)[..., :3].astype(np.float32)
|
||||
|
||||
|
||||
def smooth_pose_rotations(pose_seq: torch.Tensor, window: int) -> torch.Tensor:
|
||||
"""Temporally smooth MANO axis-angle joints by averaging in quaternion space."""
|
||||
if window <= 1:
|
||||
return pose_seq
|
||||
if window % 2 == 0:
|
||||
raise ValueError("--rotation_smooth_window must be an odd positive integer")
|
||||
|
||||
n_frames, pose_dim = pose_seq.shape
|
||||
if pose_dim % 3:
|
||||
raise ValueError(f"Expected axis-angle pose dimension divisible by 3, got {pose_dim}")
|
||||
quats = axis_angle_to_quaternion(pose_seq.reshape(n_frames, -1, 3))
|
||||
radius = window // 2
|
||||
weights = (radius + 1 - torch.arange(-radius, radius + 1, device=pose_seq.device).abs()).to(pose_seq)
|
||||
weights = weights / weights.sum()
|
||||
smoothed = []
|
||||
for frame_idx in range(n_frames):
|
||||
indices = torch.arange(frame_idx - radius, frame_idx + radius + 1, device=pose_seq.device)
|
||||
indices = indices.clamp(0, n_frames - 1)
|
||||
neighbors = quats[indices]
|
||||
reference = quats[frame_idx].unsqueeze(0)
|
||||
signs = torch.where((neighbors * reference).sum(dim=-1, keepdim=True) < 0, -1.0, 1.0)
|
||||
averaged = (neighbors * signs * weights[:, None, None]).sum(dim=0)
|
||||
smoothed.append(averaged / averaged.norm(dim=-1, keepdim=True).clamp_min(1e-8))
|
||||
return quaternion_to_axis_angle(torch.stack(smoothed)).reshape_as(pose_seq)
|
||||
|
||||
|
||||
def get_camera_geometry(frames_bgr, fix_camera: bool, intr, vipe_estimate: bool, T: int):
|
||||
"""Return (intr(4,), c2w(T,4,4), vipe_slam_time_s).
|
||||
|
||||
fix_camera=True: c2w = identity (no ViPE); intr is the provided/generic value.
|
||||
fix_camera=False: run ViPE — gt_intr (c2w only) when intr is known, else default
|
||||
(ViPE estimates intr + c2w).
|
||||
"""
|
||||
if fix_camera:
|
||||
if vipe_estimate:
|
||||
print("[demo] --fix_camera ignores --intrinsics auto; using generic 600,600,W/2,H/2")
|
||||
intr = np.array([600.0, 600.0, frames_bgr[0].shape[1] / 2.0,
|
||||
frames_bgr[0].shape[0] / 2.0], dtype=np.float32)
|
||||
c2w = np.broadcast_to(np.eye(4, dtype=np.float32), (T, 4, 4)).copy()
|
||||
return intr.astype(np.float32), c2w, 0.0
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
img_dir = str(Path(tmp) / "images")
|
||||
extract_frames_to_dir(frames_bgr, img_dir)
|
||||
if vipe_estimate:
|
||||
print("[demo] ViPE SLAM (default, estimating intrinsics + c2w) ...")
|
||||
slam = run_vipe_slam(img_dir, tmp, variant="default")
|
||||
intr_arr = slam["intrinsics"]
|
||||
intr = intr_arr[0] if intr_arr.ndim == 2 else intr_arr
|
||||
else:
|
||||
print("[demo] ViPE SLAM (gt_intr, estimating c2w only) ...")
|
||||
slam = run_vipe_slam(img_dir, tmp, variant="gt_intr", gt_intrinsics_4d=intr)
|
||||
c2w = slam["poses"]
|
||||
print(f" c2w shape={c2w.shape}, intrinsics={intr}, slam_time={slam['slam_time_s']:.2f}s")
|
||||
return intr.astype(np.float32), c2w.astype(np.float32), float(slam["slam_time_s"])
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser(description="HandFlow demo: cam overlay + orthographic view")
|
||||
ap.add_argument("--input", required=True, help="Input video path")
|
||||
ap.add_argument("--fm_ckpt", required=True, help="HandFlow checkpoint (.pt or DeepSpeed directory)")
|
||||
ap.add_argument("--config", default="configs/inference.yaml")
|
||||
ap.add_argument("--intrinsics", default=None,
|
||||
help="fx,fy,cx,cy | auto | (unset -> default 600,600,W/2,H/2)")
|
||||
ap.add_argument("--fix_camera", action="store_true",
|
||||
help="Fixed camera: skip ViPE (c2w=identity), ortho view assumes a level fixed camera")
|
||||
ap.add_argument("--camera_npz", default=None,
|
||||
help="Reuse intrinsics and c2w from a previous result npz instead of rerunning ViPE")
|
||||
ap.add_argument("--rotation_smooth_window", type=int, default=0,
|
||||
help="Odd temporal window for quaternion-space MANO rotation smoothing (0 disables)")
|
||||
ap.add_argument("--side", default="right", choices=["right", "left"])
|
||||
ap.add_argument("--view", default=None, choices=["topdown", "side", "third_person"],
|
||||
help="Ortho view; default side for --fix_camera else third_person")
|
||||
ap.add_argument("--output_dir", default="output/demo")
|
||||
ap.add_argument("--device", default="cuda")
|
||||
ap.add_argument("--save_npz", default=None, help="Optional: path to save inference results npz")
|
||||
args = ap.parse_args()
|
||||
|
||||
device = torch.device(args.device)
|
||||
t0 = time.perf_counter()
|
||||
|
||||
# ── config ──
|
||||
cfg = OmegaConf.load(str(_PROJECT_ROOT / args.config))
|
||||
model_cfg = OmegaConf.load(str(_PROJECT_ROOT / cfg.model_yaml))
|
||||
|
||||
# ── read video ──
|
||||
print(f"[demo] Reading video {args.input} ...")
|
||||
frames_bgr, fps, (H, W) = read_video_frames(args.input)
|
||||
T = len(frames_bgr)
|
||||
print(f" {T} frames, {W}x{H}, {fps:.1f} fps")
|
||||
t1 = time.perf_counter()
|
||||
|
||||
# ── intrinsics + camera geometry (c2w) ──
|
||||
camera_cache = Path(args.output_dir) / "camera_slam.npz"
|
||||
camera_source = Path(args.camera_npz) if args.camera_npz else camera_cache
|
||||
if camera_source.exists():
|
||||
camera_data = np.load(camera_source)
|
||||
intr = np.asarray(camera_data["intrinsics"], dtype=np.float32).reshape(-1)[:4]
|
||||
c2w = np.asarray(camera_data["c2w"], dtype=np.float32)
|
||||
if c2w.shape[0] < T:
|
||||
raise ValueError(f"Camera npz has {c2w.shape[0]} poses for {T} video frames")
|
||||
c2w = c2w[:T]
|
||||
vipe_t = 0.0
|
||||
print(f"[demo] Reusing camera geometry from {camera_source}")
|
||||
else:
|
||||
intr, vipe_estimate = resolve_intrinsics(args.intrinsics, W, H)
|
||||
intr, c2w, vipe_t = get_camera_geometry(frames_bgr, args.fix_camera, intr, vipe_estimate, T)
|
||||
if not args.fix_camera:
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
np.savez(camera_cache, intrinsics=intr, c2w=c2w)
|
||||
print(f"[demo] Cached camera geometry -> {camera_cache}")
|
||||
intr_list = [intr] * T
|
||||
t2 = time.perf_counter()
|
||||
|
||||
# ── load models ──
|
||||
print("[demo] Loading denoiser / HaMeR / MANO ...")
|
||||
mano_root = cfg.eval.mano_root or os.environ.get("MANO_ROOT")
|
||||
if not mano_root:
|
||||
raise RuntimeError("MANO_ROOT not set: please export MANO_ROOT=/path/to/mano")
|
||||
denoiser = load_denoiser_from_ckpt(OmegaConf.merge(model_cfg, cfg), args.fm_ckpt, device)
|
||||
denoiser.hamer_backbone.to("cpu")
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
print("[demo] Offloaded unused frozen HaMeR backbone to CPU")
|
||||
online_cache = Path(args.output_dir) / "online_hamer.pt"
|
||||
online = None if online_cache.exists() else OnlineHaMeRPipeline(device=str(device))
|
||||
mano_fk = MANOForwardKinematics(str(mano_root), device)
|
||||
t3 = time.perf_counter()
|
||||
|
||||
# ── Online HaMeR ──
|
||||
if online_cache.exists():
|
||||
print(f"[demo] Reusing Online HaMeR features from {online_cache}")
|
||||
online_result = torch.load(online_cache, map_location="cpu", weights_only=False)
|
||||
else:
|
||||
print("[demo] Online HaMeR inference ...")
|
||||
online_result = online.process_sequence(frames_bgr, intr_list, target_side=args.side)
|
||||
torch.save(online_result, online_cache)
|
||||
print(f"[demo] Cached Online HaMeR features -> {online_cache}")
|
||||
det_valid = online_result["detection_valid"].numpy()
|
||||
det_sides = online_result["sides"]
|
||||
bbox_xyxy = online_result["bbox_xyxy"].numpy()
|
||||
bbox_conf = online_result["bbox_conf"].numpy()
|
||||
dominant_side = max(set(det_sides), key=det_sides.count) if det_sides else args.side
|
||||
print(f" Detections: {int(det_valid.sum())}/{T} frames, dominant hand side={dominant_side}")
|
||||
t4 = time.perf_counter()
|
||||
|
||||
bf = online_result["backbone_features"]
|
||||
if online is not None:
|
||||
print("[demo] HaMeR cache ready; restart to release third-party CUDA hooks")
|
||||
return
|
||||
token_chunks = []
|
||||
compressor_batch_size = 1
|
||||
with torch.no_grad():
|
||||
for s in range(0, T, compressor_batch_size):
|
||||
e = min(s + compressor_batch_size, T)
|
||||
chunk = denoiser.frame_compressor(bf[s:e].to(device).unsqueeze(0)).squeeze(0)
|
||||
token_chunks.append(chunk.detach().cpu())
|
||||
image_tokens = torch.cat(token_chunks, dim=0)
|
||||
del bf, token_chunks
|
||||
|
||||
seq_batch = {
|
||||
"mano_params": torch.zeros((1, T, 48), dtype=torch.float32),
|
||||
"mano_trans": torch.zeros((1, T, 3), dtype=torch.float32),
|
||||
"mano_betas": torch.zeros((1, 10), dtype=torch.float32),
|
||||
"padding_mask": torch.zeros((1, T), dtype=torch.bool),
|
||||
"images": online_result["crop_images"].unsqueeze(0),
|
||||
"hamer_landmarks": online_result["hamer_landmarks"].unsqueeze(0),
|
||||
"crop_intrinsics": online_result["crop_intrinsics"].unsqueeze(0),
|
||||
"hamer_confidence": online_result["hamer_confidence"].unsqueeze(0),
|
||||
"image_tokens": image_tokens.unsqueeze(0),
|
||||
"side": [dominant_side],
|
||||
"source": ["custom"],
|
||||
}
|
||||
|
||||
# ── FM inference (overlapping windows) ──
|
||||
win = int(cfg.inference.window_size)
|
||||
overlap = int(cfg.inference.overlap_size)
|
||||
ode_steps = int(cfg.inference.ode_steps)
|
||||
stride = win - overlap
|
||||
win_batch, _ = split_sequence_into_windows(seq_batch, win, stride, device)
|
||||
win_batch = {k: (v.to(device) if isinstance(v, torch.Tensor) else v) for k, v in win_batch.items()}
|
||||
|
||||
print("[demo] FM Denoiser inference ...")
|
||||
pose_seq, trans_seq, betas_pred = run_fm_inference_with_overlap(
|
||||
denoiser, denoiser, win_batch, win, overlap, ode_steps, device,
|
||||
overlap_method=cfg.inference.get("overlap_method", "vblend"),
|
||||
)
|
||||
if args.rotation_smooth_window > 1:
|
||||
pose_seq = smooth_pose_rotations(pose_seq, args.rotation_smooth_window)
|
||||
print(f" Applied quaternion rotation smoothing: window={args.rotation_smooth_window}")
|
||||
nf = pose_seq.shape[0]
|
||||
print(f" Output {nf} frames of pose/trans")
|
||||
t5 = time.perf_counter()
|
||||
|
||||
# ── MANO FK -> verts_cam (meters, per-frame hand side from detection) ──
|
||||
sides = (det_sides[:nf] if len(det_sides) >= nf
|
||||
else det_sides + [det_sides[-1]] * (nf - len(det_sides)))
|
||||
verts_cam_m = (mano_fk.verts(pose_seq, betas_pred, trans_seq, sides) / 1000.0
|
||||
).cpu().numpy().astype(np.float32) # (nf,778,3) m
|
||||
faces = mano_fk.get_faces(args.side)
|
||||
t6 = time.perf_counter()
|
||||
|
||||
# ── world verts + camera positions for the ortho view ──
|
||||
n = min(nf, T, c2w.shape[0])
|
||||
verts_world = transform_verts_to_world(verts_cam_m[:n], c2w[:n]) # identity if fix_camera
|
||||
cam_positions = c2w[:n, :3, 3] # zeros if fix_camera
|
||||
|
||||
# ── Output 1: cam overlay (clean verts_cam projection onto RGB) ──
|
||||
print("[demo] Rendering overlay ...")
|
||||
renderer = PhongRenderer(device)
|
||||
n_render = min(nf, T)
|
||||
out_frames = []
|
||||
for fi in range(T):
|
||||
frame = frames_bgr[fi]
|
||||
if fi < n_render and det_valid[fi]:
|
||||
x1, y1, x2, y2 = bbox_xyxy[fi].astype(int)
|
||||
cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
|
||||
label = f"{det_sides[fi]} {bbox_conf[fi]:.2f}"
|
||||
cv2.putText(frame, label, (x1, max(8, y1 - 8)),
|
||||
cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2, cv2.LINE_AA)
|
||||
frame = renderer.render_overlay(frame, verts_cam_m[fi], faces, intr, sides[fi])
|
||||
out_frames.append(frame)
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
overlay_path = str(Path(args.output_dir) / "overlay.mp4")
|
||||
write_video_ffmpeg(out_frames, overlay_path, fps)
|
||||
print(f"[demo] ✅ {overlay_path}")
|
||||
t7 = time.perf_counter()
|
||||
|
||||
# ── Output 2: orthographic view (side for fix_camera / level; else world topdown) ──
|
||||
view = args.view or ("side" if args.fix_camera else "third_person")
|
||||
print(f"[demo] Rendering ortho ({view}) ...")
|
||||
traj_frames = renderer.render_ortho_video(verts_world, faces, c2w[:n], view=view, side=args.side)
|
||||
ortho_path = str(Path(args.output_dir) / f"ortho_{view}.mp4")
|
||||
write_video_ffmpeg(traj_frames, ortho_path, fps)
|
||||
print(f"[demo] ✅ {ortho_path}")
|
||||
t8 = time.perf_counter()
|
||||
|
||||
# ── timing breakdown ──
|
||||
det_t = float(online_result.get("det_time_s", 0.0))
|
||||
hamer_t = float(online_result.get("hamer_time_s", 0.0))
|
||||
print(
|
||||
f"[timing] read={t1-t0:.1f}s | "
|
||||
f"vipe={vipe_t:.1f}s(+frames/subprocess {t2-t1-vipe_t:.1f}s) | "
|
||||
f"load={t3-t2:.1f}s | "
|
||||
f"hamer={t4-t3:.1f}s (det={det_t:.1f} infer={hamer_t:.1f}) | "
|
||||
f"fm={t5-t4:.1f}s | mano={t6-t5:.1f}s | render={t7-t6:.1f}s | ortho={t8-t7:.1f}s | "
|
||||
f"total={t8-t0:.1f}s"
|
||||
)
|
||||
|
||||
# ── optional: save inference npz ──
|
||||
if args.save_npz:
|
||||
np.savez(
|
||||
args.save_npz,
|
||||
verts_cam=verts_cam_m[:n_render], verts_world=verts_world, faces=faces,
|
||||
intrinsics=intr, c2w=c2w[:n], fix_camera=bool(args.fix_camera),
|
||||
pose=pose_seq.cpu().numpy(), trans=trans_seq.cpu().numpy(),
|
||||
betas=betas_pred.cpu().numpy(), pred_valid=det_valid[:n_render],
|
||||
side=args.side, source=args.input, fps=fps,
|
||||
)
|
||||
print(f"[demo] npz -> {args.save_npz}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,54 @@
|
||||
"""Independently evaluate SPIDER's zero-assistance rollout and a control-replay baseline."""
|
||||
import os
|
||||
os.environ.setdefault('MUJOCO_GL','osmesa')
|
||||
from pathlib import Path
|
||||
import json,re
|
||||
import numpy as np
|
||||
import mujoco,imageio.v2 as imageio
|
||||
from PIL import Image,ImageDraw
|
||||
ROOT=Path(__file__).resolve().parents[1];OUT=ROOT/'output/spider_l20_contact';TASK=OUT/'datasets/processed/current/l20/right/bottle_hold';TRIAL=TASK/'0'
|
||||
DT=json.loads((OUT/'config.json').read_text())['sim_dt']
|
||||
ref=np.load(TRIAL/'trajectory_kinematic_act.npz');result=np.load(TRIAL/'trajectory_mjwp_act.npz');q=result['qpos'].reshape(-1,result['qpos'].shape[-1]);ctrl=result['ctrl'].reshape(-1,result['ctrl'].shape[-1]);v=result['qvel'].reshape(-1,result['qvel'].shape[-1]);time=result['time'].ravel();N=len(q)
|
||||
model=mujoco.MjModel.from_xml_path(str(TASK/'scene_act.xml'));model.opt.iterations=20;model.opt.ls_iterations=50;model.opt.o_solref[:]=[.02,1];model.opt.o_solimp[:]=[0,.95,.03,.5,2]
|
||||
obj=model.body('right_object').id;hand=model.body('hand_base_link').id
|
||||
assert np.all(model.actuator_gainprm[-6:]==0) and np.all(model.actuator_biasprm[-6:]==0)
|
||||
def replay(controls):
|
||||
d=mujoco.MjData(model);d.qpos[:]=ref['qpos'][0];d.qvel[:]=ref['qvel'][0];d.ctrl[:]=ref['ctrl'][0];mujoco.mj_step(model,d);d.time=0
|
||||
rows=[]
|
||||
for u in controls:
|
||||
d.ctrl[:]=u;mujoco.mj_step(model,d);rows.append(d.qpos.copy())
|
||||
return np.array(rows)
|
||||
baseline=replay(ref['ctrl'][1:N+1]);cpu=replay(ctrl)
|
||||
def metrics(rows):
|
||||
d=mujoco.MjData(model);dr=mujoco.MjData(model);gaps=[];errors=[];relerrors=[];counts=[];wrist=[];objheight=[]
|
||||
for t,row in enumerate(rows):
|
||||
d.qpos[:]=row;mujoco.mj_forward(model,d);dr.qpos[:]=ref['qpos'][min(t+1,len(ref['qpos'])-1)];mujoco.mj_forward(model,dr)
|
||||
errors.append(np.linalg.norm(d.xpos[obj]-dr.xpos[obj]));relerrors.append(np.linalg.norm((d.xpos[obj]-d.xpos[hand])-(dr.xpos[obj]-dr.xpos[hand])))
|
||||
selected=[c for c in d.contact if {int(model.geom_contype[c.geom[0]]),int(model.geom_contype[c.geom[1]])}=={1,2}]
|
||||
counts.append(len(selected));gaps.append(max([max(0,-c.dist) for c in selected],default=0));wrist.append(d.xpos[hand].copy());objheight.append(d.xpos[obj,2])
|
||||
e=np.array(errors);r=np.array(relerrors)
|
||||
return dict(finite=bool(np.isfinite(rows).all()),max_object_position_error_mm=float(e.max()*1000),final_object_position_error_mm=float(e[-1]*1000),mean_object_position_error_mm=float(e.mean()*1000),max_object_hand_relative_error_mm=float(r.max()*1000),contact_frames=int(np.sum(np.array(counts)>0)),contact_duration_s=float(np.sum(np.array(counts)>0)*DT),last_hand_contact_time_s=float((np.flatnonzero(np.array(counts)>0)[-1]+1)*DT) if np.any(np.array(counts)>0) else 0.,frames=len(rows),max_hand_object_penetration_mm=float(np.max(gaps)*1000),wrist_range_m=np.ptp(wrist,axis=0).tolist(),final_bottle_base_height_m=float(objheight[-1]),retained=bool(r.max()<.02 and counts[-1]>0)),np.array(counts),e
|
||||
mr,counts,errors=metrics(q);mb,_,_=metrics(baseline);mc,_,_=metrics(cpu)
|
||||
report=dict(spider=mr,reference_control_baseline=mb,independent_cpu_control_replay=mc,zero_object_actuator_gains=True,zero_object_actuator_bias=True,source_frames=[600,660],simulation_seconds=float(time[-1]),nonlinear_json_coupling=False,linear_mimic=True,source_object_pose='Estimated target, not measured video object motion',cpu_gpu_max_qpos_difference=float(np.abs(cpu-q).max()))
|
||||
log=(OUT/'run.log').read_text();report['nonfinite_sample_warnings']=len(re.findall('NaNs or infs in rews',log));report['upstream_completion_marker']='SPIDER_RUN_COMPLETE' in log
|
||||
if (OUT/'independent_gpu_replay.npz').exists():
|
||||
independent=np.load(OUT/'independent_gpu_replay.npz')
|
||||
report['independent_gpu_control_replay']=metrics(independent['optimized_controls'])[0]
|
||||
report['reference_control_gpu_baseline']=metrics(independent['baseline'])[0]
|
||||
(OUT/'validation.json').write_text(json.dumps(report,indent=2));np.savez_compressed(OUT/'comparison_motion.npz',spider=q,baseline=baseline,cpu_replay=cpu,reference=ref['qpos'][1:N+1],ctrl=ctrl,qvel=v,time=time,contact_counts=counts,object_error_m=errors)
|
||||
np.savetxt(OUT/'optimized_trajectory.csv',np.c_[time,q,ctrl],delimiter=',',header=','.join(['time_s']+[model.joint(i).name for i in range(model.njnt)]+['ctrl_'+model.actuator(i).name for i in range(model.nu)]),comments='')
|
||||
model.vis.headlight.ambient[:]=.8;model.vis.quality.offsamples=1
|
||||
renderer=mujoco.Renderer(model,height=400,width=400);cam=mujoco.MjvCamera();cam.lookat[:]=[-.02,.045,.34];cam.distance=.65;cam.azimuth=50;cam.elevation=12
|
||||
opt=mujoco.MjvOption();opt.geomgroup[3]=0;opt.sitegroup[:]=0;data=mujoco.MjData(model)
|
||||
frames=np.minimum(np.round(np.arange(60)/30/DT).astype(int),N-1)
|
||||
writer=imageio.get_writer(OUT/'spider_contact_comparison.mp4',fps=30,codec='libx264',quality=8)
|
||||
for i,t in enumerate(frames):
|
||||
panels=[]
|
||||
for row in [ref['qpos'][t+1],baseline[t],q[t],cpu[t]]:
|
||||
data.qpos[:]=row;mujoco.mj_forward(model,data);renderer.update_scene(data,camera=cam,scene_option=opt);panels.append(renderer.render().copy())
|
||||
img=Image.fromarray(np.concatenate(panels,axis=1));draw=ImageDraw.Draw(img);draw.rectangle([0,0,1600,54],fill=(20,25,35))
|
||||
for j,label in enumerate(['Estimated reference','Reference controls, free bottle (CPU)','SPIDER contact guidance (GPU)','Same controls, independent CPU replay']):draw.text((j*400+8,8),label,fill='white')
|
||||
draw.text((8,30),f't={time[t]:.2f}s | Moving wrist + URDF linear mimic | SPIDER error={errors[t]*1000:.1f} mm, contacts={counts[t]} | 20mm tracking gate={mr["retained"]}',fill='white')
|
||||
writer.append_data(np.asarray(img))
|
||||
if i in [0,30,59]:img.save(OUT/f'comparison_{i:02d}.png')
|
||||
writer.close();renderer.close();print(json.dumps(report,indent=2))
|
||||
@@ -0,0 +1,50 @@
|
||||
"""Export Dyn-HaMR MANO joints, preserving the exact reconstruction FK convention."""
|
||||
import os
|
||||
os.environ.setdefault('OMP_NUM_THREADS', '4')
|
||||
import sys
|
||||
from pathlib import Path
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
sys.path.insert(0, str(ROOT / 'third_party/Dyn-HaMR/dyn-hamr'))
|
||||
from body_model.mano_wrapper import MANO
|
||||
|
||||
def main():
|
||||
torch.set_num_threads(4)
|
||||
source = ROOT / 'output/results_2047635068/world_results.npz'
|
||||
out = ROOT / 'output/dex_2047635068'
|
||||
out.mkdir(exist_ok=True)
|
||||
data = np.load(source)
|
||||
candidates = np.flatnonzero(np.all(data['is_right'] > .5, axis=1))
|
||||
assert len(candidates) == 1, 'Expected exactly one consistent right-hand track'
|
||||
track = int(candidates[0])
|
||||
pose = data['pose_body'][track].reshape(-1, 45)
|
||||
model = MANO(model_path=str(ROOT / 'third_party/Dyn-HaMR/_DATA/data/mano'),
|
||||
batch_size=128, pose2rot=True)
|
||||
joints = []
|
||||
wrists = []
|
||||
with torch.no_grad():
|
||||
for start in range(0, len(pose), 128):
|
||||
p = torch.from_numpy(pose[start:start+128]).float()
|
||||
n = len(p)
|
||||
result = model(hand_pose=p, global_orient=torch.zeros(n, 3),
|
||||
transl=torch.zeros(n, 3),
|
||||
betas=torch.from_numpy(data['betas'][track]).float().expand(n, -1))
|
||||
j = result.joints.numpy()
|
||||
joints.append(j - j[:, :1])
|
||||
world = model(hand_pose=p,
|
||||
global_orient=torch.from_numpy(data['root_orient'][track,start:start+n]).float(),
|
||||
transl=torch.from_numpy(data['trans'][track,start:start+n]).float(),
|
||||
betas=torch.from_numpy(data['betas'][track]).float().expand(n,-1))
|
||||
wrists.append(world.joints[:,0].numpy())
|
||||
joints = np.concatenate(joints)
|
||||
assert joints.shape == (len(pose), 21, 3) and np.isfinite(joints).all()
|
||||
np.savez_compressed(out / 'human_joints.npz', joints=joints, fps=30,
|
||||
source_track=track, source=str(source.resolve()),
|
||||
wrist_world=np.concatenate(wrists),
|
||||
root_orient=data['root_orient'][track], trans=data['trans'][track])
|
||||
print(f'Exported {joints.shape}, right track {track}, to {out}', flush=True)
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,100 @@
|
||||
"""Export HandFlow with its own MANO convention and camera-to-world transform."""
|
||||
import argparse
|
||||
import json
|
||||
from pathlib import Path
|
||||
import sys
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from scipy.spatial.transform import Rotation, Slerp
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
sys.path.insert(0, str(ROOT))
|
||||
from utils.mano_utils import MANOForwardKinematics
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--input', type=Path, required=True)
|
||||
parser.add_argument('--output', type=Path, required=True)
|
||||
parser.add_argument('--gap-policy', choices=['interpolate', 'raw'], default='interpolate')
|
||||
parser.add_argument('--min-detection-run', type=int, default=3,
|
||||
help='Reject isolated detection runs shorter than this before gap filling')
|
||||
args = parser.parse_args()
|
||||
torch.set_num_threads(4)
|
||||
source = np.load(args.input)
|
||||
assert str(source['side']) == 'right', 'This exporter targets the L20 right hand'
|
||||
pose = torch.tensor(source['pose'], dtype=torch.float32)
|
||||
count = len(pose)
|
||||
betas = torch.tensor(source['betas'], dtype=torch.float32)
|
||||
if betas.ndim == 1:
|
||||
betas = betas[None].expand(count, -1)
|
||||
elif len(betas) == 1:
|
||||
betas = betas.expand(count, -1)
|
||||
assert betas.shape == (count, 10)
|
||||
trans = torch.tensor(source['trans'], dtype=torch.float32)
|
||||
c2w = source['c2w']
|
||||
assert len(c2w) == count and len(source['verts_cam']) == count
|
||||
fk = MANOForwardKinematics(str(ROOT / 'third_party/hamer/_DATA/data/mano'), torch.device('cpu'))
|
||||
local_pose = pose.clone()
|
||||
local_pose[:, :3] = 0
|
||||
local, camera, verts = [], [], []
|
||||
for start in range(0, count, 64):
|
||||
end = min(count, start + 64)
|
||||
sides = ['right'] * (end - start)
|
||||
local.append(fk.joints(local_pose[start:end], betas[start:end], torch.zeros_like(trans[start:end]), sides).numpy() / 1000)
|
||||
camera.append(fk.joints(pose[start:end], betas[start:end], trans[start:end], sides).numpy() / 1000)
|
||||
verts.append(fk.verts(pose[start:end], betas[start:end], trans[start:end], sides).numpy() / 1000)
|
||||
local = np.concatenate(local)
|
||||
camera = np.concatenate(camera)
|
||||
verts = np.concatenate(verts)
|
||||
local -= local[:, :1].copy()
|
||||
world = np.einsum('tij,tkj->tki', c2w[:, :3, :3], camera) + c2w[:, None, :3, 3]
|
||||
root_R = c2w[:, :3, :3] @ Rotation.from_rotvec(pose[:, :3].numpy()).as_matrix()
|
||||
reconstructed = np.einsum('tij,tkj->tki', root_R, local) + world[:, :1]
|
||||
vertex_error = float(np.abs(verts - source['verts_cam']).max())
|
||||
rigid_error = float(np.abs(reconstructed - world).max())
|
||||
world_vertex_error = float(np.abs(np.einsum('tij,tkj->tki', c2w[:, :3, :3], verts) + c2w[:, None, :3, 3] - source['verts_world']).max())
|
||||
assert np.isfinite(world).all() and np.isfinite(local).all()
|
||||
assert max(vertex_error, rigid_error, world_vertex_error) < 1e-5
|
||||
valid = np.asarray(source['pred_valid'], dtype=bool)
|
||||
observed_valid = valid.copy()
|
||||
valid_ids = np.flatnonzero(valid)
|
||||
for group in np.split(valid_ids, np.flatnonzero(np.diff(valid_ids) > 1) + 1):
|
||||
if len(group) < args.min_detection_run:
|
||||
valid[group] = False
|
||||
local_raw, wrist_raw, root_raw = local.copy(), world[:, 0].copy(), root_R.copy()
|
||||
wrist = wrist_raw.copy()
|
||||
if a_policy := (args.gap_policy == 'interpolate' and not valid.all()):
|
||||
observed = np.flatnonzero(valid)
|
||||
assert len(observed) >= 2, 'Too few detected frames for motion export'
|
||||
times = np.arange(count)
|
||||
local = np.stack([np.interp(times, observed, local[:, i, j][observed])
|
||||
for i in range(21) for j in range(3)], axis=1).reshape(count,21,3)
|
||||
wrist = np.stack([np.interp(times, observed, wrist[:, j][observed]) for j in range(3)],axis=1)
|
||||
root_R = Slerp(observed, Rotation.from_matrix(root_R[observed]))(np.clip(times,observed[0],observed[-1])).as_matrix()
|
||||
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||||
np.savez_compressed(args.output, joints=local, wrist_world=wrist,
|
||||
root_orient=Rotation.from_matrix(root_R).as_rotvec(), fps=source['fps'],
|
||||
detection_valid=valid, detection_valid_raw=observed_valid, source=str(args.input.resolve()),
|
||||
source_model='HandFlow/manopth flat_hand_mean=True', joints_world_raw=world,
|
||||
joints_raw=local_raw, wrist_world_raw=wrist_raw, root_world_R_raw=root_raw,
|
||||
gap_policy=args.gap_policy)
|
||||
report = dict(frames=count, fps=float(source['fps']), finite=True,
|
||||
detected_frames=int(source['pred_valid'].sum()),
|
||||
usable_detection_frames=int(valid.sum()),
|
||||
rejected_isolated_detection_frames=np.flatnonzero(observed_valid & ~valid).tolist(),
|
||||
min_detection_run=args.min_detection_run,
|
||||
camera_vertex_reproduction_max_error_m=vertex_error,
|
||||
world_vertex_reproduction_max_error_m=world_vertex_error,
|
||||
local_to_world_joint_max_error_m=rigid_error,
|
||||
source=str(args.input.resolve()),
|
||||
gap_policy=args.gap_policy, filled_frames=int((~valid).sum()) if a_policy else 0,
|
||||
gap_note='Retargeting holds endpoint detections and interpolates internal gaps; raw inference preserved',
|
||||
convention='HandFlow manopth, flat_hand_mean=True; mm converted to m; c2w applied once')
|
||||
args.output.with_suffix('.validation.json').write_text(json.dumps(report, indent=2))
|
||||
print(json.dumps(report, indent=2))
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,201 @@
|
||||
"""Export the existing right-hand replay without changing its model or motion.
|
||||
|
||||
Uses the dataset layout of HDF5_REQUIREMENTS.md, with explicit right-hand
|
||||
model/coordinate exceptions. This is not a left-hand l20_tracking_v1 PASS.
|
||||
"""
|
||||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
from pathlib import Path
|
||||
import xml.etree.ElementTree as ET
|
||||
|
||||
import h5py
|
||||
import numpy as np
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
NAMES = [f'{f}_{j}' for f in ('index', 'middle', 'pinky', 'ring')
|
||||
for j in ('dip', 'mcp_pitch', 'mcp_roll', 'pip')] + [
|
||||
'thumb_cmc_pitch', 'thumb_cmc_roll', 'thumb_cmc_yaw', 'thumb_ip', 'thumb_mcp']
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument('--source', type=Path, default=ROOT/'output/l20_15886123_stable/motion.npz')
|
||||
parser.add_argument('--output', type=Path, default=ROOT/'output/hdf5_delivery/right')
|
||||
parser.add_argument('--scene', type=Path, help='MuJoCo scene for a full-replay qpos source')
|
||||
args = parser.parse_args()
|
||||
args.output.mkdir(parents=True, exist_ok=True)
|
||||
a = np.load(args.source, allow_pickle=False)
|
||||
scene_description = None
|
||||
if args.scene:
|
||||
import mujoco
|
||||
model = mujoco.MjModel.from_xml_path(str(args.scene.resolve()))
|
||||
data = mujoco.MjData(model)
|
||||
rows = a['qpos']
|
||||
assert rows.ndim == 2 and rows.shape[1] == model.nq and np.isfinite(rows).all()
|
||||
addresses = [model.jnt_qposadr[model.joint(n).id] for n in NAMES]
|
||||
body_id = model.body('hand_base_link').id
|
||||
positions, rotations = [], []
|
||||
for row in rows:
|
||||
data.qpos[:] = row
|
||||
mujoco.mj_forward(model, data)
|
||||
positions.append(data.xpos[body_id].copy())
|
||||
rotations.append(data.xquat[body_id].copy())
|
||||
positions = np.asarray(positions)
|
||||
assert np.allclose(positions, a['wrist_pos'], atol=1e-9, rtol=0)
|
||||
a = dict(qpos=rows[:,addresses], joint_names=np.asarray(NAMES), wrist_pos=positions,
|
||||
wrist_quat_wxyz=np.asarray(rotations), time=a['time'],
|
||||
detection_valid=np.ones(len(rows), dtype=bool),
|
||||
scene_rotation=np.eye(3), scene_translation=np.zeros(3))
|
||||
scene_description = (
|
||||
'Video 2047635068, full 1333-frame right-hand replay: Dyn-HaMR -> '
|
||||
'run_dex_l20.py -> temporal smoothing -> build_l20_full_replay.py contact registration. '
|
||||
'Actual source is the final replay qpos, not the earlier raw retargeting. '
|
||||
'hand_base_link world position and active local-to-world WXYZ rotation extracted '
|
||||
'with MuJoCo mj_forward for every frame; finger columns resolved by joint name. '
|
||||
'Source world is explicitly the FINAL replay scene world; world_from_source is identity '
|
||||
'because exported FK is already in that same world, not because camera calibration is known. '
|
||||
'Axes/origin are the existing fixed right-handed replay scene with display Z up: upstream '
|
||||
'scene was aligned using estimated bottle orientation at frame 600, shifted to the display '
|
||||
'placement, then raised for floor clearance; physical gravity alignment is unverified. '
|
||||
'Original human wrist/palm mapping, temporal smoothing and per-frame contact-registration '
|
||||
'corrections are inherited unchanged. Root is robot link origin, not its center of mass. '
|
||||
'valid=true means finite kinematic reference state satisfying the actual right URDF '
|
||||
'limits and mimic; original detection validity was not supplied, so this is NOT an '
|
||||
'observed-frame mask or measurement confidence. Motion and object registration are '
|
||||
'estimated, not ground truth; no full dynamic rollout or hardware execution claim. '
|
||||
'Source scene: ' + str(args.scene.resolve()))
|
||||
asset = ROOT/'third_party/l20_assets/L20/RIGHT'
|
||||
urdf = asset/'linkerhand_g20_right.urdf'
|
||||
tree = ET.parse(urdf).getroot()
|
||||
files = {urdf}
|
||||
for mesh in tree.iter('mesh'):
|
||||
files.add((asset/mesh.attrib['filename']).resolve())
|
||||
entries = [{'path': p.relative_to(asset.resolve()).as_posix(),
|
||||
'sha256': hashlib.sha256(p.read_bytes()).hexdigest()}
|
||||
for p in sorted(files, key=lambda p: p.relative_to(asset.resolve()).as_posix())]
|
||||
# Explicit local algorithm, not the unavailable receiver's asset hash tool.
|
||||
payload = ''.join(f"{e['path']}\t{e['sha256']}\n" for e in entries).encode('utf-8')
|
||||
digest = hashlib.sha256(payload).hexdigest()
|
||||
js = {j.attrib['name']: j for j in tree.findall('joint') if j.attrib['type'] != 'fixed'}
|
||||
manifest = dict(hand_side='right', root_link='hand_base_link', asset_sha256=digest,
|
||||
hash_algorithm='sha256(UTF-8 concatenation of sorted relative_path TAB file_sha256 LF); URDF plus referenced meshes only; local export convention',
|
||||
source_urdf=str(urdf), files=entries, joints=[])
|
||||
for name in NAMES:
|
||||
j = js[name]
|
||||
item = dict(name=name, lower_rad=float(j.find('limit').get('lower')),
|
||||
upper_rad=float(j.find('limit').get('upper')))
|
||||
m = j.find('mimic')
|
||||
if m is not None:
|
||||
item['mimic'] = dict(joint=m.get('joint'), multiplier=float(m.get('multiplier', '1')),
|
||||
offset=float(m.get('offset', '0')))
|
||||
manifest['joints'].append(item)
|
||||
(args.output/'right_model_manifest.json').write_text(json.dumps(manifest, indent=2)+'\n')
|
||||
source_names = a['joint_names'].tolist()
|
||||
assert len(source_names) == len(set(source_names)) == len(NAMES)
|
||||
assert set(source_names) == set(NAMES)
|
||||
q = a['qpos'][:, [source_names.index(n) for n in NAMES]].astype(np.float32)
|
||||
pos = a['wrist_pos'].astype(np.float32)
|
||||
quat = a['wrist_quat_wxyz'].copy()
|
||||
quat /= np.linalg.norm(quat, axis=1, keepdims=True)
|
||||
for i in range(1, len(quat)):
|
||||
if quat[i] @ quat[i-1] < 0:
|
||||
quat[i] *= -1
|
||||
quat = quat.astype(np.float32)
|
||||
valid = a['detection_valid'].astype(bool)
|
||||
transform = np.eye(4, dtype=np.float64)
|
||||
transform[:3, :3] = a['scene_rotation']
|
||||
transform[:3, 3] = a['scene_translation']
|
||||
description = (
|
||||
'Video 15886123; HandFlow/manopth + ViPE camera-to-world; '
|
||||
'scripts/export_handflow_dex.py -> retarget_l20_video.py -> stabilize_l20_temporal.py. '
|
||||
'Right L20 URDF linear mimic, 16 independent and 21 total joints. '
|
||||
'hand_base_link origin assigned to estimated MANO wrist; orientation uses human palm_basis '
|
||||
'and robot neutral palm_basis, not a measured anatomical mount. '
|
||||
'Positions already use the stored fixed scene transform, followed by offline temporal smoothing. '
|
||||
'Scene origin and XYZ axes follow the original replay: first unsmoothed wrist at (0,0,0.2), '
|
||||
'first unsmoothed root orientation identity. Right-handed fixed scene with display Z up; '
|
||||
'physical gravity alignment is unknown. No new pose recentering applied by exporter. '
|
||||
'world_from_source maps the upstream estimated world to this existing display world; '
|
||||
'smoothing additionally changes individual poses. '
|
||||
'valid is upstream usable detection mask, not ground-truth tracking accuracy: '
|
||||
'first 24 frames are held/interpolated finite estimates and remain false. '
|
||||
'All frames have monocular scale/depth uncertainty, possible drift/occlusion errors, '
|
||||
'future-frame temporal smoothing; no contact, dynamics or hardware validation.')
|
||||
attrs = dict(schema_version='l20_tracking_v1', embodiment='L20', hand_side='right',
|
||||
asset_sha256=digest, root_link='hand_base_link', provenance='expert_retargeted',
|
||||
source_description=description,
|
||||
metric_scale_provenance='Upstream MANO FK millimeters divided by 1000; camera translations used in upstream meter convention. Monocular/model-estimated metric scale, no measured reference length or gravity calibration; absolute scale error unknown. Input motion.npz positions already in estimated meters, so export scale_to_meters=1.0.',
|
||||
scale_to_meters=1.0,
|
||||
contract_status='RIGHT_HAND_VARIANT_NOT_STRICT_LEFT_V1; actual URDF mimic; local asset hash algorithm; display-world gravity unverified',
|
||||
source_sha256=hashlib.sha256(args.source.read_bytes()).hexdigest(),
|
||||
asset_hash_algorithm=manifest['hash_algorithm'])
|
||||
if scene_description:
|
||||
attrs['source_description'] = scene_description
|
||||
attrs['metric_scale_provenance'] = (
|
||||
'Final replay positions already in estimated meters, inherited from Dyn-HaMR MANO '
|
||||
'world reconstruction and robot geometry; export scale_to_meters=1.0. '
|
||||
'No measured absolute scale reference is established; error unknown. '
|
||||
'Scene alignment and contact fitting are engineering estimates, not scale measurements.')
|
||||
attrs['valid_semantics'] = 'kinematic_reference_valid; source_detection_validity_unavailable'
|
||||
attrs['source_scene_sha256'] = hashlib.sha256(args.scene.read_bytes()).hexdigest()
|
||||
good = np.flatnonzero(valid)
|
||||
runs = [r for r in np.split(good, np.flatnonzero(np.diff(good) > 1)+1) if len(r) >= 2]
|
||||
assert runs
|
||||
run = max(runs, key=len)
|
||||
sample = slice(int(run[0]), int(min(run[-1]+1, run[0]+90)))
|
||||
reports = []
|
||||
for filename, selection in [('sample_right.hdf5', sample), ('demonstrations_right.hdf5', slice(0, len(q)))]:
|
||||
path = args.output/filename
|
||||
times = a['time'][selection].astype(np.float64)
|
||||
times -= times[0]
|
||||
with h5py.File(path, 'w') as f:
|
||||
for k, v in attrs.items():
|
||||
f.attrs[k] = v
|
||||
meta = f.create_group('metadata')
|
||||
meta.create_dataset('joint_names', data=NAMES, dtype=h5py.string_dtype('utf-8'))
|
||||
meta.create_dataset('world_from_source', data=transform)
|
||||
ep = f.create_group('episodes/demo_000000')
|
||||
ep.attrs['source_frame_start'] = selection.start
|
||||
ep.attrs['source_frame_end_exclusive'] = selection.stop
|
||||
for name, value in dict(time=times, wrist_position=pos[selection],
|
||||
wrist_quaternion=quat[selection], joint_position=q[selection],
|
||||
valid=valid[selection]).items():
|
||||
ep.create_dataset(name, data=value)
|
||||
# Reopen the actual serialized file and validate against the actual right URDF.
|
||||
with h5py.File(path, 'r') as f:
|
||||
ep = f['episodes/demo_000000']
|
||||
t, p, r, angles, mask = [ep[n][:] for n in
|
||||
('time','wrist_position','wrist_quaternion','joint_position','valid')]
|
||||
count = len(t)
|
||||
assert [(x.shape, x.dtype) for x in (t,p,r,angles,mask)] == [
|
||||
((count,),np.dtype('float64')),((count,3),np.dtype('float32')),
|
||||
((count,4),np.dtype('float32')),((count,21),np.dtype('float32')),((count,),np.dtype('bool'))]
|
||||
assert count >= 2 and t[0] == 0 and np.all(np.diff(t)>0) and mask.any()
|
||||
assert all(np.isfinite(x).all() for x in (t,p,r,angles))
|
||||
assert f['metadata/joint_names'].asstr()[:].tolist() == NAMES
|
||||
assert h5py.check_string_dtype(f['metadata/joint_names'].dtype).encoding == 'utf-8'
|
||||
assert np.max(np.abs(np.linalg.norm(r[mask],axis=1)-1)) < 1e-4
|
||||
assert np.all(np.sum(r[1:]*r[:-1],axis=1)[mask[1:] & mask[:-1]] >= 0)
|
||||
assert np.allclose(transform[3], [0,0,0,1])
|
||||
assert np.allclose(transform[:3,:3].T@transform[:3,:3], np.eye(3))
|
||||
assert abs(np.linalg.det(transform[:3,:3])-1)<1e-8
|
||||
residuals = {}
|
||||
for i, joint in enumerate(manifest['joints']):
|
||||
assert angles[mask,i].min() >= joint['lower_rad']-1e-6
|
||||
assert angles[mask,i].max() <= joint['upper_rad']+1e-6
|
||||
if 'mimic' in joint:
|
||||
m = joint['mimic']
|
||||
err = np.abs(angles[mask,i]-m['multiplier']*angles[mask,NAMES.index(m['joint'])]-m['offset'])
|
||||
residuals[joint['name']] = float(err.max())
|
||||
assert err.max() <= 1e-3
|
||||
reports.append(dict(file=filename, right_model_validation='PASS', strict_left_v1='NOT_COMPATIBLE',
|
||||
frames=count, valid_frames=int(mask.sum()), source_frame_start=selection.start,
|
||||
source_frame_end_exclusive=selection.stop, time_last_s=float(t[-1]),
|
||||
mimic_residual_rad=residuals, sha256=hashlib.sha256(path.read_bytes()).hexdigest()))
|
||||
(args.output/'validation.json').write_text(json.dumps(reports,indent=2)+'\n')
|
||||
print(json.dumps(reports,indent=2))
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,190 @@
|
||||
"""Actuated L20 + freely simulated metric bottle, with a prescribed moving wrist."""
|
||||
import os
|
||||
os.environ.setdefault('MUJOCO_GL','osmesa')
|
||||
import argparse,json
|
||||
import xml.etree.ElementTree as ET
|
||||
from pathlib import Path
|
||||
import numpy as np
|
||||
import mujoco
|
||||
from scipy.spatial.transform import Rotation
|
||||
from scipy.optimize import lsq_linear
|
||||
from l20_calibrated import ROOT,Calibration
|
||||
|
||||
BASE=ROOT/'output/l20_2047635068'
|
||||
OUT=BASE/'bottle_grasp'
|
||||
|
||||
def feedback_torque(model,data,cal,addr,objects,object_body,target_pos,target_velocity):
|
||||
selected=[c for c in data.contact if any(g in objects for g in c.geom)]
|
||||
if not selected:return None,None
|
||||
center=data.xipos[object_body].copy();velocity=np.zeros(6)
|
||||
mujoco.mj_objectVelocity(model,data,mujoco.mjtObj.mjOBJ_BODY,object_body,velocity,0)
|
||||
desiredforce=np.array([0,0,model.body_mass[object_body]*9.81])+100*(target_pos-center)+2*(target_velocity-velocity[3:])
|
||||
rotation=Rotation.from_matrix(data.xmat[object_body].reshape(3,3))
|
||||
desiredtorque=-.08*rotation.as_rotvec()-.008*velocity[:3]
|
||||
target=np.r_[desiredforce,desiredtorque/.03]
|
||||
columns=[];jacobians=[];rays_all=[];preferred=[]
|
||||
for c in selected:
|
||||
object_first=c.geom[0] in objects
|
||||
normal=-c.frame[:3]*(1 if object_first else -1)
|
||||
tangent=np.cross(normal,[0,0,1])
|
||||
if np.linalg.norm(tangent)<.01:tangent=np.cross(normal,[1,0,0])
|
||||
tangent/=np.linalg.norm(tangent);second=np.cross(normal,tangent)
|
||||
rays=np.array([normal+.6*(np.cos(a)*tangent+np.sin(a)*second) for a in np.arange(4)*np.pi/2]).T
|
||||
rays_all.append(rays);columns.append(np.vstack([rays,np.cross((c.pos-center)[None,:],rays.T).T/.03]))
|
||||
body=int(model.geom_bodyid[c.geom[1] if object_first else c.geom[0]])
|
||||
preferred.append(6 if model.body(body).name.startswith('thumb') else 1.5)
|
||||
jac=np.zeros((3,model.nv));mujoco.mj_jac(model,data,jac,None,c.pos,body)
|
||||
ja=np.stack([jac[:,model.jnt_dofadr[model.joint(n).id]] for n in cal.active],axis=1)
|
||||
for n,(parent,_,_) in cal.curves.items():
|
||||
slope=cal.passive_value(n,data.qpos[addr[parent]])[1]
|
||||
ja[:,cal.active.index(parent)]+=slope*jac[:,model.jnt_dofadr[model.joint(n).id]]
|
||||
jacobians.append(ja)
|
||||
A=np.concatenate(columns,axis=1);reg=np.zeros((len(selected),len(selected)*4))
|
||||
for i in range(len(selected)):reg[i,i*4:i*4+4]=.025
|
||||
sol=lsq_linear(np.vstack([A,reg]),np.r_[target,.025*np.array(preferred)],bounds=(0,10),tol=1e-6,max_iter=50)
|
||||
tau=sum(j.T@(r@sol.x[i*4:i*4+4]) for i,(j,r) in enumerate(zip(jacobians,rays_all)))
|
||||
for i,n in enumerate(cal.active):tau[i]+=data.qfrc_bias[model.jnt_dofadr[model.joint(n).id]]
|
||||
return tau,float(np.linalg.norm(A@sol.x-target))
|
||||
|
||||
def build_physics(kp=30):
|
||||
tree=ET.parse(BASE/'contact_object/contact_scene.xml');root=tree.getroot()
|
||||
hand=root.find("worldbody/body[@name='hand_base_link']")
|
||||
hand.remove(hand.find('freejoint'));hand.set('mocap','true')
|
||||
option=root.find('option');option.set('timestep','.001');option.set('integrator','implicitfast');option.set('iterations','80')
|
||||
root.find('default/joint').set('damping','.08')
|
||||
eq=root.find('equality')
|
||||
if eq is None:eq=ET.SubElement(root,'equality')
|
||||
actuator=ET.SubElement(root,'actuator')
|
||||
cal=Calibration()
|
||||
for n in cal.active:
|
||||
ET.SubElement(actuator,'position',name='motor_'+n,joint=n,kp=str(kp),kv='.3',
|
||||
forcelimited='true',forcerange='-1 1')
|
||||
for n,(parent,x,y) in cal.curves.items():
|
||||
ET.SubElement(eq,'joint',name='lut_'+n,joint1=n,joint2=parent,
|
||||
polycoef='0 1 0 0 0',solref='.003 1',solimp='.95 .99 .001')
|
||||
path=OUT/'physics_scene.xml';tree.write(path)
|
||||
return path
|
||||
|
||||
def run(seconds=3.,lift=False,kp=30.,preload=0.,force_preload=False,force_control=False,feedback=False):
|
||||
path=build_physics(kp);model=mujoco.MjModel.from_xml_path(str(path));data=mujoco.MjData(model)
|
||||
cal=Calibration();old=mujoco.MjModel.from_xml_path(str(BASE/'contact_object/contact_scene.xml'))
|
||||
fitted=np.load(BASE/'contact_object/contact_result.npz')['qpos']
|
||||
names=list(cal.tables)
|
||||
addr={n:model.jnt_qposadr[mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_JOINT,n)] for n in names}
|
||||
for n in names:
|
||||
oldaddr=old.jnt_qposadr[mujoco.mj_name2id(old,mujoco.mjtObj.mjOBJ_JOINT,n)]
|
||||
data.qpos[addr[n]]=fitted[oldaddr]
|
||||
ow=old.jnt_qposadr[mujoco.mj_name2id(old,mujoco.mjtObj.mjOBJ_JOINT,'wrist_free')]
|
||||
oo=old.jnt_qposadr[mujoco.mj_name2id(old,mujoco.mjtObj.mjOBJ_JOINT,'object_free')]
|
||||
obj=model.jnt_qposadr[mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_JOINT,'object_free')]
|
||||
# Rigidly reorient the entire registration so bottle +Z is world up.
|
||||
rotation=Rotation.from_quat(fitted[oo+3:oo+7][[1,2,3,0]]).as_matrix().T
|
||||
handpos=rotation@(fitted[ow:ow+3]-fitted[oo:oo+3])+[0,0,.3]
|
||||
handrot=rotation@Rotation.from_quat(fitted[ow+3:ow+7][[1,2,3,0]]).as_matrix()
|
||||
data.mocap_pos[0]=handpos;data.mocap_quat[0]=Rotation.from_matrix(handrot).as_quat()[[3,0,1,2]]
|
||||
data.qpos[obj:obj+7]=[0,0,.3,1,0,0,0]
|
||||
controls=np.array([data.qpos[addr[n]] for n in cal.active])
|
||||
reference_controls=controls.copy()
|
||||
for i,n in enumerate(cal.active):
|
||||
if n.endswith(('pitch','pip','mcp')):controls[i]+=preload
|
||||
controls[i]=np.clip(controls[i],*cal.bounds[n])
|
||||
data.ctrl[:]=controls
|
||||
objects=np.flatnonzero(model.geom_contype==2)
|
||||
object_body=mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_BODY,'grasp_object')
|
||||
preload_report={}
|
||||
if force_preload:
|
||||
for n,(parent,x,y) in cal.curves.items():
|
||||
a=data.qpos[addr[parent]];v,s=cal.passive_value(n,a)
|
||||
ei=mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_EQUALITY,'lut_'+n)
|
||||
model.eq_data[ei,:5]=[v-s*a,s,0,0,0]
|
||||
mujoco.mj_forward(model,data)
|
||||
selected=[c for c in data.contact if any(g in objects for g in c.geom)]
|
||||
columns=[];jacobians=[];cone=[]
|
||||
center=data.xipos[object_body].copy()
|
||||
for c in selected:
|
||||
object_first=c.geom[0] in objects
|
||||
outward=c.frame[:3]*(1 if object_first else -1)
|
||||
normal=-outward
|
||||
tangent=np.cross(normal,[0,0,1])
|
||||
if np.linalg.norm(tangent)<.01:tangent=np.cross(normal,[1,0,0])
|
||||
tangent/=np.linalg.norm(tangent);second=np.cross(normal,tangent)
|
||||
rays=np.array([normal+.6*(np.cos(a)*tangent+np.sin(a)*second) for a in np.arange(4)*np.pi/2]).T
|
||||
cone.append(rays)
|
||||
columns.append(np.vstack([rays,np.cross((c.pos-center)[None,:],rays.T).T/.03]))
|
||||
body=model.geom_bodyid[c.geom[1] if object_first else c.geom[0]]
|
||||
jac=np.zeros((3,model.nv));mujoco.mj_jac(model,data,jac,None,c.pos,int(body))
|
||||
activejac=np.stack([jac[:,model.jnt_dofadr[mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_JOINT,n)]] for n in cal.active],axis=1)
|
||||
for n,(parent,x,y) in cal.curves.items():
|
||||
slope=cal.passive_value(n,data.qpos[addr[parent]])[1]
|
||||
di=model.jnt_dofadr[mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_JOINT,n)]
|
||||
activejac[:,cal.active.index(parent)]+=slope*jac[:,di]
|
||||
jacobians.append(activejac)
|
||||
if not selected:raise RuntimeError('No contacts for force preload')
|
||||
A=np.concatenate(columns,axis=1)
|
||||
regular=np.zeros((len(selected),len(selected)*4))
|
||||
for i in range(len(selected)):regular[i,i*4:i*4+4]=.025
|
||||
target=np.r_[0,0,model.body_mass[object_body]*9.81,0,0,0]
|
||||
preferred=np.array([6.]+[1.5]*(len(selected)-1))
|
||||
solution=lsq_linear(np.vstack([A,regular]),np.r_[target,.025*preferred],bounds=(0,10),tol=1e-10)
|
||||
tau=np.zeros(16)
|
||||
for i,jac in enumerate(jacobians):tau+=jac.T@(cone[i]@solution.x[i*4:i*4+4])
|
||||
for i,n in enumerate(cal.active):
|
||||
dof=model.jnt_dofadr[mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_JOINT,n)]
|
||||
tau[i]+=data.qfrc_bias[dof]
|
||||
controls+=np.clip(tau,-1,1)/kp
|
||||
for i,n in enumerate(cal.active):controls[i]=np.clip(controls[i],*cal.bounds[n])
|
||||
data.ctrl[:]=controls
|
||||
preload_report=dict(wrench_residual=float(np.linalg.norm(A@solution.x-target)),desired_motor_torque_Nm=tau.tolist(),contact_normal_loads_N=(regular@solution.x/.025).tolist())
|
||||
rows=[];forces=[];contacts=[];relative=[];torques=[];wrist_path=[]
|
||||
initial_rel=np.array([0,0,.3])-handpos
|
||||
feedback_errors=[]
|
||||
for k in range(int(seconds/model.opt.timestep)):
|
||||
t=k*model.opt.timestep
|
||||
for n,(parent,x,y) in cal.curves.items():
|
||||
active_angle=data.qpos[addr[parent]];passive,slope=cal.passive_value(n,active_angle)
|
||||
eid=mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_EQUALITY,'lut_'+n)
|
||||
model.eq_data[eid,:5]=[passive-slope*active_angle,slope,0,0,0]
|
||||
displacement=np.zeros(3)
|
||||
if lift:
|
||||
phase=np.clip((t-.5)/1.,0,1);smooth=phase*phase*(3-2*phase)
|
||||
displacement=np.array([.02*smooth,0,.10*smooth])
|
||||
data.mocap_pos[0]=handpos+displacement
|
||||
if feedback and k%10==0:
|
||||
phase=np.clip((t-.5),0,1)
|
||||
target_velocity=np.array([.02,0,.10])*6*phase*(1-phase) if lift else np.zeros(3)
|
||||
target_center=np.array([0,0,.3663])+displacement
|
||||
newtau,residual=feedback_torque(model,data,cal,addr,objects,object_body,target_center,target_velocity)
|
||||
if newtau is not None:tau=newtau;feedback_errors.append(residual)
|
||||
if force_control:
|
||||
if not force_preload:raise ValueError('force-control requires force-preload')
|
||||
angles=np.array([data.qpos[addr[n]] for n in cal.active])
|
||||
data.ctrl[:]=angles+(tau+1.0*(reference_controls-angles))/kp
|
||||
mujoco.mj_step(model,data)
|
||||
if not np.isfinite(data.qpos).all():raise RuntimeError('Nonfinite simulation')
|
||||
if k%33==0:
|
||||
rows.append(data.qpos.copy());wrist_path.append(np.r_[data.mocap_pos[0],data.mocap_quat[0]])
|
||||
selected=[i for i,c in enumerate(data.contact) if any(g in objects for g in c.geom)]
|
||||
normalforce=0
|
||||
for ci in selected:
|
||||
f=np.zeros(6);mujoco.mj_contactForce(model,data,ci,f);normalforce+=max(0,f[0])
|
||||
forces.append(normalforce);contacts.append(len(selected));torques.append(data.actuator_force.copy())
|
||||
relative.append(np.linalg.norm(data.qpos[obj:obj+3]-data.mocap_pos[0]-initial_rel))
|
||||
rows=np.array(rows);relative=np.array(relative)
|
||||
report=dict(seconds=seconds,lift=lift,kp=kp,preload_rad=preload,force_preload=preload_report,force_control=force_control,feedback=feedback,
|
||||
feedback_wrench_residual_mean=float(np.mean(feedback_errors)) if feedback_errors else None,
|
||||
max_relative_drift_m=float(relative.max()),final_relative_drift_m=float(relative[-1]),
|
||||
object_vertical_displacement_m=float(data.qpos[obj+2]-.3),
|
||||
contact_frames=int(np.sum(np.array(contacts)>0)),total_frames=len(rows),
|
||||
normal_force_mean_N=float(np.mean(forces)),peak_actuator_torque_Nm=float(np.max(np.abs(torques))),
|
||||
retained=bool(relative.max()<.02),
|
||||
mass_kg=float(model.body_mass[object_body]),friction=.6,
|
||||
boundaries=['Prescribed moving wrist, no arm dynamics','16 force-limited active joints; 5 nonlinear LUT equality constraints','No object weld or position reset','Object registration is estimated from hand pose, not tri-camera object tracking'])
|
||||
np.savez_compressed(OUT/'physics_motion.npz',qpos=rows,wrist=np.array(wrist_path),fps=1/(33*model.opt.timestep),
|
||||
object_qpos_address=obj,contact_count=contacts,normal_force=forces,relative_drift=relative,actuator_force=torques,
|
||||
source_to_upright_rotation=rotation,controls=controls)
|
||||
(OUT/'physics_validation.json').write_text(json.dumps(report,indent=2));print(json.dumps(report,indent=2),flush=True)
|
||||
return report
|
||||
|
||||
if __name__=='__main__':
|
||||
p=argparse.ArgumentParser();p.add_argument('--seconds',type=float,default=3);p.add_argument('--lift',action='store_true');p.add_argument('--kp',type=float,default=30);p.add_argument('--preload',type=float,default=0);p.add_argument('--force-preload',action='store_true');p.add_argument('--force-control',action='store_true');p.add_argument('--feedback',action='store_true');a=p.parse_args()
|
||||
run(a.seconds,a.lift,a.kp,a.preload,a.force_preload,a.force_control,a.feedback)
|
||||
@@ -0,0 +1,73 @@
|
||||
"""Full-length contact registration; object poses are inferred targets, not physics."""
|
||||
import json
|
||||
from pathlib import Path
|
||||
import numpy as np
|
||||
from scipy.spatial.transform import Rotation,Slerp
|
||||
from l20_contact import ContactRefiner
|
||||
from l20_calibrated import ROOT
|
||||
|
||||
BASE=ROOT/'output/l20_2047635068';OUT=BASE/'bottle_grasp'
|
||||
def main():
|
||||
motion=np.load(BASE/'motion.npz');spec=json.loads((OUT/'object_spec.json').read_text())
|
||||
solver=ContactRefiner(spec,OUT)
|
||||
count=len(motion['qpos']);times=np.arange(count)
|
||||
rotation=Rotation.from_quat(motion['wrist_quat_wxyz'][:,[1,2,3,0]])
|
||||
reference=666
|
||||
# Infer a constant bottle-to-palm registration over the holding phase.
|
||||
rel_pos=rotation[reference].inv().apply(np.array(spec['position'])-motion['wrist_pos'][reference])
|
||||
rel_rot=rotation[reference].inv()*Rotation.from_quat(np.array(spec['quat_wxyz'])[[1,2,3,0]])
|
||||
bottle_pos=rotation.apply(np.tile(rel_pos,(count,1)))+motion['wrist_pos']
|
||||
bottle_rot=rotation*rel_rot
|
||||
# Contact onset bracket comes from visual inspection: open at 100, grasped by 200.
|
||||
onset=160
|
||||
bottle_pos[:onset]=bottle_pos[onset]
|
||||
bottle_quat=bottle_rot.as_quat()[:,[3,0,1,2]]
|
||||
bottle_quat[:onset]=bottle_quat[onset]
|
||||
samples=np.unique(np.r_[np.arange(onset,count,8),count-1])
|
||||
fits=[];reports=[]
|
||||
for t in samples:
|
||||
active=np.array([motion['qpos'][t,list(motion['joint_names']).index(n)] for n in solver.cal.active])
|
||||
x,q,report=solver.solve(active,motion['wrist_pos'][t],motion['wrist_quat_wxyz'][t],bottle_pos[t],bottle_quat[t])
|
||||
fits.append(x);reports.append(report)
|
||||
if len(fits)%20==0:print(f'contact keyframes {len(fits)}/{len(samples)}',flush=True)
|
||||
fits=np.array(fits)
|
||||
interp=np.stack([np.interp(times,samples,fits[:,j]) for j in range(22)],axis=1)
|
||||
source_active=np.stack([motion['qpos'][:,list(motion['joint_names']).index(n)] for n in solver.cal.active],axis=1)
|
||||
# Blend acquisition only; preserve pre-contact source motion.
|
||||
weight=np.clip((times-(onset-30))/30,0,1);weight=weight*weight*(3-2*weight)
|
||||
interp[:,:16]=source_active*(1-weight[:,None])+interp[:,:16]*weight[:,None]
|
||||
interp[:,16:]*=weight[:,None]
|
||||
repairs=[]
|
||||
for t in range(onset,count):
|
||||
solver.set_pose(interp[t],motion['wrist_pos'][t],motion['wrist_quat_wxyz'][t],bottle_pos[t],bottle_quat[t])
|
||||
if -solver.distances().min()>.0015:
|
||||
x,_,r=solver.solve(source_active[t],motion['wrist_pos'][t],motion['wrist_quat_wxyz'][t],bottle_pos[t],bottle_quat[t],initial=interp[t])
|
||||
if r['max_penetration_after_mm']>1.5 and t>onset:
|
||||
alternative,_,retry=solver.solve(source_active[t],motion['wrist_pos'][t],motion['wrist_quat_wxyz'][t],bottle_pos[t],bottle_quat[t],initial=interp[t-1])
|
||||
if retry['max_penetration_after_mm']<r['max_penetration_after_mm']:x,r=alternative,retry
|
||||
interp[t]=x;repairs.append(dict(frame=t,**r))
|
||||
if t%100==0:print(f'frame collision repair {t}/{count}: {len(repairs)} repaired',flush=True)
|
||||
qpos=[];wp=[];wq=[];penetration=[]
|
||||
for t,x in enumerate(interp):
|
||||
solver.set_pose(x,motion['wrist_pos'][t],motion['wrist_quat_wxyz'][t],bottle_pos[t],bottle_quat[t])
|
||||
values=solver.expand(x[:16]);qpos.append([values[str(n)] for n in motion['joint_names']])
|
||||
wp.append(solver.data.qpos[solver.wrist:solver.wrist+3].copy());wq.append(solver.data.qpos[solver.wrist+3:solver.wrist+7].copy())
|
||||
penetration.append(max(0,-solver.distances().min()))
|
||||
qpos=np.array(qpos);wp=np.array(wp);wq=np.array(wq)
|
||||
commands=np.array([solver.cal.command(dict(zip(motion['joint_names'],q))) for q in qpos])
|
||||
np.savez_compressed(OUT/'registered_motion.npz',qpos=qpos,joint_names=motion['joint_names'],wrist_pos=wp,wrist_quat_wxyz=wq,
|
||||
bottle_pos=bottle_pos,bottle_quat_wxyz=bottle_quat,command_u8=commands,fps=30,
|
||||
keyframes=samples,optimization_variables=interp,penetration_m=penetration,contact_onset_frame=onset)
|
||||
np.savetxt(OUT/'registered_trajectory.csv',np.c_[times/30,wp,wq,qpos,bottle_pos,bottle_quat],delimiter=',',
|
||||
header=','.join(['time_s','wrist_x','wrist_y','wrist_z','wrist_qw','wrist_qx','wrist_qy','wrist_qz']+list(motion['joint_names'])+['bottle_x','bottle_y','bottle_z','bottle_qw','bottle_qx','bottle_qy','bottle_qz']),comments='')
|
||||
report=dict(frames=count,fps=30,contact_keyframes=len(samples),converged_keyframes=sum(r['converged'] for r in reports),
|
||||
mean_keyframe_gap_before_mm=float(np.mean([r['mean_selected_surface_gap_before_mm'] for r in reports])),
|
||||
mean_keyframe_gap_after_mm=float(np.mean([r['mean_selected_surface_gap_after_mm'] for r in reports])),
|
||||
max_holding_penetration_mm=float(np.max(penetration[onset:])*1000),max_all_frame_penetration_mm=float(np.max(penetration)*1000),
|
||||
contact_onset_frame=onset,source='current 2047635068 video + supplied tri-camera bottle mesh',
|
||||
repaired_frames=len(repairs),repair_reports=repairs,
|
||||
alignment='Inferred fixed bottle-to-palm transform during holding, with contact pose corrections; no measured object trajectory or camera extrinsic alignment',
|
||||
boundary='Object poses in this file are prescribed registration targets. Dynamic grasp tests are separate and have not passed.',keyframe_reports=reports)
|
||||
(OUT/'sequence_validation.json').write_text(json.dumps(report,indent=2));print({k:v for k,v in report.items() if not k.endswith('_reports')},flush=True)
|
||||
|
||||
if __name__=='__main__':main()
|
||||
@@ -0,0 +1,104 @@
|
||||
"""L20 assets and JSON-calibrated passive coupling for dex-retargeting."""
|
||||
from pathlib import Path
|
||||
import json
|
||||
import xml.etree.ElementTree as ET
|
||||
import numpy as np
|
||||
from dex_retargeting.kinematics_adaptor import MimicJointKinematicAdaptor
|
||||
import l20_model_source as source
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
source.REPO_ROOT = ROOT / 'third_party/l20_assets'
|
||||
ASSET = source.REPO_ROOT / 'L20/RIGHT'
|
||||
FINGERS = ['thumb','index','middle','ring','pinky']
|
||||
TIPS = np.array([4,8,12,16,20])
|
||||
CHAINS = [(1,2,3,4),(5,6,7,8),(9,10,11,12),(13,14,15,16),(17,18,19,20)]
|
||||
|
||||
def palm_basis(p):
|
||||
z=p[9]-p[0]; z=z/np.linalg.norm(z)
|
||||
y=p[5]-p[17]; y=y-y.dot(z)*z; y=y/np.linalg.norm(y)
|
||||
return np.column_stack([np.cross(y,z),y,z])
|
||||
|
||||
class Calibration:
|
||||
def __init__(self):
|
||||
self.raw=json.loads((ASSET/'g20_right_G20.json').read_text())
|
||||
self.tables=self.raw['joints']
|
||||
urdf=ET.parse(ASSET/'linkerhand_g20_right.urdf').getroot()
|
||||
self.urdf_limits={j.get('name'):np.array([float(j.find('limit').get(k)) for k in ['lower','upper']]) for j in urdf.findall('joint')}
|
||||
self.active=sorted([n for n,x in self.tables.items() if not x.get('passive',False)],key=lambda n:self.tables[n]['motor_index'])
|
||||
self.passive=[n for n in self.tables if n not in self.active]
|
||||
self.values={n:np.array(x['angle_rad'])+x.get('zero_angles',{}).get('urdf_zero_offset_rad',0) for n,x in self.tables.items()}
|
||||
self.valid_commands={}
|
||||
self.bounds={}
|
||||
for n in self.active:
|
||||
motor=self.tables[n]['motor_index']
|
||||
valid=np.ones(256,dtype=bool)
|
||||
for name,x in self.tables.items():
|
||||
if x['motor_index']==motor:
|
||||
low,high=self.urdf_limits[name]
|
||||
valid &= (self.values[name]>=low)&(self.values[name]<=high)
|
||||
commands=np.flatnonzero(valid)
|
||||
assert len(commands)>1 and np.all(np.diff(commands)==1),n
|
||||
self.valid_commands[n]=commands
|
||||
self.bounds[n]=[self.values[n][commands].min(),self.values[n][commands].max()]
|
||||
self.curves={}
|
||||
for n in self.passive:
|
||||
motor=self.tables[n]['motor_index']
|
||||
active=next(a for a in self.active if self.tables[a]['motor_index']==motor)
|
||||
order=np.argsort(self.values[active])
|
||||
x=self.values[active][order]; y=self.values[n][order]
|
||||
x,idx=np.unique(x,return_index=True)
|
||||
self.curves[n]=(active,x,y[idx])
|
||||
|
||||
def passive_value(self,name,q):
|
||||
_,x,y=self.curves[name]
|
||||
i=np.clip(np.searchsorted(x,q,side='right')-1,0,len(x)-2)
|
||||
slope=(y[i+1]-y[i])/(x[i+1]-x[i])
|
||||
return float(np.interp(q,x,y)),float(slope)
|
||||
|
||||
def command(self,q_by_name):
|
||||
cmd=np.array(self.raw['baseline_command_u8'],dtype=np.uint8)
|
||||
for n in self.active:
|
||||
allowed=self.valid_commands[n]
|
||||
cmd[self.tables[n]['motor_index']]=allowed[np.argmin(np.abs(self.values[n][allowed]-q_by_name[n]))]
|
||||
return cmd
|
||||
|
||||
def decode(self,commands,joint_names):
|
||||
return np.array([self.values[n][commands[self.tables[n]['motor_index']]] for n in joint_names])
|
||||
|
||||
class CalibratedAdaptor(MimicJointKinematicAdaptor):
|
||||
def __init__(self,robot,calibration):
|
||||
c=calibration
|
||||
parents=[c.curves[n][0] for n in c.passive]
|
||||
super().__init__(robot,c.active,parents,c.passive,[1.]*len(parents),[0.]*len(parents))
|
||||
self.calibration=c
|
||||
|
||||
def forward_qpos(self,pin_qpos):
|
||||
for i,n in enumerate(self.calibration.passive):
|
||||
value,slope=self.calibration.passive_value(n,pin_qpos[self.idx_pin2source[i]])
|
||||
pin_qpos[self.idx_pin2mimic[i]]=value
|
||||
self.multipliers[i]=slope
|
||||
return pin_qpos
|
||||
|
||||
def build_assets(out):
|
||||
out.mkdir(parents=True,exist_ok=True)
|
||||
fixed=source.build_model('right',out)
|
||||
mj=ET.parse(fixed)
|
||||
# Preserve both the original URDF and JSON. Adapt only generated files.
|
||||
eq=mj.getroot().find('equality')
|
||||
if eq is not None: mj.getroot().remove(eq)
|
||||
mj.write(fixed)
|
||||
tree=ET.parse(ASSET/'linkerhand_g20_right.urdf')
|
||||
root=tree.getroot()
|
||||
for mesh in root.iter('mesh'): mesh.set('filename',str(ASSET/mesh.get('filename')))
|
||||
for j in root.findall('joint'):
|
||||
mimic=j.find('mimic')
|
||||
if mimic is not None:j.remove(mimic)
|
||||
for body in mj.getroot().iter('body'):
|
||||
for site in body.findall('site'):
|
||||
name=site.get('name')
|
||||
ET.SubElement(root,'link',name=name)
|
||||
j=ET.SubElement(root,'joint',name=name+'_fixed',type='fixed')
|
||||
ET.SubElement(j,'parent',link=body.get('name'));ET.SubElement(j,'child',link=name)
|
||||
ET.SubElement(j,'origin',xyz=site.get('pos'),rpy='0 0 0')
|
||||
path=out/'l20_dex.urdf';tree.write(path)
|
||||
return fixed,path
|
||||
@@ -0,0 +1,11 @@
|
||||
"""Local correction for dex 0.5.0 PositionOptimizer's inconsistent loss/gradient."""
|
||||
import numpy as np
|
||||
from dex_retargeting.optimizer import PositionOptimizer
|
||||
|
||||
class ConsistentPositionOptimizer(PositionOptimizer):
|
||||
def get_objective_function(self,target_pos,fixed_qpos,last_qpos):
|
||||
original=super().get_objective_function(target_pos,fixed_qpos,last_qpos)
|
||||
def objective(x,grad):
|
||||
# Upstream already includes this term's gradient, but omits its value.
|
||||
return original(x,grad)+self.norm_delta*float(np.sum((x-last_qpos)**2))
|
||||
return objective
|
||||
@@ -0,0 +1,176 @@
|
||||
"""Contact-aware L20 pose refinement with a required, explicit object specification.
|
||||
|
||||
Optimizes calibrated finger coordinates and all six wrist correction coordinates.
|
||||
The --fixture path is an analytic sphere test, NOT the object in the source video.
|
||||
"""
|
||||
import os
|
||||
os.environ.setdefault('MUJOCO_GL','osmesa')
|
||||
import argparse,json
|
||||
from pathlib import Path
|
||||
import xml.etree.ElementTree as ET
|
||||
import copy
|
||||
import numpy as np
|
||||
import mujoco
|
||||
from scipy.optimize import least_squares
|
||||
from scipy.spatial.transform import Rotation
|
||||
from l20_calibrated import ROOT,Calibration,FINGERS,TIPS
|
||||
|
||||
OUT=ROOT/'output/l20_2047635068'
|
||||
|
||||
def scene_for_object(spec,out):
|
||||
tree=ET.parse(OUT/'l20_moving.xml');root=tree.getroot()
|
||||
root.find('option').set('gravity','0 0 -9.81')
|
||||
root.find('option').set('timestep','0.001')
|
||||
root.find('option').set('cone','elliptic')
|
||||
# Resolve all mesh references relative to the existing scene.
|
||||
for mesh in root.findall('asset/mesh'):
|
||||
mesh.set('file',str((OUT/mesh.get('file')).resolve()))
|
||||
hand=root.find("worldbody/body[@name='hand_base_link']")
|
||||
for i,geom in enumerate(hand.iter('geom')):
|
||||
geom.set('name',f'hand_collision_{i}')
|
||||
geom.set('contype','1');geom.set('conaffinity','2')
|
||||
geom.set('friction',f"{spec.get('friction',.8)} .005 .0001")
|
||||
geom.set('condim','4')
|
||||
obj=ET.SubElement(root.find('worldbody'),'body',name='grasp_object',pos='0 0 0')
|
||||
ET.SubElement(obj,'freejoint',name='object_free')
|
||||
if spec['type']=='mjcf':
|
||||
source=Path(spec['mjcf_path']).resolve();bottle=ET.parse(source).getroot()
|
||||
for asset in bottle.find('asset'):
|
||||
item=copy.deepcopy(asset)
|
||||
if item.get('file'):item.set('file',str((source.parent/item.get('file')).resolve()))
|
||||
root.find('asset').append(item)
|
||||
for item in bottle.find('worldbody/body'):
|
||||
if item.tag=='freejoint':continue
|
||||
child=copy.deepcopy(item)
|
||||
if child.tag=='geom':
|
||||
child.set('name','bottle_part_'+str(len(obj)))
|
||||
if child.get('contype','1')!='0':
|
||||
child.set('contype','2');child.set('conaffinity','1')
|
||||
child.set('friction',f"{spec.get('friction',.6)} .005 .001")
|
||||
child.set('solref','.004 1');child.set('solimp','.95 .99 .001')
|
||||
obj.append(child)
|
||||
path=out/'contact_scene.xml';tree.write(path)
|
||||
return path
|
||||
attrs=dict(name='object_geom',type=spec['type'],mass=str(spec.get('mass_kg',.1)),
|
||||
contype='2',conaffinity='1',friction=f"{spec.get('friction',.8)} .005 .0001",
|
||||
condim='4',rgba='.9 .65 .2 1',solref='.005 1')
|
||||
if spec['type']=='mesh':
|
||||
mesh_path=Path(spec['mesh_path']).resolve()
|
||||
ET.SubElement(root.find('asset'),'mesh',name='object_mesh',file=str(mesh_path),
|
||||
scale=' '.join(map(str,spec.get('mesh_scale',[1,1,1]))))
|
||||
attrs['mesh']='object_mesh'
|
||||
else:attrs['size']=' '.join(map(str,spec['size']))
|
||||
ET.SubElement(obj,'geom',**attrs)
|
||||
path=out/'contact_scene.xml';tree.write(path)
|
||||
return path
|
||||
|
||||
class ContactRefiner:
|
||||
def __init__(self,spec,out):
|
||||
self.spec=spec;self.cal=Calibration()
|
||||
self.model=mujoco.MjModel.from_xml_path(str(scene_for_object(spec,out)))
|
||||
self.data=mujoco.MjData(self.model)
|
||||
self.names=[n for n in self.cal.tables]
|
||||
self.addr={n:self.model.jnt_qposadr[mujoco.mj_name2id(self.model,mujoco.mjtObj.mjOBJ_JOINT,n)] for n in self.names}
|
||||
self.wrist=self.model.jnt_qposadr[mujoco.mj_name2id(self.model,mujoco.mjtObj.mjOBJ_JOINT,'wrist_free')]
|
||||
self.obj=self.model.jnt_qposadr[mujoco.mj_name2id(self.model,mujoco.mjtObj.mjOBJ_JOINT,'object_free')]
|
||||
self.object_geoms=np.flatnonzero(self.model.geom_contype==2)
|
||||
self.object_geom=int(self.object_geoms[0])
|
||||
self.hand_geoms=np.flatnonzero(self.model.geom_contype==1)
|
||||
self.tip_geoms=[]
|
||||
contact_links=spec.get('contact_links',[f+'_distal' for f in spec.get('contact_fingers',['thumb','index','middle'])])
|
||||
for link in contact_links:
|
||||
bid=mujoco.mj_name2id(self.model,mujoco.mjtObj.mjOBJ_BODY,link)
|
||||
self.tip_geoms.append(next(g for g in self.hand_geoms if self.model.geom_bodyid[g]==bid))
|
||||
self.lower=np.array([self.cal.bounds[n][0] for n in self.cal.active]+[-.06]*3+[-.5]*3)
|
||||
self.upper=np.array([self.cal.bounds[n][1] for n in self.cal.active]+[.06]*3+[.5]*3)
|
||||
|
||||
def expand(self,active):
|
||||
q=dict(zip(self.cal.active,active))
|
||||
for n in self.cal.passive:q[n]=self.cal.passive_value(n,q[self.cal.curves[n][0]])[0]
|
||||
return q
|
||||
|
||||
def set_pose(self,x,wrist_pos,wrist_quat,obj_pos,obj_quat):
|
||||
q=self.expand(x[:16])
|
||||
for n,v in q.items():self.data.qpos[self.addr[n]]=v
|
||||
self.data.qpos[self.wrist:self.wrist+3]=wrist_pos+x[16:19]
|
||||
rotation=Rotation.from_quat(wrist_quat[[1,2,3,0]])*Rotation.from_rotvec(x[19:22])
|
||||
self.data.qpos[self.wrist+3:self.wrist+7]=rotation.as_quat()[[3,0,1,2]]
|
||||
self.data.qpos[self.obj:self.obj+3]=obj_pos
|
||||
self.data.qpos[self.obj+3:self.obj+7]=obj_quat
|
||||
mujoco.mj_kinematics(self.model,self.data)
|
||||
|
||||
def distances(self):
|
||||
return np.array([min(mujoco.mj_geomDistance(self.model,self.data,int(g),int(o),.5,None) for o in self.object_geoms) for g in self.hand_geoms])
|
||||
|
||||
def solve(self,reference,wrist_pos,wrist_quat,obj_pos,obj_quat,initial=None):
|
||||
initial=np.r_[reference,np.zeros(6)] if initial is None else np.asarray(initial).copy()
|
||||
tip_index=[list(self.hand_geoms).index(g) for g in self.tip_geoms]
|
||||
def residual(x):
|
||||
self.set_pose(x,wrist_pos,wrist_quat,obj_pos,obj_quat)
|
||||
dist=self.distances()
|
||||
compression=self.spec.get('compression_m',.0002)
|
||||
return np.r_[(dist[tip_index]+compression)/.004,
|
||||
np.maximum(-dist-max(.0003,compression),0)/.0015,
|
||||
.15*(x[:16]-reference),.2*x[16:19]/.03,.2*x[19:22]/.25]
|
||||
before=residual(initial);before_distance=self.distances()
|
||||
result=least_squares(residual,np.clip(initial,self.lower+1e-9,self.upper-1e-9),
|
||||
bounds=(self.lower,self.upper),diff_step=1e-4,max_nfev=100,ftol=1e-5,xtol=1e-5,gtol=1e-5)
|
||||
after=residual(result.x);after_distance=self.distances()
|
||||
report=dict(converged=bool(result.success),nfev=result.nfev,
|
||||
cost_before=float(before@before),cost_after=float(after@after),
|
||||
mean_selected_surface_gap_before_mm=float(np.abs(before_distance[tip_index]).mean()*1000),
|
||||
mean_selected_surface_gap_after_mm=float(np.abs(after_distance[tip_index]).mean()*1000),
|
||||
max_penetration_before_mm=float(max(0,-before_distance.min())*1000),
|
||||
max_penetration_after_mm=float(max(0,-after_distance.min())*1000),
|
||||
wrist_correction_m=result.x[16:19].tolist(),wrist_correction_rotvec=result.x[19:22].tolist())
|
||||
return result.x,self.data.qpos.copy(),report
|
||||
|
||||
def dynamic_hold(self,qpos,seconds=2.):
|
||||
# A diagnostic with prescribed hand coordinates and a freely simulated object.
|
||||
# It does not validate actuator effort, compliant tracking, or force closure.
|
||||
mujoco.mj_resetData(self.model,self.data);self.data.qpos[:]=qpos
|
||||
object_initial=qpos[self.obj:self.obj+3].copy()
|
||||
hand_addr=[self.addr[n] for n in self.names]+list(range(self.wrist,self.wrist+7))
|
||||
object_dof=self.model.jnt_dofadr[mujoco.mj_name2id(self.model,mujoco.mjtObj.mjOBJ_JOINT,'object_free')]
|
||||
controlled_dof=[i for i in range(self.model.nv) if not object_dof<=i<object_dof+6]
|
||||
frames=[];contacts=[]
|
||||
steps=int(seconds/self.model.opt.timestep)
|
||||
for k in range(steps):
|
||||
self.data.qpos[hand_addr]=qpos[hand_addr];self.data.qvel[controlled_dof]=0
|
||||
mujoco.mj_step(self.model,self.data)
|
||||
if not np.isfinite(self.data.qpos).all():raise RuntimeError('Non-finite physics state')
|
||||
if k%33==0:
|
||||
frames.append(self.data.qpos.copy())
|
||||
contacts.append(sum(any(g in self.object_geoms for g in c.geom) for c in self.data.contact))
|
||||
displacement=float(np.linalg.norm(self.data.qpos[self.obj:self.obj+3]-object_initial))
|
||||
return np.array(frames),dict(seconds=seconds,object_displacement_m=displacement,
|
||||
contact_frames=int(np.sum(np.array(contacts)>0)),total_frames=len(frames),
|
||||
retained_within_2cm=displacement<.02,
|
||||
boundary='Prescribed hand, free object under gravity. No object weld or pose resets. Hand actuator/arm dynamics not validated.')
|
||||
|
||||
def main():
|
||||
parser=argparse.ArgumentParser();parser.add_argument('--spec',type=Path);parser.add_argument('--fixture',action='store_true')
|
||||
parser.add_argument('--frame',type=int,default=666);args=parser.parse_args()
|
||||
motion=np.load(OUT/'motion.npz');t=args.frame
|
||||
if not args.fixture and args.spec is None:parser.error('Provide an object specification, or explicitly select the synthetic --fixture test.')
|
||||
out=OUT/('contact_fixture' if args.fixture else 'contact_object');out.mkdir(exist_ok=True)
|
||||
if args.fixture:
|
||||
pts=motion['actual'][t]
|
||||
center=(pts[4]+pts[12])*.5
|
||||
radius=np.linalg.norm(pts[4]-pts[12])*.42
|
||||
r=Rotation.from_quat(motion['wrist_quat_wxyz'][t,[1,2,3,0]]).as_matrix()
|
||||
spec=dict(type='sphere',size=[float(radius)],mass_kg=.05,friction=.8,
|
||||
position=(r@center+motion['wrist_pos'][t]).tolist(),quat_wxyz=[1,0,0,0],
|
||||
contact_fingers=['thumb','index','middle'],purpose='Synthetic interface/physics fixture; not source-video reconstruction')
|
||||
else:spec=json.loads(args.spec.read_text())
|
||||
(out/'object_spec.json').write_text(json.dumps(spec,indent=2))
|
||||
solver=ContactRefiner(spec,out)
|
||||
active=np.array([motion['qpos'][t,list(motion['joint_names']).index(n)] for n in solver.cal.active])
|
||||
x,q,report=solver.solve(active,motion['wrist_pos'][t],motion['wrist_quat_wxyz'][t],np.array(spec['position']),np.array(spec['quat_wxyz']))
|
||||
path,dynamic=solver.dynamic_hold(q)
|
||||
report['dynamic_hold']=dynamic;report['frame']=t;report['fixture']=args.fixture
|
||||
np.savez_compressed(out/'contact_result.npz',optimization_variables=x,qpos=q,dynamic_qpos=path)
|
||||
(out/'validation.json').write_text(json.dumps(report,indent=2))
|
||||
print(json.dumps(report,indent=2),flush=True)
|
||||
|
||||
if __name__=='__main__':main()
|
||||
@@ -0,0 +1,43 @@
|
||||
"""Isolate contact geometry using prescribed link poses and a free bottle."""
|
||||
from pathlib import Path
|
||||
import xml.etree.ElementTree as E
|
||||
import copy,json
|
||||
import numpy as np
|
||||
import mujoco
|
||||
|
||||
p=Path(__file__).resolve().parents[1]/'output/l20_2047635068/bottle_grasp'
|
||||
tree=E.parse(p/'physics_scene.xml');root=tree.getroot();world=root.find('worldbody')
|
||||
hand=world.find("body[@name='hand_base_link']")
|
||||
original=mujoco.MjModel.from_xml_path(str(p/'physics_scene.xml'));d=mujoco.MjData(original)
|
||||
a=np.load(p/'physics_motion.npz');d.qpos[:]=a['qpos'][0]
|
||||
d.mocap_pos[0]=a['wrist'][0,:3];d.mocap_quat[0]=a['wrist'][0,3:];mujoco.mj_forward(original,d)
|
||||
poses=[]
|
||||
for b in hand.iter('body'):
|
||||
bid=original.body(b.get('name')).id
|
||||
new=E.SubElement(world,'body',name=b.get('name'),mocap='true')
|
||||
poses.append((d.xpos[bid].copy(),d.xquat[bid].copy()))
|
||||
for g in b.findall('geom'):new.append(copy.deepcopy(g))
|
||||
world.remove(hand)
|
||||
for tag in ['equality','actuator']:
|
||||
e=root.find(tag)
|
||||
if e is not None:root.remove(e)
|
||||
path=p/'kinematic_grasp_scene.xml';tree.write(path)
|
||||
m=mujoco.MjModel.from_xml_path(str(path));s=mujoco.MjData(m)
|
||||
oi=original.jnt_qposadr[original.joint('object_free').id]
|
||||
s.qpos[:]=d.qpos[oi:oi+7]
|
||||
pos=np.array([x[0] for x in poses]);quat=np.array([x[1] for x in poses])
|
||||
s.mocap_pos[:]=pos;s.mocap_quat[:]=quat
|
||||
start=s.qpos[:3].copy();rows=[];mp=[];counts=[];drift=[]
|
||||
for k in range(3000):
|
||||
t=k*.001;f=np.clip(t-.5,0,1);lift=.1*f*f*(3-2*f)
|
||||
s.mocap_pos[:]=pos+[0,0,lift];mujoco.mj_step(m,s)
|
||||
if k%33==0:
|
||||
rows.append(s.qpos.copy());mp.append(s.mocap_pos.copy());counts.append(s.ncon)
|
||||
drift.append(np.linalg.norm(s.qpos[:3]-start-[0,0,lift]))
|
||||
np.savez_compressed(p/'kinematic_grasp_motion.npz',qpos=rows,mocap_pos=mp,mocap_quat=quat,fps=1/.033)
|
||||
r={'free_object':True,'hand_prescribed_kinematics':True,
|
||||
'final_object_displacement':(s.qpos[:3]-start).tolist(),'expected_lift_m':.1,
|
||||
'relative_drift_m':float(np.linalg.norm(s.qpos[:3]-start-[0,0,.1])),
|
||||
'maximum_drift_m':float(max(drift)),'retained':bool(max(drift)<.02),
|
||||
'contact_frames':int(np.sum(np.array(counts)>0))}
|
||||
(p/'kinematic_grasp_validation.json').write_text(json.dumps(r,indent=2));print(json.dumps(r,indent=2))
|
||||
@@ -0,0 +1,276 @@
|
||||
"""Build an L20 visual MJCF and provide coupled 21-landmark kinematics.
|
||||
|
||||
The supplied L20 URDFs have 21 revolute joints and five passive mimic joints.
|
||||
The accompanying G20 calibration also has only 16 active motor channels; its
|
||||
20-byte command slots 11--14 are unassigned. This module preserves the URDF
|
||||
linear mimic model. Calibration tables are metadata, not a hardware command
|
||||
conversion: their nonlinear passive curves differ from the URDF mimic ratios.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections import defaultdict
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import xml.etree.ElementTree as ET
|
||||
|
||||
import mujoco
|
||||
import numpy as np
|
||||
from scipy.spatial.transform import Rotation
|
||||
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[1]
|
||||
FINGERS = ("thumb", "index", "middle", "ring", "pinky")
|
||||
LANDMARK_NAMES = (
|
||||
"wrist", "thumb_cmc", "thumb_mcp", "thumb_ip", "thumb_tip",
|
||||
"index_mcp", "index_pip", "index_dip", "index_tip",
|
||||
"middle_mcp", "middle_pip", "middle_dip", "middle_tip",
|
||||
"ring_mcp", "ring_pip", "ring_dip", "ring_tip",
|
||||
"pinky_mcp", "pinky_pip", "pinky_dip", "pinky_tip",
|
||||
)
|
||||
|
||||
|
||||
def source_urdf(side: str = "right") -> Path:
|
||||
side = side.lower()
|
||||
if side not in ("right", "left"):
|
||||
raise ValueError(f"Expected right or left hand, got {side!r}")
|
||||
return REPO_ROOT / "L20" / side.upper() / f"linkerhand_g20_{side}.urdf"
|
||||
|
||||
|
||||
def _fmt(values) -> str:
|
||||
return " ".join(f"{float(value):.12g}" for value in values)
|
||||
|
||||
|
||||
def _origin(element: ET.Element | None) -> dict[str, str]:
|
||||
if element is None:
|
||||
return {"pos": "0 0 0", "quat": "1 0 0 0"}
|
||||
rpy = np.fromstring(element.get("rpy", "0 0 0"), sep=" ")
|
||||
quat = Rotation.from_euler("xyz", rpy).as_quat()
|
||||
return {"pos": element.get("xyz", "0 0 0"), "quat": _fmt(quat[[3, 0, 1, 2]])}
|
||||
|
||||
|
||||
def _mesh_tip(path: Path) -> np.ndarray:
|
||||
"""Approximate distal surface landmark from the terminal 2 mm STL cap.
|
||||
|
||||
Distal link mesh coordinates have their fingertip in local +Z. Averaging
|
||||
unique cap vertices avoids triangle tessellation multiplicity bias. This
|
||||
is a geometric approximation, not a measured tactile contact center.
|
||||
"""
|
||||
dtype = np.dtype([("normal", "<f4", (3,)), ("vertices", "<f4", (3, 3)), ("attr", "<u2")])
|
||||
with path.open("rb") as stream:
|
||||
stream.seek(80)
|
||||
count = int(np.fromfile(stream, dtype="<u4", count=1)[0])
|
||||
triangles = np.fromfile(stream, dtype=dtype, count=count)
|
||||
if len(triangles) != count:
|
||||
raise ValueError(f"Invalid binary STL: {path}")
|
||||
vertices = np.unique(triangles["vertices"].reshape(-1, 3), axis=0)
|
||||
return vertices[vertices[:, 2] >= vertices[:, 2].max() - .002].mean(axis=0)
|
||||
|
||||
|
||||
def build_model(side: str = "right", output_dir: str | Path | None = None) -> Path:
|
||||
"""Create a fixed-wrist, visual-only MJCF without modifying source assets.
|
||||
|
||||
Mesh references are relative to the generated XML. The output therefore
|
||||
remains relocatable with its enclosing repository, but is not an asset
|
||||
bundle independent of the repository. All URDF joint limits and all five
|
||||
mimic equalities are preserved. Contacts are disabled for kinematic replay.
|
||||
"""
|
||||
urdf_path = source_urdf(side)
|
||||
side = side.lower()
|
||||
output_dir = Path(output_dir or (REPO_ROOT / "retarget_l20" / "models")).resolve()
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
tree = ET.parse(urdf_path).getroot()
|
||||
links = {link.get("name"): link for link in tree.findall("link")}
|
||||
joints = tree.findall("joint")
|
||||
children = defaultdict(list)
|
||||
for joint in joints:
|
||||
children[joint.find("parent").get("link")].append(joint)
|
||||
|
||||
root = ET.Element("mujoco", model=f"l20_{side}_retarget")
|
||||
root.append(ET.Comment("Kinematic preview; meshes and joint limits from the original L20 URDF."))
|
||||
ET.SubElement(root, "compiler", angle="radian", balanceinertia="true", autolimits="true", strippath="false")
|
||||
ET.SubElement(root, "option", timestep="0.002", gravity="0 0 0")
|
||||
ET.SubElement(root, "statistic", center="0 0 0.135", extent="0.32")
|
||||
visual = ET.SubElement(root, "visual")
|
||||
ET.SubElement(visual, "global", offwidth="1600", offheight="1200", azimuth="20", elevation="-8")
|
||||
ET.SubElement(visual, "quality", shadowsize="2048", offsamples="4")
|
||||
ET.SubElement(visual, "headlight", ambient="0.22 0.22 0.26", diffuse="0.35 0.35 0.35", specular="0.15 0.15 0.15")
|
||||
ET.SubElement(visual, "rgba", haze="0.08 0.10 0.14 1")
|
||||
default = ET.SubElement(root, "default")
|
||||
ET.SubElement(default, "joint", damping="0.1", armature="0.00001")
|
||||
ET.SubElement(default, "geom", contype="0", conaffinity="0", group="1")
|
||||
ET.SubElement(default, "site", size="0.0018", rgba="0.1 0.9 0.8 0", group="3")
|
||||
asset = ET.SubElement(root, "asset")
|
||||
ET.SubElement(asset, "material", name="shell", rgba="0.76 0.82 0.90 1", specular="0.4", shininess="0.5")
|
||||
ET.SubElement(asset, "material", name="joints", rgba="0.15 0.24 0.34 1", specular="0.35", shininess="0.4")
|
||||
ET.SubElement(asset, "material", name="tips", rgba="0.10 0.68 0.72 1", specular="0.3", shininess="0.3")
|
||||
ET.SubElement(asset, "material", name="palm", rgba="0.36 0.44 0.55 1", specular="0.45", shininess="0.5")
|
||||
for name, link in links.items():
|
||||
for index, visual_link in enumerate(link.findall("visual")):
|
||||
mesh = visual_link.find("geometry/mesh")
|
||||
if mesh is None:
|
||||
raise ValueError(f"Expected visual mesh for {name}")
|
||||
path = (urdf_path.parent / mesh.get("filename")).resolve()
|
||||
attributes = {"name": f"{name}_mesh_{index}", "file": os.path.relpath(path, output_dir)}
|
||||
if mesh.get("scale"):
|
||||
attributes["scale"] = mesh.get("scale")
|
||||
ET.SubElement(asset, "mesh", **attributes)
|
||||
|
||||
world = ET.SubElement(root, "worldbody")
|
||||
ET.SubElement(world, "light", name="key", pos="0.4 0.25 0.5", dir="-0.7 -0.3 -0.65", diffuse="0.55 0.6 0.65", castshadow="true")
|
||||
ET.SubElement(world, "light", name="fill", pos="-0.3 -0.25 0.3", dir="0.8 0.4 -0.3", diffuse="0.25 0.3 0.35", castshadow="false")
|
||||
ET.SubElement(world, "camera", name="palm", pos="0.56 0 0.14", xyaxes="0 1 0 0 0 1", fovy="36")
|
||||
ET.SubElement(world, "geom", name="floor", type="plane", pos="0 0 -0.003", size="0.7 0.7 0.01", rgba="0.07 0.09 0.13 1", group="0")
|
||||
body_elements = {}
|
||||
|
||||
def add_link(parent: ET.Element, name: str, joint: ET.Element | None = None):
|
||||
body = ET.SubElement(parent, "body", name=name, **_origin(None if joint is None else joint.find("origin")))
|
||||
body_elements[name] = body
|
||||
link = links[name]
|
||||
inertial = link.find("inertial")
|
||||
if inertial is not None:
|
||||
inertia = inertial.find("inertia")
|
||||
# MuJoCo fullinertia uses the body-aligned frame. All current L20
|
||||
# inertial origins have zero orientation; rotate defensively.
|
||||
matrix = np.array([[float(inertia.get("ixx")), float(inertia.get("ixy")), float(inertia.get("ixz"))],
|
||||
[float(inertia.get("ixy")), float(inertia.get("iyy")), float(inertia.get("iyz"))],
|
||||
[float(inertia.get("ixz")), float(inertia.get("iyz")), float(inertia.get("izz"))]])
|
||||
attributes = _origin(inertial.find("origin"))
|
||||
quat = np.fromstring(attributes.pop("quat"), sep=" ")
|
||||
rotation = Rotation.from_quat(quat[[1, 2, 3, 0]]).as_matrix()
|
||||
matrix = rotation @ matrix @ rotation.T
|
||||
ET.SubElement(body, "inertial", **attributes, mass=inertial.find("mass").get("value"),
|
||||
fullinertia=_fmt([matrix[0, 0], matrix[1, 1], matrix[2, 2], matrix[0, 1], matrix[0, 2], matrix[1, 2]]))
|
||||
if joint is not None and joint.get("type") != "fixed":
|
||||
if joint.get("type") != "revolute":
|
||||
raise ValueError(f"Unsupported joint type {joint.get('type')}")
|
||||
limit = joint.find("limit")
|
||||
ET.SubElement(body, "joint", name=joint.get("name"), type="hinge", axis=joint.find("axis").get("xyz"),
|
||||
range=f"{limit.get('lower')} {limit.get('upper')}")
|
||||
material = "palm" if name == "hand_base_link" else "tips" if name.endswith("distal") else "joints" if "metacarpals" in name else "shell"
|
||||
for index, link_visual in enumerate(link.findall("visual")):
|
||||
ET.SubElement(body, "geom", name=f"{name}_visual_{index}", type="mesh", mesh=f"{name}_mesh_{index}", material=material, **_origin(link_visual.find("origin")))
|
||||
for child_joint in children[name]:
|
||||
add_link(body, child_joint.find("child").get("link"), child_joint)
|
||||
|
||||
add_link(world, "hand_base_link")
|
||||
placements = [("hand_base_link", "0 0 0"), ("thumb_metacarpals", "0 0 0"),
|
||||
("thumb_proximal", "0 0 0"), ("thumb_distal", "0 0 0"),
|
||||
("thumb_distal", _fmt(_mesh_tip(urdf_path.parent / "meshes/thumb_distal.STL")))]
|
||||
for finger in FINGERS[1:]:
|
||||
placements.extend([(f"{finger}_proximal", "0 0 0"), (f"{finger}_middle", "0 0 0"),
|
||||
(f"{finger}_distal", "0 0 0"), (f"{finger}_distal", _fmt(_mesh_tip(urdf_path.parent / f"meshes/{finger}_distal.STL")))])
|
||||
for index, (body_name, pos) in enumerate(placements):
|
||||
ET.SubElement(body_elements[body_name], "site", name=f"landmark_{index:02d}", pos=pos)
|
||||
equality = ET.SubElement(root, "equality")
|
||||
for joint in joints:
|
||||
mimic = joint.find("mimic")
|
||||
if mimic is not None:
|
||||
ET.SubElement(equality, "joint", name=f"mimic_{joint.get('name')}", joint1=joint.get("name"), joint2=mimic.get("joint"),
|
||||
polycoef=f"{mimic.get('offset', '0')} {mimic.get('multiplier', '1')} 0 0 0")
|
||||
ET.indent(root, space=" ")
|
||||
output_path = output_dir / f"l20_{side}.xml"
|
||||
ET.ElementTree(root).write(output_path, encoding="utf-8", xml_declaration=True)
|
||||
return output_path
|
||||
|
||||
|
||||
class HandKinematics:
|
||||
"""FK in native URDF coordinates; optimization variables are in radians.
|
||||
|
||||
+Z points from wrist to extended fingers. +X is palm-facing and positive
|
||||
flexion brings fingers toward +X. Radial direction (pinky toward index) is
|
||||
+Y for the right hand and -Y for the left. Wrist pose is fixed at identity.
|
||||
|
||||
`landmarks` and `jacobian` accept independent coordinates (16). `landmarks`
|
||||
additionally accepts full URDF coordinates (21) for playback. `expand`
|
||||
enforces exact linear mimic coupling. Returned arrays are fresh copies.
|
||||
"""
|
||||
|
||||
def __init__(self, model_path: str | Path, side: str = "right"):
|
||||
self.side = side.lower()
|
||||
self.model_path = Path(model_path).resolve()
|
||||
self.model = mujoco.MjModel.from_xml_path(str(self.model_path))
|
||||
self.data = mujoco.MjData(self.model)
|
||||
urdf = ET.parse(source_urdf(side)).getroot()
|
||||
xml_joints = [j for j in urdf.findall("joint") if j.get("type") != "fixed"]
|
||||
by_name = {joint.get("name"): joint for joint in xml_joints}
|
||||
self.joint_names = [mujoco.mj_id2name(self.model, mujoco.mjtObj.mjOBJ_JOINT, j) for j in range(self.model.njnt)]
|
||||
self.independent_joint_names = [name for name in self.joint_names if by_name[name].find("mimic") is None]
|
||||
self.independent_indices = np.array([self.joint_names.index(name) for name in self.independent_joint_names], dtype=int)
|
||||
self.expansion = np.zeros((len(self.joint_names), len(self.independent_joint_names)))
|
||||
self.offset = np.zeros(len(self.joint_names))
|
||||
self.mimic = {}
|
||||
for index, name in enumerate(self.joint_names):
|
||||
mimic = by_name[name].find("mimic")
|
||||
if mimic is None:
|
||||
self.expansion[index, self.independent_joint_names.index(name)] = 1.
|
||||
else:
|
||||
parent = mimic.get("joint")
|
||||
multiplier = float(mimic.get("multiplier", "1"))
|
||||
offset = float(mimic.get("offset", "0"))
|
||||
if parent not in self.independent_joint_names:
|
||||
raise ValueError("Nested mimic coupling is not supported")
|
||||
self.expansion[index, self.independent_joint_names.index(parent)] = multiplier
|
||||
self.offset[index] = offset
|
||||
self.mimic[name] = {"joint": parent, "multiplier": multiplier, "offset": offset}
|
||||
self.full_lower = self.model.jnt_range[:, 0].copy()
|
||||
self.full_upper = self.model.jnt_range[:, 1].copy()
|
||||
self.lower = self.full_lower[self.independent_indices].copy()
|
||||
self.upper = self.full_upper[self.independent_indices].copy()
|
||||
# Shrink active limits as needed so passive limits are respected too.
|
||||
for index, name in enumerate(self.joint_names):
|
||||
if name in self.mimic:
|
||||
column = np.flatnonzero(self.expansion[index])[0]
|
||||
limits = (np.array([self.full_lower[index], self.full_upper[index]]) - self.offset[index]) / self.expansion[index, column]
|
||||
self.lower[column] = max(self.lower[column], min(limits))
|
||||
self.upper[column] = min(self.upper[column], max(limits))
|
||||
self.site_ids = np.array([mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_SITE, f"landmark_{i:02d}") for i in range(21)])
|
||||
if np.any(self.site_ids < 0):
|
||||
raise ValueError("MJCF does not contain all 21 hand landmark sites")
|
||||
self.calibration = json.loads((source_urdf(side).parent / f"g20_{self.side}_G20.json").read_text())
|
||||
self.motor_indices = np.array([self.calibration["joints"][name]["motor_index"] for name in self.independent_joint_names], dtype=int)
|
||||
self.palm_basis = np.array([[0., 0., 1.], [1. if self.side == "right" else -1., 0., 0.], [0., 1., 0.]])
|
||||
self.neutral_landmarks = self.landmarks(np.zeros(len(self.independent_joint_names)))
|
||||
|
||||
def expand(self, q: np.ndarray) -> np.ndarray:
|
||||
q = np.asarray(q, dtype=float)
|
||||
if q.shape[-1] != len(self.independent_joint_names):
|
||||
raise ValueError(f"Expected {len(self.independent_joint_names)} independent coordinates, got {q.shape}")
|
||||
return q @ self.expansion.T + self.offset
|
||||
|
||||
def landmarks(self, q: np.ndarray) -> np.ndarray:
|
||||
q = np.asarray(q, dtype=float)
|
||||
if q.shape == (len(self.independent_joint_names),):
|
||||
q = self.expand(q)
|
||||
if q.shape != (self.model.nq,):
|
||||
raise ValueError(f"Expected independent or full hand qpos, got {q.shape}")
|
||||
self.data.qpos[:] = q
|
||||
mujoco.mj_kinematics(self.model, self.data)
|
||||
return self.data.site_xpos[self.site_ids].copy()
|
||||
|
||||
def jacobian(self, q: np.ndarray) -> np.ndarray:
|
||||
self.landmarks(q)
|
||||
mujoco.mj_comPos(self.model, self.data)
|
||||
jacobian = np.empty((21, 3, len(self.independent_joint_names)))
|
||||
jacp = np.zeros((3, self.model.nv))
|
||||
for index, site_id in enumerate(self.site_ids):
|
||||
mujoco.mj_jacSite(self.model, self.data, jacp, None, int(site_id))
|
||||
jacobian[index] = jacp @ self.expansion
|
||||
return jacobian
|
||||
|
||||
def metadata(self) -> dict:
|
||||
return {
|
||||
"source_urdf": str(source_urdf(self.side).relative_to(REPO_ROOT)),
|
||||
"side": self.side,
|
||||
"joint_names": self.joint_names,
|
||||
"independent_joint_names": self.independent_joint_names,
|
||||
"motor_indices": self.motor_indices.tolist(),
|
||||
"unused_command_indices": [11, 12, 13, 14],
|
||||
"independent_limits_rad": np.column_stack([self.lower, self.upper]).tolist(),
|
||||
"mimic": self.mimic,
|
||||
"landmark_names": list(LANDMARK_NAMES),
|
||||
"landmark_definition": "Joint origins; wrist at base origin; thumb CMC at thumb_metacarpals origin; fingertips approximate distal STL terminal-cap centers.",
|
||||
"model_note": "URDF kinematics and linear mimic coupling; calibration nonlinear passive tables differ. No hardware-equivalence claim or motor command conversion.",
|
||||
"contact_note": "Collision geoms and contact disabled for kinematic visual replay; no object interaction or collision avoidance.",
|
||||
}
|
||||
@@ -0,0 +1,151 @@
|
||||
"""Build a source-only, history-free snapshot for an internal Git platform."""
|
||||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
from pathlib import Path
|
||||
import re
|
||||
import shutil
|
||||
import subprocess
|
||||
import tarfile
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
EXTENSIONS = {'.py','.pyi','.sh','.bash','.md','.rst','.txt','.json','.yaml','.yml',
|
||||
'.toml','.cfg','.ini','.xml','.urdf','.cpp','.cc','.c','.h','.hpp','.cu',
|
||||
'.cuh','.cmake','.in','.lock','.gitignore','.gitmodules','.bib','.css',
|
||||
'.html','.js','.ipynb','.pxd','.pyx','.bat','.ps1','.m','.glsl','.vert','.frag'}
|
||||
SPECIAL = {'LICENSE','LICENCE','COPYING','NOTICE','AUTHORS','Makefile','CMakeLists.txt',
|
||||
'Dockerfile','.gitignore','.gitattributes','environment.yml'}
|
||||
EXCLUDE_PARTS = {'.git','__pycache__','.dynhamr','.venv','venv','outputs','output',
|
||||
'_DATA','checkpoints','weights','.pytest_cache','.mypy_cache'}
|
||||
|
||||
|
||||
def git(repo, *args):
|
||||
return subprocess.check_output(['git','-C',str(repo),*args], stderr=subprocess.DEVNULL).decode().strip()
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument('--output', type=Path, default=ROOT/'dist/hand-motion-pipeline')
|
||||
args = parser.parse_args()
|
||||
dest = args.output.resolve()
|
||||
if dest.exists():
|
||||
raise SystemExit(f'Refusing to overwrite existing release: {dest}')
|
||||
dest.mkdir(parents=True)
|
||||
excluded, repos, portability = [], [], []
|
||||
|
||||
def copy_file(src, relative):
|
||||
reason = None
|
||||
if src.is_symlink():
|
||||
reason = 'symlink: restore dependency explicitly'
|
||||
elif any(p in EXCLUDE_PARTS for p in relative.parts):
|
||||
reason = 'runtime/data/cache directory'
|
||||
elif relative.name == '.gitmodules' or (relative.parts[0] == 'third_party' and relative.name in {'.gitignore','.gitattributes'}):
|
||||
reason = 'vendored snapshot: omit nested Git filters and metadata'
|
||||
elif src.suffix not in EXTENSIONS and src.name not in SPECIAL and not src.name.startswith(('LICENSE','COPYING','NOTICE')):
|
||||
reason = 'non-source file; restore separately when needed'
|
||||
elif src.stat().st_size > 5*1024*1024:
|
||||
reason = 'file exceeds 5 MiB source limit'
|
||||
if reason:
|
||||
excluded.append({'path':relative.as_posix(),'reason':reason})
|
||||
return
|
||||
raw = src.read_bytes()
|
||||
if b'\0' in raw:
|
||||
excluded.append({'path':relative.as_posix(),'reason':'binary contents'})
|
||||
return
|
||||
target = dest/relative
|
||||
target.parent.mkdir(parents=True, exist_ok=True)
|
||||
shutil.copy2(src,target)
|
||||
if src.suffix in {'.py','.sh','.yaml','.yml','.json','.toml'}:
|
||||
for i,line in enumerate(raw.decode('utf-8',errors='replace').splitlines(),1):
|
||||
if re.search(r'/home/|/tmp/|2047635068|15886123',line):
|
||||
portability.append({'path':relative.as_posix(),'line':i,
|
||||
'kind':'machine path or fixed example identifier'})
|
||||
|
||||
def snapshot(repo, relative, ancestors=()):
|
||||
real = repo.resolve()
|
||||
if real in ancestors:
|
||||
raise RuntimeError(f'Recursive repository: {repo}')
|
||||
records = subprocess.check_output(['git','-C',str(repo),'ls-files','--stage','-z']).split(b'\0')
|
||||
repos.append({'path':relative.as_posix(),'commit':git(repo,'rev-parse','HEAD'),
|
||||
'origin':git(repo,'remote','get-url','origin'),
|
||||
'local_changes':git(repo,'status','--porcelain'),
|
||||
'snapshot':'working tree content; tracked files plus source additions'})
|
||||
visited = set()
|
||||
for record in records:
|
||||
if not record:
|
||||
continue
|
||||
metadata,name = record.decode().split('\t',1)
|
||||
visited.add(name)
|
||||
src = repo/name
|
||||
if metadata.startswith('160000'):
|
||||
if src.exists() and (src/'.git').exists():
|
||||
snapshot(src,relative/name,(*ancestors,real))
|
||||
else:
|
||||
excluded.append({'path':(relative/name).as_posix(),'reason':'unavailable submodule; restore pinned dependency', 'commit':metadata.split()[1]})
|
||||
elif src.is_file() or src.is_symlink():
|
||||
copy_file(src,relative/name)
|
||||
additions = subprocess.check_output(['git','-C',str(repo),'ls-files','--others','--exclude-standard','-z']).split(b'\0')
|
||||
for item in additions:
|
||||
if item:
|
||||
name=item.decode()
|
||||
if name not in visited and (repo/name).is_file():
|
||||
copy_file(repo/name,relative/name)
|
||||
|
||||
for folder in ['scripts','docs','model','utils','preprocessing','visualization','configs','requirements']:
|
||||
for src in sorted((ROOT/folder).rglob('*')):
|
||||
if src.is_file():
|
||||
copy_file(src,src.relative_to(ROOT))
|
||||
for name in ['README.md','LICENSE','.gitignore','requirements.txt','setup_env.sh','setup_vipe_env.sh']:
|
||||
copy_file(ROOT/name,Path(name))
|
||||
for name in ['hamer','vipe','Dyn-HaMR','spider']:
|
||||
snapshot(ROOT/'third_party'/name,Path('third_party')/name)
|
||||
repos.insert(0,{'path':'.','commit':git(ROOT,'rev-parse','HEAD'),
|
||||
'origin':git(ROOT,'remote','get-url','origin'),
|
||||
'snapshot':'selected integration source and documentation from current working tree'})
|
||||
# Broad upstream patterns such as data/ would hide vendored Python source.
|
||||
(dest/'.gitignore').write_text(
|
||||
'# Source snapshot: runtime exclusions anchored to repository root.\n'
|
||||
'/output/\n/results/\n/logs/\n/weights/\n/dist/\n/.venv/\n/venv/\n'
|
||||
'/.dex/\n/.spider/\n/.cuda/\n/.hf_cache/\n/.torch_cache/\n/.spider_cache/\n'
|
||||
'/.uv_cache_spider/\n/.env\n__pycache__/\n*.pyc\n')
|
||||
assets=[]
|
||||
for name in ['l20_assets','bottle_model_parametric','dex-assets']:
|
||||
folder=ROOT/'third_party'/name
|
||||
files=[{'path':p.relative_to(ROOT).as_posix(),'bytes':p.stat().st_size,
|
||||
'sha256':hashlib.sha256(p.read_bytes()).hexdigest()}
|
||||
for p in sorted(folder.rglob('*')) if p.is_file() and not p.is_symlink()]
|
||||
assets.append({'path':folder.relative_to(ROOT).as_posix(),'bundled':False,'files':files})
|
||||
manifests={'UPSTREAM_SOURCES.json':repos,'EXTERNAL_ASSETS.json':assets,
|
||||
'EXCLUDED_FILES.json':excluded,'PORTABILITY_REPORT.json':portability}
|
||||
for name,content in manifests.items():
|
||||
(dest/name).write_text(json.dumps(content,indent=2,ensure_ascii=False)+'\n')
|
||||
# Validate local Python syntax without importing scripts (some execute on import).
|
||||
checked=[]
|
||||
for folder in ['scripts','model','utils','preprocessing','visualization']:
|
||||
for p in (dest/folder).rglob('*.py'):
|
||||
compile(p.read_bytes(),str(p),'exec')
|
||||
checked.append(p.relative_to(dest).as_posix())
|
||||
for p in (dest/'scripts').glob('*.sh'):
|
||||
subprocess.run(['bash','-n',str(p)],check=True)
|
||||
allfiles=[p for p in dest.rglob('*') if p.is_file()]
|
||||
assert not any(p.is_symlink() or '.git' in p.relative_to(dest).parts for p in allfiles)
|
||||
report={'local_python_syntax_checked':len(checked),'shell_syntax':'PASS',
|
||||
'source_files':len(allfiles),'bytes':sum(p.stat().st_size for p in allfiles),
|
||||
'portability_findings':len(portability),'excluded_files':len(excluded),
|
||||
'end_to_end_reconstruction_rerun':False,'remote_upload_performed':False}
|
||||
(dest/'RELEASE_VALIDATION.json').write_text(json.dumps(report,indent=2)+'\n')
|
||||
hashes=[f'{hashlib.sha256(p.read_bytes()).hexdigest()} {p.relative_to(dest).as_posix()}'
|
||||
for p in sorted(dest.rglob('*')) if p.is_file()]
|
||||
(dest/'SHA256SUMS').write_text('\n'.join(hashes)+'\n')
|
||||
archive=dest.with_suffix('.tar.gz')
|
||||
with tarfile.open(archive,'w:gz') as tar:
|
||||
tar.add(dest,arcname=dest.name)
|
||||
with tarfile.open(archive) as tar:
|
||||
members=tar.getmembers()
|
||||
assert all(not m.issym() and not m.islnk() and '..' not in Path(m.name).parts for m in members)
|
||||
print(json.dumps({**report,'directory':str(dest),'archive':str(archive),
|
||||
'archive_bytes':archive.stat().st_size},indent=2))
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,58 @@
|
||||
"""Interactive MuJoCo full replay, or export the exact same poses as video."""
|
||||
import argparse,os,time,json
|
||||
from pathlib import Path
|
||||
p=argparse.ArgumentParser(description=__doc__)
|
||||
p.add_argument('--directory',type=Path,default=Path(__file__).resolve().parents[1]/'output/l20_full_replay')
|
||||
p.add_argument('--start',type=int,default=0);p.add_argument('--render',action='store_true');p.add_argument('--frames',type=int,default=0);p.add_argument('--check',action='store_true')
|
||||
a=p.parse_args()
|
||||
if a.render or a.check:os.environ.setdefault('MUJOCO_GL','osmesa')
|
||||
import numpy as np
|
||||
import mujoco
|
||||
motion=np.load(a.directory/'motion.npz');q=motion['qpos'];fps=float(motion['fps']);N=len(q)
|
||||
m=mujoco.MjModel.from_xml_path(str(a.directory/'scene.xml'));d=mujoco.MjData(m)
|
||||
option=mujoco.MjvOption();option.sitegroup[:]=0;option.geomgroup[3:]=0
|
||||
cam=mujoco.MjvCamera();cam.lookat[:]=motion['camera_center'];cam.distance=1.05;cam.azimuth=135;cam.elevation=-23
|
||||
if a.check:
|
||||
for row in q:d.qpos[:]=row;mujoco.mj_forward(m,d)
|
||||
print(json.dumps(dict(frames=N,nq=m.nq,actuators=m.nu,finite=bool(np.isfinite(q).all()),mode='kinematic')))
|
||||
elif a.render:
|
||||
import imageio.v2 as imageio
|
||||
from PIL import Image,ImageDraw,ImageFont
|
||||
total=min(a.frames or N,N-a.start)
|
||||
renderer=mujoco.Renderer(m,height=720,width=640)
|
||||
close=mujoco.MjvCamera();close.distance=.62;close.azimuth=35;close.elevation=-12
|
||||
font=ImageFont.truetype('/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf',16)
|
||||
path=a.directory/('preview.mp4' if a.frames else 'full_replay.mp4')
|
||||
writer=imageio.get_writer(path,fps=fps,codec='libx264',quality=8,macro_block_size=1)
|
||||
try:
|
||||
for t in range(a.start,a.start+total):
|
||||
d.qpos[:]=q[t];mujoco.mj_forward(m,d)
|
||||
close.lookat[:]=(d.xpos[m.body('hand_base_link').id]+d.xpos[m.body('right_object').id])/2;close.lookat[2]+=.06
|
||||
panels=[]
|
||||
for c in [cam,close]:renderer.update_scene(d,camera=c,scene_option=option);panels.append(renderer.render().copy())
|
||||
img=Image.fromarray(np.concatenate(panels,1));draw=ImageDraw.Draw(img)
|
||||
draw.rectangle((0,0,1280,64),fill=(28,38,49))
|
||||
draw.text((18,9),f'L20 + bottle | full video motion | {t/fps:05.2f} / {N/fps:.2f} s',font=font,fill='white')
|
||||
draw.text((18,36),'World view - moving wrist',font=font,fill=(195,220,235));draw.text((660,36),'Grasp detail - linear URDF coupling',font=font,fill=(195,220,235))
|
||||
draw.rectangle((0,690,1280,720),fill=(28,38,49));draw.text((18,696),'Kinematic contact replay | estimated bottle trajectory | floor + ambient lighting',font=font,fill=(210,220,230))
|
||||
writer.append_data(np.asarray(img))
|
||||
if t in [a.start,0,160,600,N//2,N-1]:img.save(a.directory/f'preview_{t:04d}.png')
|
||||
if t%100==0:print(f'Render {t}/{total}',flush=True)
|
||||
finally:writer.close();renderer.close()
|
||||
print('RENDER_COMPLETE',path,flush=True)
|
||||
else:
|
||||
import mujoco.viewer
|
||||
state={'paused':False,'frame':0.,'last':time.perf_counter()}
|
||||
def key(k):
|
||||
if k==32:state['paused']=not state['paused']
|
||||
elif k in [82,114]:state['frame']=0.
|
||||
elif k==262:state['frame']=(int(state['frame'])+1)%N;state['paused']=True
|
||||
elif k==263:state['frame']=(int(state['frame'])-1)%N;state['paused']=True
|
||||
with mujoco.viewer.launch_passive(m,d,key_callback=key) as viewer:
|
||||
viewer.cam.lookat[:]=cam.lookat;viewer.cam.distance=cam.distance;viewer.cam.azimuth=cam.azimuth;viewer.cam.elevation=cam.elevation
|
||||
viewer.opt.sitegroup[:]=0;viewer.opt.geomgroup[3:]=0
|
||||
while viewer.is_running():
|
||||
now=time.perf_counter();elapsed=now-state['last'];state['last']=now
|
||||
if not state['paused']:state['frame']=(state['frame']+elapsed*fps)%N
|
||||
with viewer.lock():d.qpos[:]=q[int(state['frame'])];d.time=int(state['frame'])/fps;mujoco.mj_forward(m,d)
|
||||
viewer.sync();time.sleep(.005)
|
||||
@@ -0,0 +1,51 @@
|
||||
"""Play a hand-only L20 motion file with a free wrist in MuJoCo."""
|
||||
import argparse
|
||||
import time
|
||||
import mujoco
|
||||
import mujoco.viewer
|
||||
import numpy as np
|
||||
|
||||
p = argparse.ArgumentParser()
|
||||
p.add_argument('--directory', required=True)
|
||||
a = p.parse_args()
|
||||
from pathlib import Path
|
||||
directory = Path(a.directory)
|
||||
with np.load(directory / 'motion.npz') as data:
|
||||
motion = {name:data[name] for name in data.files}
|
||||
model = mujoco.MjModel.from_xml_path(str(directory / 'l20_moving.xml'))
|
||||
data = mujoco.MjData(model)
|
||||
addr = [model.jnt_qposadr[model.joint(str(n)).id] for n in motion['joint_names']]
|
||||
wrist = model.jnt_qposadr[model.joint('wrist_free').id]
|
||||
fps = float(motion['fps'])
|
||||
count = len(motion['qpos'])
|
||||
state = dict(frame=0., paused=False)
|
||||
|
||||
def key(k):
|
||||
if k == 32:
|
||||
state['paused'] = not state['paused']
|
||||
elif k in [82,114]:
|
||||
state['frame'] = 0.
|
||||
elif k in [262,263]:
|
||||
state['frame'] = (int(state['frame']) + (1 if k == 262 else -1)) % count
|
||||
state['paused'] = True
|
||||
|
||||
with mujoco.viewer.launch_passive(model, data, key_callback=key) as viewer:
|
||||
viewer.cam.lookat[:] = motion['wrist_pos'].mean(0)
|
||||
viewer.cam.distance = .9
|
||||
viewer.cam.azimuth, viewer.cam.elevation = 135, 10
|
||||
viewer.opt.sitegroup[:] = 0
|
||||
last = time.perf_counter()
|
||||
while viewer.is_running():
|
||||
now = time.perf_counter()
|
||||
if not state['paused']:
|
||||
state['frame'] = (state['frame'] + (now-last)*fps) % count
|
||||
last = now
|
||||
t = int(state['frame'])
|
||||
with viewer.lock():
|
||||
data.qpos[addr] = motion['qpos'][t]
|
||||
data.qpos[wrist:wrist+3] = motion['wrist_pos'][t]
|
||||
data.qpos[wrist+3:wrist+7] = motion['wrist_quat_wxyz'][t]
|
||||
data.time = t/fps
|
||||
mujoco.mj_forward(model,data)
|
||||
viewer.sync()
|
||||
time.sleep(.005)
|
||||
@@ -0,0 +1,118 @@
|
||||
"""Prepare an L20 linear-mimic bottle task for the actual SPIDER contact-guided runner."""
|
||||
import os
|
||||
os.environ.setdefault('MUJOCO_GL','osmesa')
|
||||
from pathlib import Path
|
||||
import json,copy,xml.etree.ElementTree as E
|
||||
import numpy as np
|
||||
import mujoco
|
||||
from scipy.spatial.transform import Rotation
|
||||
from scipy.optimize import least_squares
|
||||
import l20_model_source as source
|
||||
ROOT=Path(__file__).resolve().parents[1];OUT=ROOT/'output/spider_l20_contact'
|
||||
TASK=OUT/'datasets/processed/current/l20/right/bottle_hold';TRIAL=TASK/'0';TRIAL.mkdir(parents=True,exist_ok=True)
|
||||
source.REPO_ROOT=ROOT/'third_party/l20_assets'
|
||||
p=source.build_model('right',OUT);tree=E.parse(p);root=tree.getroot();world=root.find('worldbody');hand=world.find("body[@name='hand_base_link']")
|
||||
for mesh in root.findall('asset/mesh'):
|
||||
mesh.set('file',str((OUT/mesh.get('file')).resolve()));mesh.set('maxhullvert','64')
|
||||
root.find('option').attrib.update(timestep='.0025',gravity='0 0 -9.81',integrator='implicitfast',cone='elliptic')
|
||||
world.find("geom[@name='floor']").attrib.update(pos='0 0 0',contype='4',conaffinity='2')
|
||||
for g in hand.iter('geom'):g.attrib.update(contype='1',conaffinity='2',friction='.6 .005 .001',condim='3',solref='.01 1')
|
||||
root.find('default/joint').attrib.update(damping='.05',armature='.002')
|
||||
# Scalar wrist coordinates are SPIDER's dexterous-hand control convention.
|
||||
axes=['1 0 0','0 1 0','0 0 1'];wristnames=[]
|
||||
for i in range(6):
|
||||
name='right_hand_'+('pos_'+'xyz'[i] if i<3 else 'rot_'+'xyz'[i-3]);wristnames.append(name)
|
||||
hand.insert(i,E.Element('joint',name=name,type='slide' if i<3 else 'hinge',axis=axes[i%3],limited='false',damping='0',armature='.02'))
|
||||
urdf=E.parse(source.source_urdf()).getroot();joints=urdf.findall('joint');names=[j.get('name') for j in joints];active=[j.get('name') for j in joints if j.find('mimic') is None]
|
||||
limits={j.get('name'):np.array([float(j.find('limit').get('lower')),float(j.find('limit').get('upper'))]) for j in joints};mimic={}
|
||||
for j in joints:
|
||||
mi=j.find('mimic')
|
||||
if mi is not None:
|
||||
name=j.get('name');a=mi.get('joint');scale=float(mi.get('multiplier','1'));offset=float(mi.get('offset','0'));mimic[name]=(a,scale,offset)
|
||||
lo,hi=np.sort((limits[name]-offset)/scale);limits[a]=np.array([max(limits[a][0],lo),min(limits[a][1],hi)])
|
||||
for eq in root.findall('equality/joint'):eq.attrib.update(solref='.01 1')
|
||||
act=E.SubElement(root,'actuator')
|
||||
for i,name in enumerate(wristnames):E.SubElement(act,'position',name=name+'_act',joint=name,kp='2000' if i<3 else '40',kv='60' if i<3 else '2',forcerange='-200 200' if i<3 else '-20 20')
|
||||
for name in active:E.SubElement(act,'position',name=name+'_act',joint=name,kp='8',kv='.15',forcerange='-1 1',ctrlrange=source._fmt(limits[name]))
|
||||
bottlepath=ROOT/'third_party/bottle_model_parametric/bottle.xml';bottle=E.parse(bottlepath).getroot()
|
||||
for item in bottle.find('asset'):
|
||||
item=copy.deepcopy(item)
|
||||
if item.get('file'):item.set('file',str((bottlepath.parent/item.get('file')).resolve()))
|
||||
if item.tag=='mesh':item.set('maxhullvert','64')
|
||||
root.find('asset').append(item)
|
||||
obj=E.SubElement(world,'body',name='right_object')
|
||||
objectnames=[]
|
||||
for i in range(6):
|
||||
name='right_object_'+('pos_'+'xyz'[i] if i<3 else 'rot_'+'xyz'[i-3]);objectnames.append(name)
|
||||
E.SubElement(obj,'joint',name=name,type='slide' if i<3 else 'hinge',axis=axes[i%3],limited='false',damping='0',armature='0')
|
||||
E.SubElement(act,'position',name=name,joint=name,kp='0',kv='0')
|
||||
for j,item in enumerate(bottle.find('worldbody/body')):
|
||||
if item.tag=='freejoint':continue
|
||||
item=copy.deepcopy(item)
|
||||
if item.tag=='geom':
|
||||
item.set('name',f'bottle_{j}')
|
||||
if item.get('contype','1')!='0':item.attrib.update(contype='2',conaffinity='5',friction='.6 .005 .001',condim='3',solref='.01 1')
|
||||
obj.append(item)
|
||||
E.SubElement(obj,'site',name='right_object_track',pos='0 0 .07')
|
||||
scene=TASK/'scene_act.xml';tree.write(scene)
|
||||
model=mujoco.MjModel.from_xml_path(str(scene));data=mujoco.MjData(model)
|
||||
addr={n:int(model.jnt_qposadr[model.joint(n).id]) for n in names+wristnames+objectnames}
|
||||
actnames=[wristnames[i] for i in range(6)]+active+objectnames
|
||||
source_motion=np.load(ROOT/'output/l20_2047635068/jitter_audit/temporal/motion.npz');start=600;stop=660
|
||||
sourcenames=list(source_motion['joint_names']);wp=source_motion['wrist_pos'][start:stop];wr=Rotation.from_quat(source_motion['wrist_quat_wxyz'][start:stop][:,[1,2,3,0]])
|
||||
# Estimated bottle registration, not an observed video object track.
|
||||
relpos=np.array([.12,-.065,.15]);relrot=Rotation.from_euler('x',-np.pi/2)
|
||||
op=wp+wr.apply(np.tile(relpos,(len(wp),1)));orr=wr*relrot
|
||||
align=orr[0].inv();shift=np.array([0,0,.3])-align.apply(op[0]);wp=align.apply(wp)+shift;wr=align*wr;op=align.apply(op)+shift;orr=align*orr
|
||||
baseq=np.stack([source_motion['qpos'][start:stop,sourcenames.index(n)] for n in active],axis=1)
|
||||
low=np.array([limits[n][0] for n in active]);high=np.array([limits[n][1] for n in active]);baseq=np.clip(baseq,low,high)
|
||||
handgeoms=np.flatnonzero(model.geom_contype==1);objectgeoms=np.flatnonzero(model.geom_contype==2)
|
||||
tips=[next(g for g in handgeoms if model.geom_bodyid[g]==model.body(f+'_distal').id) for f in source.FINGERS]
|
||||
def pose(a,pos,rot,objpos,objrot):
|
||||
q=dict(zip(active,a));q.update({n:sc*q[parent]+off for n,(parent,sc,off) in mimic.items()})
|
||||
for n,v in q.items():data.qpos[addr[n]]=v
|
||||
data.qpos[[addr[n] for n in wristnames]]=np.r_[pos,rot.as_euler('XYZ')]
|
||||
data.qpos[[addr[n] for n in objectnames]]=np.r_[objpos,objrot.as_euler('XYZ')]
|
||||
mujoco.mj_kinematics(model,data)
|
||||
def distances():return np.array([min(mujoco.mj_geomDistance(model,data,int(g),int(o),.5,None) for o in objectgeoms) for g in handgeoms])
|
||||
ti=[list(handgeoms).index(g) for g in tips]
|
||||
def residual(x):
|
||||
pose(x[:16],wp[0]+x[16:19],wr[0]*Rotation.from_rotvec(x[19:]),op[0],orr[0]);d=distances()
|
||||
return np.r_[(d[ti]+.0004)/.004,np.maximum(-d-.0005,0)/.001,.08*(x[:16]-baseq[0]),.2*x[16:19]/.03,.15*x[19:]/.3]
|
||||
x0=np.r_[baseq[0],np.zeros(6)];residual(x0);before=distances()
|
||||
fits=[]
|
||||
for shift_x,shift_z in [(0,0),(.025,0),(-.025,0),(0,.025),(0,-.025),(.025,.025),(-.025,-.025)]:
|
||||
seed=x0.copy();seed[16]=shift_x;seed[18]=shift_z
|
||||
trial=least_squares(residual,seed,bounds=(np.r_[low,[-.06]*3,[-.5]*3],np.r_[high,[.06]*3,[.5]*3]),max_nfev=250,diff_step=1e-4)
|
||||
residual(trial.x);ds=distances();score=float(np.abs(ds[ti]).mean()+3*max(0,-ds.min()-.001))
|
||||
fits.append((score,trial));print('initialization seed',shift_x,shift_z,'gap mm',np.abs(ds[ti]).mean()*1000,'penetration mm',max(0,-ds.min())*1000,flush=True)
|
||||
fit=min(fits,key=lambda x:x[0])[1]
|
||||
residual(fit.x);after=distances();contactlocal=[]
|
||||
# Sites mark the actual closest point on each distal link at initialization.
|
||||
for finger,g in zip(source.FINGERS,tips):
|
||||
candidates=[]
|
||||
for o in objectgeoms:
|
||||
segment=np.zeros(6);dist=mujoco.mj_geomDistance(model,data,int(g),int(o),.5,segment);candidates.append((dist,segment.copy()))
|
||||
_,segment=min(candidates,key=lambda x:x[0]);bid=model.geom_bodyid[g];local=data.xmat[bid].reshape(3,3).T@(segment[:3]-data.xpos[bid])
|
||||
link=hand.find(f".//body[@name='{finger}_distal']");E.SubElement(link,'site',name=f'right_hand_{finger}_track',pos=source._fmt(local))
|
||||
contactlocal.append(orr[0].inv().apply(segment[3:]-op[0]))
|
||||
tree.write(scene);model=mujoco.MjModel.from_xml_path(str(scene));data=mujoco.MjData(model)
|
||||
contacts=[model.site(f'right_hand_{f}_track').id for f in source.FINGERS]
|
||||
# Preserve source movement, applying one smooth initial registration correction.
|
||||
source_times=np.arange(stop-start)/30;times=np.arange(801)*.0025;rows=[];controls=[];contactpos=[]
|
||||
from scipy.spatial.transform import Slerp
|
||||
wpi=np.stack([np.interp(times,source_times,wp[:,k]) for k in range(3)],axis=1)
|
||||
opi=np.stack([np.interp(times,source_times,op[:,k]) for k in range(3)],axis=1)
|
||||
wri=Slerp(source_times,wr)(np.minimum(times,source_times[-1]));ori=Slerp(source_times,orr)(np.minimum(times,source_times[-1]))
|
||||
ai=np.stack([np.interp(times,source_times,baseq[:,k]) for k in range(16)],axis=1);ai=np.clip(ai+fit.x[:16]-baseq[0],low,high)
|
||||
for t in range(len(times)):
|
||||
pose(ai[t],wpi[t]+fit.x[16:19],wri[t]*Rotation.from_rotvec(fit.x[19:]),opi[t],ori[t])
|
||||
rows.append(data.qpos.copy());controls.append([data.qpos[addr[n]] for n in actnames]);contactpos.append(ori[t].apply(contactlocal)+opi[t])
|
||||
q=np.array(rows);ctrl=np.array(controls);vel=np.gradient(q,.0025,axis=0);vel[0]=0
|
||||
np.savez_compressed(TRIAL/'trajectory_kinematic_act.npz',qpos=q,qvel=vel,ctrl=ctrl,contact=np.ones((len(q),5)),contact_pos=contactpos,time=times)
|
||||
(TASK/'task_info.json').write_text(json.dumps(dict(ref_dt=.0025,contact_site_ids=contacts),indent=2))
|
||||
# The final real rollout has zero object actuator gains; virtual assistance only during search.
|
||||
config=dict(dataset_dir=str(OUT/'datasets'),dataset_name='current',robot_type='l20',embodiment_type='right',task='bottle_hold',data_id=0,contact_guidance=True,contact_rew_scale=1.,device='cuda:0',sim_dt=.0025,ref_dt=.0025,ctrl_dt=.1,knot_dt=.1,horizon=.6,trace_dt=.01,render_dt=.02,num_samples=512,max_num_iterations=6,num_dr=1,num_dyn=1,pair_margin_range=[0.,0.],xy_offset_range=[0.,0.],guidance_decay_ratio=.4,init_pos_actuator_gain=10.,init_pos_actuator_bias=1.,init_rot_actuator_gain=.1,init_rot_actuator_bias=.02,max_sim_steps=800,show_viewer=False,viewer='',wait_on_finish=False,save_video=False,save_info=True,save_metrics=False,save_config=True,sanity_check_seconds=0.,use_torch_compile=False,improvement_threshold=0.,pos_noise_scale=.008,rot_noise_scale=.05,joint_noise_scale=.2,exploit_ratio=.02,base_pos_rew_scale=.3,base_rot_rew_scale=.1,joint_rew_scale=.003,pos_rew_scale=5.,rot_rew_scale=.3,vel_rew_scale=.0001,temperature=.1,nconmax_per_env=128,njmax_per_env=512)
|
||||
(OUT/'config.json').write_text(json.dumps(config,indent=2))
|
||||
report=dict(source_frames=[start,stop],duration_s=2.,source='Current video trajectory, offline smoothed',bottle_alignment='Estimated constant bottle-to-palm relation, not tracked from video',coupling='Original URDF linear mimic, no nonlinear JSON adaptor',active_finger_joints=16,passive_finger_joints=5,controlled_wrist_dofs=6,object_virtual_actuators=6,object_real_rollout_gains_zero=True,initial_fit_converged=bool(fit.success),initial_mean_surface_gap_before_mm=float(np.abs(before[ti]).mean()*1000),initial_mean_surface_gap_after_mm=float(np.abs(after[ti]).mean()*1000),initial_max_penetration_mm=float(max(0,-after.min())*1000),nq=model.nq,nv=model.nv,nu=model.nu,contact_sites=contacts)
|
||||
(OUT/'preparation.json').write_text(json.dumps(report,indent=2));print(json.dumps(report,indent=2))
|
||||
@@ -0,0 +1,42 @@
|
||||
"""Prepare project-local Hugging Face cache metadata and resumable downloads."""
|
||||
import json
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
CACHE = ROOT / '.hf_cache' / 'hub'
|
||||
|
||||
def fetch(url, target):
|
||||
target.parent.mkdir(parents=True, exist_ok=True)
|
||||
subprocess.run(['curl', '-fsSL', '--retry', '3', '--retry-all-errors', '--max-time', '45', url, '-o', str(target)], check=True)
|
||||
|
||||
manifest = []
|
||||
for repo, files, aliases in [
|
||||
('google-bert/bert-base-uncased', ['config.json', 'tokenizer.json', 'tokenizer_config.json', 'vocab.txt', 'model.safetensors'], ['bert-base-uncased']),
|
||||
('lpiccinelli/unidepth-v2-vitl14', ['config.json', 'model.safetensors'], []),
|
||||
('Rain729/Prior-Depth-Anything', ['depth_anything_v2_vitb.pth', 'prior_depth_anything_vitb.pth'], []),
|
||||
]:
|
||||
meta = ROOT / 'output' / (repo.replace('/', '_') + '_metadata.json')
|
||||
fetch('https://hf-mirror.com/api/models/' + repo, meta)
|
||||
revision = json.loads(meta.read_text())['sha']
|
||||
cache = CACHE / ('models--' + repo.replace('/', '--'))
|
||||
snapshot = cache / 'snapshots' / revision
|
||||
snapshot.mkdir(parents=True, exist_ok=True)
|
||||
(cache / 'refs').mkdir(exist_ok=True)
|
||||
(cache / 'refs' / 'main').write_text(revision)
|
||||
for name in files:
|
||||
target = snapshot / name
|
||||
url = 'https://hf-mirror.com/' + repo + '/resolve/' + revision + '/' + name
|
||||
if repo.startswith('lpiccinelli/'):
|
||||
local = ROOT / 'weights' / 'unidepth-v2-vitl14' / name
|
||||
if not target.exists() and not target.is_symlink():
|
||||
target.symlink_to(local)
|
||||
elif name.endswith(('.json', '.txt')):
|
||||
fetch(url, target)
|
||||
else:
|
||||
manifest.extend([url, ' dir=' + str(snapshot), ' out=' + name])
|
||||
for alias in aliases:
|
||||
alias_path = CACHE / ('models--' + alias.replace('/', '--'))
|
||||
if not alias_path.exists():
|
||||
alias_path.symlink_to(cache, target_is_directory=True)
|
||||
(ROOT / 'output' / 'vipe_weights_downloads.txt').write_text('\n'.join(manifest) + '\n')
|
||||
print('Prepared', len(manifest) // 3, 'weight downloads')
|
||||
@@ -0,0 +1,70 @@
|
||||
"""Render fixed-world and wrist-following views of a moving L20 trajectory."""
|
||||
import os
|
||||
os.environ.setdefault('MUJOCO_GL','osmesa')
|
||||
import argparse,json
|
||||
from pathlib import Path
|
||||
import numpy as np
|
||||
import mujoco,imageio.v2 as imageio
|
||||
from PIL import Image,ImageDraw
|
||||
from scipy.spatial.transform import Rotation
|
||||
|
||||
def main():
|
||||
parser=argparse.ArgumentParser()
|
||||
parser.add_argument('--directory',type=Path,default=Path(__file__).resolve().parents[1]/'output/l20_2047635068')
|
||||
parser.add_argument('--frames',type=int,default=0)
|
||||
args=parser.parse_args();out=args.directory
|
||||
with np.load(out/'motion.npz') as archive:
|
||||
motion={name:archive[name] for name in archive.files}
|
||||
q=motion['qpos'][:args.frames or None];positions=motion['wrist_pos'][:len(q)]
|
||||
model=mujoco.MjModel.from_xml_path(str(out/'l20_moving.xml'));data=mujoco.MjData(model)
|
||||
addr=[model.jnt_qposadr[mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_JOINT,str(n))] for n in motion['joint_names']]
|
||||
wrist=model.jnt_qposadr[mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_JOINT,'wrist_free')]
|
||||
rotations=Rotation.from_quat(motion['wrist_quat_wxyz'][:,[1,2,3,0]]).as_matrix()
|
||||
world=np.einsum('tij,tkj->tki',rotations,motion['actual'])+motion['wrist_pos'][:,None]
|
||||
center=(world.min((0,1))+world.max((0,1)))/2
|
||||
span=np.ptp(world,axis=(0,1)).max()
|
||||
floor=mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_GEOM,'floor')
|
||||
model.geom_pos[floor,2]=world[:,:,2].min()-.035
|
||||
model.vis.quality.shadowsize=1024;model.vis.quality.offsamples=1
|
||||
model.vis.headlight.ambient[:]=.65;model.vis.headlight.diffuse[:]=.7
|
||||
renderer=mujoco.Renderer(model,height=480,width=640)
|
||||
cameras=[]
|
||||
for azimuth in [135,25]:
|
||||
cam=mujoco.MjvCamera();cam.lookat[:]=center;cam.distance=max(.65,span*2.3);cam.azimuth=azimuth;cam.elevation=10
|
||||
cameras.append(cam)
|
||||
option=mujoco.MjvOption();option.sitegroup[:]=0
|
||||
fps=float(motion['fps'])
|
||||
coupling=str(motion['coupling']) if 'coupling' in motion else 'JSON-calibrated'
|
||||
writer=imageio.get_writer(out/'l20_moving_replay.mp4',fps=fps,codec='libx264',quality=8)
|
||||
max_wrist_error=0.
|
||||
try:
|
||||
for t,pose in enumerate(q):
|
||||
data.qpos[addr]=pose;data.qpos[wrist:wrist+3]=positions[t];data.qpos[wrist+3:wrist+7]=motion['wrist_quat_wxyz'][t]
|
||||
mujoco.mj_forward(model,data)
|
||||
base=mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_BODY,'hand_base_link')
|
||||
max_wrist_error=max(max_wrist_error,float(np.abs(data.xpos[base]-positions[t]).max()))
|
||||
panels=[]
|
||||
cameras[1].lookat[:]=world[t].mean(0);cameras[1].distance=.58
|
||||
for vi,cam in enumerate(cameras):
|
||||
renderer.update_scene(data,camera=cam,scene_option=option)
|
||||
if vi==0:
|
||||
points=positions[np.unique(np.linspace(0,t,min(t+1,60),dtype=int))]
|
||||
for a,b in zip(points[:-1],points[1:]):
|
||||
scene=renderer.scene
|
||||
mujoco.mjv_initGeom(scene.geoms[scene.ngeom],mujoco.mjtGeom.mjGEOM_CAPSULE,np.ones(3)*.001,a,np.eye(3).ravel(),np.array([.2,.8,1,1]))
|
||||
mujoco.mjv_connector(scene.geoms[scene.ngeom],mujoco.mjtGeom.mjGEOM_CAPSULE,.0015,a,b)
|
||||
scene.ngeom+=1
|
||||
panels.append(renderer.render().copy())
|
||||
img=Image.fromarray(np.concatenate(panels,axis=1));draw=ImageDraw.Draw(img)
|
||||
draw.rectangle([0,0,1280,45],fill=(15,22,30))
|
||||
draw.text((10,8),f'L20 right | frame {t+1}/{len(q)} | moving 6D wrist + {coupling} fingers',fill='white')
|
||||
draw.text((10,27),'Fixed world camera / blue wrist trail',fill='white')
|
||||
draw.text((650,27),'Following camera / finger detail',fill='white')
|
||||
writer.append_data(np.asarray(img))
|
||||
if t in [0,len(q)//2,len(q)-1]:img.save(out/f'preview_{t:04d}.png')
|
||||
if t%200==0:print(f'render L20 {t}/{len(q)}',flush=True)
|
||||
finally:writer.close();renderer.close()
|
||||
(out/'render_validation.json').write_text(json.dumps(dict(frames=len(q),fps=fps,max_wrist_position_error_m=max_wrist_error),indent=2))
|
||||
print('COMPLETE',flush=True)
|
||||
|
||||
if __name__=='__main__':main()
|
||||
@@ -0,0 +1,73 @@
|
||||
"""Render contact registration separately from the free-object physics test."""
|
||||
import os
|
||||
os.environ.setdefault('MUJOCO_GL','osmesa')
|
||||
import argparse,json
|
||||
from pathlib import Path
|
||||
import numpy as np
|
||||
import mujoco,imageio.v2 as imageio
|
||||
from PIL import Image,ImageDraw
|
||||
from scipy.spatial.transform import Rotation
|
||||
|
||||
OUT=Path(__file__).resolve().parents[1]/'output/l20_2047635068/bottle_grasp'
|
||||
def main():
|
||||
p=argparse.ArgumentParser();p.add_argument('--physics',action='store_true');p.add_argument('--frames',type=int,default=0);args=p.parse_args()
|
||||
physics=args.physics
|
||||
path=OUT/('physics_scene.xml' if physics else 'contact_scene.xml')
|
||||
motion=np.load(OUT/('physics_motion.npz' if physics else 'registered_motion.npz'))
|
||||
model=mujoco.MjModel.from_xml_path(str(path));data=mujoco.MjData(model)
|
||||
model.vis.headlight.ambient[:]=.65;model.vis.headlight.diffuse[:]=.7;model.vis.quality.offsamples=1
|
||||
renderer=mujoco.Renderer(model,height=480,width=640)
|
||||
option=mujoco.MjvOption();option.geomgroup[3]=0;option.sitegroup[:]=0
|
||||
count=len(motion['qpos']);count=min(count,args.frames) if args.frames else count
|
||||
fps=float(motion['fps'])
|
||||
obj=model.jnt_qposadr[model.joint('object_free').id]
|
||||
if not physics:
|
||||
addresses=[model.jnt_qposadr[model.joint(str(n)).id] for n in motion['joint_names']]
|
||||
wrist=model.jnt_qposadr[model.joint('wrist_free').id]
|
||||
reference=int(motion['contact_onset_frame'])
|
||||
rotation=Rotation.from_quat(motion['bottle_quat_wxyz'][reference,[1,2,3,0]]).inv()
|
||||
translation=np.array([0,0,.3])-rotation.apply(motion['bottle_pos'][reference])
|
||||
bottle_pos=rotation.apply(motion['bottle_pos'])+translation
|
||||
wrist_pos=rotation.apply(motion['wrist_pos'])+translation
|
||||
bottle_quat=(rotation*Rotation.from_quat(motion['bottle_quat_wxyz'][:,[1,2,3,0]])).as_quat()[:,[3,0,1,2]]
|
||||
wrist_quat=(rotation*Rotation.from_quat(motion['wrist_quat_wxyz'][:,[1,2,3,0]])).as_quat()[:,[3,0,1,2]]
|
||||
both=np.concatenate([bottle_pos+[0,0,.08],wrist_pos])
|
||||
center=(both.min(0)+both.max(0))/2;distance=max(.75,np.ptp(both,axis=0).max()*2.5)
|
||||
floor=model.geom('floor').id;model.geom_pos[floor,2]=min(both[:,2])-.06
|
||||
else:center=np.array([0,0,.38]);distance=.65
|
||||
cameras=[]
|
||||
for azimuth in [45,225]:
|
||||
cam=mujoco.MjvCamera();cam.lookat[:]=center;cam.distance=distance;cam.azimuth=azimuth;cam.elevation=12;cameras.append(cam)
|
||||
filename='physics_test.mp4' if physics else 'contact_registration.mp4'
|
||||
writer=imageio.get_writer(OUT/filename,fps=fps,codec='libx264',quality=8)
|
||||
try:
|
||||
for t in range(count):
|
||||
if physics:
|
||||
data.qpos[:]=motion['qpos'][t];data.mocap_pos[0]=motion['wrist'][t,:3];data.mocap_quat[0]=motion['wrist'][t,3:]
|
||||
else:
|
||||
data.qpos[addresses]=motion['qpos'][t]
|
||||
data.qpos[wrist:wrist+7]=np.r_[wrist_pos[t],wrist_quat[t]]
|
||||
data.qpos[obj:obj+7]=np.r_[bottle_pos[t],bottle_quat[t]]
|
||||
cameras[1].lookat[:]=(bottle_pos[t]+wrist_pos[t])/2+[0,0,.04];cameras[1].distance=.65
|
||||
mujoco.mj_forward(model,data)
|
||||
panels=[]
|
||||
for camera in cameras:
|
||||
renderer.update_scene(data,camera=camera,scene_option=option);panels.append(renderer.render().copy())
|
||||
img=Image.fromarray(np.concatenate(panels,axis=1));draw=ImageDraw.Draw(img);draw.rectangle([0,0,1280,48],fill=(15,22,30))
|
||||
if physics:
|
||||
title='FREE OBJECT PHYSICS TEST | 16 actuators + nonlinear coupling | moving wrist'
|
||||
note=f"t={t/fps:.2f}s | slip={motion['relative_drift'][t]*1000:.1f} mm | contacts={motion['contact_count'][t]} | result: NOT RETAINED"
|
||||
color=(255,120,100)
|
||||
else:
|
||||
title='L20 + tri-camera bottle | CONTACT REGISTRATION | object poses prescribed, not a physical grasp'
|
||||
note=f"frame {t+1}/{count} | moving wrist | maximum penetration {motion['penetration_m'][t]*1000:.2f} mm | estimated object-to-hand alignment"
|
||||
color='white'
|
||||
draw.text((10,8),title,fill='white');draw.text((10,28),note,fill=color)
|
||||
writer.append_data(np.asarray(img))
|
||||
if t in [0,count//2,count-1]:img.save(OUT/(('physics' if physics else 'registration')+f'_{t:04d}.png'))
|
||||
if t%200==0:print(f'{filename} {t}/{count}',flush=True)
|
||||
finally:writer.close();renderer.close()
|
||||
(OUT/(filename+'.json')).write_text(json.dumps({'frames':count,'fps':fps,'physics':physics},indent=2))
|
||||
print('COMPLETE '+filename,flush=True)
|
||||
|
||||
if __name__=='__main__':main()
|
||||
@@ -0,0 +1,32 @@
|
||||
"""Matched wrist-local diagnostic view of the worst contact-jump interval."""
|
||||
import os
|
||||
os.environ.setdefault('MUJOCO_GL','osmesa')
|
||||
import json
|
||||
import numpy as np
|
||||
import mujoco,imageio.v2 as imageio
|
||||
from PIL import Image,ImageDraw
|
||||
from l20_calibrated import ROOT
|
||||
base=ROOT/'output/l20_2047635068';out=base/'jitter_audit'
|
||||
sources=[base/'motion.npz',out/'temporal/motion.npz',base/'bottle_grasp/registered_motion.npz']
|
||||
motion=[np.load(p) for p in sources]
|
||||
model=mujoco.MjModel.from_xml_path(str(base/'l20_moving.xml'));data=mujoco.MjData(model)
|
||||
addr=[model.jnt_qposadr[model.joint(str(n)).id] for n in motion[0]['joint_names']]
|
||||
model.geom_rgba[model.geom('floor').id,3]=0;model.vis.headlight.ambient[:]=.8
|
||||
model.vis.quality.offsamples=1;renderer=mujoco.Renderer(model,height=400,width=400)
|
||||
cam=mujoco.MjvCamera();cam.lookat[:]=[.025,0,.105];cam.distance=.48;cam.azimuth=130;cam.elevation=5
|
||||
opt=mujoco.MjvOption();opt.sitegroup[:]=0
|
||||
writer=imageio.get_writer(out/'comparison_12s_18s.mp4',fps=30,codec='libx264',quality=8)
|
||||
labels=['Original L20 (hand only)','Corrected + offline temporal smoothing','Previous contact fit (known jumps)']
|
||||
for t in range(360,540):
|
||||
panels=[]
|
||||
for m in motion:
|
||||
data.qpos[addr]=m['qpos'][t];mujoco.mj_forward(model,data);renderer.update_scene(data,camera=cam,scene_option=opt);panels.append(renderer.render().copy())
|
||||
img=Image.fromarray(np.concatenate(panels,axis=1));draw=ImageDraw.Draw(img);draw.rectangle([0,0,1200,64],fill=(20,25,35))
|
||||
for j,label in enumerate(labels):
|
||||
step=np.rad2deg(np.abs(motion[j]['qpos'][t]-motion[j]['qpos'][t-1]).max())
|
||||
draw.text((j*400+8,7),label,fill='white');draw.text((j*400+8,24),f'frame {t} | maximum joint step {step:.2f} deg',fill='white')
|
||||
draw.text((8,45),'Wrist-local diagnostic view only; saved trajectories retain moving wrists. No physical grasp claim.',fill='white')
|
||||
writer.append_data(np.asarray(img))
|
||||
if t==461:img.save(out/'comparison_frame461.png')
|
||||
writer.close();renderer.close()
|
||||
print('COMPLETE 180 frames, source interval [360,540), 30 fps')
|
||||
@@ -0,0 +1,60 @@
|
||||
"""Compare two trajectories with exactly the same fixed world camera."""
|
||||
import os
|
||||
os.environ.setdefault('MUJOCO_GL','osmesa')
|
||||
from pathlib import Path
|
||||
import json
|
||||
import mujoco
|
||||
import numpy as np
|
||||
import imageio.v2 as imageio
|
||||
from PIL import Image,ImageDraw,ImageFont
|
||||
from scipy.spatial.transform import Rotation
|
||||
|
||||
ROOT=Path(__file__).resolve().parents[1]
|
||||
out=ROOT/'output/l20_15886123_stable'
|
||||
motions=[]
|
||||
for folder in [ROOT/'output/l20_15886123',out]:
|
||||
with np.load(folder/'motion.npz') as f:
|
||||
motions.append({k:f[k] for k in f.files})
|
||||
count=len(motions[0]['qpos']);fps=float(motions[0]['fps'])
|
||||
assert count==len(motions[1]['qpos'])
|
||||
model=mujoco.MjModel.from_xml_path(str(out/'l20_moving.xml'))
|
||||
data=mujoco.MjData(model)
|
||||
addr=[model.jnt_qposadr[model.joint(str(n)).id] for n in motions[0]['joint_names']]
|
||||
wrist=model.jnt_qposadr[model.joint('wrist_free').id]
|
||||
world=[]
|
||||
for m in motions:
|
||||
r=Rotation.from_quat(m['wrist_quat_wxyz'][:,[1,2,3,0]]).as_matrix()
|
||||
world.append(np.einsum('tij,tkj->tki',r,m['actual'])+m['wrist_pos'][:,None])
|
||||
points=np.concatenate(world).reshape(-1,3)
|
||||
center=(points.min(0)+points.max(0))/2
|
||||
model.geom_pos[model.geom('floor').id,2]=points[:,2].min()-.035
|
||||
model.vis.quality.offsamples=1
|
||||
model.vis.headlight.ambient[:]=.65
|
||||
model.vis.headlight.diffuse[:]=.7
|
||||
cam=mujoco.MjvCamera();cam.lookat[:]=center;cam.distance=max(.65,float(np.ptp(points,axis=0).max())*2.1);cam.azimuth=100;cam.elevation=5
|
||||
option=mujoco.MjvOption();option.sitegroup[:]=0
|
||||
renderer=mujoco.Renderer(model,height=480,width=640)
|
||||
font=ImageFont.truetype('/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf',18)
|
||||
writer=imageio.get_writer(out/'comparison.mp4',fps=fps,codec='libx264',quality=8)
|
||||
try:
|
||||
for t in range(count):
|
||||
panels=[]
|
||||
for m in motions:
|
||||
data.qpos[addr]=m['qpos'][t]
|
||||
data.qpos[wrist:wrist+3]=m['wrist_pos'][t]
|
||||
data.qpos[wrist+3:wrist+7]=m['wrist_quat_wxyz'][t]
|
||||
mujoco.mj_forward(model,data)
|
||||
renderer.update_scene(data,camera=cam,scene_option=option)
|
||||
panels.append(renderer.render().copy())
|
||||
im=Image.fromarray(np.concatenate(panels,axis=1));d=ImageDraw.Draw(im)
|
||||
d.rectangle([0,0,1280,50],fill=(15,22,30))
|
||||
d.text((12,8),'Before - previous smoothing',font=font,fill='white')
|
||||
d.text((652,8),'After - temporal regularization',font=font,fill='white')
|
||||
d.text((12,29),f'Same fixed camera | {t/fps:.2f} s | moving wrist preserved',font=font,fill=(190,210,225))
|
||||
writer.append_data(np.asarray(im))
|
||||
if t==280:im.save(out/'comparison_preview.png')
|
||||
if t%200==0:print('AB render',t,count,flush=True)
|
||||
finally:
|
||||
writer.close();renderer.close()
|
||||
(out/'comparison_camera.json').write_text(json.dumps(dict(same_camera=True,fixed_world_camera=True,frames=count,fps=fps,lookat=center.tolist(),distance=cam.distance,azimuth=cam.azimuth,elevation=cam.elevation),indent=2))
|
||||
print('COMPLETE',flush=True)
|
||||
@@ -0,0 +1,28 @@
|
||||
"""Retry residual collision failures from the preceding valid pose."""
|
||||
import json
|
||||
import numpy as np
|
||||
from l20_bottle_sequence import BASE,OUT
|
||||
from l20_contact import ContactRefiner
|
||||
|
||||
src=np.load(BASE/'motion.npz');m=dict(np.load(OUT/'registered_motion.npz'))
|
||||
s=ContactRefiner(json.loads((OUT/'object_spec.json').read_text()),OUT)
|
||||
names=list(m['joint_names']);reports=[]
|
||||
for t in np.where(m['penetration_m']>.0015)[0]:
|
||||
active=np.array([src['qpos'][t,names.index(n)] for n in s.cal.active])
|
||||
x,_,r=s.solve(active,src['wrist_pos'][t],src['wrist_quat_wxyz'][t],m['bottle_pos'][t],m['bottle_quat_wxyz'][t],initial=m['optimization_variables'][t-1])
|
||||
if r['max_penetration_after_mm']<m['penetration_m'][t]*1000:
|
||||
m['optimization_variables'][t]=x
|
||||
m['qpos'][t]=[s.expand(x[:16])[str(n)] for n in names]
|
||||
m['wrist_pos'][t]=s.data.qpos[s.wrist:s.wrist+3]
|
||||
m['wrist_quat_wxyz'][t]=s.data.qpos[s.wrist+3:s.wrist+7]
|
||||
m['penetration_m'][t]=r['max_penetration_after_mm']/1000
|
||||
m['command_u8'][t]=s.cal.command(dict(zip(names,m['qpos'][t])))
|
||||
reports.append(dict(frame=int(t),**r))
|
||||
np.savez_compressed(OUT/'registered_motion.npz',**m)
|
||||
header=(OUT/'registered_trajectory.csv').open().readline().strip()
|
||||
np.savetxt(OUT/'registered_trajectory.csv',np.c_[np.arange(len(m['qpos']))/30,m['wrist_pos'],m['wrist_quat_wxyz'],m['qpos'],m['bottle_pos'],m['bottle_quat_wxyz']],delimiter=',',header=header,comments='')
|
||||
r=json.loads((OUT/'sequence_validation.json').read_text());r['neighbor_retry_reports']=reports
|
||||
r['max_all_frame_penetration_mm']=float(m['penetration_m'].max()*1000)
|
||||
r['max_holding_penetration_mm']=float(m['penetration_m'][160:].max()*1000)
|
||||
(OUT/'sequence_validation.json').write_text(json.dumps(r,indent=2))
|
||||
print({'retries':len(reports),'max_penetration_mm':r['max_all_frame_penetration_mm'],'remaining_bad_frames':np.where(m['penetration_m']>.0015)[0].tolist()})
|
||||
@@ -0,0 +1,20 @@
|
||||
"""Fresh GPU rollout of saved controls, without optimization or virtual assistance."""
|
||||
import os
|
||||
from pathlib import Path
|
||||
ROOT=Path(__file__).resolve().parents[1];OUT=ROOT/'output/spider_l20_contact'
|
||||
os.environ['WARP_CACHE_PATH']=str(ROOT/'.spider_cache/warp');os.environ.setdefault('MUJOCO_GL','osmesa')
|
||||
import json,numpy as np,mujoco,mujoco_warp as mw,warp as wp
|
||||
TASK=OUT/'datasets/processed/current/l20/right/bottle_hold';P=TASK/'0'
|
||||
r=np.load(P/'trajectory_kinematic_act.npz');a=np.load(P/'trajectory_mjwp_act.npz');ctrl=a['ctrl'].reshape(-1,a['ctrl'].shape[-1]);expected=a['qpos'].reshape(-1,a['qpos'].shape[-1])
|
||||
m=mujoco.MjModel.from_xml_path(str(TASK/'scene_act.xml'));m.opt.iterations=20;m.opt.ls_iterations=50;m.opt.o_solref[:]=[.02,1];m.opt.o_solimp[:]=[0,.95,.03,.5,2]
|
||||
assert not m.actuator_gainprm[-6:].any() and not m.actuator_biasprm[-6:].any()
|
||||
d=mujoco.MjData(m);d.qpos[:]=r['qpos'][0];d.qvel[:]=r['qvel'][0];d.ctrl[:]=r['ctrl'][0];mujoco.mj_step(m,d);d.time=0
|
||||
wp.init();rows=[]
|
||||
with wp.ScopedDevice('cuda:0'):
|
||||
wm=mw.put_model(m);wd=mw.put_data(m,d,nworld=2,nconmax=128,njmax=512)
|
||||
with wp.ScopedCapture() as capture:mw.step(wm,wd)
|
||||
for i,u in enumerate(ctrl):
|
||||
wd.ctrl.assign(np.stack([r['ctrl'][i+1],u]).astype(np.float32));wp.capture_launch(capture.graph);wp.synchronize();rows.append(wd.qpos.numpy())
|
||||
rows=np.array(rows);np.savez_compressed(OUT/'independent_gpu_replay.npz',baseline=rows[:,0],optimized_controls=rows[:,1])
|
||||
report=dict(all_finite=bool(np.isfinite(rows).all()),max_difference_from_native_spider_qpos=float(np.abs(rows[:,1]-expected).max()),optimized_final_object_pos=rows[-1,1,-6:-3].tolist(),native_final_object_pos=expected[-1,-6:-3].tolist(),zero_assistance=True)
|
||||
(OUT/'independent_gpu_replay.json').write_text(json.dumps(report,indent=2));print(json.dumps(report,indent=2),flush=True)
|
||||
@@ -0,0 +1,4 @@
|
||||
# Install into HandFlow .venv with --no-deps; reuse existing Torch/CUDA.
|
||||
transformers==4.48.3
|
||||
tokenizers==0.21.0
|
||||
huggingface-hub==0.28.1
|
||||
@@ -0,0 +1,55 @@
|
||||
"""Resume the existing 1333-frame reconstruction with project-local model caches."""
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import subprocess
|
||||
import time
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
OUT = ROOT / 'output'
|
||||
STATUS = OUT / 'dynhamr_auto_status.json'
|
||||
def status(stage, **extra):
|
||||
payload = dict(stage=stage, time=time.strftime('%Y-%m-%d %H:%M:%S'), **extra)
|
||||
STATUS.write_text(json.dumps(payload, indent=2) + '\n')
|
||||
print(json.dumps(payload), flush=True)
|
||||
|
||||
def required_files():
|
||||
hub = ROOT / '.torch_cache/hub'
|
||||
files = [hub/'sam/sam_vit_b_01ec64.pth', hub/'aot/R50_DeAOTL_PRE_YTB_DAV.pth',
|
||||
hub/'checkpoints/groundingdino_swint_ogc.pth', hub/'checkpoints/video_depth_anything_vits.pth',
|
||||
ROOT/'weights/unidepth-v2-vitl14/model.safetensors']
|
||||
manifest = (OUT/'vipe_weights_downloads.txt').read_text().splitlines()
|
||||
for i in range(0, len(manifest), 3):
|
||||
files.append(Path(manifest[i+1].strip().removeprefix('dir='))/manifest[i+2].strip().removeprefix('out='))
|
||||
return files
|
||||
|
||||
files = required_files()
|
||||
deadline = time.monotonic() + 5400
|
||||
while True:
|
||||
missing = [str(p) for p in files if not p.is_file() or p.stat().st_size < 1000000 or Path(str(p)+'.aria2').exists()]
|
||||
if not missing:
|
||||
break
|
||||
status('WAITING_WEIGHTS', pending=missing)
|
||||
if time.monotonic() > deadline:
|
||||
status('FAILED', reason='Weight preparation timeout')
|
||||
raise SystemExit(1)
|
||||
time.sleep(30)
|
||||
|
||||
env = os.environ.copy()
|
||||
env.update(TORCH_HOME=str(ROOT/'.torch_cache'), HF_HOME=str(ROOT/'.hf_cache'),
|
||||
HF_HUB_OFFLINE='1', TRANSFORMERS_OFFLINE='1', PYTHONUNBUFFERED='1',
|
||||
VIPE_CACHE_DIR=str(OUT/'vipe_frame_cache'),
|
||||
MPLCONFIGDIR='/tmp/matplotlib-dynhamr', HYDRA_FULL_ERROR='1', PYOPENGL_PLATFORM='egl')
|
||||
dyn = ROOT/'third_party/Dyn-HaMR'
|
||||
env['PATH'] = str(dyn/'.dynhamr/bin') + ':' + env['PATH']
|
||||
env['PYTHONPATH'] = ':'.join(str(p) for p in [dyn/'dyn-hamr', dyn/'third-party/hamer', dyn/'third-party/hamer/third-party/ViTPose'])
|
||||
cmd = [str(dyn/'.dynhamr/bin/python'), '-u', 'run_opt.py', 'data=video_vipe',
|
||||
'data.root='+str(dyn/'test'), 'data.seq=2047635068',
|
||||
'data.vipe_dir='+str(OUT/'vipe_2047635068'), 'is_static=False',
|
||||
'run_opt=True', 'run_vis=True', 'run_prior=True', '+data.vipe_pipeline=dynhamr_cameras',
|
||||
'HMP.vid_path='+str(dyn/'test/videos/2047635068.mp4')]
|
||||
status('RUNNING', command=cmd)
|
||||
with (OUT/'dynhamr_auto_run.log').open('a') as log:
|
||||
result = subprocess.run(cmd, cwd=dyn/'dyn-hamr', env=env, stdout=log, stderr=subprocess.STDOUT)
|
||||
status('PROCESS_EXITED' if result.returncode == 0 else 'FAILED', returncode=result.returncode,
|
||||
log=str(OUT/'dynhamr_auto_run.log'))
|
||||
raise SystemExit(result.returncode)
|
||||
@@ -0,0 +1,165 @@
|
||||
"""Retarget exported video joints to L20 using original URDF linear coupling."""
|
||||
import os
|
||||
os.environ.setdefault('MUJOCO_GL', 'osmesa')
|
||||
import argparse
|
||||
import copy
|
||||
import json
|
||||
from pathlib import Path
|
||||
import xml.etree.ElementTree as ET
|
||||
|
||||
import mujoco
|
||||
import numpy as np
|
||||
import torch
|
||||
from scipy.ndimage import gaussian_filter1d
|
||||
from scipy.spatial.transform import Rotation
|
||||
from dex_retargeting.retargeting_config import RetargetingConfig
|
||||
from dex_retargeting.kinematics_adaptor import MimicJointKinematicAdaptor
|
||||
from l20_calibrated import ASSET, build_assets, palm_basis, CHAINS, TIPS
|
||||
from l20_consistent_optimizer import ConsistentPositionOptimizer
|
||||
|
||||
|
||||
def main():
|
||||
p = argparse.ArgumentParser()
|
||||
p.add_argument('--input', type=Path, required=True)
|
||||
p.add_argument('--output-dir', type=Path, required=True)
|
||||
p.add_argument('--sigma', type=float, default=2.5)
|
||||
a = p.parse_args()
|
||||
torch.set_num_threads(4)
|
||||
out = a.output_dir.resolve()
|
||||
fixed, urdf = build_assets(out)
|
||||
original = ET.parse(ASSET / 'linkerhand_g20_right.urdf').getroot()
|
||||
joints = [j for j in original.findall('joint') if j.get('type') != 'fixed']
|
||||
active = [j.get('name') for j in joints if j.find('mimic') is None]
|
||||
limits = {j.get('name'):np.array([float(j.find('limit').get(k)) for k in ['lower','upper']]) for j in joints}
|
||||
bounds_by_name = {n: limits[n].copy() for n in active}
|
||||
mimic = {}
|
||||
for j in joints:
|
||||
m = j.find('mimic')
|
||||
if m is None:
|
||||
continue
|
||||
parent = m.get('joint')
|
||||
scale, offset = float(m.get('multiplier', '1')), float(m.get('offset', '0'))
|
||||
assert parent in active and scale != 0
|
||||
mimic[j.get('name')] = (parent, scale, offset)
|
||||
lo, hi = np.sort((limits[j.get('name')] - offset) / scale)
|
||||
bounds_by_name[parent] = np.array([max(bounds_by_name[parent][0], lo), min(bounds_by_name[parent][1], hi)])
|
||||
bounds = np.array([bounds_by_name[n] for n in active])
|
||||
config = dict(type='position', urdf_path=str(urdf), target_joint_names=active,
|
||||
target_link_names=[f'landmark_{i:02d}' for i in range(21)],
|
||||
target_link_human_indices=list(range(21)), low_pass_alpha=-1, normal_delta=.002)
|
||||
retarget = RetargetingConfig.from_dict(copy.deepcopy(config)).build()
|
||||
retarget.optimizer.__class__ = ConsistentPositionOptimizer
|
||||
robot = retarget.optimizer.robot
|
||||
adaptor = MimicJointKinematicAdaptor(robot, active, [m[0] for m in mimic.values()],
|
||||
list(mimic), [m[1] for m in mimic.values()], [m[2] for m in mimic.values()])
|
||||
retarget.optimizer.set_kinematic_adaptor(adaptor)
|
||||
retarget.joint_limits = bounds
|
||||
retarget.optimizer.set_joint_limit(bounds, epsilon=0)
|
||||
neutral_q = np.zeros(robot.dof)
|
||||
neutral_q[adaptor.idx_pin2target] = np.clip(np.zeros(len(active)), bounds[:, 0], bounds[:, 1])
|
||||
neutral_q = adaptor.forward_qpos(neutral_q)
|
||||
retarget.set_qpos(neutral_q)
|
||||
sites = [robot.get_link_index(f'landmark_{i:02d}') for i in range(21)]
|
||||
|
||||
def fk(q):
|
||||
robot.compute_forward_kinematics(q)
|
||||
return np.array([robot.get_link_pose(i)[:3, 3] for i in sites])
|
||||
|
||||
neutral = fk(neutral_q)
|
||||
robot_basis = palm_basis(neutral)
|
||||
human = np.load(a.input)
|
||||
hp = human['joints']
|
||||
count, fps = len(hp), float(human['fps'])
|
||||
assert hp.shape == (count, 21, 3) and np.isfinite(hp).all() and fps > 0
|
||||
targets, world_R, qraw, codes = [], [], [], []
|
||||
for t, pts in enumerate(hp):
|
||||
align = palm_basis(pts) @ robot_basis.T
|
||||
aligned = (pts - pts[0]) @ align
|
||||
target = neutral.copy()
|
||||
for chain in CHAINS:
|
||||
for i, j in zip(chain[:-1], chain[1:]):
|
||||
d = aligned[j] - aligned[i]
|
||||
assert np.linalg.norm(d) > 1e-8
|
||||
target[j] = target[i] + d / np.linalg.norm(d) * np.linalg.norm(neutral[j] - neutral[i])
|
||||
targets.append(target)
|
||||
world_R.append(Rotation.from_rotvec(human['root_orient'][t]).as_matrix() @ align)
|
||||
qraw.append(retarget.retarget(target).copy())
|
||||
codes.append(retarget.optimizer.opt.last_optimize_result())
|
||||
qraw, targets, world_R = np.array(qraw), np.array(targets), np.array(world_R)
|
||||
active_q = gaussian_filter1d(qraw[:, adaptor.idx_pin2target], a.sigma, axis=0) if a.sigma > 0 else qraw[:, adaptor.idx_pin2target]
|
||||
active_q = np.clip(active_q, bounds[:, 0], bounds[:, 1])
|
||||
qpos = []
|
||||
for row in active_q:
|
||||
q = np.zeros(robot.dof)
|
||||
q[adaptor.idx_pin2target] = row
|
||||
qpos.append(adaptor.forward_qpos(q).copy())
|
||||
qpos = np.array(qpos)
|
||||
actual = np.array([fk(q) for q in qpos])
|
||||
scene_R = world_R[0].T
|
||||
scene_t = np.array([0, 0, .2]) - scene_R @ human['wrist_world'][0]
|
||||
wrist_raw = human['wrist_world'] @ scene_R.T + scene_t
|
||||
wrist = gaussian_filter1d(wrist_raw, a.sigma, axis=0) if a.sigma > 0 else wrist_raw.copy()
|
||||
rot = Rotation.from_matrix(scene_R @ world_R)
|
||||
quat_raw = rot.as_quat()
|
||||
quat = quat_raw.copy()
|
||||
for t in range(1, count):
|
||||
if quat[t] @ quat[t-1] < 0:
|
||||
quat[t] *= -1
|
||||
if a.sigma > 0:
|
||||
quat = gaussian_filter1d(quat, a.sigma, axis=0)
|
||||
quat /= np.linalg.norm(quat, axis=1, keepdims=True)
|
||||
quat = quat[:, [3, 0, 1, 2]]
|
||||
times = np.arange(count) / fps
|
||||
np.savez_compressed(out / 'motion.npz', qpos=qpos, qpos_raw=qraw, joint_names=robot.dof_joint_names,
|
||||
active_qpos=active_q, active_joint_names=active, targets=targets, actual=actual,
|
||||
wrist_pos=wrist, wrist_pos_unsmoothed=wrist_raw, wrist_quat_wxyz=quat,
|
||||
wrist_quat_unsmoothed_wxyz=quat_raw[:, [3, 0, 1, 2]],
|
||||
wrist_world=human['wrist_world'], wrist_world_R=world_R,
|
||||
scene_rotation=scene_R, scene_translation=scene_t, time=times, fps=fps,
|
||||
coupling='linear_urdf', offline_sigma_frames=a.sigma,
|
||||
source=str(a.input.resolve()), detection_valid=human['detection_valid'])
|
||||
np.savetxt(out / 'trajectory.csv', np.c_[times, wrist, quat, qpos], delimiter=',',
|
||||
header=','.join(['time_s','wrist_x','wrist_y','wrist_z','qw','qx','qy','qz'] + robot.dof_joint_names), comments='')
|
||||
model = mujoco.MjModel.from_xml_path(str(fixed))
|
||||
data = mujoco.MjData(model)
|
||||
addresses = [model.jnt_qposadr[model.joint(n).id] for n in robot.dof_joint_names]
|
||||
mids = [model.site(f'landmark_{i:02d}').id for i in range(21)]
|
||||
fk_error = 0.
|
||||
for t, row in enumerate(qpos):
|
||||
data.qpos[addresses] = row
|
||||
mujoco.mj_forward(model, data)
|
||||
fk_error = max(fk_error, float(np.abs(data.site_xpos[mids] - actual[t]).max()))
|
||||
lim = np.array([limits[n] for n in robot.dof_joint_names])
|
||||
violation = float(np.maximum(lim[:, 0] - qpos, 0).max() + np.maximum(qpos - lim[:, 1], 0).max())
|
||||
names = list(robot.dof_joint_names)
|
||||
mimic_error = max(float(np.abs(qpos[:, names.index(n)] - (sc*qpos[:, names.index(parent)] + off)).max()) for n,(parent,sc,off) in mimic.items())
|
||||
assert violation < 1e-8 and mimic_error < 1e-10 and fk_error < 1e-5
|
||||
assert all(np.isfinite(x).all() for x in [qpos, wrist, quat, actual])
|
||||
error = np.linalg.norm(actual[:, TIPS] - targets[:, TIPS], axis=-1) * 1000
|
||||
report = dict(frames=count, fps=fps, duration_seconds=count/fps, source=str(a.input.resolve()),
|
||||
coupling='linear_urdf', independent_joints=len(active), urdf_joints=len(names),
|
||||
tip_error_mean_mm=float(error.mean()), tip_error_p95_mm=float(np.percentile(error,95)),
|
||||
tip_error_finger_mm=dict(zip(['thumb','index','middle','ring','pinky'],error.mean(0).tolist())),
|
||||
max_joint_step_deg=float(np.rad2deg(np.abs(np.diff(qpos,axis=0)).max())),
|
||||
joint_second_difference_rms_deg=float(np.rad2deg(np.sqrt(np.mean(np.diff(qpos,n=2,axis=0)**2)))),
|
||||
wrist_translation_range_m=np.ptp(wrist,axis=0).tolist(),
|
||||
wrist_smoothing_max_displacement_mm=float(np.linalg.norm(wrist-wrist_raw,axis=1).max()*1000),
|
||||
urdf_limit_violation_rad=violation, linear_mimic_max_error_rad=mimic_error,
|
||||
pinocchio_mujoco_fk_error_m=fk_error, fk_checked_frames=count, all_finite=True,
|
||||
quaternion_norm_max_error=float(np.abs(np.linalg.norm(quat,axis=1)-1).max()),
|
||||
offline_sigma_frames=a.sigma, uses_future_frames=a.sigma>0,
|
||||
optimizer_status_counts={str(k):int(v) for k,v in zip(*np.unique(codes,return_counts=True))},
|
||||
boundary='Hand-only kinematic retargeting; no object tracking, contact optimization, dynamics or hardware commands')
|
||||
(out / 'retarget_validation.json').write_text(json.dumps(report,indent=2))
|
||||
(out / 'dex_config.json').write_text(json.dumps(config,indent=2))
|
||||
tree = ET.parse(fixed)
|
||||
root = tree.getroot()
|
||||
ET.SubElement(root.find("worldbody/body[@name='hand_base_link']"), 'freejoint', name='wrist_free')
|
||||
root.find('option').set('gravity','0 0 0')
|
||||
# Replay supplies all angles; physical mimic equality is not needed for mj_forward.
|
||||
tree.write(out / 'l20_moving.xml')
|
||||
print(json.dumps(report,indent=2))
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,5 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
cd "$(dirname "$0")/.."
|
||||
env PYTHONPATH= LD_LIBRARY_PATH= third_party/Dyn-HaMR/.dynhamr/bin/python scripts/export_dynhamr_dex.py
|
||||
env PYTHONPATH= LD_LIBRARY_PATH= MUJOCO_GL=osmesa .dex/bin/python scripts/run_dex_leap.py "$@"
|
||||
@@ -0,0 +1,127 @@
|
||||
"""Retarget Dyn-HaMR to the supplied L20 with JSON coupling and moving wrist."""
|
||||
import os
|
||||
os.environ.setdefault('MUJOCO_GL','osmesa')
|
||||
os.environ.setdefault('OMP_NUM_THREADS','4')
|
||||
import argparse,json,time
|
||||
from pathlib import Path
|
||||
import xml.etree.ElementTree as ET
|
||||
import numpy as np
|
||||
import torch,mujoco
|
||||
from scipy.spatial.transform import Rotation
|
||||
from scipy.ndimage import gaussian_filter1d
|
||||
from dex_retargeting.retargeting_config import RetargetingConfig
|
||||
from l20_calibrated import ROOT,Calibration,CalibratedAdaptor,build_assets,palm_basis,CHAINS,TIPS
|
||||
|
||||
OUT=ROOT/'output/l20_2047635068'
|
||||
|
||||
def main():
|
||||
global OUT
|
||||
parser=argparse.ArgumentParser();parser.add_argument('--frames',type=int,default=0)
|
||||
parser.add_argument('--output-dir',type=Path,default=OUT)
|
||||
parser.add_argument('--consistent-objective',action=argparse.BooleanOptionalAction,default=True)
|
||||
args=parser.parse_args();OUT=args.output_dir.resolve()
|
||||
torch.set_num_threads(4)
|
||||
fixed,urdf=build_assets(OUT)
|
||||
cal=Calibration()
|
||||
config=dict(type='position',urdf_path=str(urdf),target_joint_names=cal.active,
|
||||
target_link_names=[f'landmark_{i:02d}' for i in range(21)],
|
||||
target_link_human_indices=list(range(21)),low_pass_alpha=-1,normal_delta=.002)
|
||||
retarget=RetargetingConfig.from_dict(config.copy()).build()
|
||||
if args.consistent_objective:
|
||||
from l20_consistent_optimizer import ConsistentPositionOptimizer
|
||||
retarget.optimizer.__class__=ConsistentPositionOptimizer
|
||||
robot=retarget.optimizer.robot
|
||||
adaptor=CalibratedAdaptor(robot,cal);retarget.optimizer.set_kinematic_adaptor(adaptor)
|
||||
bounds=np.array([cal.bounds[n] for n in cal.active]);retarget.joint_limits=bounds
|
||||
retarget.optimizer.set_joint_limit(bounds,epsilon=0)
|
||||
neutral_q=np.zeros(robot.dof)
|
||||
for i,n in enumerate(cal.active):neutral_q[robot.get_joint_index(n)]=np.clip(cal.values[n][cal.raw['baseline_command_u8'][cal.tables[n]['motor_index']]],*bounds[i])
|
||||
neutral_q=adaptor.forward_qpos(neutral_q)
|
||||
retarget.set_qpos(neutral_q)
|
||||
sites=[robot.get_link_index(f'landmark_{i:02d}') for i in range(21)]
|
||||
def fk(q):
|
||||
robot.compute_forward_kinematics(q)
|
||||
return np.array([robot.get_link_pose(k)[:3,3] for k in sites])
|
||||
neutral=fk(neutral_q);basis_robot=palm_basis(neutral)
|
||||
human=np.load(ROOT/'output/dex_2047635068/human_joints.npz')
|
||||
joints=human['joints'][:args.frames or None]
|
||||
targets=[];wrist_rot=[]
|
||||
for t,p in enumerate(joints):
|
||||
basis=palm_basis(p)
|
||||
align=basis@basis_robot.T
|
||||
aligned=(p-p[0])@align
|
||||
target=neutral.copy()
|
||||
for chain in CHAINS:
|
||||
for a,b in zip(chain[:-1],chain[1:]):
|
||||
direction=aligned[b]-aligned[a]
|
||||
target[b]=target[a]+direction/np.linalg.norm(direction)*np.linalg.norm(neutral[b]-neutral[a])
|
||||
targets.append(target)
|
||||
wrist_rot.append(Rotation.from_rotvec(human['root_orient'][t]).as_matrix()@align)
|
||||
targets=np.array(targets);wrist_rot=np.array(wrist_rot)
|
||||
qraw=[];codes=[];start=time.monotonic()
|
||||
for t,target in enumerate(targets):
|
||||
q=retarget.retarget(target)
|
||||
qraw.append(q.copy());codes.append(retarget.optimizer.opt.last_optimize_result())
|
||||
if t%200==0: print(f'L20 dex {t}/{len(targets)}',flush=True)
|
||||
qraw=np.array(qraw)
|
||||
active=qraw[:,adaptor.idx_pin2target]
|
||||
active=np.clip(gaussian_filter1d(active,1,axis=0),bounds[:,0],bounds[:,1])
|
||||
qpos=[]
|
||||
for a in active:
|
||||
q=np.zeros(robot.dof);q[adaptor.idx_pin2target]=a;qpos.append(adaptor.forward_qpos(q))
|
||||
qpos=np.array(qpos);actual=np.array([fk(q) for q in qpos])
|
||||
commands=np.array([cal.command(dict(zip(robot.dof_joint_names,q))) for q in qpos])
|
||||
decoded=np.array([cal.decode(c,robot.dof_joint_names) for c in commands])
|
||||
wrist_world=human['wrist_world'][:len(qpos)]
|
||||
# A single rigid transform makes the initial palm upright for visualization.
|
||||
scene_rot=wrist_rot[0].T
|
||||
wrist_pos=(wrist_world-wrist_world[0])@scene_rot.T+np.array([0,0,.2])
|
||||
wrist_scene=np.einsum('ij,tjk->tik',scene_rot,wrist_rot)
|
||||
quat=Rotation.from_matrix(wrist_scene).as_quat()[:,[3,0,1,2]]
|
||||
for i in range(1,len(quat)):
|
||||
if quat[i]@quat[i-1]<0:quat[i]*=-1
|
||||
np.savez_compressed(OUT/'motion.npz',qpos=qpos,qpos_raw=qraw,joint_names=robot.dof_joint_names,
|
||||
active_qpos=active,active_joint_names=cal.active,command_u8=commands,decoded_qpos=decoded,
|
||||
wrist_pos=wrist_pos,wrist_quat_wxyz=quat,wrist_world=wrist_world,wrist_world_R=wrist_rot,
|
||||
scene_rotation=scene_rot,scene_translation=np.array([0,0,.2])-scene_rot@wrist_world[0],
|
||||
targets=targets,actual=actual,fps=30,time=np.arange(len(qpos))/30)
|
||||
np.savetxt(OUT/'trajectory.csv',np.c_[np.arange(len(qpos))/30,wrist_pos,quat,qpos],delimiter=',',
|
||||
header=','.join(['time_s','wrist_x','wrist_y','wrist_z','qw','qx','qy','qz']+robot.dof_joint_names),comments='')
|
||||
np.savetxt(OUT/'commands_u8.csv',np.c_[np.arange(len(qpos)),commands],delimiter=',',fmt='%d',
|
||||
header=','.join(['frame']+[f'slot_{i}' for i in range(20)]),comments='')
|
||||
limits=np.array([cal.urdf_limits[n] for n in robot.dof_joint_names])
|
||||
violation=float(np.maximum(limits[:,0]-qpos,0).max()+np.maximum(qpos-limits[:,1],0).max())
|
||||
assert violation<1e-6 and np.isfinite(qpos).all()
|
||||
err=np.linalg.norm(actual[:,TIPS]-targets[:,TIPS],axis=-1)*1000
|
||||
# Compare independent FK libraries, with all 21 nonlinear-coupled angles supplied.
|
||||
model=mujoco.MjModel.from_xml_path(str(fixed));data=mujoco.MjData(model)
|
||||
addresses=[model.jnt_qposadr[mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_JOINT,n)] for n in robot.dof_joint_names]
|
||||
ids=[mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_SITE,f'landmark_{i:02d}') for i in range(21)]
|
||||
fk_error=0
|
||||
for t in np.linspace(0,len(qpos)-1,min(20,len(qpos)),dtype=int):
|
||||
data.qpos[addresses]=qpos[t];mujoco.mj_forward(model,data)
|
||||
fk_error=max(fk_error,float(np.abs(data.site_xpos[ids]-actual[t]).max()))
|
||||
assert fk_error<1e-5
|
||||
report=dict(frames=len(qpos),fps=30,independent_joints=16,urdf_joints=21,
|
||||
calibration_model=cal.raw['model'],source_directory_name='L20',
|
||||
calibration_interpretation='q_urdf=angle_rad[command]+urdf_zero_offset_rad; passive joints share motor command; not hardware verified',
|
||||
reserved_command_slots=[11,12,13,14],tip_error_mean_mm=float(err.mean()),tip_error_p95_mm=float(np.percentile(err,95)),
|
||||
tip_error_finger_mm=dict(zip(['thumb','index','middle','ring','pinky'],err.mean(0).tolist())),
|
||||
wrist_translation_range_m=np.ptp(wrist_pos,axis=0).tolist(),
|
||||
wrist_rotation_max_from_initial_rad=float(np.linalg.norm(Rotation.from_matrix(wrist_scene).as_rotvec(),axis=1).max()),
|
||||
urdf_limit_violation_rad=violation,quantization_max_error_rad=float(np.abs(decoded-qpos).max()),
|
||||
pinocchio_mujoco_fk_error_m=fk_error,seconds=time.monotonic()-start,
|
||||
optimizer_status_counts={str(k):int(v) for k,v in zip(*np.unique(codes,return_counts=True))},
|
||||
consistent_objective=args.consistent_objective,
|
||||
object_contact_stage='not applied in this hand-only retargeting run')
|
||||
(OUT/'retarget_validation.json').write_text(json.dumps(report,indent=2))
|
||||
(OUT/'dex_config.json').write_text(json.dumps(config,indent=2))
|
||||
# Moving wrist replay model; contact scene construction is a separate stage.
|
||||
tree=ET.parse(fixed);root=tree.getroot()
|
||||
body=root.find("worldbody/body[@name='hand_base_link']")
|
||||
ET.SubElement(body,'freejoint',name='wrist_free')
|
||||
root.find('option').set('gravity','0 0 0')
|
||||
tree.write(OUT/'l20_moving.xml')
|
||||
print(json.dumps(report,indent=2),flush=True)
|
||||
|
||||
if __name__=='__main__':main()
|
||||
@@ -0,0 +1,165 @@
|
||||
"""Offline right-hand vector retargeting and MuJoCo kinematic replay."""
|
||||
import os
|
||||
os.environ.setdefault('MUJOCO_GL', 'osmesa')
|
||||
os.environ.setdefault('OMP_NUM_THREADS', '4')
|
||||
import argparse
|
||||
import json
|
||||
import time
|
||||
from pathlib import Path
|
||||
import xml.etree.ElementTree as ET
|
||||
import numpy as np
|
||||
import torch
|
||||
import mujoco
|
||||
import imageio.v2 as imageio
|
||||
from PIL import Image, ImageDraw
|
||||
from dex_retargeting.constants import RobotName, RetargetingType, HandType, get_default_config_path, OPERATOR2MANO_RIGHT
|
||||
from dex_retargeting.retargeting_config import RetargetingConfig
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
OUT = ROOT / 'output/dex_2047635068'
|
||||
ASSETS = ROOT / 'third_party/dex-assets/robots/hands'
|
||||
|
||||
def operator_frame(joints):
|
||||
# Same palm frame construction as upstream SingleHandDetector, without MediaPipe.
|
||||
points = joints[[0, 5, 9]]
|
||||
x = points[0] - points[2]
|
||||
points = points - points.mean(axis=0)
|
||||
_, _, v = np.linalg.svd(points)
|
||||
normal = v[2]
|
||||
x -= x.dot(normal) * normal
|
||||
x /= np.linalg.norm(x)
|
||||
z = np.cross(x, normal)
|
||||
if z.dot(points[1] - points[2]) < 0:
|
||||
normal *= -1
|
||||
z *= -1
|
||||
frame = np.stack([x, normal, z], axis=1)
|
||||
assert np.allclose(frame.T @ frame, np.eye(3), atol=1e-5)
|
||||
return frame @ OPERATOR2MANO_RIGHT
|
||||
|
||||
def load_scene(urdf):
|
||||
tree = ET.parse(urdf)
|
||||
root = tree.getroot()
|
||||
ext = ET.SubElement(root, 'mujoco')
|
||||
ET.SubElement(ext, 'compiler', discardvisual='false', fusestatic='false', strippath='false')
|
||||
for mesh in root.iter('mesh'):
|
||||
mesh.set('filename', str(urdf.parent / mesh.get('filename')))
|
||||
adapted = OUT / 'leap_replay.urdf'
|
||||
tree.write(adapted)
|
||||
model = mujoco.MjModel.from_xml_path(str(adapted))
|
||||
scene = OUT / 'scene.xml'
|
||||
mujoco.mj_saveLastXML(str(scene), model)
|
||||
tree = ET.parse(scene)
|
||||
root = tree.getroot()
|
||||
visual = ET.SubElement(root, 'visual')
|
||||
ET.SubElement(visual, 'global', offwidth='1280', offheight='480')
|
||||
ET.SubElement(visual, 'headlight', ambient='0.5 0.5 0.5', diffuse='0.7 0.7 0.7')
|
||||
ET.SubElement(root.find('worldbody'), 'light', pos='0 -1 1', dir='0 1 -1', directional='true')
|
||||
tree.write(scene)
|
||||
model = mujoco.MjModel.from_xml_path(str(scene))
|
||||
# Replay is kinematic. Retain only visual geometry in the viewer.
|
||||
model.geom_rgba[:, 3] = 1
|
||||
# LEAP fingertip pads are supplied only as collision meshes in this asset.
|
||||
# Display these meshes too, while keeping collision primitives hidden.
|
||||
model.geom_group[model.geom_type == mujoco.mjtGeom.mjGEOM_MESH] = 1
|
||||
return model
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--frames', type=int, default=0)
|
||||
parser.add_argument('--no-render', action='store_true')
|
||||
args = parser.parse_args()
|
||||
torch.set_num_threads(4)
|
||||
human = np.load(OUT / 'human_joints.npz')
|
||||
joints = human['joints'][:args.frames or None]
|
||||
RetargetingConfig.set_default_urdf_dir(ASSETS)
|
||||
config_path = get_default_config_path(RobotName.leap, RetargetingType.vector, HandType.right)
|
||||
cfg = RetargetingConfig.load_from_file(config_path)
|
||||
(OUT / 'retargeting_config.yml').write_text(Path(config_path).read_text())
|
||||
retarget = cfg.build()
|
||||
robot = retarget.optimizer.robot
|
||||
index = retarget.optimizer.target_link_human_indices
|
||||
tip_ids = [robot.get_link_index(x) for x in cfg.target_task_link_names]
|
||||
base_id = robot.get_link_index('base')
|
||||
qs, targets, actual, codes = [], [], [], []
|
||||
start = time.monotonic()
|
||||
for i, j in enumerate(joints):
|
||||
local = (j - j[:1]) @ operator_frame(j)
|
||||
ref = local[index[1]] - local[index[0]]
|
||||
q = retarget.retarget(ref)
|
||||
robot.compute_forward_kinematics(q)
|
||||
base = robot.get_link_pose(base_id)[:3, 3]
|
||||
tips = np.array([robot.get_link_pose(k)[:3, 3] for k in tip_ids])
|
||||
qs.append(q); targets.append(ref * cfg.scaling_factor + base); actual.append(tips)
|
||||
codes.append(retarget.optimizer.opt.last_optimize_result())
|
||||
if i % 200 == 0:
|
||||
print(f'retarget {i}/{len(joints)}', flush=True)
|
||||
qpos, targets, actual = map(np.array, (qs, targets, actual))
|
||||
elapsed = time.monotonic() - start
|
||||
limits = robot.joint_limits
|
||||
assert np.isfinite(qpos).all()
|
||||
assert np.all(qpos >= limits[:, 0] - 1e-6) and np.all(qpos <= limits[:, 1] + 1e-6)
|
||||
np.savez_compressed(OUT / 'leap_right_qpos.npz', qpos=qpos, joint_names=retarget.joint_names,
|
||||
fps=30, time=np.arange(len(qpos))/30, target_tips=targets,
|
||||
actual_tips=actual, joint_limits=limits, source_track=human['source_track'])
|
||||
np.savetxt(OUT / 'leap_right_qpos.csv', np.c_[np.arange(len(qpos))/30, qpos], delimiter=',',
|
||||
header=','.join(['time_s'] + retarget.joint_names), comments='')
|
||||
errors = np.linalg.norm(actual-targets, axis=2)*1000
|
||||
stats = dict(frames=len(qpos), fps=30, dof=qpos.shape[1], finite=True, joint_limits_pass=True,
|
||||
retarget_seconds=elapsed, tip_error_mean_mm=float(errors.mean()),
|
||||
tip_error_p95_mm=float(np.percentile(errors,95)),
|
||||
tip_error_by_finger_mm=dict(zip(['thumb','index','middle','ring'],errors.mean(0).tolist())),
|
||||
max_frame_joint_step_rad=float(np.max(np.abs(np.diff(qpos,axis=0)))),
|
||||
optimizer_status_counts={str(k):int(v) for k,v in zip(*np.unique(codes,return_counts=True))},
|
||||
scope='Fixed-wrist kinematic replay; no object, contact, dynamics, or real-robot validation.')
|
||||
(OUT / 'validation.json').write_text(json.dumps(stats, indent=2))
|
||||
print(json.dumps(stats, indent=2), flush=True)
|
||||
if args.no_render:
|
||||
return
|
||||
model = load_scene(ASSETS / 'leap_hand/leap_hand_right.urdf')
|
||||
data = mujoco.MjData(model)
|
||||
addresses = [model.jnt_qposadr[mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_JOINT,n)] for n in retarget.joint_names]
|
||||
assert len(set(addresses)) == len(addresses) == model.nq
|
||||
tip_bodies = [mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_BODY,n) for n in cfg.target_task_link_names]
|
||||
assert min(tip_bodies) >= 0
|
||||
fk_errors = []
|
||||
for i in np.linspace(0,len(qpos)-1,min(20,len(qpos)),dtype=int):
|
||||
data.qpos[addresses] = qpos[i]
|
||||
mujoco.mj_forward(model,data)
|
||||
fk_errors.append(np.max(np.abs(data.xpos[tip_bodies]-actual[i])))
|
||||
assert max(fk_errors) < 1e-5, f'Pinocchio/MuJoCo FK mismatch: {max(fk_errors)}'
|
||||
stats['pinocchio_mujoco_max_fk_error_m'] = float(max(fk_errors))
|
||||
renderer = mujoco.Renderer(model, height=480, width=640)
|
||||
cameras=[]
|
||||
for azimuth in [135, 225]:
|
||||
cam=mujoco.MjvCamera(); cam.lookat[:]=model.stat.center;cam.distance=.65;cam.azimuth=azimuth;cam.elevation=15
|
||||
cameras.append(cam)
|
||||
option=mujoco.MjvOption()
|
||||
option.geomgroup[0]=0 # URDF collision geometry; show visual group 1 only.
|
||||
writer=imageio.get_writer(OUT/'leap_right_replay.mp4',fps=30,codec='libx264',quality=8)
|
||||
try:
|
||||
for i,q in enumerate(qpos):
|
||||
data.qpos[addresses]=q; mujoco.mj_forward(model,data)
|
||||
panels=[]
|
||||
for cam in cameras:
|
||||
renderer.update_scene(data,camera=cam,scene_option=option)
|
||||
for p in targets[i]:
|
||||
scene=renderer.scene
|
||||
mujoco.mjv_initGeom(scene.geoms[scene.ngeom],mujoco.mjtGeom.mjGEOM_SPHERE,
|
||||
np.array([.004]*3),p,np.eye(3).ravel(),np.array([1,.25,.1,1]))
|
||||
scene.ngeom+=1
|
||||
panels.append(renderer.render().copy())
|
||||
img=Image.fromarray(np.concatenate(panels,axis=1)); draw=ImageDraw.Draw(img)
|
||||
draw.rectangle([0,0,1280,46],fill=(20,25,32))
|
||||
draw.text((12,8),f'LEAP right | frame {i+1}/{len(qpos)} | red dots: scaled human fingertips | fixed wrist',fill='white')
|
||||
draw.text((12,26),f'Mean fingertip error: {errors[i].mean():.1f} mm | kinematic replay',fill='white')
|
||||
writer.append_data(np.asarray(img))
|
||||
if i in [0,len(qpos)//2,len(qpos)-1]: img.save(OUT/f'preview_{i:04d}.png')
|
||||
if i % 200==0: print(f'render {i}/{len(qpos)}',flush=True)
|
||||
finally:
|
||||
writer.close(); renderer.close()
|
||||
stats['rendered_frames']=len(qpos)
|
||||
(OUT/'validation.json').write_text(json.dumps(stats,indent=2))
|
||||
print('COMPLETE',flush=True)
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,44 @@
|
||||
"""Resume the requested egocentric video inference with existing local models."""
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import subprocess
|
||||
import time
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
OUT = ROOT / 'output/15886123_handflow'
|
||||
OUT.mkdir(parents=True, exist_ok=True)
|
||||
env = os.environ.copy()
|
||||
env.update(PYTHONPATH=str(ROOT), LD_LIBRARY_PATH='', OMP_NUM_THREADS='4',
|
||||
MKL_NUM_THREADS='4', OPENBLAS_NUM_THREADS='4',
|
||||
TORCH_HOME=str(ROOT / '.torch_cache'), HF_HOME=str(ROOT / '.hf_cache'),
|
||||
HF_HUB_OFFLINE='1', TRANSFORMERS_OFFLINE='1',
|
||||
TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD='1', PYTHONUNBUFFERED='1',
|
||||
VIPE_PYTHON=str(ROOT / '.venv/bin/python'),
|
||||
HAMER_CKPT=str(ROOT / 'third_party/hamer/_DATA/hamer_ckpts/checkpoints/hamer.ckpt'),
|
||||
DETECTOR_CKPT=str(ROOT / 'weights/detector.pt'),
|
||||
MANO_ROOT=str(ROOT / 'third_party/hamer/_DATA/data/mano'),
|
||||
HANDFLOW_NORMALIZATION_STATS=str(ROOT / 'weights/normalization_stats.npz'))
|
||||
cmd = [str(ROOT / '.venv/bin/python'), '-u', 'scripts/demo.py',
|
||||
'--input', '/home/timessage/下载/15886123.mp4',
|
||||
'--fm_ckpt', str(ROOT / 'weights/handflow_denoiser.pt'),
|
||||
'--intrinsics', 'auto', '--side', 'right', '--rotation_smooth_window', '5',
|
||||
'--output_dir', str(OUT), '--save_npz', str(OUT / 'handflow_results.npz')]
|
||||
|
||||
def status(stage, **extra):
|
||||
(OUT / 'status.json').write_text(json.dumps(dict(stage=stage, time=time.strftime('%Y-%m-%d %H:%M:%S'), **extra), indent=2))
|
||||
print(stage, flush=True)
|
||||
|
||||
try:
|
||||
for attempt in range(2):
|
||||
if (OUT / 'handflow_results.npz').exists():
|
||||
break
|
||||
status('INFERENCE', pass_index=attempt + 1, command=cmd)
|
||||
with (OUT / f'inference_pass_{attempt + 1}.log').open('a') as log:
|
||||
subprocess.run(cmd, cwd=ROOT, env=env, stdout=log, stderr=subprocess.STDOUT, check=True)
|
||||
if not (OUT / 'handflow_results.npz').exists():
|
||||
raise RuntimeError('Inference exited without final arrays after cache restart')
|
||||
status('INFERENCE_FINISHED_PENDING_VALIDATION')
|
||||
except Exception as exc:
|
||||
status('FAILED', error=str(exc))
|
||||
raise
|
||||
@@ -0,0 +1,10 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
SPIDER_WORKSPACE="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)"
|
||||
export PYTHONPATH=""
|
||||
export LD_LIBRARY_PATH=""
|
||||
export WARP_CACHE_PATH="$SPIDER_WORKSPACE/.spider_cache/warp"
|
||||
export TORCHINDUCTOR_CACHE_DIR="$SPIDER_WORKSPACE/.spider_cache/torchinductor"
|
||||
export MUJOCO_GL="${MUJOCO_GL:-osmesa}"
|
||||
cd "$SPIDER_WORKSPACE/third_party/spider"
|
||||
exec "$SPIDER_WORKSPACE/.spider/bin/python" examples/run_mjwp_fast.py "$@"
|
||||
@@ -0,0 +1,18 @@
|
||||
"""Run upstream SPIDER contact guidance on the prepared L20 task."""
|
||||
import os
|
||||
from pathlib import Path
|
||||
ROOT=Path(__file__).resolve().parents[1]
|
||||
os.environ['WARP_CACHE_PATH']=str(ROOT/'.spider_cache/warp')
|
||||
os.environ['TORCHINDUCTOR_CACHE_DIR']=str(ROOT/'.spider_cache/torchinductor')
|
||||
os.environ.setdefault('MUJOCO_GL','osmesa')
|
||||
os.environ.setdefault('OMP_NUM_THREADS','4')
|
||||
import sys,json,importlib.util
|
||||
sys.path.insert(0,str(ROOT/'third_party/spider'))
|
||||
import torch
|
||||
torch.set_num_threads(4)
|
||||
from spider.config import Config
|
||||
spec=importlib.util.spec_from_file_location('spider_original_runner',ROOT/'third_party/spider/examples/run_mjwp.py')
|
||||
module=importlib.util.module_from_spec(spec);sys.modules[spec.name]=module;spec.loader.exec_module(module)
|
||||
config=Config(**json.loads((ROOT/'output/spider_l20_contact/config.json').read_text()))
|
||||
module.main(config)
|
||||
print('SPIDER_RUN_COMPLETE',flush=True)
|
||||
@@ -0,0 +1,39 @@
|
||||
"""Explicit offline smoothing ablation, preserving calibration and moving wrist."""
|
||||
import json,shutil
|
||||
import numpy as np
|
||||
import mujoco
|
||||
from scipy.ndimage import gaussian_filter1d
|
||||
from scipy.spatial.transform import Rotation
|
||||
from l20_calibrated import ROOT,Calibration
|
||||
base=ROOT/'output/l20_2047635068/jitter_audit';out=base/'temporal';out.mkdir(exist_ok=True)
|
||||
m=dict(np.load(base/'corrected/motion.npz'));cal=Calibration();names=list(m['joint_names'])
|
||||
sigma=2.5
|
||||
a=gaussian_filter1d(m['active_qpos'],sigma,axis=0)
|
||||
q=[]
|
||||
for row in a:
|
||||
d=dict(zip(cal.active,row))
|
||||
for n in cal.passive:d[n]=cal.passive_value(n,d[cal.curves[n][0]])[0]
|
||||
q.append([d[n] for n in names])
|
||||
m['qpos']=np.array(q);m['active_qpos']=a
|
||||
m['wrist_pos_unsmoothed']=m['wrist_pos'].copy();m['wrist_quat_unsmoothed_wxyz']=m['wrist_quat_wxyz'].copy()
|
||||
m['wrist_pos']=gaussian_filter1d(m['wrist_pos'],sigma,axis=0)
|
||||
rot=Rotation.from_quat(m['wrist_quat_wxyz'][:,[1,2,3,0]])
|
||||
relative=(rot[0].inv()*rot).as_rotvec();assert np.linalg.norm(relative,axis=1).max()<np.pi/2
|
||||
m['wrist_quat_wxyz']=(rot[0]*Rotation.from_rotvec(gaussian_filter1d(relative,sigma,axis=0))).as_quat()[:,[3,0,1,2]]
|
||||
for n in ['l20_moving.xml','l20_right.xml']:shutil.copy2(base/'corrected'/n,out/n)
|
||||
model=mujoco.MjModel.from_xml_path(str(out/'l20_moving.xml'));data=mujoco.MjData(model)
|
||||
addr=[model.jnt_qposadr[model.joint(str(n)).id] for n in names];sites=[model.site(f'landmark_{i:02d}').id for i in range(21)]
|
||||
actual=[]
|
||||
for q in m['qpos']:
|
||||
data.qpos[addr]=q;mujoco.mj_forward(model,data);actual.append(data.site_xpos[sites].copy())
|
||||
m['actual']=np.array(actual);m['command_u8']=np.array([cal.command(dict(zip(names,q))) for q in m['qpos']]);m['decoded_qpos']=np.array([cal.decode(c,names) for c in m['command_u8']])
|
||||
m['offline_smoothing_sigma_frames']=sigma
|
||||
np.savez_compressed(out/'motion.npz',**m)
|
||||
np.savetxt(out/'trajectory.csv',np.c_[m['time'],m['wrist_pos'],m['wrist_quat_wxyz'],m['qpos']],delimiter=',',header=','.join(['time_s','wrist_x','wrist_y','wrist_z','qw','qx','qy','qz']+names),comments='')
|
||||
limits=np.array([cal.urdf_limits[n] for n in names]);r=dict(frames=len(a),offline_sigma_frames=sigma,uses_future_frames=True,
|
||||
tip_error_mean_mm=float(np.linalg.norm(m['actual'][:,[4,8,12,16,20]]-m['targets'][:,[4,8,12,16,20]],axis=-1).mean()*1000),
|
||||
max_joint_step_deg=float(np.abs(np.diff(m['qpos'],axis=0)).max()*180/np.pi),joint_second_difference_rms_deg=float(np.sqrt(np.mean(np.diff(m['qpos'],n=2,axis=0)**2))*180/np.pi),
|
||||
wrist_second_difference_rms_mm=float(np.sqrt(np.mean(np.diff(m['wrist_pos'],n=2,axis=0)**2))*1000),wrist_max_change_from_source_mm=float(np.linalg.norm(m['wrist_pos']-m['wrist_pos_unsmoothed'],axis=1).max()*1000),
|
||||
wrist_range_m=np.ptp(m['wrist_pos'],axis=0).tolist(),limit_violations=int(np.sum((m['qpos']<limits[:,0]-1e-8)|(m['qpos']>limits[:,1]+1e-8))),
|
||||
boundary='Offline hand-only smoothing ablation, not contact or hardware validation')
|
||||
(out/'validation.json').write_text(json.dumps(r,indent=2));print(json.dumps(r,indent=2))
|
||||
@@ -0,0 +1,149 @@
|
||||
"""Offline whole-sequence temporal regularization with explicit distortion budgets."""
|
||||
import argparse
|
||||
import json
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
import xml.etree.ElementTree as ET
|
||||
|
||||
import mujoco
|
||||
import numpy as np
|
||||
from scipy import sparse
|
||||
from scipy.sparse.linalg import splu
|
||||
from scipy.spatial.transform import Rotation
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
|
||||
|
||||
def main():
|
||||
p = argparse.ArgumentParser()
|
||||
p.add_argument('--source', type=Path, default=ROOT/'output/l20_15886123')
|
||||
p.add_argument('--output', type=Path, default=ROOT/'output/l20_15886123_stable')
|
||||
a = p.parse_args()
|
||||
a.output.mkdir(parents=True, exist_ok=True)
|
||||
with np.load(a.source/'motion.npz') as f:
|
||||
m = {k:f[k] for k in f.files}
|
||||
n = len(m['qpos'])
|
||||
D2 = sparse.diags([np.ones(n-2), -2*np.ones(n-2), np.ones(n-2)], [0,1,2], shape=(n-2,n),format='csc')
|
||||
D3 = sparse.diags([-np.ones(n-3),3*np.ones(n-3),-3*np.ones(n-3),np.ones(n-3)], [0,1,2,3],shape=(n-3,n),format='csc')
|
||||
eye = sparse.eye(n,format='csc')
|
||||
names, active = list(m['joint_names']), list(m['active_joint_names'])
|
||||
ai = [names.index(k) for k in active]
|
||||
original = ET.parse(ROOT/'third_party/l20_assets/L20/RIGHT/linkerhand_g20_right.urdf').getroot()
|
||||
limits = {j.get('name'):np.array([float(j.find('limit').get(k)) for k in ['lower','upper']]) for j in original.findall('joint') if j.get('type')!='fixed'}
|
||||
low, high = np.array([limits[k] for k in active]).T.copy()
|
||||
mimic = {}
|
||||
for j in original.findall('joint'):
|
||||
mi = j.find('mimic')
|
||||
if mi is None:
|
||||
continue
|
||||
par, sc, off = mi.get('joint'),float(mi.get('multiplier','1')),float(mi.get('offset','0'))
|
||||
mimic[j.get('name')] = par,sc,off
|
||||
lo,hi = np.sort((limits[j.get('name')]-off)/sc)
|
||||
k = active.index(par)
|
||||
low[k],high[k] = max(low[k],lo),min(high[k],hi)
|
||||
|
||||
def smooth(x, weight, bounds=None):
|
||||
# Sample-domain objective ||x-y||² + ||D²x||² + weight*||D³x||².
|
||||
A = eye + D2.T@D2 + weight*(D3.T@D3)
|
||||
result = splu(A).solve(x)
|
||||
if bounds is not None:
|
||||
lo,hi = bounds
|
||||
# Explicit post-solve projection; evaluate smoothness after projection.
|
||||
# This is not claimed to solve the box-constrained temporal QP exactly.
|
||||
result = np.clip(result,lo,hi)
|
||||
return result
|
||||
|
||||
def d2(x):
|
||||
return float(np.sqrt(np.mean(np.diff(x[30:],n=2,axis=0)**2)))
|
||||
|
||||
rot = Rotation.from_quat(m['wrist_quat_wxyz'][:,[1,2,3,0]])
|
||||
rv = (rot[0].inv()*rot).as_rotvec()
|
||||
assert np.linalg.norm(rv,axis=1).max() < np.pi/2
|
||||
choices, sweep = {}, {}
|
||||
# Budgets are engineering choices relative to prior output, not true-motion tolerances.
|
||||
components = dict(fingers=(m['qpos'][:,ai],(low,high),np.deg2rad(8)),
|
||||
wrist_position=(m['wrist_pos'],None,.020),
|
||||
wrist_rotation=(rv,None,np.deg2rad(6)))
|
||||
for key,(x,bounds,budget) in components.items():
|
||||
rows, candidates = [], []
|
||||
for weight in [100.,300.,1000.,3000.,10000.,30000.]:
|
||||
y = smooth(x,weight,bounds)
|
||||
if key=='fingers':
|
||||
dev = float(np.abs(y-x).max())
|
||||
elif key=='wrist_rotation':
|
||||
dev = float(np.linalg.norm((rot.inv()*(rot[0]*Rotation.from_rotvec(y))).as_rotvec(),axis=1).max())
|
||||
else:
|
||||
dev = float(np.linalg.norm(y-x,axis=1).max())
|
||||
ratio = d2(y)/d2(x)
|
||||
row = dict(weight_d2=1.,weight_d3=weight,max_deviation=dev,budget=float(budget),
|
||||
second_difference_ratio=ratio,within_budget=bool(dev<=budget))
|
||||
rows.append(row)
|
||||
if dev<=budget:
|
||||
candidates.append((ratio,y,row))
|
||||
assert candidates, key
|
||||
_,chosen,row = min(candidates,key=lambda v:v[0])
|
||||
choices[key] = chosen
|
||||
sweep[key] = dict(selected=row,candidates=rows)
|
||||
print(key,row,flush=True)
|
||||
before = {k:m[k].copy() for k in ['qpos','wrist_pos','wrist_quat_wxyz','actual']}
|
||||
q = m['qpos'].copy()
|
||||
q[:,ai] = choices['fingers']
|
||||
for name,(parent,sc,off) in mimic.items():
|
||||
q[:,names.index(name)] = sc*q[:,names.index(parent)]+off
|
||||
m['qpos'] = q
|
||||
m['active_qpos'] = q[:,ai]
|
||||
m['wrist_pos'] = choices['wrist_position']
|
||||
newrot = rot[0]*Rotation.from_rotvec(choices['wrist_rotation'])
|
||||
m['wrist_quat_wxyz'] = newrot.as_quat()[:,[3,0,1,2]]
|
||||
for k,v in before.items():
|
||||
m['before_stabilization_'+k] = v
|
||||
m['stabilization_method'] = 'whole-sequence data fidelity + second/third difference regularization; component distortion budgets'
|
||||
m['stabilization_source'] = str(a.source.resolve())
|
||||
for filename in ['l20_moving.xml','l20_right.xml','l20_dex.urdf','dex_config.json']:
|
||||
shutil.copy2(a.source/filename,a.output/filename)
|
||||
model = mujoco.MjModel.from_xml_path(str(a.output/'l20_right.xml'))
|
||||
data = mujoco.MjData(model)
|
||||
addr = [model.jnt_qposadr[model.joint(k).id] for k in names]
|
||||
sites = [model.site(f'landmark_{i:02d}').id for i in range(21)]
|
||||
actual = []
|
||||
for row in q:
|
||||
data.qpos[addr] = row
|
||||
mujoco.mj_forward(model,data)
|
||||
actual.append(data.site_xpos[sites].copy())
|
||||
m['actual'] = np.array(actual)
|
||||
np.savez_compressed(a.output/'motion.npz',**m)
|
||||
np.savetxt(a.output/'trajectory.csv',np.c_[m['time'],m['wrist_pos'],m['wrist_quat_wxyz'],q],delimiter=',',
|
||||
header=','.join(['time_s','wrist_x','wrist_y','wrist_z','qw','qx','qy','qz']+names),comments='')
|
||||
lim = np.array([limits[k] for k in names])
|
||||
violation = float(max(np.maximum(lim[:,0]-q,0).max(),np.maximum(q-lim[:,1],0).max()))
|
||||
assert violation < 1e-9
|
||||
mimic_error = max(float(np.abs(q[:,names.index(k)]-sc*q[:,names.index(par)]-off).max()) for k,(par,sc,off) in mimic.items())
|
||||
assert mimic_error < 1e-10
|
||||
|
||||
def metrics(q,pos,quat,actual):
|
||||
r = Rotation.from_quat(quat[:,[1,2,3,0]])[30:]
|
||||
w = (r[:-1].inv()*r[1:]).as_rotvec()
|
||||
tips = np.linalg.norm(actual[:,[4,8,12,16,20]]-m['targets'][:,[4,8,12,16,20]],axis=-1)*1000
|
||||
return dict(joint_second_difference_rms_deg=np.rad2deg(d2(q)),
|
||||
joint_max_step_deg=float(np.rad2deg(np.abs(np.diff(q[30:],axis=0)).max())),
|
||||
wrist_second_difference_rms_mm=d2(pos)*1000,
|
||||
wrist_max_step_mm=float(np.linalg.norm(np.diff(pos[30:],axis=0),axis=1).max()*1000),
|
||||
wrist_rotation_velocity_difference_rms_deg=float(np.rad2deg(np.sqrt(np.mean(np.diff(w,axis=0)**2)))),
|
||||
wrist_rotation_max_step_deg=float(np.rad2deg(np.linalg.norm(w,axis=1).max())),
|
||||
wrist_range_m=np.ptp(pos,axis=0).tolist(),
|
||||
tip_error_mean_mm=float(tips.mean()),tip_error_p95_mm=float(np.percentile(tips,95)))
|
||||
report = dict(frames=n,fps=float(m['fps']),evaluation_start_frame=30,
|
||||
before=metrics(**dict(q=before['qpos'],pos=before['wrist_pos'],quat=before['wrist_quat_wxyz'],actual=before['actual'])),
|
||||
after=metrics(q,m['wrist_pos'],m['wrist_quat_wxyz'],m['actual']),
|
||||
selection=sweep,urdf_limit_violation_rad=violation,linear_mimic_error_rad=mimic_error,
|
||||
max_joint_change_deg=float(np.rad2deg(np.abs(q-before['qpos']).max())),
|
||||
max_wrist_change_mm=float(np.linalg.norm(m['wrist_pos']-before['wrist_pos'],axis=1).max()*1000),
|
||||
max_wrist_rotation_change_deg=float(np.rad2deg(np.linalg.norm((rot.inv()*newrot).as_rotvec(),axis=1).max())),
|
||||
uses_future_frames=True,boundary='Lower temporal derivatives do not prove true-motion accuracy; no contact/dynamics validation')
|
||||
assert all(np.isfinite(m[k]).all() for k in ['qpos','wrist_pos','wrist_quat_wxyz','actual'])
|
||||
(a.output/'stabilization_validation.json').write_text(json.dumps(report,indent=2))
|
||||
print(json.dumps({k:v for k,v in report.items() if k!='selection'},indent=2))
|
||||
|
||||
|
||||
if __name__=='__main__':
|
||||
main()
|
||||
@@ -0,0 +1,18 @@
|
||||
"""Decode and verify the complete exported LEAP replay."""
|
||||
import json
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
|
||||
root = Path(__file__).resolve().parents[1] / 'output/dex_2047635068'
|
||||
video = root / 'leap_right_replay.mp4'
|
||||
result = subprocess.run(['ffprobe', '-v', 'error', '-select_streams', 'v:0',
|
||||
'-count_frames', '-show_entries',
|
||||
'stream=width,height,avg_frame_rate,nb_read_frames',
|
||||
'-of', 'json', str(video)], check=True, capture_output=True, text=True)
|
||||
stream = json.loads(result.stdout)['streams'][0]
|
||||
assert int(stream['nb_read_frames']) == 1333
|
||||
assert stream['avg_frame_rate'] == '30/1'
|
||||
assert (stream['width'], stream['height']) == (1280, 480)
|
||||
subprocess.run(['ffmpeg', '-v', 'error', '-xerror', '-i', str(video), '-f', 'null', '-'], check=True)
|
||||
(root / 'video_validation.json').write_text(json.dumps(dict(stream=stream, full_decode_pass=True), indent=2))
|
||||
print('PASS: 1333 frames, 30 fps, 1280x480, full decode')
|
||||
@@ -0,0 +1,32 @@
|
||||
"""Validate full registered trajectory and decode both deliverable videos."""
|
||||
import json, subprocess
|
||||
import numpy as np
|
||||
from l20_calibrated import Calibration, ROOT
|
||||
|
||||
out=ROOT/'output/l20_2047635068/bottle_grasp'
|
||||
m=np.load(out/'registered_motion.npz'); c=Calibration()
|
||||
names=list(m['joint_names']); q=m['qpos']
|
||||
limits=np.array([c.urdf_limits[n] for n in names])
|
||||
coupling=max(float(np.max(np.abs(q[:,names.index(n)]-np.interp(q[:,names.index(a)],x,y)))) for n,(a,x,y) in c.curves.items())
|
||||
report=dict(frames=len(q),finite=all(bool(np.isfinite(m[k]).all()) for k in ['qpos','wrist_pos','wrist_quat_wxyz','bottle_pos','bottle_quat_wxyz']),
|
||||
limit_violations=int(np.sum((q<limits[:,0]-1e-8)|(q>limits[:,1]+1e-8))),
|
||||
max_passive_coupling_error_rad=coupling,
|
||||
max_penetration_mm=float(m['penetration_m'].max()*1000),
|
||||
frames_above_1_5_mm=int(np.sum(m['penetration_m']>.0015)),
|
||||
wrist_axis_range_m=np.ptp(m['wrist_pos'],axis=0).tolist(),
|
||||
max_joint_step_rad=float(np.abs(np.diff(q,axis=0)).max()),
|
||||
frames_with_joint_step_above_0_2_rad=int(np.sum(np.max(np.abs(np.diff(q,axis=0)),axis=1)>.2)),
|
||||
max_wrist_step_m=float(np.linalg.norm(np.diff(m['wrist_pos'],axis=0),axis=1).max()),
|
||||
max_command_quantization_error_rad=float(np.max(np.abs(np.array([c.decode(cmd,names) for cmd in m['command_u8']])-q))),videos={})
|
||||
for filename,expected in [('contact_registration.mp4',len(q)),('physics_test.mp4',len(np.load(out/'physics_motion.npz')['qpos']))]:
|
||||
path=out/filename
|
||||
probe=json.loads(subprocess.check_output(['ffprobe','-v','error','-count_frames','-select_streams','v:0','-show_entries','stream=width,height,nb_read_frames,r_frame_rate','-of','json',str(path)]))['streams'][0]
|
||||
subprocess.run(['ffmpeg','-v','error','-i',str(path),'-f','null','-'],check=True)
|
||||
assert int(probe['nb_read_frames'])==expected
|
||||
report['videos'][filename]=dict(probe,full_decode_passed=True)
|
||||
report['geometry_checks_passed']=report['finite'] and report['limit_violations']==0 and coupling<1e-8 and report['frames_above_1_5_mm']==0
|
||||
report['physical_retention_passed']=json.loads((out/'physics_validation.json').read_text())['retained']
|
||||
report['execution_ready']=False
|
||||
report['execution_limitations']=['Local pose jumps remain; no joint velocity constrained optimization','Free object retention test failed','Object trajectory is an inferred registration target']
|
||||
(out/'delivery_validation.json').write_text(json.dumps(report,indent=2))
|
||||
print(json.dumps(report,indent=2))
|
||||
@@ -0,0 +1,29 @@
|
||||
"""Independent narrow-phase contact, coupling, floor and pose verification."""
|
||||
import os
|
||||
os.environ.setdefault('MUJOCO_GL','osmesa')
|
||||
from pathlib import Path
|
||||
import json,xml.etree.ElementTree as E
|
||||
import numpy as np,mujoco
|
||||
OUT=Path(__file__).resolve().parents[1]/'output/l20_full_replay'
|
||||
x=np.load(OUT/'motion.npz');m=mujoco.MjModel.from_xml_path(str(OUT/'scene.xml'));d=mujoco.MjData(m)
|
||||
penetration=[];counts=[];floor_count=0;coupling=0.;limit=0.
|
||||
for row in x['qpos']:
|
||||
d.qpos[:]=row;mujoco.mj_forward(m,d)
|
||||
ds=[]
|
||||
for c in d.contact:
|
||||
types={int(m.geom_contype[c.geom1]),int(m.geom_contype[c.geom2])}
|
||||
if types=={1,2}:ds.append(float(c.dist))
|
||||
if 4 in types:floor_count+=1
|
||||
penetration.append(max(0,-min(ds,default=0)));counts.append(len(ds))
|
||||
for eq in range(m.neq):
|
||||
j1,j2=m.eq_obj1id[eq],m.eq_obj2id[eq];v=row[m.jnt_qposadr[j2]]
|
||||
expected=sum(m.eq_data[eq,k]*v**k for k in range(5))
|
||||
coupling=max(coupling,abs(row[m.jnt_qposadr[j1]]-expected))
|
||||
for j in range(m.njnt):
|
||||
if m.jnt_limited[j]:
|
||||
v=row[m.jnt_qposadr[j]];limit=max(limit,m.jnt_range[j,0]-v,v-m.jnt_range[j,1])
|
||||
r=json.loads((OUT/'validation.json').read_text());r['geom_distance_max_penetration_mm']=float(np.max(x['penetration_m'])*1000)
|
||||
r.update(max_all_penetration_mm=float(max(penetration)*1000),max_holding_penetration_mm=float(max(penetration[160:])*1000),holding_frames_with_contact=int(np.count_nonzero(counts[160:])),holding_frames=len(counts)-160,floor_contacts=floor_count,linear_mimic_max_error_rad=float(coupling),joint_limit_violation_rad=float(limit),contact_verification='mj_forward contact manifold independently checked every frame; geomDistance has isolated inconsistent spikes')
|
||||
(OUT/'validation.json').write_text(json.dumps(r,indent=2));np.savez_compressed(OUT/'contact_verification.npz',penetration_m=penetration,contact_count=counts)
|
||||
assert np.isfinite(x['qpos']).all() and coupling<1e-8 and limit<1e-8 and floor_count==0
|
||||
print(json.dumps(r,indent=2))
|
||||
@@ -0,0 +1,36 @@
|
||||
"""Check SPIDER imports, dependency versions and a two-world GPU contact simulation."""
|
||||
import os
|
||||
from pathlib import Path
|
||||
ROOT=Path(__file__).resolve().parents[1]
|
||||
os.environ['WARP_CACHE_PATH']=str(ROOT/'.spider_cache/warp')
|
||||
os.environ['TORCHINDUCTOR_CACHE_DIR']=str(ROOT/'.spider_cache/torchinductor')
|
||||
os.environ.setdefault('MUJOCO_GL','osmesa')
|
||||
import json,sys,importlib,importlib.metadata as metadata,tomllib
|
||||
import numpy as np
|
||||
from packaging.requirements import Requirement
|
||||
import torch,mujoco,mujoco_warp as mw,warp as wp
|
||||
|
||||
report={'python':sys.version,'compatibility_mode':'Python 3.11 with reused Torch; upstream declares Python >=3.12','versions':{},'imports':{},'requirements':{}}
|
||||
reqs=tomllib.loads((ROOT/'third_party/spider/pyproject.toml').read_text())['project']['dependencies']
|
||||
for text in reqs:
|
||||
r=Requirement(text);v=metadata.version(r.name);report['versions'][r.name]=v;report['requirements'][r.name]=v in r.specifier
|
||||
for name in ['spider.config','spider.simulators.mjwp','spider.optimizers.sampling_fast','spider.viewers','spider.preprocess.ik','open3d','pymeshlab','mink']:
|
||||
importlib.import_module(name);report['imports'][name]=True
|
||||
report['torch_path']=torch.__file__;report['cuda_available']=torch.cuda.is_available()
|
||||
assert report['cuda_available'],'GPU test requires access outside the device-restricted sandbox'
|
||||
report['gpu']=torch.cuda.get_device_name(0)
|
||||
assert float((torch.ones(8,device='cuda')*2).sum())==16
|
||||
wp.init()
|
||||
xml='<mujoco><option timestep="0.002" iterations="5"/><worldbody><geom type="plane" size="1 1 .1"/><body pos="0 0 .105"><freejoint/><geom type="sphere" size=".1" mass=".1"/></body></worldbody></mujoco>'
|
||||
m=mujoco.MjModel.from_xml_string(xml);d=mujoco.MjData(m);mujoco.mj_forward(m,d)
|
||||
with wp.ScopedDevice('cuda:0'):
|
||||
wm=mw.put_model(m);wd=mw.put_data(m,d,nworld=2,nconmax=32,njmax=128)
|
||||
for _ in range(60):mw.step(wm,wd)
|
||||
wp.synchronize()
|
||||
q=wd.qpos.numpy();time=wd.time.numpy();contacts=wd.nacon.numpy()
|
||||
assert np.isfinite(q).all() and np.all(q[:,2]>.085) and np.all(q[:,2]<.105),q
|
||||
assert np.allclose(time,.12,atol=1e-5),time
|
||||
report['gpu_physics']={'worlds':2,'steps':60,'time_s':time.tolist(),'sphere_height_m':q[:,2].tolist(),'contacts':contacts.tolist(),'passed':True}
|
||||
assert all(report['requirements'].values())
|
||||
out=ROOT/'output/spider_install';out.mkdir(exist_ok=True)
|
||||
(out/'validation.json').write_text(json.dumps(report,indent=2));print(json.dumps(report,indent=2),flush=True)
|
||||
@@ -0,0 +1,81 @@
|
||||
"""Independently validate saved video/robot outputs, including full video decode."""
|
||||
import argparse
|
||||
import json
|
||||
from pathlib import Path
|
||||
import xml.etree.ElementTree as ET
|
||||
|
||||
import cv2
|
||||
import mujoco
|
||||
import numpy as np
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
|
||||
|
||||
def main():
|
||||
p = argparse.ArgumentParser()
|
||||
p.add_argument('--inference', type=Path, required=True)
|
||||
p.add_argument('--retarget', type=Path, required=True)
|
||||
a = p.parse_args()
|
||||
h = np.load(a.inference / 'handflow_results.npz')
|
||||
m = np.load(a.retarget / 'motion.npz')
|
||||
count, fps = len(h['pose']), float(h['fps'])
|
||||
assert len(m['qpos']) == count and float(m['fps']) == fps
|
||||
names = list(m['joint_names'])
|
||||
model = mujoco.MjModel.from_xml_path(str(a.retarget / 'l20_moving.xml'))
|
||||
data = mujoco.MjData(model)
|
||||
addr = [model.jnt_qposadr[model.joint(n).id] for n in names]
|
||||
wrist = model.jnt_qposadr[model.joint('wrist_free').id]
|
||||
sites = [model.site(f'landmark_{i:02d}').id for i in range(21)]
|
||||
limits = []
|
||||
mimic_error = 0.
|
||||
urdf = ET.parse(ROOT / 'third_party/l20_assets/L20/RIGHT/linkerhand_g20_right.urdf').getroot()
|
||||
for n in names:
|
||||
j = urdf.find(f"joint[@name='{n}']")
|
||||
limits.append([float(j.find('limit').get(k)) for k in ['lower','upper']])
|
||||
mi = j.find('mimic')
|
||||
if mi is not None:
|
||||
expected = m['qpos'][:,names.index(mi.get('joint'))] * float(mi.get('multiplier','1')) + float(mi.get('offset','0'))
|
||||
mimic_error = max(mimic_error, float(np.abs(expected-m['qpos'][:,names.index(n)]).max()))
|
||||
limits = np.array(limits)
|
||||
violation = float(max(np.maximum(limits[:,0]-m['qpos'],0).max(),np.maximum(m['qpos']-limits[:,1],0).max()))
|
||||
fk_error = 0.
|
||||
from scipy.spatial.transform import Rotation
|
||||
rotations = Rotation.from_quat(m['wrist_quat_wxyz'][:,[1,2,3,0]]).as_matrix()
|
||||
for t in range(count):
|
||||
data.qpos[addr] = m['qpos'][t]
|
||||
data.qpos[wrist:wrist+3] = m['wrist_pos'][t]
|
||||
data.qpos[wrist+3:wrist+7] = m['wrist_quat_wxyz'][t]
|
||||
mujoco.mj_forward(model,data)
|
||||
expected = m['actual'][t] @ rotations[t].T + m['wrist_pos'][t]
|
||||
fk_error = max(fk_error,float(np.abs(data.site_xpos[sites]-expected).max()))
|
||||
csv = np.loadtxt(a.retarget/'trajectory.csv',delimiter=',',skiprows=1)
|
||||
expected_csv = np.c_[m['time'],m['wrist_pos'],m['wrist_quat_wxyz'],m['qpos']]
|
||||
assert np.allclose(csv,expected_csv,atol=1e-12)
|
||||
assert np.isfinite(expected_csv).all() and violation < 1e-8 and mimic_error < 1e-10 and fk_error < 1e-5
|
||||
video_paths = [a.inference/'overlay.mp4',a.inference/'ortho_third_person.mp4',a.retarget/'l20_moving_replay.mp4']
|
||||
if (a.retarget/'comparison.mp4').exists():
|
||||
video_paths.append(a.retarget/'comparison.mp4')
|
||||
videos = []
|
||||
for path in video_paths:
|
||||
cap = cv2.VideoCapture(str(path))
|
||||
vf = cap.get(cv2.CAP_PROP_FPS)
|
||||
decoded = 0
|
||||
while True:
|
||||
ok, image = cap.read()
|
||||
if not ok:
|
||||
break
|
||||
assert image is not None
|
||||
decoded += 1
|
||||
cap.release()
|
||||
assert decoded == count and abs(vf-fps)<1e-3, (path,decoded,vf)
|
||||
videos.append(dict(path=str(path.resolve()),decoded_frames=decoded,fps=vf))
|
||||
report = dict(status='COMPLETE',frames=count,fps=fps,duration_seconds=count/fps,
|
||||
finite=True,urdf_limit_violation_rad=violation,linear_mimic_error_rad=mimic_error,
|
||||
independent_world_fk_max_error_m=fk_error,checked_frames=count,csv_matches_npz=True,
|
||||
videos=videos,boundary='Kinematic hand motion only; monocular estimates; no object or hardware validation')
|
||||
(a.retarget/'delivery_validation.json').write_text(json.dumps(report,indent=2))
|
||||
print(json.dumps(report,indent=2))
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
Executable
+49
@@ -0,0 +1,49 @@
|
||||
#!/bin/bash
|
||||
# HandFlow inference + render environment (training deps excluded)
|
||||
# Usage: bash setup_env.sh [env_name=handflow] [python=3.10]
|
||||
set -e
|
||||
|
||||
ENV_NAME=${1:-handflow}
|
||||
PYTHON_VERSION=${2:-3.10}
|
||||
|
||||
echo "=========================================="
|
||||
echo "HandFlow inference environment setup: $ENV_NAME (Python $PYTHON_VERSION)"
|
||||
echo "=========================================="
|
||||
|
||||
# Create conda environment
|
||||
eval "$(conda shell.bash hook)"
|
||||
if ! conda env list | grep -qE "^${ENV_NAME} "; then
|
||||
conda create -y -n "$ENV_NAME" "python=$PYTHON_VERSION"
|
||||
fi
|
||||
conda activate "$ENV_NAME"
|
||||
|
||||
echo ""
|
||||
echo "Installing PyTorch 2.7.0 + CUDA 12.8..."
|
||||
pip install torch==2.7.0+cu128 torchvision==0.22.0+cu128 \
|
||||
--index-url https://download.pytorch.org/whl/cu128
|
||||
|
||||
echo ""
|
||||
echo "Installing HaMeR backbone (submodule, includes its deps smplx/yacs etc.)..."
|
||||
pip install -e third_party/hamer
|
||||
|
||||
echo ""
|
||||
echo "Installing HandFlow inference + render deps..."
|
||||
pip install numpy scipy omegaconf tqdm opencv-python "imageio[ffmpeg]" matplotlib
|
||||
pip install ultralytics manopth
|
||||
|
||||
echo ""
|
||||
echo "Installing pytorch3d (Phong renderer; build from source — needs nvcc, matches torch 2.7+cu128)..."
|
||||
pip install "git+https://github.com/facebookresearch/pytorch3d.git@v0.7.8"
|
||||
|
||||
echo ""
|
||||
echo "=========================================="
|
||||
echo "Inference environment ready: conda activate $ENV_NAME"
|
||||
echo "=========================================="
|
||||
echo ""
|
||||
echo "Next — set the weight-path environment variables (see README for weight downloads):"
|
||||
echo " export HAMER_CKPT=/path/to/hamer.ckpt"
|
||||
echo " export DETECTOR_CKPT=/path/to/detector.pt # from WiLoR"
|
||||
echo " export MANO_ROOT=/path/to/mano # MANO model"
|
||||
echo " export HANDFLOW_NORMALIZATION_STATS=/path/to/normalization_stats.npz"
|
||||
echo ""
|
||||
echo "Optional — for moving-camera world trajectory (demo.py without --fix_camera): bash setup_vipe_env.sh"
|
||||
Executable
+41
@@ -0,0 +1,41 @@
|
||||
#!/bin/bash
|
||||
# ViPE SLAM standalone environment (invoked by demo.py via conda run -n vipe when --fix_camera is unset)
|
||||
# Install only if you need the world-view trajectory (demo.py without --fix_camera); --fix_camera mode skips ViPE.
|
||||
#
|
||||
# ViPE's CUDA extensions are compiled against torch 2.7.0+cu128, matching the handflow environment.
|
||||
# Usage: bash setup_vipe_env.sh [env_name=vipe]
|
||||
set -e
|
||||
|
||||
ENV_NAME=${1:-vipe}
|
||||
|
||||
echo "=========================================="
|
||||
echo "ViPE SLAM environment setup: $ENV_NAME"
|
||||
echo "=========================================="
|
||||
|
||||
eval "$(conda shell.bash hook)"
|
||||
if ! conda env list | grep -qE "^${ENV_NAME} "; then
|
||||
conda create -y -n "$ENV_NAME" python=3.10
|
||||
fi
|
||||
conda activate "$ENV_NAME"
|
||||
|
||||
echo ""
|
||||
echo "Installing PyTorch 2.7.0 + CUDA 12.8 (matching handflow)..."
|
||||
pip install torch==2.7.0+cu128 torchvision==0.22.0+cu128 \
|
||||
--index-url https://download.pytorch.org/whl/cu128
|
||||
|
||||
echo ""
|
||||
echo "Installing CUDA 12.8 build tools (nvcc, required to compile the ViPE extension)..."
|
||||
conda install -y -c "nvidia/label/cuda-12.8.0" cuda-nvcc cuda-cudart-dev
|
||||
conda install -y -c conda-forge gxx_linux-64=11 # gcc 11
|
||||
|
||||
echo ""
|
||||
echo "Installing ViPE (submodule; compiles the vipe_ext CUDA extension + downloads Eigen)..."
|
||||
pip install -e third_party/vipe
|
||||
|
||||
echo ""
|
||||
echo "=========================================="
|
||||
echo "ViPE environment ready: conda activate $ENV_NAME"
|
||||
echo "=========================================="
|
||||
echo ""
|
||||
echo "demo.py will invoke this environment automatically via 'conda run -n ${ENV_NAME}' (moving-camera mode only)."
|
||||
echo "The first build of vipe_ext downloads Eigen 3.4 online; ensure network access."
|
||||
Vendored
+21
@@ -0,0 +1,21 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2024 Zhengdi Yu
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
Vendored
+241
@@ -0,0 +1,241 @@
|
||||
<p align="center">
|
||||
|
||||
<h1 align="center">Dyn-HaMR: Recovering 4D Interacting Hand Motion from a Dynamic Camera</h1>
|
||||
<p align="center">
|
||||
<a href="https://github.com/ZhengdiYu"><strong>Zhengdi Yu</strong></a>
|
||||
·
|
||||
<a href="https://scholar.google.com/citations?user=QKOH5iYAAAAJ&hl=en&oi=ao"><strong>Stefanos Zafeiriou</strong></a>
|
||||
·
|
||||
<a href="https://tolgabirdal.github.io/"><strong>Tolga Birdal</strong></a>
|
||||
</p>
|
||||
<p align="center">
|
||||
<strong>Imperial College London</strong>
|
||||
</p>
|
||||
</p>
|
||||
<p align="center">
|
||||
<strong>CVPR 2025 (Highlight)</strong>
|
||||
</p>
|
||||
|
||||
<div align="center" style="display: flex; justify-content: center; gap: 20px;">
|
||||
<img src="./assets/wild.gif" alt="Wild GIF" width="45%" style="border-radius: 10px;">
|
||||
<img src="./assets/global.gif" alt="Global GIF" width="45%" style="border-radius: 10px;">
|
||||
</div>
|
||||
|
||||
<p align="center">
|
||||
<a href='https://arxiv.org/abs/2412.12861'>
|
||||
<img src='https://img.shields.io/badge/Arxiv-2412.12861-A42C25?style=flat&logo=arXiv&logoColor=A42C25'>
|
||||
</a>
|
||||
<a href='https://arxiv.org/pdf/2412.12861'>
|
||||
<img src='https://img.shields.io/badge/Paper-PDF-green?style=flat&logo=arXiv&logoColor=green'>
|
||||
</a>
|
||||
<a href='https://dyn-hamr.github.io/'>
|
||||
<img src='https://img.shields.io/badge/Project-Page-blue?style=flat&logo=Google%20chrome&logoColor=blue'></a>
|
||||
<a href='https://youtu.be/n25NGIWiA7M'>
|
||||
<img src='https://img.shields.io/badge/YouTube-Video-EA3323?style=flat&logo=youtube&logoColor=EA3323'></a>
|
||||
<img src="https://visitor-badge.laobi.icu/badge?page_id=ZhengdiYu.Dyn-HaMR&left_color=gray&right_color=orange">
|
||||
</a>
|
||||
</p>
|
||||
</p>
|
||||
|
||||
|
||||
## Introduction
|
||||
<img src="./assets/teaser.png">
|
||||
|
||||
We propose **Dyn-HaMR** to reconstruct 4D global hand motion from monocular videos recorded by dynamic cameras in the wild, as a remedy for the motion entanglement in the wild.
|
||||
|
||||
<!-- TABLE OF CONTENTS -->
|
||||
<details open="open" style='padding: 10px; border-radius:5px 30px 30px 5px; border-style: solid; border-width: 1px;'>
|
||||
<summary>Table of Contents</summary>
|
||||
<ol>
|
||||
<li>
|
||||
<a href="#installation">Installation</a>
|
||||
</li>
|
||||
<li>
|
||||
<a href="#get-started">Get Started</a>
|
||||
</li>
|
||||
<li>
|
||||
<a href="#citation">Citation</a>
|
||||
</li>
|
||||
</ol>
|
||||
</details>
|
||||
<br/>
|
||||
|
||||
## News :triangular_flag_on_post:
|
||||
- [2025/11/20] 🚀 **Major Update**:
|
||||
- **Integrated [VIPE](https://github.com/nv-tlabs/vipe)** for camera estimation, significantly improving reconstruction quality over DROID-SLAM
|
||||
- **Enhanced Hand Tracker** with robust hallucination prevention and handedness correction for better hand tracking and **significantly** Improved temporal consistency. Please `pip install ultralytics==8.1.34` since YOLO is using in this version (Thanks to [WiloR](https://github.com/rolpotamias/WiLoR)). Please download the checkpoint from [here](https://huggingface.co/spaces/rolpotamias/WiLoR/resolve/main/pretrained_models/detector.pt) and put it under `third-party/hamer/pretrained_models`.
|
||||
|
||||
See comparison below:
|
||||
|
||||
<table>
|
||||
<tr>
|
||||
<th>Before: Jitter from DROID-SLAM</th>
|
||||
<th>New: VIPE + enhanced HaMeR (Recommended)</th>
|
||||
</tr>
|
||||
<tr>
|
||||
<td align="center">
|
||||
<img src="./assets/droid_result.gif" width="100%">
|
||||
</td>
|
||||
<td align="center">
|
||||
<img src="./assets/vipe_result.gif" width="100%">
|
||||
</td>
|
||||
</tr>
|
||||
</table>
|
||||
|
||||
<table>
|
||||
<tr>
|
||||
<th>Before: Handedness shifting with original HaMeR</th>
|
||||
<th>New: Enhanced hand tracker (Recommended)</th>
|
||||
</tr>
|
||||
<tr>
|
||||
<td align="center">
|
||||
<img src="./assets/handedness1.gif" width="100%">
|
||||
</td>
|
||||
<td align="center">
|
||||
<img src="./assets/handedness2.gif" width="100%">
|
||||
</td>
|
||||
</tr>
|
||||
</table>
|
||||
|
||||
<table>
|
||||
<tr>
|
||||
<th>Before: Handedness shifting with original HaMeR</th>
|
||||
<th>New: Enhanced hand tracker (Recommended)</th>
|
||||
</tr>
|
||||
<tr>
|
||||
<td align="center">
|
||||
<img src="./assets/handedness3.gif" width="100%">
|
||||
</td>
|
||||
<td align="center">
|
||||
<img src="./assets/handedness4.gif" width="100%">
|
||||
</td>
|
||||
</tr>
|
||||
</table>
|
||||
|
||||
- [2025/06/04] Code released.
|
||||
- [2024/12/18] [Paper](https://arxiv.org/abs/2412.12861) is now available on arXiv. ⭐
|
||||
|
||||
## Installation
|
||||
|
||||
### Environment setup
|
||||
1. Clone the repository with submodules with the following command:
|
||||
```bash
|
||||
git clone --recursive https://github.com/ZhengdiYu/Dyn-HaMR.git
|
||||
cd Dyn-HaMR
|
||||
```
|
||||
You can also run the following command to fetch the submodules:
|
||||
```bash
|
||||
git submodule update --init --recursive .
|
||||
```
|
||||
|
||||
2. To set up the virtual environment for Dyn-HaMR, we provide the integrated commands in `\scripts`. You can create the environment from
|
||||
```bash
|
||||
source install_pip.sh
|
||||
```
|
||||
|
||||
Or, alternatively, create the environment from conda:
|
||||
```bash
|
||||
source install_conda.sh
|
||||
```
|
||||
|
||||
### Model checkpoints download
|
||||
Please run the following command to fetch the data dependencies. This will create a folder in `_DATA`:
|
||||
```bash
|
||||
source prepare.sh
|
||||
```
|
||||
After processing, the folder layout should be:
|
||||
```
|
||||
|-- _DATA
|
||||
| |-- data/
|
||||
| | |-- mano/
|
||||
| | | |-- MANO_RIGHT.pkl
|
||||
| | |-- mano_mean_params.npz
|
||||
| |-- BMC/
|
||||
| |-- hamer_ckpts/
|
||||
| |-- vitpose_ckpts/
|
||||
| |-- <SLAM model .pkl>
|
||||
```
|
||||
|
||||
### Prerequisites
|
||||
We use [MANO](https://mano.is.tue.mpg.de) model for hand mesh representation. Please visit the [MANO website](https://mano.is.tue.mpg.de) for registration and the model downloading. Please download `MANO_RIGHT.pkl` and put under the `_DATA/data/mano` folder.
|
||||
|
||||
## Get Started🚀
|
||||
|
||||
### Preparation
|
||||
Please follow the instructions [here](https://github.com/MengHao666/Hand-BMC-pytorch) to calculate the below `.npz` files in order `dyn-hamr/optim/BMC/`:
|
||||
```
|
||||
|-- BMC
|
||||
| |-- bone_len_max.npy
|
||||
| |-- bone_len_min.npy
|
||||
| |-- CONVEX_HULLS.npy
|
||||
| |-- curvatures_max.npy
|
||||
| |-- curvatures_min.npy
|
||||
| |-- joint_angles.npy
|
||||
| |-- PHI_max.npy
|
||||
| |-- PHI_min.npy
|
||||
```
|
||||
|
||||
> [!NOTE]
|
||||
> If accurate camera parameters are available, please follow the format of `Dyn-HaMR/test/dynhamr/cameras/demo/shot-0/cameras.npz` to prepare the camera parameters for loading. Similarly, you can use Dyn-HaMR to refine and recover the hand mesh in the world coordinate system initializing from your own 2D & 3D motion data.
|
||||
|
||||
### Customize configurations
|
||||
| Config | Operation |
|
||||
|--------|-----------------|
|
||||
| GPU | Edit in [`<CONFIG_GPU>`](https://github.com/ZhengdiYu/Dyn-HaMR/blob/main/dyn-hamr/confs/config.yaml#L56) |
|
||||
| Video info | Edit in [`<VIDEO_SEQ>`](https://github.com/ZhengdiYu/Dyn-HaMR/blob/main/dyn-hamr/confs/data/video.yaml#L5) |
|
||||
| Interval | Edit in [`<VIDEO_START_END>`](https://github.com/ZhengdiYu/Dyn-HaMR/blob/main/dyn-hamr/confs/data/video.yaml#L16-L17) |
|
||||
| Optimization configurations | Edit in [`<OPT_WEIGHTS>`](https://github.com/ZhengdiYu/Dyn-HaMR/blob/main/dyn-hamr/confs/optim.yaml#L29-L49) |
|
||||
| General configurations | Edit in [`<GENERAL_CONFIG>`](https://github.com/ZhengdiYu/Dyn-HaMR/blob/main/dyn-hamr/confs/config.yaml) |
|
||||
|
||||
### Fitting on RGB-(D) videos 🎮
|
||||
To run the optimization pipeline for fitting on arbitrary RGB-(D) videos, please first edit the path information here in `dyn-hamr/confs/data/video.yaml`, where `root` is the root folder to all of your datasets. `video_dir` is the corresponding folder that contains the videos. The key `seq` represents the video name you wanted to process. For example, you can run the following command to recover the global motion for `test/videos/demo1.mp4`:
|
||||
|
||||
|
||||
#### 🌟 Using VIPE for Camera Estimation (Recommended)
|
||||
For significantly better camera estimation quality, use VIPE instead of DROID-SLAM:
|
||||
|
||||
```bash
|
||||
python run_opt.py data=video_vipe run_opt=True data.seq=demo1 is_static=False
|
||||
```
|
||||
|
||||
#### 🌟 Using original DROID-SLAM for Camera Estimation
|
||||
```bash
|
||||
python run_opt.py data=video run_opt=True data.seq=demo1 is_static=<True or False>
|
||||
```
|
||||
|
||||
VIPE will automatically run if results are not found. Make sure you have:
|
||||
1. Installed VIPE in `third-party/vipe/` with conda environment named `vipe`
|
||||
2. Set `src_path` in `dyn-hamr/confs/data/video_vipe.yaml` to your video file
|
||||
|
||||
By default, the camera parameters will be predicted during the process and assumes a moving camera (`is_static=False`). If your video is recorded with a static camera, you can add `is_static=True` for more stable optimization. The result will be saved to `outputs/logs/video-custom/<DATE>/<VIDEO_NAME>-<tracklet>-shot-<shot_id>-<start_frame_id>-<end_frame_id>`. After optimization, you can specify the output log dir and visualize the results by running the following command:
|
||||
```
|
||||
python run_vis.py --log_root <LOG_ROOT>
|
||||
```
|
||||
This will visualize all log subdirectories and save the rendered videos and images, as well as saved 3D meshes in the world space in `<LOG_ROOT>`. Please visit `run_vis.py` for further details. Alternatively, you can also use the following command to run and visualize the results in one-stage:
|
||||
```
|
||||
python -u run_opt.py data=video_vipe run_opt=True run_vis=True is_static=<True of False>
|
||||
```
|
||||
As a multi-stage pipeline, you can customize the optimization process. Add `is_static=True` for static camera videos. Adding `run_prior=True` can activate the motion prior in stage III. Please note that in the current version, each motion chunk size needs to be set to 128 to be compatible with the original setting of HMP only when the prior module is activated.
|
||||
|
||||
### Blender Addon
|
||||
Coming soon.
|
||||
|
||||
## Acknowledgements
|
||||
The PyTorch implementation of MANO is based on [manopth](https://github.com/hassony2/manopth). Part of the fitting and optimization code of this repository is borrowed from [SLAHMR](https://github.com/vye16/slahmr). For data preprocessing and observation, [ViTPose](https://github.com/ViTAE-Transformer/ViTPose) and [HaMeR](https://github.com/geopavlakos/hamer/) is used for 2D keypoints detection and MANO parameter initilization. For camera motion estimation, we support [VIPE](https://github.com/facebookresearch/vipe) (recommended), [DPVO](https://github.com/princeton-vl/DPVO), and [DROID-SLAM](https://github.com/princeton-vl/DROID-SLAM). For biomechanical constraints and motion prior, we use the code from [here](https://github.com/MengHao666/Hand-BMC-pytorch) and [HMP](https://hmp.is.tue.mpg.de/). We thank all the authors for their impressive work!
|
||||
|
||||
## License
|
||||
Please see [License](https://github.com/ZhengdiYu/Dyn-HaMR/blob/main/LICENSE) for details of Dyn-HaMR. This code and model are available only for non-commercial research purposes as defined in the LICENSE (i.e., MIT LICENSE). Note that, for MANO you must agree with the LICENSE of it. You can check the LICENSE of MANO from https://mano.is.tue.mpg.de/license.html.
|
||||
|
||||
## Citation
|
||||
```bibtex
|
||||
@inproceedings{yu2025dynhamr,
|
||||
title={Dyn-HaMR: Recovering 4D Interacting Hand Motion from a Dynamic Camera},
|
||||
author={Yu, Zhengdi and Zafeiriou, Stefanos and Birdal, Tolga},
|
||||
booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
|
||||
month = {June},
|
||||
year = {2025},
|
||||
}
|
||||
```
|
||||
|
||||
## Contact
|
||||
For any technical questions, please contact z.yu23@imperial.ac.uk or ZhengdiYu@hotmail.com.
|
||||
+11
@@ -0,0 +1,11 @@
|
||||
dataset_dir: ./data/amass/generative
|
||||
data:
|
||||
fps: 30
|
||||
clip_length: 128
|
||||
gender: neutral
|
||||
up: z
|
||||
|
||||
canonical: False
|
||||
unified_orientation: True
|
||||
single_motion: False
|
||||
normalize: True
|
||||
+36
@@ -0,0 +1,36 @@
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
import yaml
|
||||
|
||||
|
||||
class Struct:
|
||||
def __init__(self, **entries):
|
||||
self.__dict__.update(entries)
|
||||
|
||||
|
||||
class Arguments:
|
||||
def __init__(self, base_dir, config_path, filename='default.yaml'):
|
||||
with open(os.path.join(config_path, 'mano.yaml'), 'r') as f:
|
||||
smpl = yaml.safe_load(f)
|
||||
|
||||
self.smpl = Struct(**smpl)
|
||||
self.smpl.offsets['right'] = np.array(self.smpl.offsets['right'])
|
||||
self.smpl.offsets['left'] = np.array(self.smpl.offsets['left'])
|
||||
self.smpl.parents = np.array(self.smpl.parents).astype(np.int32)
|
||||
self.smpl.joint_num = 16 # len(self.smpl.joints_to_use)
|
||||
self.smpl.joints_to_use = np.array(self.smpl.joints_to_use)
|
||||
self.smpl.joints_to_use = np.arange(0, 63).reshape((-1, 3))[self.smpl.joints_to_use].reshape(-1)
|
||||
self.smpl.smpl_body_model = os.path.join(base_dir, '_DATA/data/mano')
|
||||
|
||||
self.filename = os.path.splitext(filename)[0]
|
||||
with open(os.path.join(config_path, filename), 'r') as f:
|
||||
config = yaml.safe_load(f)
|
||||
|
||||
for key, value in config.items():
|
||||
if isinstance(value, dict):
|
||||
setattr(self, key, Struct(**value))
|
||||
else:
|
||||
setattr(self, key, value)
|
||||
|
||||
self.json = config
|
||||
+639
@@ -0,0 +1,639 @@
|
||||
import os
|
||||
import glob
|
||||
import json
|
||||
import torch
|
||||
import joblib
|
||||
import random
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from scipy import linalg
|
||||
from arguments import Arguments
|
||||
from nemf.losses import GeodesicLoss
|
||||
from nemf.generative import Architecture
|
||||
from nemf.fk import ForwardKinematicsLayer
|
||||
from scripts.train_gmm import MaxMixturePrior
|
||||
from human_body_prior.tools.omni_tools import copy2cpu as c2c
|
||||
from utils import FOOT_IDX, estimate_angular_velocity, estimate_linear_velocity
|
||||
from rotations import axis_angle_to_matrix, matrix_to_rotation_6d, rotation_6d_to_matrix
|
||||
|
||||
|
||||
def padding_prior_data(data):
|
||||
N, T, J = data['pos'].shape[:3]
|
||||
# pad inputs to match the training sequence length
|
||||
zero_padding = torch.zeros(N, args.data.clip_length - T, J, 3).to(prior.device)
|
||||
identity_padding = torch.zeros(N, args.data.clip_length - T, 3, 3).to(prior.device)
|
||||
identity_padding[:, :, 0, 0] = 1
|
||||
identity_padding[:, :, 1, 1] = 1
|
||||
identity_padding[:, :, 2, 2] = 1
|
||||
rotation_6d_padding = torch.zeros(N, args.data.clip_length - T, 6).to(prior.device)
|
||||
rotation_6d_padding[..., 0] = 1
|
||||
rotation_6d_padding[..., 4] = 1
|
||||
|
||||
pos = torch.cat((data['pos'], zero_padding), dim=1)
|
||||
velocity = torch.cat((data['velocity'], zero_padding), dim=1)
|
||||
angular = torch.cat((data['angular'], zero_padding), dim=1)
|
||||
root_orient = torch.cat((data['root_orient'], rotation_6d_padding), dim=1)
|
||||
root_vel = torch.cat((data['root_vel'], zero_padding[:, :, 0]), dim=1)
|
||||
rotmat_padding = rotation_6d_padding.unsqueeze(2).expand(-1, -1, J, -1)
|
||||
global_xform = torch.cat((data['global_xform'], rotmat_padding), dim=1)
|
||||
|
||||
data['pos'] = pos
|
||||
data['velocity'] = velocity
|
||||
data['global_xform'] = global_xform
|
||||
data['angular'] = angular
|
||||
data['root_orient'] = root_orient
|
||||
data['root_vel'] = root_vel
|
||||
|
||||
return data
|
||||
|
||||
|
||||
def load_data(files, padding=False, max_samples=400):
|
||||
poses = []
|
||||
trans = []
|
||||
visibility = []
|
||||
assert len(files) != 0, 'files not found'
|
||||
|
||||
max_samples = min(max_samples, len(files))
|
||||
|
||||
# import ipdb; ipdb.set_trace()
|
||||
|
||||
for f in files[:max_samples]:
|
||||
bdata = dict(np.load(f, allow_pickle=True))
|
||||
if 'r_wrist_orient' in bdata.keys() and 'pose_right_hand' in bdata.keys():
|
||||
|
||||
coef = len(bdata['r_wrist_orient']) // args.data.clip_length
|
||||
|
||||
for k_ in range(coef):
|
||||
root_orient = bdata['r_wrist_orient'][k_*args.data.clip_length:(k_+1)*args.data.clip_length]
|
||||
pose_body = bdata['pose_right_hand'][k_*args.data.clip_length:(k_+1)*args.data.clip_length]
|
||||
trans.append(bdata['r_wrist_location'][k_*args.data.clip_length:(k_+1)*args.data.clip_length])
|
||||
|
||||
poses.append(np.concatenate((root_orient, pose_body), axis=-1))
|
||||
|
||||
elif 'pose_right_hand' in bdata.keys():
|
||||
pass
|
||||
# poses.append(bdata['pose_right_hand'][:, :72])
|
||||
else:
|
||||
raise RuntimeError(f'missing pose parameters in the file: {f}')
|
||||
|
||||
|
||||
if 'visibility' in bdata.keys():
|
||||
visibility.append(bdata['visibility'])
|
||||
|
||||
# if 'poses' in bdata.keys():
|
||||
# poses.append(bdata['poses'][:, :72])
|
||||
# elif 'root_orient' in bdata.keys() and 'pose_body' in bdata.keys():
|
||||
# root_orient = bdata['root_orient']
|
||||
# pose_body = bdata['pose_body']
|
||||
# poses.append(np.concatenate((root_orient, pose_body), axis=-1))
|
||||
# else:
|
||||
# raise RuntimeError(f'missing pose parameters in the file: {f}')
|
||||
# trans.append(bdata['trans'])
|
||||
# if 'visibility' in bdata.keys():
|
||||
# visibility.append(bdata['visibility'])
|
||||
|
||||
trans = torch.from_numpy(np.asarray(trans, np.float32)).cuda() # global translation (N, T, 3)
|
||||
N, T = trans.shape[:2]
|
||||
poses = torch.from_numpy(np.asarray(poses, np.float32)).cuda()
|
||||
poses = poses.view(N, T, -1, 3) # axis-angle (N, T, J, 3)
|
||||
|
||||
|
||||
# poses = poses.view(N, T, 24, 3) # axis-angle (N, T, J, 3)
|
||||
|
||||
root_orient = poses[:, :, 0].clone()
|
||||
root_rotation = axis_angle_to_matrix(root_orient) # (N, T, 3, 3)
|
||||
poses[:, :, 0] = 0
|
||||
|
||||
rotmat = axis_angle_to_matrix(poses) # (N, T, J, 3, 3)
|
||||
angular = estimate_angular_velocity(rotmat.clone(), dt=1.0 / args.data.fps) # angular velocity of all the joints (N, T, J, 3)
|
||||
pos, global_xform = fk(rotmat.view(-1, 16, 3, 3))
|
||||
pos = pos.contiguous().view(N, T, 16, 3) # local joint positions (N, T, J, 3)
|
||||
global_xform = global_xform.view(N, T, 16, 4, 4)
|
||||
global_xform = global_xform[:, :, :, :3, :3] # global transformation matrix for each joint (N, T, J, 3, 3)
|
||||
velocity = estimate_linear_velocity(pos, dt=1.0 / args.data.fps) # linear velocity of all the joints (N, T, J, 3)
|
||||
|
||||
root_vel = estimate_linear_velocity(trans, dt=1.0 / args.data.fps) # linear velocity of the root joint (N, T, 3)
|
||||
|
||||
global_pos = torch.matmul(root_rotation.unsqueeze(2), pos.unsqueeze(-1)).squeeze(-1) # (N, T, J, 3)
|
||||
global_pos = global_pos + trans.unsqueeze(2)
|
||||
|
||||
data = {
|
||||
'pos': pos,
|
||||
'velocity': velocity,
|
||||
'global_xform': matrix_to_rotation_6d(global_xform),
|
||||
'angular': angular,
|
||||
'root_orient': matrix_to_rotation_6d(root_rotation),
|
||||
'root_vel': root_vel,
|
||||
'global_pos': global_pos,
|
||||
'rotmat': rotmat,
|
||||
'trans': trans
|
||||
}
|
||||
|
||||
if len(visibility) != 0:
|
||||
visibility = torch.from_numpy(np.asarray(visibility, bool)).cuda()
|
||||
padding = torch.zeros(visibility.shape[:2] + torch.Size([2]), dtype=torch.bool).cuda()
|
||||
visibility = torch.cat((visibility, padding), dim=-1) # (T, 24)
|
||||
print(f'visibility: {visibility.shape}')
|
||||
data['visibility'] = visibility
|
||||
|
||||
if T < args.data.clip_length:
|
||||
padding = True
|
||||
|
||||
if padding:
|
||||
if T == args.data.clip_length:
|
||||
data['pos'] = data['pos'].view(2 * N, T // 2, 24, 3)
|
||||
data['velocity'] = data['velocity'].view(2 * N, T // 2, 24, 3)
|
||||
data['global_xform'] = data['global_xform'].view(2 * N, T // 2, 24, 6)
|
||||
data['angular'] = data['angular'].view(2 * N, T // 2, 24, 3)
|
||||
data['root_orient'] = data['root_orient'].view(2 * N, T // 2, 6)
|
||||
data['root_vel'] = data['root_vel'].view(2 * N, T // 2, 3)
|
||||
data['global_pos'] = data['global_pos'].view(2 * N, T // 2, 24, 3)
|
||||
data['rotmat'] = data['rotmat'].view(2 * N, T // 2, 24, 3, 3)
|
||||
data['trans'] = data['trans'].view(2 * N, T // 2, 3)
|
||||
|
||||
padding_data = padding_prior_data(data)
|
||||
for key in padding_data.keys():
|
||||
data[key] = padding_data[key]
|
||||
|
||||
return data
|
||||
|
||||
|
||||
def calculate_feature_statistics(feature):
|
||||
feature = feature.cpu().numpy()
|
||||
mu = np.mean(feature, axis=0)
|
||||
sigma = np.cov(feature, rowvar=False)
|
||||
|
||||
return mu, sigma
|
||||
|
||||
|
||||
def calculate_frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6):
|
||||
"""Numpy implementation of the Frechet Distance.
|
||||
The Frechet distance between two multivariate Gaussians X_1 ~ N(mu_1, C_1)
|
||||
and X_2 ~ N(mu_2, C_2) is
|
||||
d^2 = ||mu_1 - mu_2||^2 + Tr(C_1 + C_2 - 2*sqrt(C_1*C_2)).
|
||||
Stable version by Dougal J. Sutherland.
|
||||
Params:
|
||||
-- mu1 : Numpy array containing the activations of a layer of the
|
||||
inception net (like returned by the function 'get_predictions')
|
||||
for generated samples.
|
||||
-- mu2 : The sample mean over activations, precalculated on an
|
||||
representative data set.
|
||||
-- sigma1: The covariance matrix over activations for generated samples.
|
||||
-- sigma2: The covariance matrix over activations, precalculated on an
|
||||
representative data set.
|
||||
Returns:
|
||||
-- : The Frechet Distance.
|
||||
"""
|
||||
|
||||
mu1 = np.atleast_1d(mu1)
|
||||
mu2 = np.atleast_1d(mu2)
|
||||
|
||||
sigma1 = np.atleast_2d(sigma1)
|
||||
sigma2 = np.atleast_2d(sigma2)
|
||||
|
||||
assert mu1.shape == mu2.shape, \
|
||||
'Training and test mean vectors have different lengths'
|
||||
assert sigma1.shape == sigma2.shape, \
|
||||
'Training and test covariances have different dimensions'
|
||||
|
||||
diff = mu1 - mu2
|
||||
|
||||
# Product might be almost singular
|
||||
covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)
|
||||
if not np.isfinite(covmean).all():
|
||||
msg = ('fid calculation produces singular product; '
|
||||
'adding %s to diagonal of cov estimates') % eps
|
||||
print(msg)
|
||||
offset = np.eye(sigma1.shape[0]) * eps
|
||||
covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))
|
||||
|
||||
# Numerical error might give slight imaginary component
|
||||
if np.iscomplexobj(covmean):
|
||||
if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):
|
||||
m = np.max(np.abs(covmean.imag))
|
||||
raise ValueError('Imaginary component {}'.format(m))
|
||||
covmean = covmean.real
|
||||
|
||||
tr_covmean = np.trace(covmean)
|
||||
|
||||
return (diff.dot(diff) + np.trace(sigma1) + np.trace(sigma2) - 2 * tr_covmean)
|
||||
|
||||
|
||||
def evaluate_fid(data, gt_data, num_motions=None, mask=None):
|
||||
if num_motions is None:
|
||||
index = list(range(data['pos'].shape[0]))
|
||||
index_gt = list(range(gt_data['pos'].shape[0]))
|
||||
else:
|
||||
# randomly pick [num_motions] motions from the data
|
||||
assert num_motions <= data['pos'].shape[0]
|
||||
assert num_motions <= gt_data['pos'].shape[0]
|
||||
|
||||
index = list(range(data['pos'].shape[0]))
|
||||
index = random.sample(index, num_motions)
|
||||
index_gt = list(range(gt_data['pos'].shape[0]))
|
||||
index_gt = random.sample(index_gt, num_motions)
|
||||
|
||||
data_selected = {}
|
||||
gt_data_selected = {}
|
||||
for key in data.keys():
|
||||
data_selected[key] = data[key][index].clone()
|
||||
gt_data_selected[key] = gt_data[key][index_gt].clone()
|
||||
|
||||
if mask is not None:
|
||||
# replace with ground truth
|
||||
for key in data.keys():
|
||||
data_selected[key][:, mask] = gt_data[key][:, mask].clone()
|
||||
|
||||
with torch.no_grad():
|
||||
prior.set_input(data_selected)
|
||||
f_l, _, _ = prior.encode_local()
|
||||
f_g, _, _ = prior.encode_global()
|
||||
feature = torch.cat((f_l, f_g), dim=-1)
|
||||
|
||||
prior.set_input(gt_data_selected)
|
||||
f_l_gt, _, _ = prior.encode_local()
|
||||
f_g_gt, _, _ = prior.encode_global()
|
||||
feature_gt = torch.cat((f_l_gt, f_g_gt), dim=-1)
|
||||
|
||||
statistics = calculate_feature_statistics(feature.view(feature.shape[0], -1))
|
||||
statistics_gt = calculate_feature_statistics(feature_gt.view(feature_gt.shape[0], -1))
|
||||
|
||||
fid = calculate_frechet_distance(statistics[0], statistics[1], statistics_gt[0], statistics_gt[1])
|
||||
print(f'FID: {fid:.4f}')
|
||||
|
||||
return fid
|
||||
|
||||
|
||||
def evaluate_diversity(data):
|
||||
index = list(range(data['pos'].shape[0]))
|
||||
random.shuffle(index)
|
||||
index1 = index[:data['pos'].shape[0] // 2]
|
||||
index2 = index[data['pos'].shape[0] // 2:]
|
||||
|
||||
with torch.no_grad():
|
||||
prior.set_input(data)
|
||||
f_l, _, _ = prior.encode_local()
|
||||
# f_g, _, _ = prior.encode_global()
|
||||
# feature = torch.cat((f_l, f_g), dim=-1)
|
||||
|
||||
feature = f_l.clone()
|
||||
|
||||
|
||||
diversity = 0
|
||||
for first_idx, second_idx in zip(index1, index2):
|
||||
diversity += torch.dist(feature[first_idx], feature[second_idx])
|
||||
diversity /= len(index1)
|
||||
|
||||
print(f'Diversity: {diversity.item():.4f}')
|
||||
|
||||
return diversity.item()
|
||||
|
||||
|
||||
def evaluate_recon_error(data, gt_data):
|
||||
criterion_geo = GeodesicLoss()
|
||||
|
||||
root_orient = rotation_6d_to_matrix(data['root_orient'])
|
||||
root_orient_gt = rotation_6d_to_matrix(gt_data['root_orient'])
|
||||
orientation_error = criterion_geo(root_orient.view(-1, 3, 3), root_orient_gt.view(-1, 3, 3))
|
||||
|
||||
rotation_error = criterion_geo(data['rotmat'][:, :, 1:].reshape(-1, 3, 3), gt_data['rotmat'][:, :, 1:].reshape(-1, 3, 3))
|
||||
position_error = torch.linalg.norm((data['pos'] - gt_data['pos']), dim=-1).mean()
|
||||
translation_error = torch.linalg.norm((data['trans'] - gt_data['trans']), dim=-1).mean()
|
||||
final_displacement_error = torch.linalg.norm(data['global_pos'][:, -1] - gt_data['global_pos'][:, -1], dim=-1).mean()
|
||||
|
||||
return {
|
||||
'rotation': c2c(rotation_error) * 180.0 / np.pi,
|
||||
'position': c2c(position_error) * 100.0,
|
||||
'orientation': c2c(orientation_error) * 180.0 / np.pi,
|
||||
'translation': c2c(translation_error) * 100.0,
|
||||
'final_displacement': c2c(final_displacement_error) * 100.0
|
||||
}
|
||||
|
||||
|
||||
def synthetic_trajectory(N_motions):
|
||||
|
||||
z_l = torch.normal(mean=0, std=1, size=(N_motions, args.local_prior.z_dim)).to(prior.device)
|
||||
z_g = torch.normal(mean=0, std=1, size=(N_motions, args.global_prior.z_dim)).to(prior.device)
|
||||
# z_g = torch.zeros((N_motions, args.global_prior.z_dim)).to(prior.device)
|
||||
|
||||
|
||||
with torch.no_grad():
|
||||
output = prior.decode(z_l, z_g, length=args.data.clip_length, step=1)
|
||||
|
||||
rotmat = output["rotmat"] # (N, T, J, 3, 3)
|
||||
angular = estimate_angular_velocity(rotmat.clone(), dt=1.0 / args.data.fps) # angular velocity of all the joints (N, T, J, 3)
|
||||
pos, global_xform = fk(rotmat.view(-1, 16, 3, 3))
|
||||
pos = pos.contiguous().view(N_motions, args.data.clip_length, 16, 3) # local joint positions (N, T, J, 3)
|
||||
global_xform = global_xform.view(N_motions, args.data.clip_length, 16, 4, 4)
|
||||
|
||||
global_xform = global_xform[:, :, :, :3, :3] # global transformation matrix for each joint (N, T, J, 3, 3)
|
||||
velocity = estimate_linear_velocity(pos, dt=1.0 / args.data.fps) # linear velocity of all the joints (N, T, J, 3)
|
||||
|
||||
root_rotation = rotmat[:, :, 0] # (N, T, 3, 3)
|
||||
global_pos = torch.matmul(root_rotation.unsqueeze(2), pos.unsqueeze(-1)).squeeze(-1) # (N, T, J, 3)
|
||||
|
||||
output_ = {'pos': pos,
|
||||
'velocity': velocity,
|
||||
'global_xform': matrix_to_rotation_6d(global_xform),
|
||||
'angular': angular,
|
||||
'root_orient': matrix_to_rotation_6d(root_rotation),
|
||||
'rotmat': rotmat}
|
||||
|
||||
return output_
|
||||
|
||||
|
||||
def ablation_study(path, subtask):
|
||||
|
||||
num_motions = 100
|
||||
print(os.path.join(path, 'TCDHands', '*.npz'))
|
||||
|
||||
test_data = load_data(sorted(glob.glob(os.path.join('./outputs/generative_copy/TCDHands', '*.npz')))) # test data TCDHands
|
||||
val_data = load_data(sorted(glob.glob(os.path.join('./outputs/generative_copy/GRAB', '*.npz')))) # test data TCDHands
|
||||
train_data = load_data(sorted(glob.glob(os.path.join('./outputs/generative_copy/SAMP', '*.npz')))) # test data TCDHands
|
||||
synthetic_data = synthetic_trajectory(N_motions=num_motions)
|
||||
|
||||
pca_aa = joblib.load("data/pca.pkl")
|
||||
pca_mean = torch.tensor(pca_aa.mean_).to("cuda")
|
||||
pca_cov = torch.tensor(pca_aa.get_covariance()).to("cuda")
|
||||
pca_sv = torch.tensor(pca_aa.singular_values_).to("cuda")
|
||||
pca_pc = torch.tensor(pca_aa.components_).to("cuda")
|
||||
|
||||
random_samples_low_dim = np.random.randn(num_motions, pca_aa.n_components_)
|
||||
pca_data = pca_aa.inverse_transform(random_samples_low_dim)
|
||||
|
||||
gmm_aa = MaxMixturePrior()
|
||||
gmm_data = gmm_aa.gmm.sample(num_motions)
|
||||
|
||||
identity = torch.eye(3)
|
||||
identity = identity.view(1, 1, 3, 3).repeat(num_motions, args.data.clip_length, 1, 1)
|
||||
|
||||
# convert to tensor
|
||||
gmm_data = torch.tensor(gmm_data[0].reshape(num_motions, args.data.clip_length, -1, 3))
|
||||
gmm_data = torch.cat((identity.unsqueeze(2), axis_angle_to_matrix(gmm_data)), dim=2)
|
||||
gmm_rot6d = matrix_to_rotation_6d(gmm_data).reshape(num_motions * args.data.clip_length, -1, 6)
|
||||
|
||||
# convert to tensor
|
||||
pca_data = torch.tensor(pca_data.reshape(num_motions, args.data.clip_length, -1, 3))
|
||||
pca_data = torch.cat((identity.unsqueeze(2), axis_angle_to_matrix(pca_data)), dim=2)
|
||||
pca_rot6d = matrix_to_rotation_6d(pca_data).reshape(num_motions * args.data.clip_length, -1, 6)
|
||||
|
||||
gmm_pos, gmm_xform = fk(gmm_rot6d.to(fk.device))
|
||||
pca_pos, pca_xform = fk(pca_rot6d.to(fk.device))
|
||||
|
||||
gmm_xform = matrix_to_rotation_6d(gmm_xform[:, :, :3, :3])
|
||||
pca_xform = matrix_to_rotation_6d(pca_xform[:, :, :3, :3])
|
||||
|
||||
gmm_velocity = estimate_linear_velocity(gmm_pos.unsqueeze(0), dt=1.0 / args.data.fps).squeeze(0)
|
||||
pca_velocity = estimate_linear_velocity(pca_pos.unsqueeze(0), dt=1.0 / args.data.fps).squeeze(0)
|
||||
|
||||
gmm_angular = estimate_angular_velocity(gmm_data.unsqueeze(0), dt=1.0 / args.data.fps).squeeze(0)
|
||||
pca_angular = estimate_angular_velocity(pca_data.unsqueeze(0), dt=1.0 / args.data.fps).squeeze(0)
|
||||
|
||||
gmm_dict = {'pos': gmm_pos.reshape(num_motions, args.data.clip_length, -1, 3),
|
||||
'global_xform': gmm_xform.reshape(num_motions, args.data.clip_length, -1, 6),
|
||||
'velocity': gmm_velocity.reshape(num_motions, args.data.clip_length, -1, 3),
|
||||
'angular': gmm_angular}
|
||||
|
||||
pca_dict = {'pos': pca_pos.reshape(num_motions, args.data.clip_length, -1, 3),
|
||||
'global_xform': pca_xform.reshape(num_motions, args.data.clip_length, -1, 6),
|
||||
'velocity': pca_velocity.reshape(num_motions, args.data.clip_length, -1, 3),
|
||||
'angular': pca_angular}
|
||||
|
||||
metrics = dict()
|
||||
|
||||
def f(data, num_motions, key):
|
||||
div = evaluate_diversity(data)
|
||||
if key in metrics.keys():
|
||||
pass
|
||||
else:
|
||||
metrics[key] = {'DIV': [div]}
|
||||
|
||||
print('Motion Synthesis:')
|
||||
div_test = f(test_data, num_motions, 'TEST_DATA (TDCHANDS)')
|
||||
div_val = f(val_data, num_motions, 'VAL_DATA (SAMP)')
|
||||
div_train = f(train_data, num_motions, 'TRAIN_DATA (GRAB)')
|
||||
div_synt = f(synthetic_data, num_motions, 'SYNTHETIC')
|
||||
div_gmm = f(gmm_dict, num_motions, 'GMM')
|
||||
div_pca = f(pca_dict, num_motions, 'PCA')
|
||||
|
||||
df = pd.DataFrame.from_dict(metrics)
|
||||
df.to_csv(os.path.join(path, f'diversity.csv'), index=False)
|
||||
|
||||
|
||||
|
||||
def comparison(path, subtask):
|
||||
gt_data = load_data(sorted(glob.glob(os.path.join(path, 'gt', '*.npz'))))
|
||||
|
||||
file_path = os.path.join(path, subtask)
|
||||
|
||||
metrics = dict()
|
||||
if subtask == 'recon':
|
||||
humor_data = load_data(sorted(glob.glob(os.path.join(file_path, 'humor', '*.npz'))))
|
||||
hm_vae_data = load_data(sorted(glob.glob(os.path.join(file_path, 'hm-vae', '*.npz'))))
|
||||
nemf_data = load_data(sorted(glob.glob(os.path.join(file_path, 'nemf', '*.npz'))))
|
||||
|
||||
metrics['HUMOR'] = evaluate_recon_error(humor_data, gt_data)
|
||||
metrics['HMVAE'] = evaluate_recon_error(hm_vae_data, gt_data)
|
||||
metrics['NEMF'] = evaluate_recon_error(nemf_data, gt_data)
|
||||
print('Motion Reconstruction:')
|
||||
print(json.dumps(metrics, sort_keys=True, indent=4))
|
||||
|
||||
df = pd.DataFrame.from_dict(metrics)
|
||||
df.to_csv(os.path.join(file_path, f'{subtask}.csv'))
|
||||
elif subtask == 'synth':
|
||||
humor_data = load_data(sorted(glob.glob(os.path.join(file_path, 'humor', '.*.npz'))), padding=True)
|
||||
hm_vae_data = load_data(sorted(glob.glob(os.path.join(file_path, 'hm-vae', '*.npz'))), padding=True)
|
||||
nemf_data = load_data(sorted(glob.glob(os.path.join(file_path, 'nemf', '*.npz'))), padding=True)
|
||||
|
||||
num_motions = 100
|
||||
|
||||
def f(data, num_motions, key):
|
||||
fid = evaluate_fid(data, gt_data, num_motions)
|
||||
div = evaluate_diversity(data)
|
||||
index = list(range(data['global_pos'].shape[0]))
|
||||
index = random.sample(index, num_motions)
|
||||
s = foot_skating(data['global_pos'][index])
|
||||
if key in metrics.keys():
|
||||
metrics[key]['FID'].append(fid)
|
||||
metrics[key]['DIV'].append(div)
|
||||
metrics[key]['FOOT'].append(s)
|
||||
else:
|
||||
metrics[key] = {'FID': [fid], 'DIV': [div], 'FOOT': [s]}
|
||||
|
||||
print('Motion Synthesis:')
|
||||
for replication in range(20):
|
||||
print(f'==================== Replication {replication} ====================')
|
||||
f(humor_data, num_motions, 'HUMOR')
|
||||
f(hm_vae_data, num_motions, 'HMVAE')
|
||||
f(nemf_data, num_motions, 'NEMF')
|
||||
|
||||
for key in metrics.keys():
|
||||
df = pd.DataFrame.from_dict(metrics[key])
|
||||
df.to_csv(os.path.join(file_path, f'{subtask}_{key}.csv'), index=False)
|
||||
else:
|
||||
raise RuntimeError(f'unsupported evaluation task for comparison: {subtask}')
|
||||
|
||||
|
||||
def inbetween_benchmarks(path, keyframes=False):
|
||||
trans_lengths = [5, 10, 15, 20] if keyframes else [10, 20, 30]
|
||||
|
||||
gt_data = load_data(sorted(glob.glob(os.path.join(path, 'gt', '*.npz'))))
|
||||
s = foot_skating(gt_data['global_pos'])
|
||||
print(f'average foot skating of the real motions: {s:.4f}')
|
||||
|
||||
def f(data, key, mask):
|
||||
fid = evaluate_fid(data, gt_data, mask=mask)
|
||||
s = foot_skating(data['global_pos'], mask=mask)
|
||||
if key in statistics.keys():
|
||||
statistics[key]['FID'].append(fid)
|
||||
statistics[key]['FOOT'].append(s)
|
||||
else:
|
||||
statistics[key] = {'FID': [fid], 'FOOT': [s]}
|
||||
|
||||
file_path = os.path.join(path, 'keyframe') if keyframes else os.path.join(path, 'clip')
|
||||
statistics = {}
|
||||
for t in trans_lengths:
|
||||
print(f'==================== Length {t} ====================')
|
||||
if keyframes:
|
||||
mask = None # don't mask data for keyframe inbetweening
|
||||
else:
|
||||
end_frame = (args.data.clip_length - t) // 2
|
||||
mask = list(range(end_frame)) + list(range(end_frame + t, args.data.clip_length))
|
||||
print(f'mask: {mask}')
|
||||
|
||||
slerp_data = load_data(sorted(glob.glob(os.path.join(file_path, 'slerp', f'{t}_frames', '*.npz'))))
|
||||
f(slerp_data, 'SLERP', mask)
|
||||
|
||||
if not keyframes:
|
||||
inertia_data = load_data(sorted(glob.glob(os.path.join(file_path, 'inertialization', f'{t}_frames', '*.npz'))))
|
||||
f(inertia_data, 'INERTIALIZATION', mask)
|
||||
|
||||
rmi_data = load_data(sorted(glob.glob(os.path.join(file_path, 'rmi', f'{t}_frames', '*.npz'))))
|
||||
f(rmi_data, 'RMI', mask)
|
||||
|
||||
hm_vae_data = load_data(sorted(glob.glob(os.path.join(file_path, 'hm-vae', f'{t}_frames', '*.npz'))))
|
||||
f(hm_vae_data, 'HMVAE', mask)
|
||||
|
||||
nemf_data = load_data(sorted(glob.glob(os.path.join(file_path, 'nemf', f'{t}_frames', '*.npz'))))
|
||||
f(nemf_data, 'NEMF', mask)
|
||||
|
||||
for key in statistics.keys():
|
||||
df = pd.DataFrame.from_dict(statistics[key])
|
||||
df.to_csv(os.path.join(file_path, f'inbetween_{key}.csv'))
|
||||
|
||||
|
||||
def compute_joint_accel(joint_seq, fps):
|
||||
''' Magnitude of joint accelerations for joint_seq : B x T x J x 3 '''
|
||||
joint_accel = joint_seq[:, :-2] - (2 * joint_seq[:, 1:-1]) + joint_seq[:, 2:]
|
||||
DATA_h = 1.0 / fps
|
||||
joint_accel = joint_accel / ((DATA_h**2))
|
||||
joint_accel_mag = torch.norm(joint_accel, dim=-1)
|
||||
|
||||
return joint_accel, joint_accel_mag
|
||||
|
||||
|
||||
def evaluate_super_sampling(path):
|
||||
fps_30_data = load_data(sorted(glob.glob(os.path.join(path, '30fps/*.npz'))))
|
||||
fps_60_data = load_data(sorted(glob.glob(os.path.join(path, '60fps/*.npz'))))
|
||||
fps_120_data = load_data(sorted(glob.glob(os.path.join(path, '120fps/*.npz'))))
|
||||
fps_240_data = load_data(sorted(glob.glob(os.path.join(path, '240fps/*.npz'))))
|
||||
|
||||
smoothness = dict()
|
||||
joint_vel = estimate_linear_velocity(fps_30_data['global_pos'], dt=1.0 / 30.0)
|
||||
joint_vel_mag = torch.norm(joint_vel, dim=-1).mean()
|
||||
smoothness['30FPS'] = [c2c(joint_vel_mag * 100.0)]
|
||||
joint_vel = estimate_linear_velocity(fps_60_data['global_pos'], dt=1.0 / 60.0)
|
||||
joint_vel_mag = torch.norm(joint_vel, dim=-1).mean()
|
||||
smoothness['60FPS'] = [c2c(joint_vel_mag * 100.0)]
|
||||
joint_vel = estimate_linear_velocity(fps_120_data['global_pos'], dt=1.0 / 120.0)
|
||||
joint_vel_mag = torch.norm(joint_vel, dim=-1).mean()
|
||||
smoothness['120FPS'] = [c2c(joint_vel_mag * 100.0)]
|
||||
joint_vel = estimate_linear_velocity(fps_240_data['global_pos'], dt=1.0 / 240.0)
|
||||
joint_vel_mag = torch.norm(joint_vel, dim=-1).mean()
|
||||
smoothness['240FPS'] = [c2c(joint_vel_mag * 100.0)]
|
||||
|
||||
df = pd.DataFrame.from_dict(smoothness)
|
||||
df.to_csv(os.path.join(path, f'smoothness.csv'))
|
||||
|
||||
|
||||
def evaluate_smoothness_test(path):
|
||||
smoothness = dict()
|
||||
|
||||
for fps in [30, 60]:
|
||||
smoothness[f'{fps}FPS'] = dict()
|
||||
for L in range(1, 23, 2):
|
||||
data = load_data([os.path.join(path, f'L{L}', f'test_{fps}fps.npz')])
|
||||
# _, joint_accel_mag = compute_joint_accel(data['global_pos'], fps=fps)
|
||||
# smoothness[f'{fps}FPS'][L] = c2c(joint_accel_mag.mean() * 100.0)
|
||||
joint_vel = estimate_linear_velocity(data['global_pos'], dt=1.0 / fps)
|
||||
joint_vel_mag = torch.norm(joint_vel, dim=-1).mean()
|
||||
smoothness[f'{fps}FPS'][L] = c2c(joint_vel_mag.mean() * 100.0)
|
||||
|
||||
df = pd.DataFrame.from_dict(smoothness)
|
||||
df.to_csv(os.path.join(path, f'smoothness.csv'))
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
plt.rcParams['text.usetex'] = True
|
||||
|
||||
smoothness_30fps = np.array(list(smoothness['30FPS'].values()))
|
||||
smoothness_60fps = np.array(list(smoothness['60FPS'].values()))
|
||||
|
||||
plt.plot(np.arange(1, 23, 2), smoothness_30fps, marker='o', label='30fps')
|
||||
plt.plot(np.arange(1, 23, 2), smoothness_60fps, marker='o', label='60fps')
|
||||
plt.legend()
|
||||
plt.grid(linestyle='--')
|
||||
plt.xticks(np.arange(1, 23, 2))
|
||||
plt.xlabel(r'$L$')
|
||||
# plt.yscale('log')
|
||||
plt.ylabel('Mean Per-Joint Velocity')
|
||||
plt.savefig(os.path.join(path, 'smoothness.pdf'))
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
from argparse import ArgumentParser
|
||||
|
||||
parser = ArgumentParser()
|
||||
parser.add_argument('--config', type=str, default='evaluation.yaml', help='name of the configutation file')
|
||||
parser.add_argument('--task', type=str, help='name of application tasks', required=True)
|
||||
parser.add_argument('--load_path', type=str, help='path to load results for evaluation', required=True)
|
||||
|
||||
opt = parser.parse_args()
|
||||
args = Arguments('./configs', filename=opt.config)
|
||||
|
||||
torch.set_default_dtype(torch.float32)
|
||||
|
||||
torch.manual_seed(0)
|
||||
np.random.seed(0)
|
||||
random.seed(0)
|
||||
|
||||
ngpu = 1
|
||||
if args.multi_gpu is True:
|
||||
ngpu = torch.cuda.device_count()
|
||||
if ngpu == 1:
|
||||
args.multi_gpu = False
|
||||
print(f'Number of GPUs: {ngpu}')
|
||||
|
||||
prior = Architecture(args, ngpu)
|
||||
prior.load(optimal=True)
|
||||
prior.eval()
|
||||
fk = ForwardKinematicsLayer(args)
|
||||
|
||||
if opt.task == 'ablation_study':
|
||||
print('Evaluating ablation study ...')
|
||||
ablation_study(opt.load_path, subtask='synth')
|
||||
elif opt.task == 'comparison':
|
||||
print('Evaluating comparison with other methods ...')
|
||||
comparison(opt.load_path, subtask='recon')
|
||||
comparison(opt.load_path, subtask='synth')
|
||||
elif opt.task == 'inbetween_benchmark':
|
||||
print('Evaluating motion in-betweening ...')
|
||||
inbetween_benchmarks(opt.load_path, keyframes=False)
|
||||
inbetween_benchmarks(opt.load_path, keyframes=True)
|
||||
elif opt.task == 'super_sampling':
|
||||
print('Evaluating super sampling ...')
|
||||
evaluate_super_sampling(opt.load_path)
|
||||
elif opt.task == 'smoothness_test':
|
||||
print('Evaluating smoothness testing ...')
|
||||
evaluate_smoothness_test(opt.load_path)
|
||||
else:
|
||||
raise NotImplementedError(f'unknown evaluation: {opt.task}')
|
||||
+1731
File diff suppressed because it is too large
Load Diff
+1119
File diff suppressed because it is too large
Load Diff
+150
@@ -0,0 +1,150 @@
|
||||
batch_size: 1
|
||||
bvh_viz: false
|
||||
|
||||
data:
|
||||
clip_length: 128
|
||||
fps: 30
|
||||
gender: neutral
|
||||
normalize:
|
||||
- pos
|
||||
- velocity
|
||||
- global_xform
|
||||
- angular
|
||||
- height
|
||||
- root_orient
|
||||
- root_vel
|
||||
root_transform: true
|
||||
up: z
|
||||
|
||||
root: None
|
||||
dataset_dir: _DATA/hmp_model
|
||||
dtw_loss: true
|
||||
geodesic_loss: true
|
||||
|
||||
global_prior:
|
||||
activation: tanh
|
||||
in_channels: 6
|
||||
kernel_size: 4
|
||||
temporal_scale: 8
|
||||
z_dim: 256
|
||||
|
||||
initialization: true
|
||||
is_train: false
|
||||
keypoint_source: mmpose
|
||||
|
||||
l1_loss: true
|
||||
lambda_angle: 1
|
||||
lambda_dtw: 0.5
|
||||
lambda_kl: 0.0001
|
||||
|
||||
local_prior:
|
||||
activation: tanh
|
||||
channel_base: 15
|
||||
extra_conv: 0
|
||||
kernel_size: 4
|
||||
num_layers: 4
|
||||
padding_mode: reflect
|
||||
skeleton_dist: 2
|
||||
skeleton_pool: mean
|
||||
temporal_scale: 8
|
||||
upsampling: linear
|
||||
use_residual_blocks: true
|
||||
z_dim: 1024
|
||||
|
||||
log: false
|
||||
multi_gpu: false
|
||||
|
||||
nemf:
|
||||
bandwidth: 7
|
||||
global_output: 6
|
||||
global_z: 256
|
||||
hidden_neuron: 1024
|
||||
local_output: 144
|
||||
local_z: 1024
|
||||
norm_layer: true
|
||||
siren: false
|
||||
skip_connection: true
|
||||
nsubject: 1
|
||||
num_workers: 0
|
||||
opt_betas: false
|
||||
output_trans: false
|
||||
overlap_len: 16
|
||||
pretrained_gmp: false
|
||||
save_dir: _DATA/hmp_model
|
||||
|
||||
scheduler:
|
||||
gamma: 0.7
|
||||
step_size: 200
|
||||
|
||||
stg1:
|
||||
betas_prior: 10
|
||||
keypoint_source: mmpose
|
||||
lambda_batch_cs: 0
|
||||
lambda_init_z_prior: 0
|
||||
lambda_j3d_smooth: 0
|
||||
lambda_motion_prior: 0
|
||||
lambda_pose_prior: 0
|
||||
lambda_orient: 0
|
||||
lambda_orient_smooth: 1
|
||||
lambda_pos: 0
|
||||
lambda_reproj: 0.05
|
||||
lambda_bio: 1
|
||||
lambda_consistency: 0
|
||||
lambda_rot: 0
|
||||
lambda_rot_smooth: 0
|
||||
lambda_trans: 10
|
||||
lambda_trans_smooth: 10
|
||||
lr: 0.05
|
||||
niters: 400
|
||||
opt_params:
|
||||
- betas
|
||||
- trans
|
||||
- root_orient
|
||||
|
||||
stg2:
|
||||
betas_prior: 10
|
||||
keypoint_source: mmpose
|
||||
lambda_batch_cs: 0
|
||||
lambda_init_z_prior: 0
|
||||
lambda_j3d_smooth: 0
|
||||
lambda_motion_prior: 200
|
||||
lambda_pose_prior: 0
|
||||
lambda_orient: 0
|
||||
lambda_orient_smooth: 1
|
||||
lambda_pos: 0
|
||||
lambda_reproj: 0.05
|
||||
lambda_bio: 1
|
||||
lambda_consistency: 0
|
||||
lambda_rot: 0
|
||||
lambda_rot_smooth: 0
|
||||
lambda_trans: 10
|
||||
lambda_trans_smooth: 10
|
||||
lr: 0.05
|
||||
niters: 400
|
||||
opt_params:
|
||||
- betas
|
||||
- trans
|
||||
- root_orient
|
||||
- z_l
|
||||
|
||||
stg3:
|
||||
betas_prior: 0
|
||||
lambda_batch_cs: 0
|
||||
lambda_init_z_prior: 0
|
||||
lambda_j3d_smooth: 0
|
||||
lambda_motion_prior: 0
|
||||
lambda_orient: 0
|
||||
lambda_orient_smooth: 0
|
||||
lambda_pos: 0
|
||||
lambda_reproj: 0
|
||||
lambda_bio: 0
|
||||
lambda_rot: 0
|
||||
lambda_rot_smooth: 0
|
||||
lambda_trans: 0
|
||||
lambda_trans_smooth: 0
|
||||
lr: 0.0
|
||||
niters: 0
|
||||
opt_params:
|
||||
- root_orient
|
||||
- z_l
|
||||
verbose: false
|
||||
@@ -0,0 +1,675 @@
|
||||
import operator
|
||||
|
||||
import numpy as np
|
||||
import numpy.core.umath as ut # (for numpy 1.25)
|
||||
# import numpy.core.umath_tests as ut
|
||||
|
||||
import AnimationStructure
|
||||
from Quaternions import Quaternions
|
||||
|
||||
|
||||
class Animation:
|
||||
"""
|
||||
Animation is a numpy-like wrapper for animation data
|
||||
|
||||
Animation data consists of several arrays consisting
|
||||
of F frames and J joints.
|
||||
|
||||
The animation is specified by
|
||||
|
||||
rotations : (F, J) Quaternions | Joint Rotations
|
||||
positions : (F, J, 3) ndarray | Joint Positions
|
||||
|
||||
The base pose is specified by
|
||||
|
||||
orients : (J) Quaternions | Joint Orientations
|
||||
offsets : (J, 3) ndarray | Joint Offsets
|
||||
|
||||
And the skeletal structure is specified by
|
||||
|
||||
parents : (J) ndarray | Joint Parents
|
||||
"""
|
||||
|
||||
def __init__(self, rotations, positions, orients, offsets, parents):
|
||||
|
||||
self.rotations = rotations
|
||||
self.positions = positions
|
||||
self.orients = orients
|
||||
self.offsets = offsets
|
||||
self.parents = parents
|
||||
|
||||
def __op__(self, op, other):
|
||||
return Animation(
|
||||
op(self.rotations, other.rotations),
|
||||
op(self.positions, other.positions),
|
||||
op(self.orients, other.orients),
|
||||
op(self.offsets, other.offsets),
|
||||
op(self.parents, other.parents))
|
||||
|
||||
def __iop__(self, op, other):
|
||||
self.rotations = op(self.roations, other.rotations)
|
||||
self.positions = op(self.roations, other.positions)
|
||||
self.orients = op(self.orients, other.orients)
|
||||
self.offsets = op(self.offsets, other.offsets)
|
||||
self.parents = op(self.parents, other.parents)
|
||||
return self
|
||||
|
||||
def __sop__(self, op):
|
||||
return Animation(
|
||||
op(self.rotations),
|
||||
op(self.positions),
|
||||
op(self.orients),
|
||||
op(self.offsets),
|
||||
op(self.parents))
|
||||
|
||||
def __add__(self, other): return self.__op__(operator.add, other)
|
||||
def __sub__(self, other): return self.__op__(operator.sub, other)
|
||||
def __mul__(self, other): return self.__op__(operator.mul, other)
|
||||
def __div__(self, other): return self.__op__(operator.div, other)
|
||||
|
||||
def __abs__(self): return self.__sop__(operator.abs)
|
||||
def __neg__(self): return self.__sop__(operator.neg)
|
||||
|
||||
def __iadd__(self, other): return self.__iop__(operator.iadd, other)
|
||||
def __isub__(self, other): return self.__iop__(operator.isub, other)
|
||||
def __imul__(self, other): return self.__iop__(operator.imul, other)
|
||||
def __idiv__(self, other): return self.__iop__(operator.idiv, other)
|
||||
|
||||
def __len__(self): return len(self.rotations)
|
||||
|
||||
def __getitem__(self, k):
|
||||
if isinstance(k, tuple):
|
||||
return Animation(
|
||||
self.rotations[k],
|
||||
self.positions[k],
|
||||
self.orients[k[1:]],
|
||||
self.offsets[k[1:]],
|
||||
self.parents[k[1:]])
|
||||
else:
|
||||
return Animation(
|
||||
self.rotations[k],
|
||||
self.positions[k],
|
||||
self.orients,
|
||||
self.offsets,
|
||||
self.parents)
|
||||
|
||||
def __setitem__(self, k, v):
|
||||
if isinstance(k, tuple):
|
||||
self.rotations.__setitem__(k, v.rotations)
|
||||
self.positions.__setitem__(k, v.positions)
|
||||
self.orients.__setitem__(k[1:], v.orients)
|
||||
self.offsets.__setitem__(k[1:], v.offsets)
|
||||
self.parents.__setitem__(k[1:], v.parents)
|
||||
else:
|
||||
self.rotations.__setitem__(k, v.rotations)
|
||||
self.positions.__setitem__(k, v.positions)
|
||||
self.orients.__setitem__(k, v.orients)
|
||||
self.offsets.__setitem__(k, v.offsets)
|
||||
self.parents.__setitem__(k, v.parents)
|
||||
|
||||
@property
|
||||
def shape(self): return (self.rotations.shape[0], self.rotations.shape[1])
|
||||
|
||||
def copy(self): return Animation(
|
||||
self.rotations.copy(), self.positions.copy(),
|
||||
self.orients.copy(), self.offsets.copy(),
|
||||
self.parents.copy())
|
||||
|
||||
def repeat(self, *args, **kw):
|
||||
return Animation(
|
||||
self.rotations.repeat(*args, **kw),
|
||||
self.positions.repeat(*args, **kw),
|
||||
self.orients, self.offsets, self.parents)
|
||||
|
||||
def ravel(self):
|
||||
return np.hstack([
|
||||
self.rotations.log().ravel(),
|
||||
self.positions.ravel(),
|
||||
self.orients.log().ravel(),
|
||||
self.offsets.ravel()])
|
||||
|
||||
@classmethod
|
||||
def unravel(clas, anim, shape, parents):
|
||||
nf, nj = shape
|
||||
rotations = anim[nf*nj*0:nf*nj*3]
|
||||
positions = anim[nf*nj*3:nf*nj*6]
|
||||
orients = anim[nf*nj*6+nj*0:nf*nj*6+nj*3]
|
||||
offsets = anim[nf*nj*6+nj*3:nf*nj*6+nj*6]
|
||||
return cls(
|
||||
Quaternions.exp(rotations), positions,
|
||||
Quaternions.exp(orients), offsets,
|
||||
parents.copy())
|
||||
|
||||
|
||||
""" Maya Interaction """
|
||||
|
||||
|
||||
def load_to_maya(anim, names=None, radius=0.5):
|
||||
"""
|
||||
Load Animation Object into Maya as Joint Skeleton
|
||||
loads each frame as a new keyfame in maya.
|
||||
|
||||
If the animation is too slow or too fast perhaps
|
||||
the framerate needs adjusting before being loaded
|
||||
such that it matches the maya scene framerate.
|
||||
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
anim : Animation
|
||||
Animation to load into Scene
|
||||
|
||||
names : [str]
|
||||
Optional list of Joint names for Skeleton
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
List of Maya Joint Nodes loaded into scene
|
||||
"""
|
||||
|
||||
import pymel.core as pm
|
||||
|
||||
joints = []
|
||||
frames = range(1, len(anim)+1)
|
||||
|
||||
if names is None:
|
||||
names = ["joint_" + str(i) for i in range(len(anim.parents))]
|
||||
|
||||
for i, offset, orient, parent, name in zip(range(len(anim.offsets)), anim.offsets, anim.orients, anim.parents, names):
|
||||
|
||||
if parent < 0:
|
||||
pm.select(d=True)
|
||||
else:
|
||||
pm.select(joints[parent])
|
||||
|
||||
joint = pm.joint(n=name, p=offset, relative=True, radius=radius)
|
||||
joint.setOrientation([orient[1], orient[2], orient[3], orient[0]])
|
||||
|
||||
curvex = pm.nodetypes.AnimCurveTA(n=name + "_rotateX")
|
||||
curvey = pm.nodetypes.AnimCurveTA(n=name + "_rotateY")
|
||||
curvez = pm.nodetypes.AnimCurveTA(n=name + "_rotateZ")
|
||||
|
||||
jrotations = (-Quaternions(orient[np.newaxis]) * anim.rotations[:, i]).euler()
|
||||
curvex.addKeys(frames, jrotations[:, 0])
|
||||
curvey.addKeys(frames, jrotations[:, 1])
|
||||
curvez.addKeys(frames, jrotations[:, 2])
|
||||
|
||||
pm.connectAttr(curvex.output, joint.rotateX)
|
||||
pm.connectAttr(curvey.output, joint.rotateY)
|
||||
pm.connectAttr(curvez.output, joint.rotateZ)
|
||||
|
||||
offsetx = pm.nodetypes.AnimCurveTU(n=name + "_translateX")
|
||||
offsety = pm.nodetypes.AnimCurveTU(n=name + "_translateY")
|
||||
offsetz = pm.nodetypes.AnimCurveTU(n=name + "_translateZ")
|
||||
|
||||
offsetx.addKeys(frames, anim.positions[:, i, 0])
|
||||
offsety.addKeys(frames, anim.positions[:, i, 1])
|
||||
offsetz.addKeys(frames, anim.positions[:, i, 2])
|
||||
|
||||
pm.connectAttr(offsetx.output, joint.translateX)
|
||||
pm.connectAttr(offsety.output, joint.translateY)
|
||||
pm.connectAttr(offsetz.output, joint.translateZ)
|
||||
|
||||
joints.append(joint)
|
||||
|
||||
return joints
|
||||
|
||||
|
||||
def load_from_maya(root, start, end):
|
||||
"""
|
||||
Load Animation Object from Maya Joint Skeleton
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
root : PyNode
|
||||
Root Joint of Maya Skeleton
|
||||
|
||||
start, end : int, int
|
||||
Start and End frame index of Maya Animation
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
animation : Animation
|
||||
Loaded animation from maya
|
||||
|
||||
names : [str]
|
||||
Joint names from maya
|
||||
"""
|
||||
|
||||
import pymel.core as pm
|
||||
|
||||
original_time = pm.currentTime(q=True)
|
||||
pm.currentTime(start)
|
||||
|
||||
""" Build Structure """
|
||||
|
||||
names, parents = AnimationStructure.load_from_maya(root)
|
||||
descendants = AnimationStructure.descendants_list(parents)
|
||||
orients = Quaternions.id(len(names))
|
||||
offsets = np.array([pm.xform(j, q=True, translation=True) for j in names])
|
||||
|
||||
for j, name in enumerate(names):
|
||||
scale = pm.xform(pm.PyNode(name), q=True, scale=True, relative=True)
|
||||
if len(descendants[j]) == 0:
|
||||
continue
|
||||
offsets[descendants[j]] *= scale
|
||||
|
||||
""" Load Animation """
|
||||
|
||||
eulers = np.zeros((end-start, len(names), 3))
|
||||
positions = np.zeros((end-start, len(names), 3))
|
||||
rotations = Quaternions.id((end-start, len(names)))
|
||||
|
||||
for i in range(end-start):
|
||||
|
||||
pm.currentTime(start+i+1, u=True)
|
||||
|
||||
scales = {}
|
||||
|
||||
for j, name, parent in zip(range(len(names)), names, parents):
|
||||
|
||||
node = pm.PyNode(name)
|
||||
|
||||
if i == 0 and pm.hasAttr(node, 'jointOrient'):
|
||||
ort = node.getOrientation()
|
||||
orients[j] = Quaternions(np.array([ort[3], ort[0], ort[1], ort[2]]))
|
||||
|
||||
if pm.hasAttr(node, 'rotate'):
|
||||
eulers[i, j] = np.radians(pm.xform(node, q=True, rotation=True))
|
||||
if pm.hasAttr(node, 'translate'):
|
||||
positions[i, j] = pm.xform(node, q=True, translation=True)
|
||||
if pm.hasAttr(node, 'scale'):
|
||||
scales[j] = pm.xform(node, q=True, scale=True, relative=True)
|
||||
|
||||
for j in scales:
|
||||
if len(descendants[j]) == 0:
|
||||
continue
|
||||
positions[i, descendants[j]] *= scales[j]
|
||||
|
||||
positions[i, 0] = pm.xform(root, q=True, translation=True, worldSpace=True)
|
||||
|
||||
rotations = orients[np.newaxis] * Quaternions.from_euler(eulers, order='xyz', world=True)
|
||||
|
||||
""" Done """
|
||||
|
||||
pm.currentTime(original_time)
|
||||
|
||||
return Animation(rotations, positions, orients, offsets, parents), names
|
||||
|
||||
# local transformation matrices
|
||||
|
||||
|
||||
def transforms_local(anim):
|
||||
"""
|
||||
Computes Animation Local Transforms
|
||||
|
||||
As well as a number of other uses this can
|
||||
be used to compute global joint transforms,
|
||||
which in turn can be used to compete global
|
||||
joint positions
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
anim : Animation
|
||||
Input animation
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
transforms : (F, J, 4, 4) ndarray
|
||||
|
||||
For each frame F, joint local
|
||||
transforms for each joint J
|
||||
"""
|
||||
|
||||
transforms = anim.rotations.transforms()
|
||||
transforms = np.concatenate([transforms, np.zeros(transforms.shape[:2] + (3, 1))], axis=-1)
|
||||
transforms = np.concatenate([transforms, np.zeros(transforms.shape[:2] + (1, 4))], axis=-2)
|
||||
# the last column is filled with the joint positions!
|
||||
transforms[:, :, 0:3, 3] = anim.positions
|
||||
transforms[:, :, 3:4, 3] = 1.0
|
||||
return transforms
|
||||
|
||||
|
||||
def transforms_multiply(t0s, t1s):
|
||||
"""
|
||||
Transforms Multiply
|
||||
|
||||
Multiplies two arrays of animation transforms
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
t0s, t1s : (F, J, 4, 4) ndarray
|
||||
Two arrays of transforms
|
||||
for each frame F and each
|
||||
joint J
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
transforms : (F, J, 4, 4) ndarray
|
||||
Array of transforms for each
|
||||
frame F and joint J multiplied
|
||||
together
|
||||
"""
|
||||
|
||||
return ut.matrix_multiply(t0s, t1s)
|
||||
|
||||
|
||||
def transforms_inv(ts):
|
||||
fts = ts.reshape(-1, 4, 4)
|
||||
fts = np.array(list(map(lambda x: np.linalg.inv(x), fts)))
|
||||
return fts.reshape(ts.shape)
|
||||
|
||||
|
||||
def transforms_blank(anim):
|
||||
"""
|
||||
Blank Transforms
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
anim : Animation
|
||||
Input animation
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
transforms : (F, J, 4, 4) ndarray
|
||||
Array of identity transforms for
|
||||
each frame F and joint J
|
||||
"""
|
||||
|
||||
ts = np.zeros(anim.shape + (4, 4))
|
||||
ts[:, :, 0, 0] = 1.0
|
||||
ts[:, :, 1, 1] = 1.0
|
||||
ts[:, :, 2, 2] = 1.0
|
||||
ts[:, :, 3, 3] = 1.0
|
||||
return ts
|
||||
|
||||
# global transformation matrices
|
||||
|
||||
|
||||
def transforms_global(anim):
|
||||
"""
|
||||
Global Animation Transforms
|
||||
|
||||
This relies on joint ordering
|
||||
being incremental. That means a joint
|
||||
J1 must not be a ancestor of J0 if
|
||||
J0 appears before J1 in the joint
|
||||
ordering.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
anim : Animation
|
||||
Input animation
|
||||
|
||||
Returns
|
||||
------
|
||||
|
||||
transforms : (F, J, 4, 4) ndarray
|
||||
Array of global transforms for
|
||||
each frame F and joint J
|
||||
"""
|
||||
|
||||
joints = np.arange(anim.shape[1])
|
||||
parents = np.arange(anim.shape[1])
|
||||
locals = transforms_local(anim)
|
||||
globals = transforms_blank(anim)
|
||||
|
||||
globals[:, 0] = locals[:, 0]
|
||||
|
||||
for i in range(1, anim.shape[1]):
|
||||
globals[:, i] = transforms_multiply(globals[:, anim.parents[i]], locals[:, i])
|
||||
|
||||
return globals
|
||||
|
||||
# !!! useful!
|
||||
|
||||
|
||||
def positions_global(anim):
|
||||
"""
|
||||
Global Joint Positions
|
||||
|
||||
Given an animation compute the global joint
|
||||
positions at at every frame
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
anim : Animation
|
||||
Input animation
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
positions : (F, J, 3) ndarray
|
||||
Positions for every frame F
|
||||
and joint position J
|
||||
"""
|
||||
|
||||
# get the last column -- corresponding to the coordinates
|
||||
positions = transforms_global(anim)[:, :, :, 3]
|
||||
return positions[:, :, :3] / positions[:, :, 3, np.newaxis]
|
||||
|
||||
|
||||
""" Rotations """
|
||||
|
||||
|
||||
def rotations_global(anim):
|
||||
"""
|
||||
Global Animation Rotations
|
||||
|
||||
This relies on joint ordering
|
||||
being incremental. That means a joint
|
||||
J1 must not be a ancestor of J0 if
|
||||
J0 appears before J1 in the joint
|
||||
ordering.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
anim : Animation
|
||||
Input animation
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
points : (F, J) Quaternions
|
||||
global rotations for every frame F
|
||||
and joint J
|
||||
"""
|
||||
|
||||
joints = np.arange(anim.shape[1])
|
||||
parents = np.arange(anim.shape[1])
|
||||
locals = anim.rotations
|
||||
globals = Quaternions.id(anim.shape)
|
||||
|
||||
globals[:, 0] = locals[:, 0]
|
||||
|
||||
for i in range(1, anim.shape[1]):
|
||||
globals[:, i] = globals[:, anim.parents[i]] * locals[:, i]
|
||||
|
||||
return globals
|
||||
|
||||
|
||||
def rotations_parents_global(anim):
|
||||
rotations = rotations_global(anim)
|
||||
rotations = rotations[:, anim.parents]
|
||||
rotations[:, 0] = Quaternions.id(len(anim))
|
||||
return rotations
|
||||
|
||||
|
||||
def rotations_load_to_maya(rotations, positions, names=None):
|
||||
"""
|
||||
Load Rotations into Maya
|
||||
|
||||
Loads a Quaternions array into the scene
|
||||
via the representation of axis
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
rotations : (F, J) Quaternions
|
||||
array of rotations to load
|
||||
into the scene where
|
||||
F = number of frames
|
||||
J = number of joints
|
||||
|
||||
positions : (F, J, 3) ndarray
|
||||
array of positions to load
|
||||
rotation axis at where:
|
||||
F = number of frames
|
||||
J = number of joints
|
||||
|
||||
names : [str]
|
||||
List of joint names
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
maxies : Group
|
||||
Grouped Maya Node of all Axis nodes
|
||||
"""
|
||||
|
||||
import pymel.core as pm
|
||||
|
||||
if names is None:
|
||||
names = ["joint_" + str(i) for i in range(rotations.shape[1])]
|
||||
|
||||
maxis = []
|
||||
frames = range(1, len(positions)+1)
|
||||
for i, name in enumerate(names):
|
||||
|
||||
name = name + "_axis"
|
||||
axis = pm.group(
|
||||
pm.curve(p=[(0, 0, 0), (1, 0, 0)], d=1, n=name+'_axis_x'),
|
||||
pm.curve(p=[(0, 0, 0), (0, 1, 0)], d=1, n=name+'_axis_y'),
|
||||
pm.curve(p=[(0, 0, 0), (0, 0, 1)], d=1, n=name+'_axis_z'),
|
||||
n=name)
|
||||
|
||||
axis.rotatePivot.set((0, 0, 0))
|
||||
axis.scalePivot.set((0, 0, 0))
|
||||
axis.childAtIndex(0).overrideEnabled.set(1)
|
||||
axis.childAtIndex(0).overrideColor.set(13)
|
||||
axis.childAtIndex(1).overrideEnabled.set(1)
|
||||
axis.childAtIndex(1).overrideColor.set(14)
|
||||
axis.childAtIndex(2).overrideEnabled.set(1)
|
||||
axis.childAtIndex(2).overrideColor.set(15)
|
||||
|
||||
curvex = pm.nodetypes.AnimCurveTA(n=name + "_rotateX")
|
||||
curvey = pm.nodetypes.AnimCurveTA(n=name + "_rotateY")
|
||||
curvez = pm.nodetypes.AnimCurveTA(n=name + "_rotateZ")
|
||||
|
||||
arotations = rotations[:, i].euler()
|
||||
curvex.addKeys(frames, arotations[:, 0])
|
||||
curvey.addKeys(frames, arotations[:, 1])
|
||||
curvez.addKeys(frames, arotations[:, 2])
|
||||
|
||||
pm.connectAttr(curvex.output, axis.rotateX)
|
||||
pm.connectAttr(curvey.output, axis.rotateY)
|
||||
pm.connectAttr(curvez.output, axis.rotateZ)
|
||||
|
||||
offsetx = pm.nodetypes.AnimCurveTU(n=name + "_translateX")
|
||||
offsety = pm.nodetypes.AnimCurveTU(n=name + "_translateY")
|
||||
offsetz = pm.nodetypes.AnimCurveTU(n=name + "_translateZ")
|
||||
|
||||
offsetx.addKeys(frames, positions[:, i, 0])
|
||||
offsety.addKeys(frames, positions[:, i, 1])
|
||||
offsetz.addKeys(frames, positions[:, i, 2])
|
||||
|
||||
pm.connectAttr(offsetx.output, axis.translateX)
|
||||
pm.connectAttr(offsety.output, axis.translateY)
|
||||
pm.connectAttr(offsetz.output, axis.translateZ)
|
||||
|
||||
maxis.append(axis)
|
||||
|
||||
return pm.group(*maxis, n='RotationAnimation')
|
||||
|
||||
|
||||
""" Offsets & Orients """
|
||||
|
||||
|
||||
def orients_global(anim):
|
||||
|
||||
joints = np.arange(anim.shape[1])
|
||||
parents = np.arange(anim.shape[1])
|
||||
locals = anim.orients
|
||||
globals = Quaternions.id(anim.shape[1])
|
||||
|
||||
globals[:, 0] = locals[:, 0]
|
||||
|
||||
for i in range(1, anim.shape[1]):
|
||||
globals[:, i] = globals[:, anim.parents[i]] * locals[:, i]
|
||||
|
||||
return globals
|
||||
|
||||
|
||||
def offsets_transforms_local(anim):
|
||||
|
||||
transforms = anim.orients[np.newaxis].transforms()
|
||||
transforms = np.concatenate([transforms, np.zeros(transforms.shape[:2] + (3, 1))], axis=-1)
|
||||
transforms = np.concatenate([transforms, np.zeros(transforms.shape[:2] + (1, 4))], axis=-2)
|
||||
transforms[:, :, 0:3, 3] = anim.offsets[np.newaxis]
|
||||
transforms[:, :, 3:4, 3] = 1.0
|
||||
return transforms
|
||||
|
||||
|
||||
def offsets_transforms_global(anim):
|
||||
|
||||
joints = np.arange(anim.shape[1])
|
||||
parents = np.arange(anim.shape[1])
|
||||
locals = offsets_transforms_local(anim)
|
||||
globals = transforms_blank(anim)
|
||||
|
||||
globals[:, 0] = locals[:, 0]
|
||||
|
||||
for i in range(1, anim.shape[1]):
|
||||
globals[:, i] = transforms_multiply(globals[:, anim.parents[i]], locals[:, i])
|
||||
|
||||
return globals
|
||||
|
||||
|
||||
def offsets_global(anim):
|
||||
offsets = offsets_transforms_global(anim)[:, :, :, 3]
|
||||
return offsets[0, :, :3] / offsets[0, :, 3, np.newaxis]
|
||||
|
||||
|
||||
""" Lengths """
|
||||
|
||||
|
||||
def offset_lengths(anim):
|
||||
return np.sum(anim.offsets[1:]**2.0, axis=1)**0.5
|
||||
|
||||
|
||||
def position_lengths(anim):
|
||||
return np.sum(anim.positions[:, 1:]**2.0, axis=2)**0.5
|
||||
|
||||
|
||||
""" Skinning """
|
||||
|
||||
|
||||
def skin(anim, rest, weights, mesh, maxjoints=4):
|
||||
|
||||
full_transforms = transforms_multiply(
|
||||
transforms_global(anim),
|
||||
transforms_inv(transforms_global(rest[0:1])))
|
||||
|
||||
weightids = np.argsort(-weights, axis=1)[:, :maxjoints]
|
||||
weightvls = np.array(list(map(lambda w, i: w[i], weights, weightids)))
|
||||
weightvls = weightvls / weightvls.sum(axis=1)[..., np.newaxis]
|
||||
|
||||
verts = np.hstack([mesh, np.ones((len(mesh), 1))])
|
||||
verts = verts[np.newaxis, :, np.newaxis, :, np.newaxis]
|
||||
verts = transforms_multiply(full_transforms[:, weightids], verts)
|
||||
verts = (verts[:, :, :, :3] / verts[:, :, :, 3:4])[:, :, :, :, 0]
|
||||
|
||||
return np.sum(weightvls[np.newaxis, :, :, np.newaxis] * verts, axis=2)
|
||||
@@ -0,0 +1,451 @@
|
||||
import numpy as np
|
||||
import scipy.sparse as sparse
|
||||
import Animation as Animation
|
||||
|
||||
|
||||
""" Maya Functions """
|
||||
|
||||
def load_from_maya(root):
|
||||
"""
|
||||
Load joint parents and names from maya
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
root : PyNode
|
||||
Root Maya Node
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
(names, parents) : ([str], (J) ndarray)
|
||||
List of joint names and array
|
||||
of indices representing the parent
|
||||
joint for each joint J.
|
||||
|
||||
Joint index -1 is used to represent
|
||||
that there is no parent joint
|
||||
"""
|
||||
|
||||
import pymel.core as pm
|
||||
|
||||
names = []
|
||||
parents = []
|
||||
|
||||
def unload_joint(j, parents, par):
|
||||
|
||||
id = len(names)
|
||||
names.append(j)
|
||||
parents.append(par)
|
||||
|
||||
children = [c for c in j.getChildren() if
|
||||
isinstance(c, pm.nt.Transform) and
|
||||
not isinstance(c, pm.nt.Constraint) and
|
||||
not any(pm.listRelatives(c, s=True)) and
|
||||
(any(pm.listRelatives(c, ad=True, ap=False, type='joint')) or isinstance(c, pm.nt.Joint))]
|
||||
|
||||
map(lambda c: unload_joint(c, parents, id), children)
|
||||
|
||||
unload_joint(root, parents, -1)
|
||||
|
||||
return (names, parents)
|
||||
|
||||
|
||||
""" Family Functions """
|
||||
|
||||
def joints(parents):
|
||||
"""
|
||||
Parameters
|
||||
----------
|
||||
|
||||
parents : (J) ndarray
|
||||
parents array
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
joints : (J) ndarray
|
||||
Array of joint indices
|
||||
"""
|
||||
return np.arange(len(parents), dtype=int)
|
||||
|
||||
def joints_list(parents):
|
||||
"""
|
||||
Parameters
|
||||
----------
|
||||
|
||||
parents : (J) ndarray
|
||||
parents array
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
joints : [ndarray]
|
||||
List of arrays of joint idices for
|
||||
each joint
|
||||
"""
|
||||
return list(joints(parents)[:,np.newaxis])
|
||||
|
||||
def parents_list(parents):
|
||||
"""
|
||||
Parameters
|
||||
----------
|
||||
|
||||
parents : (J) ndarray
|
||||
parents array
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
parents : [ndarray]
|
||||
List of arrays of joint idices for
|
||||
the parents of each joint
|
||||
"""
|
||||
return list(parents[:,np.newaxis])
|
||||
|
||||
|
||||
def children_list(parents):
|
||||
"""
|
||||
Parameters
|
||||
----------
|
||||
|
||||
parents : (J) ndarray
|
||||
parents array
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
children : [ndarray]
|
||||
List of arrays of joint indices for
|
||||
the children of each joint
|
||||
"""
|
||||
|
||||
def joint_children(i):
|
||||
return [j for j, p in enumerate(parents) if p == i]
|
||||
|
||||
return list(map(lambda j: np.array(joint_children(j)), joints(parents)))
|
||||
|
||||
|
||||
def descendants_list(parents):
|
||||
"""
|
||||
Parameters
|
||||
----------
|
||||
|
||||
parents : (J) ndarray
|
||||
parents array
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
descendants : [ndarray]
|
||||
List of arrays of joint idices for
|
||||
the descendants of each joint
|
||||
"""
|
||||
|
||||
children = children_list(parents)
|
||||
|
||||
def joint_descendants(i):
|
||||
return sum([joint_descendants(j) for j in children[i]], list(children[i]))
|
||||
|
||||
return list(map(lambda j: np.array(joint_descendants(j)), joints(parents)))
|
||||
|
||||
|
||||
def ancestors_list(parents):
|
||||
"""
|
||||
Parameters
|
||||
----------
|
||||
|
||||
parents : (J) ndarray
|
||||
parents array
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
ancestors : [ndarray]
|
||||
List of arrays of joint idices for
|
||||
the ancestors of each joint
|
||||
"""
|
||||
|
||||
decendants = descendants_list(parents)
|
||||
|
||||
def joint_ancestors(i):
|
||||
return [j for j in joints(parents) if i in decendants[j]]
|
||||
|
||||
return list(map(lambda j: np.array(joint_ancestors(j)), joints(parents)))
|
||||
|
||||
|
||||
""" Mask Functions """
|
||||
|
||||
def mask(parents, filter):
|
||||
"""
|
||||
Constructs a Mask for a give filter
|
||||
|
||||
A mask is a (J, J) ndarray truth table for a given
|
||||
condition over J joints. For example there
|
||||
may be a mask specifying if a joint N is a
|
||||
child of another joint M.
|
||||
|
||||
This could be constructed into a mask using
|
||||
`m = mask(parents, children_list)` and the condition
|
||||
of childhood tested using `m[N, M]`.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
parents : (J) ndarray
|
||||
parents array
|
||||
|
||||
filter : (J) ndarray -> [ndarray]
|
||||
function that outputs a list of arrays
|
||||
of joint indices for some condition
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
mask : (N, N) ndarray
|
||||
boolean truth table of given condition
|
||||
"""
|
||||
m = np.zeros((len(parents), len(parents))).astype(bool)
|
||||
jnts = joints(parents)
|
||||
fltr = filter(parents)
|
||||
for i,f in enumerate(fltr): m[i,:] = np.any(jnts[:,np.newaxis] == f[np.newaxis,:], axis=1)
|
||||
return m
|
||||
|
||||
def joints_mask(parents): return np.eye(len(parents)).astype(bool)
|
||||
def children_mask(parents): return mask(parents, children_list)
|
||||
def parents_mask(parents): return mask(parents, parents_list)
|
||||
def descendants_mask(parents): return mask(parents, descendants_list)
|
||||
def ancestors_mask(parents): return mask(parents, ancestors_list)
|
||||
|
||||
""" Search Functions """
|
||||
|
||||
def joint_chain_ascend(parents, start, end):
|
||||
chain = []
|
||||
while start != end:
|
||||
chain.append(start)
|
||||
start = parents[start]
|
||||
chain.append(end)
|
||||
return np.array(chain, dtype=int)
|
||||
|
||||
|
||||
""" Constraints """
|
||||
|
||||
def constraints(anim, **kwargs):
|
||||
"""
|
||||
Constraint list for Animation
|
||||
|
||||
This constraint list can be used in the
|
||||
VerletParticle solver to constrain
|
||||
a animation global joint positions.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
anim : Animation
|
||||
Input animation
|
||||
|
||||
masses : (F, J) ndarray
|
||||
Optional list of masses
|
||||
for joints J across frames F
|
||||
defaults to weighting by
|
||||
vertical height
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
constraints : [(int, int, (F, J) ndarray, (F, J) ndarray, (F, J) ndarray)]
|
||||
A list of constraints in the format:
|
||||
(Joint1, Joint2, Masses1, Masses2, Lengths)
|
||||
|
||||
"""
|
||||
|
||||
masses = kwargs.pop('masses', None)
|
||||
|
||||
children = children_list(anim.parents)
|
||||
constraints = []
|
||||
|
||||
points_offsets = Animation.offsets_global(anim)
|
||||
points = Animation.positions_global(anim)
|
||||
|
||||
if masses is None:
|
||||
masses = 1.0 / (0.1 + np.absolute(points_offsets[:,1]))
|
||||
masses = masses[np.newaxis].repeat(len(anim), axis=0)
|
||||
|
||||
for j in xrange(anim.shape[1]):
|
||||
|
||||
""" Add constraints between all joints and their children """
|
||||
for c0 in children[j]:
|
||||
|
||||
dists = np.sum((points[:, c0] - points[:, j])**2.0, axis=1)**0.5
|
||||
constraints.append((c0, j, masses[:,c0], masses[:,j], dists))
|
||||
|
||||
""" Add constraints between all children of joint """
|
||||
for c1 in children[j]:
|
||||
if c0 == c1: continue
|
||||
|
||||
dists = np.sum((points[:, c0] - points[:, c1])**2.0, axis=1)**0.5
|
||||
constraints.append((c0, c1, masses[:,c0], masses[:,c1], dists))
|
||||
|
||||
return constraints
|
||||
|
||||
""" Graph Functions """
|
||||
|
||||
def graph(anim):
|
||||
"""
|
||||
Generates a weighted adjacency matrix
|
||||
using local joint distances along
|
||||
the skeletal structure.
|
||||
|
||||
Joints which are not connected
|
||||
are assigned the weight `0`.
|
||||
|
||||
Joints which actually have zero distance
|
||||
between them, but are still connected, are
|
||||
perturbed by some minimal amount.
|
||||
|
||||
The output of this routine can be used
|
||||
with the `scipy.sparse.csgraph`
|
||||
routines for graph analysis.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
anim : Animation
|
||||
input animation
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
graph : (N, N) ndarray
|
||||
weight adjacency matrix using
|
||||
local distances along the
|
||||
skeletal structure from joint
|
||||
N to joint M. If joints are not
|
||||
directly connected are assigned
|
||||
the weight `0`.
|
||||
"""
|
||||
|
||||
graph = np.zeros(anim.shape[1], anim.shape[1])
|
||||
lengths = np.sum(anim.offsets**2.0, axis=1)**0.5 + 0.001
|
||||
|
||||
for i,p in enumerate(anim.parents):
|
||||
if p == -1: continue
|
||||
graph[i,p] = lengths[p]
|
||||
graph[p,i] = lengths[p]
|
||||
|
||||
return graph
|
||||
|
||||
|
||||
def distances(anim):
|
||||
"""
|
||||
Generates a distance matrix for
|
||||
pairwise joint distances along
|
||||
the skeletal structure
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
anim : Animation
|
||||
input animation
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
distances : (N, N) ndarray
|
||||
array of pairwise distances
|
||||
along skeletal structure
|
||||
from some joint N to some
|
||||
joint M
|
||||
"""
|
||||
|
||||
distances = np.zeros((anim.shape[1], anim.shape[1]))
|
||||
generated = distances.copy().astype(bool)
|
||||
|
||||
joint_lengths = np.sum(anim.offsets**2.0, axis=1)**0.5
|
||||
joint_children = children_list(anim)
|
||||
joint_parents = parents_list(anim)
|
||||
|
||||
def find_distance(distances, generated, prev, i, j):
|
||||
|
||||
""" If root, identity, or already generated, return """
|
||||
if j == -1: return (0.0, True)
|
||||
if j == i: return (0.0, True)
|
||||
if generated[i,j]: return (distances[i,j], True)
|
||||
|
||||
""" Find best distances along parents and children """
|
||||
par_dists = [(joint_lengths[j], find_distance(distances, generated, j, i, p)) for p in joint_parents[j] if p != prev]
|
||||
out_dists = [(joint_lengths[c], find_distance(distances, generated, j, i, c)) for c in joint_children[j] if c != prev]
|
||||
|
||||
""" Check valid distance and not dead end """
|
||||
par_dists = [a + d for (a, (d, f)) in par_dists if f]
|
||||
out_dists = [a + d for (a, (d, f)) in out_dists if f]
|
||||
|
||||
""" All dead ends """
|
||||
if (out_dists + par_dists) == []: return (0.0, False)
|
||||
|
||||
""" Get minimum path """
|
||||
dist = min(out_dists + par_dists)
|
||||
distances[i,j] = dist; distances[j,i] = dist
|
||||
generated[i,j] = True; generated[j,i] = True
|
||||
|
||||
for i in xrange(anim.shape[1]):
|
||||
for j in xrange(anim.shape[1]):
|
||||
find_distance(distances, generated, -1, i, j)
|
||||
|
||||
return distances
|
||||
|
||||
def edges(parents):
|
||||
"""
|
||||
Animation structure edges
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
parents : (J) ndarray
|
||||
parents array
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
edges : (M, 2) ndarray
|
||||
array of pairs where each
|
||||
pair contains two indices of a joints
|
||||
which corrisponds to an edge in the
|
||||
joint structure going from parent to child.
|
||||
"""
|
||||
|
||||
return np.array(list(zip(parents, joints(parents)))[1:])
|
||||
|
||||
|
||||
def incidence(parents):
|
||||
"""
|
||||
Incidence Matrix
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
parents : (J) ndarray
|
||||
parents array
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
incidence : (N, M) ndarray
|
||||
|
||||
Matrix of N joint positions by
|
||||
M edges which each entry is either
|
||||
1 or -1 and multiplication by the
|
||||
joint positions returns the an
|
||||
array of vectors along each edge
|
||||
of the structure
|
||||
"""
|
||||
|
||||
es = edges(parents)
|
||||
|
||||
inc = np.zeros((len(parents)-1, len(parents))).astype(np.int)
|
||||
for i, e in enumerate(es):
|
||||
inc[i,e[0]] = 1
|
||||
inc[i,e[1]] = -1
|
||||
|
||||
return inc.T
|
||||
|
||||
+471
@@ -0,0 +1,471 @@
|
||||
from Quaternions import Quaternions
|
||||
from Animation import Animation
|
||||
import re
|
||||
import numpy as np
|
||||
import sys
|
||||
|
||||
|
||||
channelmap = {
|
||||
'Xrotation': 'x',
|
||||
'Yrotation': 'y',
|
||||
'Zrotation': 'z'
|
||||
}
|
||||
|
||||
channelmap_inv = {
|
||||
'x': 'Xrotation',
|
||||
'y': 'Yrotation',
|
||||
'z': 'Zrotation',
|
||||
}
|
||||
|
||||
ordermap = {
|
||||
'x': 0,
|
||||
'y': 1,
|
||||
'z': 2,
|
||||
}
|
||||
|
||||
|
||||
def load(filename, start=None, end=None, order=None, world=False):
|
||||
"""
|
||||
Reads a BVH file and constructs an animation
|
||||
|
||||
Parameters
|
||||
----------
|
||||
filename: str
|
||||
File to be opened
|
||||
|
||||
start : int
|
||||
Optional Starting Frame
|
||||
|
||||
end : int
|
||||
Optional Ending Frame
|
||||
|
||||
order : str
|
||||
Optional Specifier for joint order.
|
||||
Given as string E.G 'xyz', 'zxy'
|
||||
|
||||
world : bool
|
||||
If set to true euler angles are applied
|
||||
together in world space rather than local
|
||||
space
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
(animation, joint_names, frametime)
|
||||
Tuple of loaded animation and joint names
|
||||
"""
|
||||
|
||||
f = open(filename, "r")
|
||||
|
||||
i = 0
|
||||
active = -1
|
||||
end_site = False
|
||||
|
||||
names = []
|
||||
orients = Quaternions.id(0)
|
||||
offsets = np.array([]).reshape((0, 3))
|
||||
parents = np.array([], dtype=int)
|
||||
|
||||
for line in f:
|
||||
|
||||
if "HIERARCHY" in line:
|
||||
continue
|
||||
if "MOTION" in line:
|
||||
continue
|
||||
|
||||
rmatch = re.match(r"ROOT (\w+)", line)
|
||||
if rmatch:
|
||||
names.append(rmatch.group(1))
|
||||
offsets = np.append(offsets, np.array([[0, 0, 0]]), axis=0)
|
||||
orients.qs = np.append(orients.qs, np.array([[1, 0, 0, 0]]), axis=0)
|
||||
parents = np.append(parents, active)
|
||||
active = (len(parents)-1)
|
||||
continue
|
||||
|
||||
if "{" in line:
|
||||
continue
|
||||
|
||||
if "}" in line:
|
||||
if end_site:
|
||||
end_site = False
|
||||
else:
|
||||
active = parents[active]
|
||||
continue
|
||||
|
||||
offmatch = re.match(r"\s*OFFSET\s+([\-\d\.e]+)\s+([\-\d\.e]+)\s+([\-\d\.e]+)", line)
|
||||
if offmatch:
|
||||
if not end_site:
|
||||
offsets[active] = np.array([list(map(float, offmatch.groups()))])
|
||||
continue
|
||||
|
||||
chanmatch = re.match(r"\s*CHANNELS\s+(\d+)", line)
|
||||
if chanmatch:
|
||||
channels = int(chanmatch.group(1))
|
||||
if order is None:
|
||||
channelis = 0 if channels == 3 else 3
|
||||
channelie = 3 if channels == 3 else 6
|
||||
parts = line.split()[2+channelis:2+channelie]
|
||||
if any([p not in channelmap for p in parts]):
|
||||
continue
|
||||
order = "".join([channelmap[p] for p in parts])
|
||||
continue
|
||||
|
||||
jmatch = re.match("\s*JOINT\s+(\w+)", line)
|
||||
if jmatch:
|
||||
names.append(jmatch.group(1))
|
||||
offsets = np.append(offsets, np.array([[0, 0, 0]]), axis=0)
|
||||
orients.qs = np.append(orients.qs, np.array([[1, 0, 0, 0]]), axis=0)
|
||||
parents = np.append(parents, active)
|
||||
active = (len(parents)-1)
|
||||
continue
|
||||
|
||||
if "End Site" in line:
|
||||
end_site = True
|
||||
continue
|
||||
|
||||
fmatch = re.match("\s*Frames:\s+(\d+)", line)
|
||||
if fmatch:
|
||||
if start and end:
|
||||
fnum = (end - start)-1
|
||||
else:
|
||||
fnum = int(fmatch.group(1))
|
||||
jnum = len(parents)
|
||||
# result: [fnum, J, 3]
|
||||
positions = offsets[np.newaxis].repeat(fnum, axis=0)
|
||||
# result: [fnum, len(orients), 3]
|
||||
rotations = np.zeros((fnum, len(orients), 3))
|
||||
continue
|
||||
|
||||
fmatch = re.match("\s*Frame Time:\s+([\d\.]+)", line)
|
||||
if fmatch:
|
||||
frametime = float(fmatch.group(1))
|
||||
continue
|
||||
|
||||
if (start and end) and (i < start or i >= end-1):
|
||||
i += 1
|
||||
continue
|
||||
|
||||
dmatch = line.strip().split()
|
||||
if dmatch:
|
||||
data_block = np.array(list(map(float, dmatch)))
|
||||
N = len(parents)
|
||||
fi = i - start if start else i
|
||||
if channels == 3:
|
||||
# This should be root positions[0:1] & all rotations
|
||||
positions[fi, 0:1] = data_block[0:3]
|
||||
rotations[fi, :] = data_block[3:].reshape(N, 3)
|
||||
elif channels == 6:
|
||||
data_block = data_block.reshape(N, 6)
|
||||
# fill in all positions
|
||||
positions[fi, :] = data_block[:, 0:3]
|
||||
rotations[fi, :] = data_block[:, 3:6]
|
||||
elif channels == 9:
|
||||
positions[fi, 0] = data_block[0:3]
|
||||
data_block = data_block[3:].reshape(N-1, 9)
|
||||
rotations[fi, 1:] = data_block[:, 3:6]
|
||||
positions[fi, 1:] += data_block[:, 0:3] * data_block[:, 6:9]
|
||||
else:
|
||||
raise Exception("Too many channels! %i" % channels)
|
||||
|
||||
i += 1
|
||||
|
||||
f.close()
|
||||
|
||||
rotations = Quaternions.from_euler(np.radians(rotations), order=order, world=world)
|
||||
|
||||
return (Animation(rotations, positions, orients, offsets, parents), names, frametime)
|
||||
|
||||
|
||||
def load_bfa(filename, start=None, end=None, order=None, world=False):
|
||||
"""
|
||||
Reads a BVH file and constructs an animation
|
||||
|
||||
!!! Read from bfa, will replace the end sites of arms by two joints (w/ unit rotation)
|
||||
|
||||
Parameters
|
||||
----------
|
||||
filename: str
|
||||
File to be opened
|
||||
|
||||
start : int
|
||||
Optional Starting Frame
|
||||
|
||||
end : int
|
||||
Optional Ending Frame
|
||||
|
||||
order : str
|
||||
Optional Specifier for joint order.
|
||||
Given as string E.G 'xyz', 'zxy'
|
||||
|
||||
world : bool
|
||||
If set to true euler angles are applied
|
||||
together in world space rather than local
|
||||
space
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
(animation, joint_names, frametime)
|
||||
Tuple of loaded animation and joint names
|
||||
"""
|
||||
|
||||
f = open(filename, "r")
|
||||
|
||||
i = 0
|
||||
active = -1
|
||||
end_site = False
|
||||
|
||||
hand_idx = [9, 14]
|
||||
|
||||
names = []
|
||||
orients = Quaternions.id(0)
|
||||
offsets = np.array([]).reshape((0, 3))
|
||||
parents = np.array([], dtype=int)
|
||||
|
||||
for line in f:
|
||||
|
||||
if "HIERARCHY" in line:
|
||||
continue
|
||||
if "MOTION" in line:
|
||||
continue
|
||||
|
||||
rmatch = re.match(r"ROOT (\w+)", line)
|
||||
if rmatch:
|
||||
names.append(rmatch.group(1))
|
||||
offsets = np.append(offsets, np.array([[0, 0, 0]]), axis=0)
|
||||
orients.qs = np.append(orients.qs, np.array([[1, 0, 0, 0]]), axis=0)
|
||||
parents = np.append(parents, active)
|
||||
active = (len(parents)-1)
|
||||
continue
|
||||
|
||||
if "{" in line:
|
||||
continue
|
||||
|
||||
if "}" in line:
|
||||
if end_site:
|
||||
end_site = False
|
||||
else:
|
||||
active = parents[active]
|
||||
continue
|
||||
|
||||
offmatch = re.match(r"\s*OFFSET\s+([\-\d\.e]+)\s+([\-\d\.e]+)\s+([\-\d\.e]+)", line)
|
||||
if offmatch:
|
||||
if not end_site:
|
||||
offsets[active] = np.array([list(map(float, offmatch.groups()))])
|
||||
"""
|
||||
else:
|
||||
print("active = ", active)
|
||||
if active in hand_idx:
|
||||
offsets[active] = np.array([list(map(float, offmatch.groups()))])
|
||||
"""
|
||||
continue
|
||||
|
||||
chanmatch = re.match(r"\s*CHANNELS\s+(\d+)", line)
|
||||
if chanmatch:
|
||||
channels = int(chanmatch.group(1))
|
||||
if order is None:
|
||||
channelis = 0 if channels == 3 else 3
|
||||
channelie = 3 if channels == 3 else 6
|
||||
parts = line.split()[2+channelis:2+channelie]
|
||||
if any([p not in channelmap for p in parts]):
|
||||
continue
|
||||
order = "".join([channelmap[p] for p in parts])
|
||||
continue
|
||||
|
||||
jmatch = re.match("\s*JOINT\s+(\w+)", line)
|
||||
if jmatch:
|
||||
names.append(jmatch.group(1))
|
||||
offsets = np.append(offsets, np.array([[0, 0, 0]]), axis=0)
|
||||
orients.qs = np.append(orients.qs, np.array([[1, 0, 0, 0]]), axis=0)
|
||||
parents = np.append(parents, active)
|
||||
active = (len(parents)-1)
|
||||
continue
|
||||
|
||||
if "End Site" in line:
|
||||
if active + 1 in hand_idx:
|
||||
print("parent:", names[-1])
|
||||
name = "LeftHandIndex" if active + 1 == hand_idx[0] else "RightHandIndex"
|
||||
names.append(name)
|
||||
offsets = np.append(offsets, np.array([[0, 0, 0]]), axis=0)
|
||||
orients.qs = np.append(orients.qs, np.array([[1, 0, 0, 0]]), axis=0)
|
||||
parents = np.append(parents, active)
|
||||
active = (len(parents)-1)
|
||||
else:
|
||||
end_site = True
|
||||
continue
|
||||
|
||||
fmatch = re.match("\s*Frames:\s+(\d+)", line)
|
||||
if fmatch:
|
||||
if start and end:
|
||||
fnum = (end - start)-1
|
||||
else:
|
||||
fnum = int(fmatch.group(1))
|
||||
jnum = len(parents)
|
||||
# result: [fnum, J, 3]
|
||||
positions = offsets[np.newaxis].repeat(fnum, axis=0)
|
||||
# result: [fnum, len(orients), 3]
|
||||
rotations = np.zeros((fnum, len(orients), 3))
|
||||
continue
|
||||
|
||||
fmatch = re.match("\s*Frame Time:\s+([\d\.]+)", line)
|
||||
if fmatch:
|
||||
frametime = float(fmatch.group(1))
|
||||
continue
|
||||
|
||||
if (start and end) and (i < start or i >= end-1):
|
||||
i += 1
|
||||
continue
|
||||
|
||||
dmatch = line.strip().split()
|
||||
if dmatch:
|
||||
data_block = np.array(list(map(float, dmatch)))
|
||||
N = len(parents)
|
||||
fi = i - start if start else i
|
||||
if channels == 3:
|
||||
# This should be root positions[0:1] & all rotations
|
||||
positions[fi, 0:1] = data_block[0:3]
|
||||
tmp = data_block[3:].reshape(N - 2, 3)
|
||||
tmp = np.concatenate([tmp[:hand_idx[0]],
|
||||
np.array([[0, 0, 0]]),
|
||||
tmp[hand_idx[0]: hand_idx[1] - 1],
|
||||
np.array([[0, 0, 0]]),
|
||||
tmp[hand_idx[1] - 1:]], axis=0)
|
||||
rotations[fi, :] = tmp.reshape(N, 3)
|
||||
elif channels == 6:
|
||||
data_block = data_block.reshape(N, 6)
|
||||
# fill in all positions
|
||||
positions[fi, :] = data_block[:, 0:3]
|
||||
rotations[fi, :] = data_block[:, 3:6]
|
||||
elif channels == 9:
|
||||
positions[fi, 0] = data_block[0:3]
|
||||
data_block = data_block[3:].reshape(N-1, 9)
|
||||
rotations[fi, 1:] = data_block[:, 3:6]
|
||||
positions[fi, 1:] += data_block[:, 0:3] * data_block[:, 6:9]
|
||||
else:
|
||||
raise Exception("Too many channels! %i" % channels)
|
||||
|
||||
i += 1
|
||||
|
||||
f.close()
|
||||
|
||||
rotations = Quaternions.from_euler(np.radians(rotations), order=order, world=world)
|
||||
|
||||
return (Animation(rotations, positions, orients, offsets, parents), names, frametime)
|
||||
|
||||
|
||||
def save(filename, anim, names=None, frametime=1.0/24.0, order='zyx', positions=False, orients=True):
|
||||
"""
|
||||
Saves an Animation to file as BVH
|
||||
|
||||
Parameters
|
||||
----------
|
||||
filename: str
|
||||
File to be saved to
|
||||
|
||||
anim : Animation
|
||||
Animation to save
|
||||
|
||||
names : [str]
|
||||
List of joint names
|
||||
|
||||
order : str
|
||||
Optional Specifier for joint order.
|
||||
Given as string E.G 'xyz', 'zxy'
|
||||
|
||||
frametime : float
|
||||
Optional Animation Frame time
|
||||
|
||||
positions : bool
|
||||
Optional specfier to save bone
|
||||
positions for each frame
|
||||
|
||||
orients : bool
|
||||
Multiply joint orients to the rotations
|
||||
before saving.
|
||||
|
||||
"""
|
||||
|
||||
if names is None:
|
||||
names = ["joint_" + str(i) for i in range(len(anim.parents))]
|
||||
|
||||
with open(filename, 'w') as f:
|
||||
|
||||
t = ""
|
||||
f.write("%sHIERARCHY\n" % t)
|
||||
f.write("%sROOT %s\n" % (t, names[0]))
|
||||
f.write("%s{\n" % t)
|
||||
t += '\t'
|
||||
|
||||
f.write("%sOFFSET %f %f %f\n" % (t, anim.offsets[0, 0], anim.offsets[0, 1], anim.offsets[0, 2]))
|
||||
f.write("%sCHANNELS 6 Xposition Yposition Zposition %s %s %s \n" %
|
||||
(t, channelmap_inv[order[0]], channelmap_inv[order[1]], channelmap_inv[order[2]]))
|
||||
|
||||
for i in range(anim.shape[1]):
|
||||
if anim.parents[i] == 0:
|
||||
t = save_joint(f, anim, names, t, i, order=order, positions=positions)
|
||||
|
||||
t = t[:-1]
|
||||
f.write("%s}\n" % t)
|
||||
|
||||
f.write("MOTION\n")
|
||||
f.write("Frames: %i\n" % anim.shape[0])
|
||||
f.write("Frame Time: %f\n" % frametime)
|
||||
|
||||
# if orients:
|
||||
# rots = np.degrees((-anim.orients[np.newaxis] * anim.rotations).euler(order=order[::-1]))
|
||||
# else:
|
||||
# rots = np.degrees(anim.rotations.euler(order=order[::-1]))
|
||||
rots = np.degrees(anim.rotations.euler(order=order[::-1]))
|
||||
poss = anim.positions
|
||||
|
||||
for i in range(anim.shape[0]):
|
||||
for j in range(anim.shape[1]):
|
||||
|
||||
if positions or j == 0:
|
||||
|
||||
f.write("%f %f %f %f %f %f " % (
|
||||
poss[i, j, 0], poss[i, j, 1], poss[i, j, 2],
|
||||
rots[i, j, ordermap[order[0]]], rots[i, j, ordermap[order[1]]], rots[i, j, ordermap[order[2]]]))
|
||||
|
||||
else:
|
||||
|
||||
f.write("%f %f %f " % (
|
||||
rots[i, j, ordermap[order[0]]], rots[i, j, ordermap[order[1]]], rots[i, j, ordermap[order[2]]]))
|
||||
|
||||
f.write("\n")
|
||||
|
||||
|
||||
def save_joint(f, anim, names, t, i, order='zyx', positions=False):
|
||||
|
||||
f.write("%sJOINT %s\n" % (t, names[i]))
|
||||
f.write("%s{\n" % t)
|
||||
t += '\t'
|
||||
|
||||
f.write("%sOFFSET %f %f %f\n" % (t, anim.offsets[i, 0], anim.offsets[i, 1], anim.offsets[i, 2]))
|
||||
|
||||
if positions:
|
||||
f.write("%sCHANNELS 6 Xposition Yposition Zposition %s %s %s \n" % (t,
|
||||
channelmap_inv[order[0]], channelmap_inv[order[1]], channelmap_inv[order[2]]))
|
||||
else:
|
||||
f.write("%sCHANNELS 3 %s %s %s\n" % (t,
|
||||
channelmap_inv[order[0]], channelmap_inv[order[1]], channelmap_inv[order[2]]))
|
||||
|
||||
end_site = True
|
||||
|
||||
for j in range(anim.shape[1]):
|
||||
if anim.parents[j] == i:
|
||||
t = save_joint(f, anim, names, t, j, order=order, positions=positions)
|
||||
end_site = False
|
||||
|
||||
if end_site:
|
||||
f.write("%sEnd Site\n" % t)
|
||||
f.write("%s{\n" % t)
|
||||
t += '\t'
|
||||
f.write("%sOFFSET %f %f %f\n" % (t, 0.0, 0.0, 0.0))
|
||||
t = t[:-1]
|
||||
f.write("%s}\n" % t)
|
||||
|
||||
t = t[:-1]
|
||||
f.write("%s}\n" % t)
|
||||
|
||||
return t
|
||||
@@ -0,0 +1,496 @@
|
||||
import numpy as np
|
||||
import scipy.linalg as linalg
|
||||
|
||||
import Animation
|
||||
import AnimationStructure
|
||||
|
||||
from Quaternions import Quaternions
|
||||
|
||||
|
||||
class BasicInverseKinematics:
|
||||
"""
|
||||
Basic Inverse Kinematics Solver
|
||||
|
||||
This is an extremely simple full body IK
|
||||
solver.
|
||||
|
||||
It works given the following conditions:
|
||||
|
||||
* All joint targets must be specified
|
||||
* All joint targets must be in reach
|
||||
* All joint targets must not differ
|
||||
extremely from the starting pose
|
||||
* No bone length constraints can be violated
|
||||
* The root translation and rotation are
|
||||
set to good initial values
|
||||
|
||||
It works under the observation that if the
|
||||
_directions_ the joints are pointing toward
|
||||
match the _directions_ of the vectors between
|
||||
the target joints then the pose should match
|
||||
that of the target pose.
|
||||
|
||||
Therefore it iterates over joints rotating
|
||||
each joint such that the vectors between it
|
||||
and it's children match that of the target
|
||||
positions.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
animation : Animation
|
||||
animation input
|
||||
|
||||
positions : (F, J, 3) ndarray
|
||||
target positions for each frame F
|
||||
and each joint J
|
||||
|
||||
iterations : int
|
||||
Optional number of iterations.
|
||||
If the above conditions are met
|
||||
1 iteration should be enough,
|
||||
therefore the default is 1
|
||||
|
||||
silent : bool
|
||||
Optional if to suppress output
|
||||
defaults to False
|
||||
"""
|
||||
|
||||
def __init__(self, animation, positions, iterations=1, silent=True):
|
||||
|
||||
self.animation = animation
|
||||
self.positions = positions
|
||||
self.iterations = iterations
|
||||
self.silent = silent
|
||||
|
||||
def __call__(self):
|
||||
|
||||
children = AnimationStructure.children_list(self.animation.parents)
|
||||
|
||||
for i in range(self.iterations):
|
||||
|
||||
for j in AnimationStructure.joints(self.animation.parents):
|
||||
|
||||
c = np.array(children[j])
|
||||
if len(c) == 0:
|
||||
continue
|
||||
|
||||
anim_transforms = Animation.transforms_global(self.animation)
|
||||
anim_positions = anim_transforms[:, :, :3, 3]
|
||||
anim_rotations = Quaternions.from_transforms(anim_transforms)
|
||||
|
||||
jdirs = anim_positions[:, c] - anim_positions[:, np.newaxis, j]
|
||||
ddirs = self.positions[:, c] - anim_positions[:, np.newaxis, j]
|
||||
|
||||
jsums = np.sqrt(np.sum(jdirs**2.0, axis=-1)) + 1e-10
|
||||
dsums = np.sqrt(np.sum(ddirs**2.0, axis=-1)) + 1e-10
|
||||
|
||||
jdirs = jdirs / jsums[:, :, np.newaxis]
|
||||
ddirs = ddirs / dsums[:, :, np.newaxis]
|
||||
|
||||
angles = np.arccos(np.sum(jdirs * ddirs, axis=2).clip(-1, 1))
|
||||
axises = np.cross(jdirs, ddirs)
|
||||
axises = -anim_rotations[:, j, np.newaxis] * axises
|
||||
|
||||
rotations = Quaternions.from_angle_axis(angles, axises)
|
||||
|
||||
if rotations.shape[1] == 1:
|
||||
averages = rotations[:, 0]
|
||||
else:
|
||||
averages = Quaternions.exp(rotations.log().mean(axis=-2))
|
||||
|
||||
self.animation.rotations[:, j] = self.animation.rotations[:, j] * averages
|
||||
|
||||
if not self.silent:
|
||||
anim_positions = Animation.positions_global(self.animation)
|
||||
error = np.mean(np.sum((anim_positions - self.positions)**2.0, axis=-1)**0.5, axis=-1)
|
||||
print('[BasicInverseKinematics] Iteration %i Error: %f' % (i+1, error))
|
||||
|
||||
return self.animation
|
||||
|
||||
|
||||
class JacobianInverseKinematics:
|
||||
"""
|
||||
Jacobian Based Full Body IK Solver
|
||||
|
||||
This is a full body IK solver which
|
||||
uses the dampened least squares inverse
|
||||
jacobian method.
|
||||
|
||||
It should remain fairly stable and effective
|
||||
even for joint positions which are out of
|
||||
reach and it can also take any number of targets
|
||||
to treat as end effectors.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
animation : Animation
|
||||
animation to solve inverse problem on
|
||||
|
||||
targets : {int : (F, 3) ndarray}
|
||||
Dictionary of target positions for each
|
||||
frame F, mapping joint index to
|
||||
a target position
|
||||
|
||||
references : (F, 3)
|
||||
Optional list of J joint position
|
||||
references for which the result
|
||||
should bias toward
|
||||
|
||||
iterations : int
|
||||
Optional number of iterations to
|
||||
compute. More iterations results in
|
||||
better accuracy but takes longer to
|
||||
compute. Default is 10.
|
||||
|
||||
recalculate : bool
|
||||
Optional if to recalcuate jacobian
|
||||
each iteration. Gives better accuracy
|
||||
but slower to compute. Defaults to True
|
||||
|
||||
damping : float
|
||||
Optional damping constant. Higher
|
||||
damping increases stability but
|
||||
requires more iterations to converge.
|
||||
Defaults to 5.0
|
||||
|
||||
secondary : float
|
||||
Force, or bias toward secondary target.
|
||||
Defaults to 0.25
|
||||
|
||||
silent : bool
|
||||
Optional if to suppress output
|
||||
defaults to False
|
||||
"""
|
||||
|
||||
def __init__(self, animation, targets,
|
||||
references=None, iterations=10,
|
||||
recalculate=True, damping=2.0,
|
||||
secondary=0.25, translate=False,
|
||||
silent=False, weights=None,
|
||||
weights_translate=None):
|
||||
|
||||
self.animation = animation
|
||||
self.targets = targets
|
||||
self.references = references
|
||||
|
||||
self.iterations = iterations
|
||||
self.recalculate = recalculate
|
||||
self.damping = damping
|
||||
self.secondary = secondary
|
||||
self.translate = translate
|
||||
self.silent = silent
|
||||
self.weights = weights
|
||||
self.weights_translate = weights_translate
|
||||
|
||||
def cross(self, a, b):
|
||||
o = np.empty(b.shape)
|
||||
o[..., 0] = a[..., 1]*b[..., 2] - a[..., 2]*b[..., 1]
|
||||
o[..., 1] = a[..., 2]*b[..., 0] - a[..., 0]*b[..., 2]
|
||||
o[..., 2] = a[..., 0]*b[..., 1] - a[..., 1]*b[..., 0]
|
||||
return o
|
||||
|
||||
def jacobian(self, x, fp, fr, ts, dsc, tdsc):
|
||||
""" Find parent rotations """
|
||||
prs = fr[:, self.animation.parents]
|
||||
prs[:, 0] = Quaternions.id((1))
|
||||
|
||||
""" Find global positions of target joints """
|
||||
tps = fp[:, np.array(list(ts.keys()))]
|
||||
|
||||
""" Get partial rotations """
|
||||
qys = Quaternions.from_angle_axis(x[:, 1:prs.shape[1]*3:3], np.array([[[0, 1, 0]]]))
|
||||
qzs = Quaternions.from_angle_axis(x[:, 2:prs.shape[1]*3:3], np.array([[[0, 0, 1]]]))
|
||||
|
||||
""" Find axis of rotations """
|
||||
es = np.empty((len(x), fr.shape[1]*3, 3))
|
||||
es[:, 0::3] = ((prs * qzs) * qys) * np.array([[[1, 0, 0]]])
|
||||
es[:, 1::3] = ((prs * qzs) * np.array([[[0, 1, 0]]]))
|
||||
es[:, 2::3] = ((prs * np.array([[[0, 0, 1]]])))
|
||||
|
||||
""" Construct Jacobian """
|
||||
j = fp.repeat(3, axis=1)
|
||||
j = dsc[np.newaxis, :, :, np.newaxis] * (tps[:, np.newaxis, :] - j[:, :, np.newaxis])
|
||||
j = self.cross(es[:, :, np.newaxis, :], j)
|
||||
j = np.swapaxes(j.reshape((len(x), fr.shape[1]*3, len(ts)*3)), 1, 2)
|
||||
|
||||
if self.translate:
|
||||
|
||||
es = np.empty((len(x), fr.shape[1]*3, 3))
|
||||
es[:, 0::3] = prs * np.array([[[1, 0, 0]]])
|
||||
es[:, 1::3] = prs * np.array([[[0, 1, 0]]])
|
||||
es[:, 2::3] = prs * np.array([[[0, 0, 1]]])
|
||||
|
||||
jt = tdsc[np.newaxis, :, :, np.newaxis] * es[:, :, np.newaxis, :].repeat(tps.shape[1], axis=2)
|
||||
jt = np.swapaxes(jt.reshape((len(x), fr.shape[1]*3, len(ts)*3)), 1, 2)
|
||||
|
||||
j = np.concatenate([j, jt], axis=-1)
|
||||
|
||||
return j
|
||||
|
||||
# @profile(immediate=True)
|
||||
def __call__(self, descendants=None, gamma=1.0):
|
||||
|
||||
self.descendants = descendants
|
||||
|
||||
""" Calculate Masses """
|
||||
if self.weights is None:
|
||||
self.weights = np.ones(self.animation.shape[1])
|
||||
|
||||
if self.weights_translate is None:
|
||||
self.weights_translate = np.ones(self.animation.shape[1])
|
||||
|
||||
""" Calculate Descendants """
|
||||
if self.descendants is None:
|
||||
self.descendants = AnimationStructure.descendants_mask(self.animation.parents)
|
||||
|
||||
self.tdescendants = np.eye(self.animation.shape[1]) + self.descendants
|
||||
|
||||
self.first_descendants = self.descendants[:, np.array(list(self.targets.keys()))].repeat(3, axis=0).astype(int)
|
||||
self.first_tdescendants = self.tdescendants[:, np.array(list(self.targets.keys()))].repeat(3, axis=0).astype(int)
|
||||
|
||||
""" Calculate End Effectors """
|
||||
self.endeff = np.array(list(self.targets.values()))
|
||||
self.endeff = np.swapaxes(self.endeff, 0, 1)
|
||||
|
||||
if not self.references is None:
|
||||
self.second_descendants = self.descendants.repeat(3, axis=0).astype(int)
|
||||
self.second_tdescendants = self.tdescendants.repeat(3, axis=0).astype(int)
|
||||
self.second_targets = dict([(i, self.references[:, i]) for i in range(self.references.shape[1])])
|
||||
|
||||
nf = len(self.animation)
|
||||
nj = self.animation.shape[1]
|
||||
|
||||
if not self.silent:
|
||||
gp = Animation.positions_global(self.animation)
|
||||
gp = gp[:, np.array(list(self.targets.keys()))]
|
||||
error = np.mean(np.sqrt(np.sum((self.endeff - gp)**2.0, axis=2)))
|
||||
print('[JacobianInverseKinematics] Start | Error: %f' % error)
|
||||
|
||||
for i in range(self.iterations):
|
||||
|
||||
""" Get Global Rotations & Positions """
|
||||
gt = Animation.transforms_global(self.animation)
|
||||
gp = gt[:, :, :, 3]
|
||||
gp = gp[:, :, :3] / gp[:, :, 3, np.newaxis]
|
||||
gr = Quaternions.from_transforms(gt)
|
||||
|
||||
x = self.animation.rotations.euler().reshape(nf, -1)
|
||||
w = self.weights.repeat(3)
|
||||
|
||||
if self.translate:
|
||||
x = np.hstack([x, self.animation.positions.reshape(nf, -1)])
|
||||
w = np.hstack([w, self.weights_translate.repeat(3)])
|
||||
|
||||
""" Generate Jacobian """
|
||||
if self.recalculate or i == 0:
|
||||
j = self.jacobian(x, gp, gr, self.targets, self.first_descendants, self.first_tdescendants)
|
||||
|
||||
""" Update Variables """
|
||||
l = self.damping * (1.0 / (w + 0.001))
|
||||
d = (l*l) * np.eye(x.shape[1])
|
||||
e = gamma * (self.endeff.reshape(nf, -1) - gp[:, np.array(list(self.targets.keys()))].reshape(nf, -1))
|
||||
|
||||
j = np.nan_to_num(j)
|
||||
e = np.nan_to_num(e)
|
||||
x += np.array(list(map(lambda jf, ef:
|
||||
linalg.lu_solve(linalg.lu_factor(jf.T.dot(jf) + d), jf.T.dot(ef)), j, e)))
|
||||
|
||||
""" Generate Secondary Jacobian """
|
||||
if self.references is not None:
|
||||
|
||||
ns = np.array(list(map(lambda jf:
|
||||
np.eye(x.shape[1]) - linalg.solve(jf.T.dot(jf) + d, jf.T.dot(jf)), j)))
|
||||
|
||||
if self.recalculate or i == 0:
|
||||
j2 = self.jacobian(x, gp, gr, self.second_targets, self.second_descendants, self.second_tdescendants)
|
||||
|
||||
e2 = self.secondary * (self.references.reshape(nf, -1) - gp.reshape(nf, -1))
|
||||
|
||||
x += np.array(list(map(lambda nsf, j2f, e2f:
|
||||
nsf.dot(linalg.lu_solve(linalg.lu_factor(j2f.T.dot(j2f) + d), j2f.T.dot(e2f))), ns, j2, e2)))
|
||||
|
||||
""" Set Back Rotations / Translations """
|
||||
self.animation.rotations = Quaternions.from_euler(
|
||||
x[:, :nj*3].reshape((nf, nj, 3)), order='xyz', world=True)
|
||||
|
||||
if self.translate:
|
||||
self.animation.positions = x[:, nj*3:].reshape((nf, nj, 3))
|
||||
|
||||
""" Generate Error """
|
||||
|
||||
if not self.silent:
|
||||
gp = Animation.positions_global(self.animation)
|
||||
gp = gp[:, np.array(list(self.targets.keys()))]
|
||||
error = np.mean(np.sum((self.endeff - gp)**2.0, axis=2)**0.5)
|
||||
print('[JacobianInverseKinematics] Iteration %i | Error: %f' % (i+1, error))
|
||||
|
||||
|
||||
class BasicJacobianIK:
|
||||
"""
|
||||
Same interface as BasicInverseKinematics
|
||||
but uses the Jacobian IK Solver Instead
|
||||
"""
|
||||
|
||||
def __init__(self, animation, positions, iterations=10, silent=True, **kw):
|
||||
|
||||
targets = dict([(i, positions[:, i]) for i in range(positions.shape[1])])
|
||||
self.ik = JacobianInverseKinematics(animation, targets, iterations=iterations, silent=silent, **kw)
|
||||
|
||||
def __call__(self, **kw):
|
||||
return self.ik(**kw)
|
||||
|
||||
|
||||
class ICP:
|
||||
|
||||
def __init__(self,
|
||||
anim, rest, weights, mesh, goal,
|
||||
find_closest=True, damping=10,
|
||||
iterations=10, silent=True,
|
||||
translate=True, recalculate=True,
|
||||
weights_translate=None):
|
||||
|
||||
self.animation = anim
|
||||
self.rest = rest
|
||||
self.vweights = weights
|
||||
self.mesh = mesh
|
||||
self.goal = goal
|
||||
self.find_closest = find_closest
|
||||
self.iterations = iterations
|
||||
self.silent = silent
|
||||
self.translate = translate
|
||||
self.damping = damping
|
||||
self.weights = None
|
||||
self.weights_translate = weights_translate
|
||||
self.recalculate = recalculate
|
||||
|
||||
def cross(self, a, b):
|
||||
o = np.empty(b.shape)
|
||||
o[..., 0] = a[..., 1]*b[..., 2] - a[..., 2]*b[..., 1]
|
||||
o[..., 1] = a[..., 2]*b[..., 0] - a[..., 0]*b[..., 2]
|
||||
o[..., 2] = a[..., 0]*b[..., 1] - a[..., 1]*b[..., 0]
|
||||
return o
|
||||
|
||||
def jacobian(self, x, fp, fr, goal, weights, des_r, des_t):
|
||||
""" Find parent rotations """
|
||||
prs = fr[:, self.animation.parents]
|
||||
prs[:, 0] = Quaternions.id((1))
|
||||
|
||||
""" Get partial rotations """
|
||||
qys = Quaternions.from_angle_axis(x[:, 1:prs.shape[1]*3:3], np.array([[[0, 1, 0]]]))
|
||||
qzs = Quaternions.from_angle_axis(x[:, 2:prs.shape[1]*3:3], np.array([[[0, 0, 1]]]))
|
||||
|
||||
""" Find axis of rotations """
|
||||
es = np.empty((len(x), fr.shape[1]*3, 3))
|
||||
es[:, 0::3] = ((prs * qzs) * qys) * np.array([[[1, 0, 0]]])
|
||||
es[:, 1::3] = ((prs * qzs) * np.array([[[0, 1, 0]]]))
|
||||
es[:, 2::3] = ((prs * np.array([[[0, 0, 1]]])))
|
||||
|
||||
""" Construct Jacobian """
|
||||
j = fp.repeat(3, axis=1)
|
||||
j = des_r[np.newaxis, :, :, :, np.newaxis] * (goal[:, np.newaxis, :, np.newaxis] - j[:, :, np.newaxis, np.newaxis])
|
||||
j = np.sum(j * weights[np.newaxis, np.newaxis, :, :, np.newaxis], 3)
|
||||
j = self.cross(es[:, :, np.newaxis, :], j)
|
||||
j = np.swapaxes(j.reshape((len(x), fr.shape[1]*3, goal.shape[1]*3)), 1, 2)
|
||||
|
||||
if self.translate:
|
||||
|
||||
es = np.empty((len(x), fr.shape[1]*3, 3))
|
||||
es[:, 0::3] = prs * np.array([[[1, 0, 0]]])
|
||||
es[:, 1::3] = prs * np.array([[[0, 1, 0]]])
|
||||
es[:, 2::3] = prs * np.array([[[0, 0, 1]]])
|
||||
|
||||
jt = des_t[np.newaxis, :, :, :, np.newaxis] * es[:, :, np.newaxis, np.newaxis, :].repeat(goal.shape[1], axis=2)
|
||||
jt = np.sum(jt * weights[np.newaxis, np.newaxis, :, :, np.newaxis], 3)
|
||||
jt = np.swapaxes(jt.reshape((len(x), fr.shape[1]*3, goal.shape[1]*3)), 1, 2)
|
||||
|
||||
j = np.concatenate([j, jt], axis=-1)
|
||||
|
||||
return j
|
||||
|
||||
# @profile(immediate=True)
|
||||
def __call__(self, descendants=None, maxjoints=4, gamma=1.0, transpose=False):
|
||||
""" Calculate Masses """
|
||||
if self.weights is None:
|
||||
self.weights = np.ones(self.animation.shape[1])
|
||||
|
||||
if self.weights_translate is None:
|
||||
self.weights_translate = np.ones(self.animation.shape[1])
|
||||
|
||||
nf = len(self.animation)
|
||||
nj = self.animation.shape[1]
|
||||
nv = self.goal.shape[1]
|
||||
|
||||
weightids = np.argsort(-self.vweights, axis=1)[:, :maxjoints]
|
||||
weightvls = np.array(list(map(lambda w, i: w[i], self.vweights, weightids)))
|
||||
weightvls = weightvls / weightvls.sum(axis=1)[..., np.newaxis]
|
||||
|
||||
if descendants is None:
|
||||
self.descendants = AnimationStructure.descendants_mask(self.animation.parents)
|
||||
else:
|
||||
self.descendants = descendants
|
||||
|
||||
des_r = np.eye(nj) + self.descendants
|
||||
des_r = des_r[:, weightids].repeat(3, axis=0)
|
||||
|
||||
des_t = np.eye(nj) + self.descendants
|
||||
des_t = des_t[:, weightids].repeat(3, axis=0)
|
||||
|
||||
if not self.silent:
|
||||
curr = Animation.skin(self.animation, self.rest, self.vweights, self.mesh, maxjoints=maxjoints)
|
||||
error = np.mean(np.sqrt(np.sum((curr - self.goal)**2.0, axis=-1)))
|
||||
print('[ICP] Start | Error: %f' % error)
|
||||
|
||||
for i in range(self.iterations):
|
||||
|
||||
""" Get Global Rotations & Positions """
|
||||
gt = Animation.transforms_global(self.animation)
|
||||
gp = gt[:, :, :, 3]
|
||||
gp = gp[:, :, :3] / gp[:, :, 3, np.newaxis]
|
||||
gr = Quaternions.from_transforms(gt)
|
||||
|
||||
x = self.animation.rotations.euler().reshape(nf, -1)
|
||||
w = self.weights.repeat(3)
|
||||
|
||||
if self.translate:
|
||||
x = np.hstack([x, self.animation.positions.reshape(nf, -1)])
|
||||
w = np.hstack([w, self.weights_translate.repeat(3)])
|
||||
|
||||
""" Get Current State """
|
||||
curr = Animation.skin(self.animation, self.rest, self.vweights, self.mesh, maxjoints=maxjoints)
|
||||
|
||||
""" Find Cloest Points """
|
||||
if self.find_closest:
|
||||
mapping = np.argmin(
|
||||
(curr[:, :, np.newaxis] -
|
||||
self.goal[:, np.newaxis, :])**2.0, axis=2)
|
||||
e = gamma * (np.array(list(map(lambda g, m: g[m], self.goal, mapping))) - curr).reshape(nf, -1)
|
||||
else:
|
||||
e = gamma * (self.goal - curr).reshape(nf, -1)
|
||||
|
||||
""" Generate Jacobian """
|
||||
if self.recalculate or i == 0:
|
||||
j = self.jacobian(x, gp, gr, self.goal, weightvls, des_r, des_t)
|
||||
|
||||
""" Update Variables """
|
||||
l = self.damping * (1.0 / (w + 1e-10))
|
||||
d = (l*l) * np.eye(x.shape[1])
|
||||
|
||||
if transpose:
|
||||
x += np.array(list(map(lambda jf, ef: jf.T.dot(ef), j, e)))
|
||||
else:
|
||||
x += np.array(list(map(lambda jf, ef:
|
||||
linalg.lu_solve(linalg.lu_factor(jf.T.dot(jf) + d), jf.T.dot(ef)), j, e)))
|
||||
|
||||
""" Set Back Rotations / Translations """
|
||||
self.animation.rotations = Quaternions.from_euler(
|
||||
x[:, :nj*3].reshape((nf, nj, 3)), order='xyz', world=True)
|
||||
|
||||
if self.translate:
|
||||
self.animation.positions = x[:, nj*3:].reshape((nf, nj, 3))
|
||||
|
||||
if not self.silent:
|
||||
curr = Animation.skin(self.animation, self.rest, self.vweights, self.mesh)
|
||||
error = np.mean(np.sqrt(np.sum((curr - self.goal)**2.0, axis=-1)))
|
||||
print('[ICP] Iteration %i | Error: %f' % (i+1, error))
|
||||
@@ -0,0 +1,91 @@
|
||||
import numpy as np
|
||||
|
||||
from Quaternions import Quaternions
|
||||
|
||||
|
||||
class Pivots:
|
||||
"""
|
||||
Pivots is an ndarray of angular rotations
|
||||
|
||||
This wrapper provides some functions for
|
||||
working with pivots.
|
||||
|
||||
These are particularly useful as a number
|
||||
of atomic operations (such as adding or
|
||||
subtracting) cannot be achieved using
|
||||
the standard arithmatic and need to be
|
||||
defined differently to work correctly
|
||||
"""
|
||||
|
||||
def __init__(self, ps): self.ps = np.array(ps)
|
||||
def __str__(self): return "Pivots(" + str(self.ps) + ")"
|
||||
def __repr__(self): return "Pivots(" + repr(self.ps) + ")"
|
||||
|
||||
def __add__(self, other): return Pivots(np.arctan2(np.sin(self.ps + other.ps), np.cos(self.ps + other.ps)))
|
||||
def __sub__(self, other): return Pivots(np.arctan2(np.sin(self.ps - other.ps), np.cos(self.ps - other.ps)))
|
||||
def __mul__(self, other): return Pivots(self.ps * other.ps)
|
||||
def __div__(self, other): return Pivots(self.ps / other.ps)
|
||||
def __mod__(self, other): return Pivots(self.ps % other.ps)
|
||||
def __pow__(self, other): return Pivots(self.ps ** other.ps)
|
||||
|
||||
def __lt__(self, other): return self.ps < other.ps
|
||||
def __le__(self, other): return self.ps <= other.ps
|
||||
def __eq__(self, other): return self.ps == other.ps
|
||||
def __ne__(self, other): return self.ps != other.ps
|
||||
def __ge__(self, other): return self.ps >= other.ps
|
||||
def __gt__(self, other): return self.ps > other.ps
|
||||
|
||||
def __abs__(self): return Pivots(abs(self.ps))
|
||||
def __neg__(self): return Pivots(-self.ps)
|
||||
|
||||
def __iter__(self): return iter(self.ps)
|
||||
def __len__(self): return len(self.ps)
|
||||
|
||||
def __getitem__(self, k): return Pivots(self.ps[k])
|
||||
def __setitem__(self, k, v): self.ps[k] = v.ps
|
||||
|
||||
def _ellipsis(self): return tuple(map(lambda x: slice(None), self.shape))
|
||||
|
||||
def quaternions(self, plane='xz'):
|
||||
fa = self._ellipsis()
|
||||
axises = np.ones(self.ps.shape + (3,))
|
||||
axises[fa + ("xyz".index(plane[0]),)] = 0.0
|
||||
axises[fa + ("xyz".index(plane[1]),)] = 0.0
|
||||
return Quaternions.from_angle_axis(self.ps, axises)
|
||||
|
||||
def directions(self, plane='xz'):
|
||||
dirs = np.zeros((len(self.ps), 3))
|
||||
dirs[..., "xyz".index(plane[0])] = np.sin(self.ps)
|
||||
dirs[..., "xyz".index(plane[1])] = np.cos(self.ps)
|
||||
return dirs
|
||||
|
||||
def normalized(self):
|
||||
xs = np.copy(self.ps)
|
||||
while np.any(xs > np.pi):
|
||||
xs[xs > np.pi] = xs[xs > np.pi] - 2 * np.pi
|
||||
while np.any(xs < -np.pi):
|
||||
xs[xs < -np.pi] = xs[xs < -np.pi] + 2 * np.pi
|
||||
return Pivots(xs)
|
||||
|
||||
def interpolate(self, ws):
|
||||
dir = np.average(self.directions, weights=ws, axis=0)
|
||||
return np.arctan2(dir[2], dir[0])
|
||||
|
||||
def copy(self):
|
||||
return Pivots(np.copy(self.ps))
|
||||
|
||||
@property
|
||||
def shape(self):
|
||||
return self.ps.shape
|
||||
|
||||
@classmethod
|
||||
def from_quaternions(cls, qs, forward='z', plane='xz'):
|
||||
ds = np.zeros(qs.shape + (3,))
|
||||
ds[..., 'xyz'.index(forward)] = 1.0
|
||||
return Pivots.from_directions(qs * ds, plane=plane)
|
||||
|
||||
@classmethod
|
||||
def from_directions(cls, ds, plane='xz'):
|
||||
ys = ds[..., 'xyz'.index(plane[0])]
|
||||
xs = ds[..., 'xyz'.index(plane[1])]
|
||||
return Pivots(np.arctan2(ys, xs))
|
||||
@@ -0,0 +1,481 @@
|
||||
import numpy as np
|
||||
|
||||
|
||||
class Quaternions:
|
||||
"""
|
||||
Quaternions is a wrapper around a numpy ndarray
|
||||
that allows it to act as if it were an narray of
|
||||
a quaternion data type.
|
||||
|
||||
Therefore addition, subtraction, multiplication,
|
||||
division, negation, absolute, are all defined
|
||||
in terms of quaternion operations such as quaternion
|
||||
multiplication.
|
||||
|
||||
This allows for much neater code and many routines
|
||||
which conceptually do the same thing to be written
|
||||
in the same way for point data and for rotation data.
|
||||
|
||||
The Quaternions class has been desgined such that it
|
||||
should support broadcasting and slicing in all of the
|
||||
usual ways.
|
||||
"""
|
||||
|
||||
def __init__(self, qs):
|
||||
if isinstance(qs, np.ndarray):
|
||||
|
||||
if len(qs.shape) == 1:
|
||||
qs = np.array([qs])
|
||||
self.qs = qs
|
||||
return
|
||||
|
||||
if isinstance(qs, Quaternions):
|
||||
self.qs = qs.qs
|
||||
return
|
||||
|
||||
raise TypeError('Quaternions must be constructed from iterable, numpy array, or Quaternions, not %s' % type(qs))
|
||||
|
||||
def __str__(self): return "Quaternions(" + str(self.qs) + ")"
|
||||
def __repr__(self): return "Quaternions(" + repr(self.qs) + ")"
|
||||
""" Helper Methods for Broadcasting and Data extraction """
|
||||
|
||||
@classmethod
|
||||
def _broadcast(cls, sqs, oqs, scalar=False):
|
||||
if isinstance(oqs, float):
|
||||
return sqs, oqs * np.ones(sqs.shape[:-1])
|
||||
|
||||
ss = np.array(sqs.shape) if not scalar else np.array(sqs.shape[:-1])
|
||||
os = np.array(oqs.shape)
|
||||
|
||||
if len(ss) != len(os):
|
||||
raise TypeError('Quaternions cannot broadcast together shapes %s and %s' % (sqs.shape, oqs.shape))
|
||||
|
||||
if np.all(ss == os):
|
||||
return sqs, oqs
|
||||
|
||||
if not np.all((ss == os) | (os == np.ones(len(os))) | (ss == np.ones(len(ss)))):
|
||||
raise TypeError('Quaternions cannot broadcast together shapes %s and %s' % (sqs.shape, oqs.shape))
|
||||
|
||||
sqsn, oqsn = sqs.copy(), oqs.copy()
|
||||
|
||||
for a in np.where(ss == 1)[0]:
|
||||
sqsn = sqsn.repeat(os[a], axis=a)
|
||||
for a in np.where(os == 1)[0]:
|
||||
oqsn = oqsn.repeat(ss[a], axis=a)
|
||||
|
||||
return sqsn, oqsn
|
||||
|
||||
""" Adding Quaterions is just Defined as Multiplication """
|
||||
|
||||
def __add__(self, other): return self * other
|
||||
def __sub__(self, other): return self / other
|
||||
""" Quaterion Multiplication """
|
||||
|
||||
def __mul__(self, other):
|
||||
"""
|
||||
Quaternion multiplication has three main methods.
|
||||
|
||||
When multiplying a Quaternions array by Quaternions
|
||||
normal quaternion multiplication is performed.
|
||||
|
||||
When multiplying a Quaternions array by a vector
|
||||
array of the same shape, where the last axis is 3,
|
||||
it is assumed to be a Quaternion by 3D-Vector
|
||||
multiplication and the 3D-Vectors are rotated
|
||||
in space by the Quaternions.
|
||||
|
||||
When multipplying a Quaternions array by a scalar
|
||||
or vector of different shape it is assumed to be
|
||||
a Quaternions by Scalars multiplication and the
|
||||
Quaternions are scaled using Slerp and the identity
|
||||
quaternions.
|
||||
"""
|
||||
|
||||
""" If Quaternions type do Quaternions * Quaternions """
|
||||
if isinstance(other, Quaternions):
|
||||
sqs, oqs = Quaternions._broadcast(self.qs, other.qs)
|
||||
|
||||
q0 = sqs[..., 0]
|
||||
q1 = sqs[..., 1]
|
||||
q2 = sqs[..., 2]
|
||||
q3 = sqs[..., 3]
|
||||
r0 = oqs[..., 0]
|
||||
r1 = oqs[..., 1]
|
||||
r2 = oqs[..., 2]
|
||||
r3 = oqs[..., 3]
|
||||
|
||||
qs = np.empty(sqs.shape)
|
||||
qs[..., 0] = r0 * q0 - r1 * q1 - r2 * q2 - r3 * q3
|
||||
qs[..., 1] = r0 * q1 + r1 * q0 - r2 * q3 + r3 * q2
|
||||
qs[..., 2] = r0 * q2 + r1 * q3 + r2 * q0 - r3 * q1
|
||||
qs[..., 3] = r0 * q3 - r1 * q2 + r2 * q1 + r3 * q0
|
||||
|
||||
return Quaternions(qs)
|
||||
|
||||
""" If array type do Quaternions * Vectors """
|
||||
if isinstance(other, np.ndarray) and other.shape[-1] == 3:
|
||||
vs = Quaternions(np.concatenate([np.zeros(other.shape[:-1] + (1,)), other], axis=-1))
|
||||
|
||||
return (self * (vs * -self)).imaginaries
|
||||
|
||||
""" If float do Quaternions * Scalars """
|
||||
if isinstance(other, np.ndarray) or isinstance(other, float):
|
||||
return Quaternions.slerp(Quaternions.id_like(self), self, other)
|
||||
|
||||
raise TypeError('Cannot multiply/add Quaternions with type %s' % str(type(other)))
|
||||
|
||||
def __div__(self, other):
|
||||
"""
|
||||
When a Quaternion type is supplied, division is defined
|
||||
as multiplication by the inverse of that Quaternion.
|
||||
|
||||
When a scalar or vector is supplied it is defined
|
||||
as multiplicaion of one over the supplied value.
|
||||
Essentially a scaling.
|
||||
"""
|
||||
|
||||
if isinstance(other, Quaternions):
|
||||
return self * (-other)
|
||||
if isinstance(other, np.ndarray):
|
||||
return self * (1.0 / other)
|
||||
if isinstance(other, float):
|
||||
return self * (1.0 / other)
|
||||
raise TypeError('Cannot divide/subtract Quaternions with type %s' + str(type(other)))
|
||||
|
||||
def __eq__(self, other): return self.qs == other.qs
|
||||
def __ne__(self, other): return self.qs != other.qs
|
||||
|
||||
def __neg__(self):
|
||||
""" Invert Quaternions """
|
||||
return Quaternions(self.qs * np.array([[1, -1, -1, -1]]))
|
||||
|
||||
def __abs__(self):
|
||||
""" Unify Quaternions To Single Pole """
|
||||
qabs = self.normalized().copy()
|
||||
top = np.sum((qabs.qs) * np.array([1, 0, 0, 0]), axis=-1)
|
||||
bot = np.sum((-qabs.qs) * np.array([1, 0, 0, 0]), axis=-1)
|
||||
qabs.qs[top < bot] = -qabs.qs[top < bot]
|
||||
return qabs
|
||||
|
||||
def __iter__(self): return iter(self.qs)
|
||||
def __len__(self): return len(self.qs)
|
||||
|
||||
def __getitem__(self, k): return Quaternions(self.qs[k])
|
||||
def __setitem__(self, k, v): self.qs[k] = v.qs
|
||||
|
||||
@property
|
||||
def lengths(self):
|
||||
return np.sum(self.qs**2.0, axis=-1)**0.5
|
||||
|
||||
@property
|
||||
def reals(self):
|
||||
return self.qs[..., 0]
|
||||
|
||||
@property
|
||||
def imaginaries(self):
|
||||
return self.qs[..., 1:4]
|
||||
|
||||
@property
|
||||
def shape(self): return self.qs.shape[:-1]
|
||||
|
||||
def repeat(self, n, **kwargs):
|
||||
return Quaternions(self.qs.repeat(n, **kwargs))
|
||||
|
||||
def normalized(self):
|
||||
return Quaternions(self.qs / self.lengths[..., np.newaxis])
|
||||
|
||||
def log(self):
|
||||
norm = abs(self.normalized())
|
||||
imgs = norm.imaginaries
|
||||
lens = np.sqrt(np.sum(imgs**2, axis=-1))
|
||||
lens = np.arctan2(lens, norm.reals) / (lens + 1e-10)
|
||||
return imgs * lens[..., np.newaxis]
|
||||
|
||||
def constrained(self, axis):
|
||||
|
||||
rl = self.reals
|
||||
im = np.sum(axis * self.imaginaries, axis=-1)
|
||||
|
||||
t1 = -2 * np.arctan2(rl, im) + np.pi
|
||||
t2 = -2 * np.arctan2(rl, im) - np.pi
|
||||
|
||||
top = Quaternions.exp(axis[np.newaxis] * (t1[:, np.newaxis] / 2.0))
|
||||
bot = Quaternions.exp(axis[np.newaxis] * (t2[:, np.newaxis] / 2.0))
|
||||
img = self.dot(top) > self.dot(bot)
|
||||
|
||||
ret = top.copy()
|
||||
ret[img] = top[img]
|
||||
ret[~img] = bot[~img]
|
||||
return ret
|
||||
|
||||
def constrained_x(self): return self.constrained(np.array([1, 0, 0]))
|
||||
def constrained_y(self): return self.constrained(np.array([0, 1, 0]))
|
||||
def constrained_z(self): return self.constrained(np.array([0, 0, 1]))
|
||||
|
||||
def dot(self, q): return np.sum(self.qs * q.qs, axis=-1)
|
||||
|
||||
def copy(self): return Quaternions(np.copy(self.qs))
|
||||
|
||||
def reshape(self, s):
|
||||
self.qs.reshape(s)
|
||||
return self
|
||||
|
||||
def interpolate(self, ws):
|
||||
return Quaternions.exp(np.average(abs(self).log, axis=0, weights=ws))
|
||||
|
||||
def euler(self, order='xyz'):
|
||||
|
||||
q = self.normalized().qs
|
||||
q0 = q[..., 0]
|
||||
q1 = q[..., 1]
|
||||
q2 = q[..., 2]
|
||||
q3 = q[..., 3]
|
||||
es = np.zeros(self.shape + (3,))
|
||||
|
||||
if order == 'xyz':
|
||||
es[..., 0] = np.arctan2(2 * (q0 * q1 + q2 * q3), 1 - 2 * (q1 * q1 + q2 * q2))
|
||||
es[..., 1] = np.arcsin((2 * (q0 * q2 - q3 * q1)).clip(-1, 1))
|
||||
es[..., 2] = np.arctan2(2 * (q0 * q3 + q1 * q2), 1 - 2 * (q2 * q2 + q3 * q3))
|
||||
elif order == 'yzx':
|
||||
es[..., 0] = np.arctan2(2 * (q1 * q0 - q2 * q3), -q1 * q1 + q2 * q2 - q3 * q3 + q0 * q0)
|
||||
es[..., 1] = np.arctan2(2 * (q2 * q0 - q1 * q3), q1 * q1 - q2 * q2 - q3 * q3 + q0 * q0)
|
||||
es[..., 2] = np.arcsin((2 * (q1 * q2 + q3 * q0)).clip(-1, 1))
|
||||
else:
|
||||
raise NotImplementedError('Cannot convert from ordering %s' % order)
|
||||
|
||||
"""
|
||||
|
||||
# These conversion don't appear to work correctly for Maya.
|
||||
# http://bediyap.com/programming/convert-quaternion-to-euler-rotations/
|
||||
|
||||
if order == 'xyz':
|
||||
es[fa + (0,)] = np.arctan2(2 * (q0 * q3 - q1 * q2), q0 * q0 + q1 * q1 - q2 * q2 - q3 * q3)
|
||||
es[fa + (1,)] = np.arcsin((2 * (q1 * q3 + q0 * q2)).clip(-1,1))
|
||||
es[fa + (2,)] = np.arctan2(2 * (q0 * q1 - q2 * q3), q0 * q0 - q1 * q1 - q2 * q2 + q3 * q3)
|
||||
elif order == 'yzx':
|
||||
es[fa + (0,)] = np.arctan2(2 * (q0 * q1 - q2 * q3), q0 * q0 - q1 * q1 + q2 * q2 - q3 * q3)
|
||||
es[fa + (1,)] = np.arcsin((2 * (q1 * q2 + q0 * q3)).clip(-1,1))
|
||||
es[fa + (2,)] = np.arctan2(2 * (q0 * q2 - q1 * q3), q0 * q0 + q1 * q1 - q2 * q2 - q3 * q3)
|
||||
elif order == 'zxy':
|
||||
es[fa + (0,)] = np.arctan2(2 * (q0 * q2 - q1 * q3), q0 * q0 - q1 * q1 - q2 * q2 + q3 * q3)
|
||||
es[fa + (1,)] = np.arcsin((2 * (q0 * q1 + q2 * q3)).clip(-1,1))
|
||||
es[fa + (2,)] = np.arctan2(2 * (q0 * q3 - q1 * q2), q0 * q0 - q1 * q1 + q2 * q2 - q3 * q3)
|
||||
elif order == 'xzy':
|
||||
es[fa + (0,)] = np.arctan2(2 * (q0 * q2 + q1 * q3), q0 * q0 + q1 * q1 - q2 * q2 - q3 * q3)
|
||||
es[fa + (1,)] = np.arcsin((2 * (q0 * q3 - q1 * q2)).clip(-1,1))
|
||||
es[fa + (2,)] = np.arctan2(2 * (q0 * q1 + q2 * q3), q0 * q0 - q1 * q1 + q2 * q2 - q3 * q3)
|
||||
elif order == 'yxz':
|
||||
es[fa + (0,)] = np.arctan2(2 * (q1 * q2 + q0 * q3), q0 * q0 - q1 * q1 + q2 * q2 - q3 * q3)
|
||||
es[fa + (1,)] = np.arcsin((2 * (q0 * q1 - q2 * q3)).clip(-1,1))
|
||||
es[fa + (2,)] = np.arctan2(2 * (q1 * q3 + q0 * q2), q0 * q0 - q1 * q1 - q2 * q2 + q3 * q3)
|
||||
elif order == 'zyx':
|
||||
es[fa + (0,)] = np.arctan2(2 * (q0 * q1 + q2 * q3), q0 * q0 - q1 * q1 - q2 * q2 + q3 * q3)
|
||||
es[fa + (1,)] = np.arcsin((2 * (q0 * q2 - q1 * q3)).clip(-1,1))
|
||||
es[fa + (2,)] = np.arctan2(2 * (q0 * q3 + q1 * q2), q0 * q0 + q1 * q1 - q2 * q2 - q3 * q3)
|
||||
else:
|
||||
raise KeyError('Unknown ordering %s' % order)
|
||||
|
||||
"""
|
||||
|
||||
# https://github.com/ehsan/ogre/blob/master/OgreMain/src/OgreMatrix3.cpp
|
||||
# Use this class and convert from matrix
|
||||
|
||||
return es
|
||||
|
||||
def average(self):
|
||||
|
||||
if len(self.shape) == 1:
|
||||
|
||||
import numpy.core.umath_tests as ut
|
||||
system = ut.matrix_multiply(self.qs[:, :, np.newaxis], self.qs[:, np.newaxis, :]).sum(axis=0)
|
||||
w, v = np.linalg.eigh(system)
|
||||
qiT_dot_qref = (self.qs[:, :, np.newaxis] * v[np.newaxis, :, :]).sum(axis=1)
|
||||
return Quaternions(v[:, np.argmin((1.-qiT_dot_qref**2).sum(axis=0))])
|
||||
|
||||
else:
|
||||
|
||||
raise NotImplementedError('Cannot average multi-dimensionsal Quaternions')
|
||||
|
||||
def angle_axis(self):
|
||||
|
||||
norm = self.normalized()
|
||||
s = np.sqrt(1 - (norm.reals**2.0))
|
||||
s[s == 0] = 0.001
|
||||
|
||||
angles = 2.0 * np.arccos(norm.reals)
|
||||
axis = norm.imaginaries / s[..., np.newaxis]
|
||||
|
||||
return angles, axis
|
||||
|
||||
def transforms(self):
|
||||
|
||||
qw = self.qs[..., 0]
|
||||
qx = self.qs[..., 1]
|
||||
qy = self.qs[..., 2]
|
||||
qz = self.qs[..., 3]
|
||||
|
||||
x2 = qx + qx
|
||||
y2 = qy + qy
|
||||
z2 = qz + qz
|
||||
xx = qx * x2
|
||||
yy = qy * y2
|
||||
wx = qw * x2
|
||||
xy = qx * y2
|
||||
yz = qy * z2
|
||||
wy = qw * y2
|
||||
xz = qx * z2
|
||||
zz = qz * z2
|
||||
wz = qw * z2
|
||||
|
||||
m = np.empty(self.shape + (3, 3))
|
||||
m[..., 0, 0] = 1.0 - (yy + zz)
|
||||
m[..., 0, 1] = xy - wz
|
||||
m[..., 0, 2] = xz + wy
|
||||
m[..., 1, 0] = xy + wz
|
||||
m[..., 1, 1] = 1.0 - (xx + zz)
|
||||
m[..., 1, 2] = yz - wx
|
||||
m[..., 2, 0] = xz - wy
|
||||
m[..., 2, 1] = yz + wx
|
||||
m[..., 2, 2] = 1.0 - (xx + yy)
|
||||
|
||||
return m
|
||||
|
||||
def ravel(self):
|
||||
return self.qs.ravel()
|
||||
|
||||
@classmethod
|
||||
def id(cls, n):
|
||||
|
||||
if isinstance(n, tuple):
|
||||
qs = np.zeros(n + (4,))
|
||||
qs[..., 0] = 1.0
|
||||
return Quaternions(qs)
|
||||
|
||||
if isinstance(n, int):
|
||||
qs = np.zeros((n, 4))
|
||||
qs[:, 0] = 1.0
|
||||
return Quaternions(qs)
|
||||
|
||||
raise TypeError('Cannot Construct Quaternion from %s type' % str(type(n)))
|
||||
|
||||
@classmethod
|
||||
def id_like(cls, a):
|
||||
qs = np.zeros(a.shape + (4,))
|
||||
qs[..., 0] = 1.0
|
||||
return Quaternions(qs)
|
||||
|
||||
@classmethod
|
||||
def exp(cls, ws):
|
||||
|
||||
ts = np.sum(ws**2.0, axis=-1)**0.5
|
||||
ts[ts == 0] = 0.001
|
||||
ls = np.sin(ts) / ts
|
||||
|
||||
qs = np.empty(ws.shape[:-1] + (4,))
|
||||
qs[..., 0] = np.cos(ts)
|
||||
qs[..., 1] = ws[..., 0] * ls
|
||||
qs[..., 2] = ws[..., 1] * ls
|
||||
qs[..., 3] = ws[..., 2] * ls
|
||||
|
||||
return Quaternions(qs).normalized()
|
||||
|
||||
@classmethod
|
||||
def slerp(cls, q0s, q1s, a):
|
||||
|
||||
fst, snd = cls._broadcast(q0s.qs, q1s.qs)
|
||||
fst, a = cls._broadcast(fst, a, scalar=True)
|
||||
snd, a = cls._broadcast(snd, a, scalar=True)
|
||||
|
||||
len = np.sum(fst * snd, axis=-1)
|
||||
|
||||
neg = len < 0.0
|
||||
len[neg] = -len[neg]
|
||||
snd[neg] = -snd[neg]
|
||||
|
||||
amount0 = np.zeros(a.shape)
|
||||
amount1 = np.zeros(a.shape)
|
||||
|
||||
linear = (1.0 - len) < 0.01
|
||||
omegas = np.arccos(len[~linear])
|
||||
sinoms = np.sin(omegas)
|
||||
|
||||
amount0[linear] = 1.0 - a[linear]
|
||||
amount1[linear] = a[linear]
|
||||
amount0[~linear] = np.sin((1.0 - a[~linear]) * omegas) / sinoms
|
||||
amount1[~linear] = np.sin(a[~linear] * omegas) / sinoms
|
||||
|
||||
return Quaternions(
|
||||
amount0[..., np.newaxis] * fst +
|
||||
amount1[..., np.newaxis] * snd)
|
||||
|
||||
@classmethod
|
||||
def between(cls, v0s, v1s):
|
||||
a = np.cross(v0s, v1s)
|
||||
w = np.sqrt((v0s**2).sum(axis=-1) * (v1s**2).sum(axis=-1)) + (v0s * v1s).sum(axis=-1)
|
||||
return Quaternions(np.concatenate([w[..., np.newaxis], a], axis=-1)).normalized()
|
||||
|
||||
@classmethod
|
||||
def from_angle_axis(cls, angles, axis):
|
||||
axis = axis / (np.sqrt(np.sum(axis**2, axis=-1)) + 1e-10)[..., np.newaxis]
|
||||
sines = np.sin(angles / 2.0)[..., np.newaxis]
|
||||
cosines = np.cos(angles / 2.0)[..., np.newaxis]
|
||||
return Quaternions(np.concatenate([cosines, axis * sines], axis=-1))
|
||||
|
||||
@classmethod
|
||||
def from_euler(cls, es, order='xyz', world=False):
|
||||
|
||||
axis = {
|
||||
'x': np.array([1, 0, 0]),
|
||||
'y': np.array([0, 1, 0]),
|
||||
'z': np.array([0, 0, 1]),
|
||||
}
|
||||
|
||||
q0s = Quaternions.from_angle_axis(es[..., 0], axis[order[0]])
|
||||
q1s = Quaternions.from_angle_axis(es[..., 1], axis[order[1]])
|
||||
q2s = Quaternions.from_angle_axis(es[..., 2], axis[order[2]])
|
||||
|
||||
return (q2s * (q1s * q0s)) if world else (q0s * (q1s * q2s))
|
||||
|
||||
@classmethod
|
||||
def from_transforms(cls, ts):
|
||||
|
||||
d0, d1, d2 = ts[..., 0, 0], ts[..., 1, 1], ts[..., 2, 2]
|
||||
|
||||
q0 = (d0 + d1 + d2 + 1.0) / 4.0
|
||||
q1 = (d0 - d1 - d2 + 1.0) / 4.0
|
||||
q2 = (-d0 + d1 - d2 + 1.0) / 4.0
|
||||
q3 = (-d0 - d1 + d2 + 1.0) / 4.0
|
||||
|
||||
q0 = np.sqrt(q0.clip(0, None))
|
||||
q1 = np.sqrt(q1.clip(0, None))
|
||||
q2 = np.sqrt(q2.clip(0, None))
|
||||
q3 = np.sqrt(q3.clip(0, None))
|
||||
|
||||
c0 = (q0 >= q1) & (q0 >= q2) & (q0 >= q3)
|
||||
c1 = (q1 >= q0) & (q1 >= q2) & (q1 >= q3)
|
||||
c2 = (q2 >= q0) & (q2 >= q1) & (q2 >= q3)
|
||||
c3 = (q3 >= q0) & (q3 >= q1) & (q3 >= q2)
|
||||
|
||||
q1[c0] *= np.sign(ts[c0, 2, 1] - ts[c0, 1, 2])
|
||||
q2[c0] *= np.sign(ts[c0, 0, 2] - ts[c0, 2, 0])
|
||||
q3[c0] *= np.sign(ts[c0, 1, 0] - ts[c0, 0, 1])
|
||||
|
||||
q0[c1] *= np.sign(ts[c1, 2, 1] - ts[c1, 1, 2])
|
||||
q2[c1] *= np.sign(ts[c1, 1, 0] + ts[c1, 0, 1])
|
||||
q3[c1] *= np.sign(ts[c1, 0, 2] + ts[c1, 2, 0])
|
||||
|
||||
q0[c2] *= np.sign(ts[c2, 0, 2] - ts[c2, 2, 0])
|
||||
q1[c2] *= np.sign(ts[c2, 1, 0] + ts[c2, 0, 1])
|
||||
q3[c2] *= np.sign(ts[c2, 2, 1] + ts[c2, 1, 2])
|
||||
|
||||
q0[c3] *= np.sign(ts[c3, 1, 0] - ts[c3, 0, 1])
|
||||
q1[c3] *= np.sign(ts[c3, 2, 0] + ts[c3, 0, 2])
|
||||
q2[c3] *= np.sign(ts[c3, 2, 1] + ts[c3, 1, 2])
|
||||
|
||||
qs = np.empty(ts.shape[:-2] + (4,))
|
||||
qs[..., 0] = q0
|
||||
qs[..., 1] = q1
|
||||
qs[..., 2] = q2
|
||||
qs[..., 3] = q3
|
||||
|
||||
return cls(qs)
|
||||
@@ -0,0 +1,4 @@
|
||||
import sys
|
||||
import os
|
||||
BASEPATH = os.path.dirname(__file__)
|
||||
sys.path.insert(0, BASEPATH)
|
||||
+60
@@ -0,0 +1,60 @@
|
||||
smpl_body_model: None # _DATA/data/mano
|
||||
|
||||
offsets:
|
||||
right: [[ 0.0957, 0.0064, 0.0062],
|
||||
[ 0.0076, 0.0012, 0.0269],
|
||||
[-0.0251, 0.0052, 0.0291],
|
||||
[-0.0473, 0.0039, 0.0290],
|
||||
[ 0.0010, 0.0049, 0.0028],
|
||||
[-0.0302, 0.0068, -0.0028],
|
||||
[-0.0531, 0.0055, -0.0067],
|
||||
[ 0.0269, -0.0036, -0.0370],
|
||||
[ 0.0099, -0.0035, -0.0495],
|
||||
[-0.0060, -0.0042, -0.0599],
|
||||
[ 0.0139, 0.0024, -0.0205],
|
||||
[-0.0144, 0.0045, -0.0256],
|
||||
[-0.0379, 0.0028, -0.0332],
|
||||
[ 0.0716, -0.0091, 0.0320],
|
||||
[ 0.0519, -0.0082, 0.0557],
|
||||
[ 0.0297, -0.0137, 0.0702]]
|
||||
|
||||
|
||||
left: [
|
||||
[-0.000876, 0.909315, 0.027821],
|
||||
[ 0.071361, -0.089584, -0.008046],
|
||||
[ 0.030669, -0.364209, -0.006689],
|
||||
[-0.011554, -0.383348, -0.043502],
|
||||
[ 0.023338, -0.054645, 0.114370],
|
||||
[-0.069012, -0.088960, -0.004796],
|
||||
[-0.036152, -0.370650, -0.009185],
|
||||
[ 0.014029, -0.383638, -0.041892],
|
||||
[-0.022043, -0.046410, 0.117900],
|
||||
[-0.002508, 0.103257, -0.022185],
|
||||
[ 0.003581, 0.127658, -0.001713],
|
||||
[ 0.002027, 0.049072, 0.027867],
|
||||
[-0.001963, 0.208243, -0.049765],
|
||||
[ 0.003517, 0.062322, 0.050205],
|
||||
[ 0.075298, 0.117780, -0.036875],
|
||||
[ 0.085317, 0.031739, -0.007293],
|
||||
[ 0.251247, -0.011967, -0.027518],
|
||||
[ 0.238019, 0.009007, 0.000044],
|
||||
[ 0.079668, -0.009683, -0.013206],
|
||||
[-0.077033, 0.115606, -0.041811],
|
||||
[-0.089203, 0.032785, -0.009802],
|
||||
[-0.245990, -0.013152, -0.020162],
|
||||
[-0.245177, 0.008622, -0.003532],
|
||||
[-0.080400, -0.007248, -0.010419]
|
||||
]
|
||||
|
||||
parents: [-1, 0, 1, 2, 0, 4, 5, 0, 7, 8, 0, 10, 11, 0, 13, 14]
|
||||
|
||||
joint_names: ['wrist', 'index1', 'index2', 'index3', 'middle1', 'middle2', 'middle3', 'pinky1', 'pinky2', 'pinky3',
|
||||
'ring1', 'ring2', 'ring3', 'thumb1', 'thumb2', 'thumb3']
|
||||
|
||||
|
||||
joints_to_use: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]
|
||||
|
||||
leaf_joints: [16, 17, 18, 19, 20]
|
||||
|
||||
wrist_index: [0]
|
||||
|
||||
@@ -0,0 +1,165 @@
|
||||
import math
|
||||
import os
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
import torch
|
||||
import torch.nn.init
|
||||
import torch.optim
|
||||
|
||||
|
||||
def weights_init(init_type='default'):
|
||||
def init_fun(m):
|
||||
classname = m.__class__.__name__
|
||||
if (classname.find('Conv') == 0 or classname.find('Linear') == 0) and hasattr(m, 'weight'):
|
||||
if init_type == 'gaussian':
|
||||
torch.nn.init.normal_(m.weight.data)
|
||||
elif init_type == 'xavier':
|
||||
torch.nn.init.xavier_normal_(m.weight.data, gain=math.sqrt(2))
|
||||
elif init_type == 'kaiming':
|
||||
torch.nn.init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
|
||||
elif init_type == 'orthogonal':
|
||||
torch.nn.init.orthogonal_(m.weight.data, gain=math.sqrt(2))
|
||||
elif init_type == 'default':
|
||||
torch.nn.init.normal_(m.weight.data)
|
||||
weight_norm = m.weight.pow(2).sum(2, keepdim=True).sum(1, keepdim=True).add(1e-8).sqrt()
|
||||
m.weight.data.div_(weight_norm)
|
||||
else:
|
||||
raise NotImplementedError(f"Unsupported initialization: {init_type}")
|
||||
if hasattr(m, 'bias') and m.bias is not None:
|
||||
torch.nn.init.constant_(m.bias.data, 0.0)
|
||||
return init_fun
|
||||
|
||||
|
||||
class BaseModel(ABC):
|
||||
"""This class is an abstract base class (ABC) for models.
|
||||
To create a subclass, you need to implement the following five functions:
|
||||
-- <__init__>: initialize the class; first call BaseModel.__init__(self, opt).
|
||||
-- <set_input>: unpack data from dataset and apply preprocessing.
|
||||
-- <forward>: produce intermediate results.
|
||||
-- <optimize_parameters>: calculate losses, gradients, and update network weights.
|
||||
"""
|
||||
|
||||
def __init__(self, args):
|
||||
self.args = args
|
||||
self.is_train = args.is_train
|
||||
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
self.model_save_dir = os.path.join(args.save_dir, 'checkpoints') # save all the checkpoints to save_dir
|
||||
|
||||
if self.is_train and args.log:
|
||||
import time
|
||||
from tensorboardX import SummaryWriter
|
||||
from .loss_record import LossRecorder
|
||||
|
||||
log_dir = os.path.join(args.save_dir, 'logs')
|
||||
if args.epoch_begin != 0:
|
||||
all = [d for d in os.listdir(log_dir)]
|
||||
if len(all) == 0:
|
||||
raise RuntimeError(f'Empty logging path {log_dir}')
|
||||
timestamp = sorted(all)[-1]
|
||||
else:
|
||||
timestamp = time.strftime("%Y-%m-%d-%H:%M:%S", time.localtime(time.time()))
|
||||
self.log_path = os.path.join(log_dir, timestamp)
|
||||
os.makedirs(self.log_path, exist_ok=True)
|
||||
self.writer = SummaryWriter(self.log_path)
|
||||
self.loss_recorder = LossRecorder(self.writer)
|
||||
|
||||
self.epoch_cnt = 0
|
||||
self.schedulers = []
|
||||
self.optimizers = []
|
||||
self.models = []
|
||||
|
||||
@abstractmethod
|
||||
def set_input(self, input):
|
||||
"""Unpack input data from the dataloader and perform necessary pre-processing steps.
|
||||
Parameters:
|
||||
input (dict): includes the data itself and its metadata information.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def compute_test_result(self):
|
||||
"""
|
||||
After forward, do something like output bvh, get error value
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def forward(self):
|
||||
"""Run forward pass; called by both functions <optimize_parameters> and <test>."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def optimize_parameters(self):
|
||||
"""Calculate losses, gradients, and update network weights; called in every training iteration"""
|
||||
pass
|
||||
|
||||
def get_scheduler(self, optimizer):
|
||||
if self.args.scheduler.name == 'linear':
|
||||
def lambda_rule(epoch):
|
||||
lr_l = 1.0 - max(0, epoch - self.args.n_epochs_origin) / float(self.args.n_epochs_decay + 1)
|
||||
return lr_l
|
||||
return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule)
|
||||
if self.args.scheduler.name == 'StepLR':
|
||||
print('StepLR scheduler set')
|
||||
return torch.optim.lr_scheduler.StepLR(optimizer, self.args.scheduler.step_size, self.args.scheduler.gamma, verbose=self.args.verbose)
|
||||
if self.args.scheduler.name == 'Plateau':
|
||||
print('ReduceLROnPlateau shceduler set')
|
||||
return torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=self.args.scheduler.factor, threshold=self.args.scheduler.threshold, patience=5, verbose=True)
|
||||
if self.args.scheduler.name == 'MultiStep':
|
||||
print('MultiStep shceduler set')
|
||||
return torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=self.args.scheduler.milestones, gamma=self.args.scheduler.gamma, verbose=self.args.verbose)
|
||||
else:
|
||||
print('No scheduler set')
|
||||
return None
|
||||
|
||||
def setup(self):
|
||||
"""Load and print networks; create schedulers
|
||||
Parameters:
|
||||
opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
|
||||
"""
|
||||
if self.is_train:
|
||||
self.schedulers = [self.get_scheduler(optimizer) for optimizer in self.optimizers]
|
||||
|
||||
def epoch(self):
|
||||
if self.args.verbose:
|
||||
self.loss_recorder.epoch()
|
||||
for scheduler in self.schedulers:
|
||||
if scheduler is not None:
|
||||
if self.args.scheduler.name == 'Plateau':
|
||||
loss = self.loss_recorder.losses['total_loss'].current()
|
||||
scheduler.step(loss[1])
|
||||
else:
|
||||
scheduler.step()
|
||||
self.epoch_cnt += 1
|
||||
|
||||
def test(self):
|
||||
"""Forward function used in test time.
|
||||
This function wraps <forward> function in no_grad() so we don't save intermediate steps for backprop
|
||||
It also calls <compute_visuals> to produce additional visualization results
|
||||
"""
|
||||
with torch.no_grad():
|
||||
self.forward()
|
||||
self.compute_test_result()
|
||||
|
||||
def train(self):
|
||||
for model in self.models:
|
||||
model.train()
|
||||
for param in model.parameters():
|
||||
param.requires_grad = True
|
||||
|
||||
def eval(self):
|
||||
for model in self.models:
|
||||
model.eval()
|
||||
for param in model.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
def print(self):
|
||||
for model in self.models:
|
||||
print(model)
|
||||
|
||||
def count_params(self):
|
||||
params = 0
|
||||
for model in self.models:
|
||||
params += sum(p.numel() for p in model.parameters() if p.requires_grad)
|
||||
|
||||
return params
|
||||
+291
@@ -0,0 +1,291 @@
|
||||
import os
|
||||
|
||||
import holden.BVH as BVH
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from holden.Animation import Animation
|
||||
from holden.Quaternions import Quaternions
|
||||
from human_body_prior.tools.omni_tools import copy2cpu as c2c
|
||||
from rotations import matrix_to_quaternion, quaternion_to_axis_angle, rotation_6d_to_matrix
|
||||
from utils import align_joints, compute_trajectory
|
||||
|
||||
from .base_model import BaseModel
|
||||
from .fk import ForwardKinematicsLayer
|
||||
from .losses import GeodesicLoss
|
||||
from .neural_motion import NeuralMotionField
|
||||
|
||||
|
||||
class Architecture(BaseModel):
|
||||
def __init__(self, args, ngpu):
|
||||
super(Architecture, self).__init__(args)
|
||||
|
||||
self.args = args
|
||||
|
||||
if args.amass_data:
|
||||
self.fk = ForwardKinematicsLayer(args)
|
||||
else:
|
||||
self.fk = ForwardKinematicsLayer(parents=args.animal.parents, positions=args.animal.offsets)
|
||||
|
||||
self.field = NeuralMotionField(args.nemf).to(self.device)
|
||||
if args.multi_gpu is True:
|
||||
self.field = nn.DataParallel(self.field, device_ids=range(ngpu))
|
||||
|
||||
self.models = [self.field]
|
||||
|
||||
ordermap = {
|
||||
'x': 0,
|
||||
'y': 1,
|
||||
'z': 2,
|
||||
}
|
||||
self.up_axis = ordermap[self.args.data.up]
|
||||
self.v_axis = [x for x in ordermap.values() if x != self.up_axis]
|
||||
|
||||
print(self.up_axis)
|
||||
print(self.v_axis)
|
||||
|
||||
self.input_data = dict()
|
||||
self.recon_data = dict()
|
||||
|
||||
if self.is_train:
|
||||
self.optimizer = torch.optim.Adam(self.field.parameters(), args.learning_rate)
|
||||
self.optimizers = [self.optimizer]
|
||||
self.criterion_rec = nn.L1Loss() if args.l1_loss else nn.MSELoss()
|
||||
|
||||
self.criterion_geo = GeodesicLoss().to(self.device)
|
||||
self.fps = args.data.fps
|
||||
if args.bvh_viz:
|
||||
self.viz_dir = os.path.join(args.save_dir, 'results', 'bvh')
|
||||
else:
|
||||
self.viz_dir = os.path.join(args.save_dir, 'results', 'smpl')
|
||||
|
||||
def set_input(self, input):
|
||||
self.input_data = {k: v.float().to(self.device) for k, v in input.items() if k in ['pos', 'global_xform', 'root_orient', 'root_vel', 'trans']}
|
||||
self.input_data['rotmat'] = rotation_6d_to_matrix(self.input_data['global_xform'])
|
||||
|
||||
def forward(self, step=1):
|
||||
b_size, t_length, n_joints = self.input_data['pos'].shape[:3]
|
||||
|
||||
t = torch.arange(start=0, end=t_length, step=step).unsqueeze(0) # (1, T)
|
||||
t = t / t_length * 2 - 1 # map t to [-1, 1]
|
||||
t = t.expand(b_size, -1).unsqueeze(-1).to(self.device) # (B, T, 1)
|
||||
f, _ = self.field(t) # (B, T, D), rot6d + root orient + root vel + root height
|
||||
# import ipdb; ipdb.set_trace()
|
||||
self.recon_data = dict()
|
||||
rot6d_recon = f[:, :, :n_joints * 6] # (B, T, J x 6)
|
||||
rot6d_recon = rot6d_recon.view(b_size, -1, n_joints, 6) # (B, T, J, 6)
|
||||
rotmat_recon = rotation_6d_to_matrix(rot6d_recon) # (B, T, J, 3, 3)
|
||||
if self.args.data.root_transform:
|
||||
identity = torch.zeros(b_size, t.shape[1], 3, 3).to(self.device)
|
||||
identity[:, :, 0, 0] = 1
|
||||
identity[:, :, 1, 1] = 1
|
||||
identity[:, :, 2, 2] = 1
|
||||
rotmat_recon[:, :, 0] = identity
|
||||
# rotmat_recon[:, :, -2] = rotmat_recon[:, :, -4]
|
||||
# rotmat_recon[:, :, -1] = rotmat_recon[:, :, -3]
|
||||
local_rotmat = self.fk.global_to_local(rotmat_recon.view(-1, n_joints, 3, 3)) # (B x T, J, 3. 3)
|
||||
pos_recon, _ = self.fk(local_rotmat) # (B x T, J, 3)
|
||||
pos_recon = pos_recon.contiguous().view(b_size, -1, n_joints, 3) # (B, T, J, 3)
|
||||
self.recon_data['rotmat'] = rotmat_recon
|
||||
self.recon_data['pos'] = pos_recon
|
||||
|
||||
if self.args.data.root_transform:
|
||||
root_orient = f[:, :, n_joints * 6:n_joints * 6 + 6] # (B, T, 6)
|
||||
self.recon_data['root_orient'] = root_orient
|
||||
else:
|
||||
self.recon_data['root_orient'] = rot6d_recon[:, :, 0] # (B, T, 6)
|
||||
|
||||
root_vel = f[:, :, -4:-1] # (B, T, 3)
|
||||
root_height = f[:, :, -1] # (B, T)
|
||||
|
||||
self.recon_data['root_vel'] = root_vel
|
||||
self.recon_data['root_height'] = root_height
|
||||
|
||||
dt = 1.0 / self.args.data.fps * step
|
||||
origin = torch.zeros(b_size, 3).to(self.device)
|
||||
trans = compute_trajectory(root_vel, root_height, origin, dt, up_axis=self.args.data.up)
|
||||
self.recon_data['trans'] = trans
|
||||
|
||||
def backward(self, validation=False):
|
||||
loss = 0
|
||||
|
||||
root_orient = rotation_6d_to_matrix(self.recon_data['root_orient'])
|
||||
root_orient_gt = rotation_6d_to_matrix(self.input_data['root_orient'])
|
||||
if self.args.geodesic_loss:
|
||||
orient_recon_loss = self.criterion_geo(root_orient.view(-1, 3, 3), root_orient_gt.view(-1, 3, 3))
|
||||
else:
|
||||
orient_recon_loss = self.criterion_rec(root_orient, root_orient_gt)
|
||||
self.loss_recorder.add_scalar('orient_recon_loss', orient_recon_loss, validation=validation)
|
||||
loss += self.args.lambda_orient * orient_recon_loss
|
||||
|
||||
if self.args.geodesic_loss:
|
||||
rotmat_recon_loss = self.criterion_geo(self.recon_data['rotmat'].view(-1, 3, 3), self.input_data['rotmat'].view(-1, 3, 3))
|
||||
else:
|
||||
rotmat_recon_loss = self.criterion_rec(self.recon_data['rotmat'], self.input_data['rotmat'])
|
||||
self.loss_recorder.add_scalar('rotmat_recon_loss', rotmat_recon_loss, validation=validation)
|
||||
loss += self.args.lambda_rotmat * rotmat_recon_loss
|
||||
|
||||
pos_recon_loss = self.criterion_rec(self.recon_data['pos'], self.input_data['pos'])
|
||||
self.loss_recorder.add_scalar('pos_recon_loss', pos_recon_loss, validation=validation)
|
||||
loss += self.args.lambda_pos * pos_recon_loss
|
||||
|
||||
v_loss = self.criterion_rec(self.recon_data['root_vel'], self.input_data['root_vel'])
|
||||
up_loss = self.criterion_rec(self.recon_data['root_height'], self.input_data['trans'][..., 'xyz'.index(self.args.data.up)])
|
||||
self.loss_recorder.add_scalar('v_loss', v_loss, validation=validation)
|
||||
self.loss_recorder.add_scalar('up_loss', up_loss, validation=validation)
|
||||
loss += self.args.lambda_v * v_loss + self.args.lambda_up * up_loss
|
||||
|
||||
origin = self.input_data['trans'][:, 0]
|
||||
trans = self.recon_data['trans']
|
||||
trans[..., self.v_axis] = trans[..., self.v_axis] + origin[..., self.v_axis].unsqueeze(1)
|
||||
trans_loss = self.criterion_rec(trans, self.input_data['trans'])
|
||||
self.loss_recorder.add_scalar('trans_loss', trans_loss, validation=validation)
|
||||
loss += self.args.lambda_trans * trans_loss
|
||||
|
||||
self.loss_recorder.add_scalar('total_loss', loss, validation=validation)
|
||||
|
||||
if not validation:
|
||||
loss.backward()
|
||||
|
||||
def optimize_parameters(self):
|
||||
self.optimizer.zero_grad()
|
||||
self.forward()
|
||||
self.backward(validation=False)
|
||||
self.optimizer.step()
|
||||
|
||||
def report_errors(self):
|
||||
root_orient = rotation_6d_to_matrix(self.recon_data['root_orient']) # (B, T, 3, 3)
|
||||
root_orient_gt = rotation_6d_to_matrix(self.input_data['root_orient']) # (B, T, 3, 3)
|
||||
orientation_error = self.criterion_geo(root_orient.view(-1, 3, 3), root_orient_gt.view(-1, 3, 3))
|
||||
|
||||
local_rotmat = self.fk.global_to_local(self.recon_data['rotmat'].view(-1, self.args.smpl.joint_num, 3, 3)) # (B x T, J, 3, 3)
|
||||
local_rotmat_gt = self.fk.global_to_local(self.input_data['rotmat'].view(-1, self.args.smpl.joint_num, 3, 3)) # (B x T, J, 3, 3)
|
||||
rotation_error = self.criterion_geo(local_rotmat[:, 1:].reshape(-1, 3, 3), local_rotmat_gt[:, 1:].reshape(-1, 3, 3))
|
||||
|
||||
pos = self.recon_data['pos'] # (B, T, J, 3)
|
||||
pos_gt = self.input_data['pos'] # (B, T, J, 3)
|
||||
position_error = torch.linalg.norm((pos - pos_gt), dim=-1).mean()
|
||||
|
||||
origin = self.input_data['trans'][:, 0]
|
||||
trans = self.recon_data['trans']
|
||||
trans[..., self.v_axis] = trans[..., self.v_axis] + origin[..., self.v_axis].unsqueeze(1)
|
||||
trans_gt = self.input_data['trans'] # (B, T, 3)
|
||||
translation_error = torch.linalg.norm((trans - trans_gt), dim=-1).mean()
|
||||
|
||||
return {
|
||||
'rotation': c2c(rotation_error),
|
||||
'position': c2c(position_error),
|
||||
'orientation': c2c(orientation_error),
|
||||
'translation': c2c(translation_error)
|
||||
}
|
||||
|
||||
def verbose(self):
|
||||
res = {}
|
||||
for loss in self.loss_recorder.losses.values():
|
||||
res[loss.name] = {'train': loss.current()[0]}
|
||||
|
||||
return res
|
||||
|
||||
def save(self, optimal=False):
|
||||
if optimal:
|
||||
path = os.path.join(self.args.save_dir, 'results', 'model')
|
||||
else:
|
||||
path = os.path.join(self.model_save_dir, f'{self.epoch_cnt:04d}')
|
||||
|
||||
os.makedirs(path, exist_ok=True)
|
||||
nemf = self.field.module if isinstance(self.field, nn.DataParallel) else self.field
|
||||
torch.save(nemf.state_dict(), os.path.join(path, 'nemf.pth'))
|
||||
torch.save(self.optimizer.state_dict(), os.path.join(path, 'optimizer.pth'))
|
||||
if self.args.scheduler.name:
|
||||
torch.save(self.schedulers[0].state_dict(), os.path.join(path, 'scheduler.pth'))
|
||||
self.loss_recorder.save(path)
|
||||
|
||||
print(f'Save at {path} succeeded')
|
||||
|
||||
def load(self, epoch=None, optimal=False):
|
||||
if optimal:
|
||||
path = os.path.join(self.args.save_dir, 'results', 'model')
|
||||
else:
|
||||
if epoch is None:
|
||||
all = [int(q) for q in os.listdir(self.model_save_dir)]
|
||||
if len(all) == 0:
|
||||
raise RuntimeError(f'Empty loading path {self.model_save_dir}')
|
||||
epoch = sorted(all)[-1]
|
||||
path = os.path.join(self.model_save_dir, f'{epoch:04d}')
|
||||
|
||||
print(f'Loading from {path}')
|
||||
nemf = self.field.module if isinstance(self.field, nn.DataParallel) else self.field
|
||||
nemf.load_state_dict(torch.load(os.path.join(path, 'nemf.pth'), map_location=self.device))
|
||||
if self.is_train:
|
||||
self.optimizer.load_state_dict(torch.load(os.path.join(path, 'optimizer.pth')))
|
||||
self.loss_recorder.load(path)
|
||||
if self.args.scheduler.name:
|
||||
self.schedulers[0].load_state_dict(torch.load(os.path.join(path, 'scheduler.pth')))
|
||||
self.epoch_cnt = epoch if not optimal else 0
|
||||
|
||||
print('Load succeeded')
|
||||
|
||||
def super_sampling(self, step=1):
|
||||
with torch.no_grad():
|
||||
self.forward(step)
|
||||
self.fps = self.args.data.fps / step
|
||||
self.compute_test_result()
|
||||
|
||||
def compute_test_result(self):
|
||||
os.makedirs(self.viz_dir, exist_ok=True)
|
||||
|
||||
rotmat = self.recon_data['rotmat'][0] # (T, J, 3, 3)
|
||||
rotmat_gt = self.input_data['rotmat'][0] # (T, J, 3, 3)
|
||||
local_rotmat = self.fk.global_to_local(rotmat) # (T, J, 3, 3)
|
||||
local_rotmat_gt = self.fk.global_to_local(rotmat_gt) # (T, J, 3, 3)
|
||||
|
||||
if self.args.data.root_transform:
|
||||
root_orient = rotation_6d_to_matrix(self.recon_data['root_orient'][0]) # (T, 3, 3)
|
||||
root_orient_gt = rotation_6d_to_matrix(self.input_data['root_orient'][0]) # (T, 3, 3)
|
||||
local_rotmat[:, 0] = root_orient
|
||||
local_rotmat_gt[:, 0] = root_orient_gt
|
||||
rotation = matrix_to_quaternion(local_rotmat) # (T, J, 4)
|
||||
rotation_gt = matrix_to_quaternion(local_rotmat_gt) # (T, J, 4)
|
||||
|
||||
origin = self.input_data['trans'][0, 0].clone() # [3]
|
||||
position = self.recon_data['trans'][0].clone() # (T, 3)
|
||||
position[:, self.v_axis] = position[:, self.v_axis] + origin[self.v_axis].unsqueeze(0)
|
||||
position_gt = self.input_data['trans'][0] # (T, 3)
|
||||
|
||||
if self.args.bvh_viz:
|
||||
if self.args.amass_data:
|
||||
rotation = align_joints(rotation)
|
||||
rotation_gt = align_joints(rotation_gt)
|
||||
|
||||
position = position.unsqueeze(1) # (T, 1, 3)
|
||||
position_gt = position_gt.unsqueeze(1) # (T, 1, 3)
|
||||
|
||||
# export data to BVH files
|
||||
if self.args.amass_data:
|
||||
offsets = self.args.smpl.offsets[self.args.data.gender]
|
||||
parents = self.args.smpl.parents
|
||||
names = self.args.smpl.joint_names
|
||||
else:
|
||||
offsets = self.args.animal.offsets
|
||||
parents = self.args.animal.parents
|
||||
names = self.args.animal.joint_names
|
||||
|
||||
anim = Animation(Quaternions(c2c(rotation)), c2c(position), None, offsets=offsets, parents=parents)
|
||||
BVH.save(os.path.join(self.viz_dir, f'test_{int(self.fps)}fps.bvh'), anim, names=names, frametime=1 / self.fps)
|
||||
|
||||
anim.rotations = Quaternions(c2c(rotation_gt))
|
||||
anim.positions = c2c(position_gt)
|
||||
BVH.save(os.path.join(self.viz_dir, 'test_gt.bvh'), anim, names=names, frametime=1 / self.args.data.fps)
|
||||
else:
|
||||
poses = c2c(quaternion_to_axis_angle(rotation)) # (T, J, 3)
|
||||
poses_gt = c2c(quaternion_to_axis_angle(rotation_gt)) # (T, J, 3)
|
||||
|
||||
poses = poses.reshape((poses.shape[0], -1)) # (T, J x 3)
|
||||
poses = np.pad(poses, [(0, 0), (0, 93)], mode='constant')
|
||||
poses_gt = poses_gt.reshape((poses_gt.shape[0], -1)) # (T, J x 3)
|
||||
poses_gt = np.pad(poses_gt, [(0, 0), (0, 93)], mode='constant')
|
||||
|
||||
np.savez(os.path.join(self.viz_dir, f'test_{int(self.fps)}fps.npz'),
|
||||
poses=poses, trans=c2c(position), betas=np.zeros(10), gender=self.args.data.gender, mocap_framerate=self.fps)
|
||||
np.savez(os.path.join(self.viz_dir, 'test_gt.npz'),
|
||||
poses=poses_gt, trans=c2c(position_gt), betas=np.zeros(10), gender=self.args.data.gender, mocap_framerate=self.args.data.fps)
|
||||
+152
@@ -0,0 +1,152 @@
|
||||
"""Based on Daniel Holden code from:
|
||||
A Deep Learning Framework for Character Motion Synthesis and Editing
|
||||
(http://www.ipab.inf.ed.ac.uk/cgvu/motionsynthesis.pdf)
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from rotations import euler_angles_to_matrix, quaternion_to_matrix, rotation_6d_to_matrix
|
||||
|
||||
|
||||
class ForwardKinematicsLayer(nn.Module):
|
||||
""" Forward Kinematics Layer Class """
|
||||
|
||||
def __init__(self, args=None, parents=None, positions=None, device=None):
|
||||
super().__init__()
|
||||
self.b_idxs = None
|
||||
if device is None:
|
||||
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
else:
|
||||
self.device = device
|
||||
if parents is None and positions is None:
|
||||
# Load SMPL skeleton (their joint order is different from the one we use for bvh export)
|
||||
# smpl_fname = os.path.join(args.smpl.smpl_body_model, args.data.gender, 'model.npz')
|
||||
smpl_fname = os.path.join(args.smpl.smpl_body_model, "MANO_RIGHT.pkl")
|
||||
if not os.path.exists(smpl_fname):
|
||||
raise FileNotFoundError(f'SMPL file {os.path.abspath(smpl_fname)} does not exist. Please check the path., pwd: {os.getcwd()}')
|
||||
smpl_data = np.load(smpl_fname, encoding='latin1', allow_pickle=True)
|
||||
self.parents = torch.from_numpy(smpl_data['kintree_table'][0].astype(np.int32)).to(self.device)
|
||||
self.parents = self.parents.long()
|
||||
self.positions = torch.from_numpy(smpl_data['J'].astype(np.float32)).to(self.device)
|
||||
self.positions[1:] -= self.positions[self.parents[1:]]
|
||||
else:
|
||||
self.parents = torch.from_numpy(parents).to(self.device)
|
||||
self.parents = self.parents.long()
|
||||
self.positions = torch.from_numpy(positions).to(self.device)
|
||||
self.positions = self.positions.float()
|
||||
self.positions[0] = 0
|
||||
|
||||
def rotate(self, t0s, t1s):
|
||||
return torch.matmul(t0s, t1s)
|
||||
|
||||
def identity_rotation(self, rotations):
|
||||
diagonal = torch.diag(torch.tensor([1.0, 1.0, 1.0, 1.0])).to(self.device)
|
||||
diagonal = torch.reshape(
|
||||
diagonal, torch.Size([1] * len(rotations.shape[:2]) + [4, 4]))
|
||||
ts = diagonal.repeat(rotations.shape[:2] + torch.Size([1, 1]))
|
||||
return ts
|
||||
|
||||
def make_fast_rotation_matrices(self, positions, rotations):
|
||||
if len(rotations.shape) == 4 and rotations.shape[-2:] == torch.Size([3, 3]):
|
||||
rot_matrices = rotations
|
||||
elif rotations.shape[-1] == 3:
|
||||
rot_matrices = euler_angles_to_matrix(rotations, convention='XYZ')
|
||||
elif rotations.shape[-1] == 4:
|
||||
rot_matrices = quaternion_to_matrix(rotations)
|
||||
elif rotations.shape[-1] == 6:
|
||||
rot_matrices = rotation_6d_to_matrix(rotations)
|
||||
else:
|
||||
raise NotImplementedError(f'Unimplemented rotation representation in FK layer, shape of {rotations.shape}')
|
||||
|
||||
rot_matrices = torch.cat([rot_matrices, positions[..., None]], dim=-1)
|
||||
zeros = torch.zeros(rot_matrices.shape[:-2] + torch.Size([1, 3])).to(self.device)
|
||||
ones = torch.ones(rot_matrices.shape[:-2] + torch.Size([1, 1])).to(self.device)
|
||||
zerosones = torch.cat([zeros, ones], dim=-1)
|
||||
rot_matrices = torch.cat([rot_matrices, zerosones], dim=-2)
|
||||
return rot_matrices
|
||||
|
||||
def rotate_global(self, parents, positions, rotations):
|
||||
locals = self.make_fast_rotation_matrices(positions, rotations)
|
||||
globals = self.identity_rotation(rotations)
|
||||
|
||||
globals = torch.cat([locals[:, 0:1], globals[:, 1:]], dim=1)
|
||||
b_size = positions.shape[0]
|
||||
if self.b_idxs is None:
|
||||
self.b_idxs = torch.LongTensor(np.arange(b_size)).to(self.device)
|
||||
elif self.b_idxs.shape[-1] != b_size:
|
||||
self.b_idxs = torch.LongTensor(np.arange(b_size)).to(self.device)
|
||||
|
||||
for i in range(1, positions.shape[1]):
|
||||
globals[:, i] = self.rotate(
|
||||
globals[self.b_idxs, parents[i]], locals[:, i])
|
||||
|
||||
return globals
|
||||
|
||||
def get_tpose_joints(self, offsets, parents):
|
||||
num_joints = len(parents)
|
||||
joints = [offsets[:, 0]]
|
||||
for j in range(1, len(parents)):
|
||||
joints.append(joints[parents[j]] + offsets[:, j])
|
||||
|
||||
return torch.stack(joints, dim=1)
|
||||
|
||||
def canonical_to_local(self, canonical_xform, global_orient=None):
|
||||
"""
|
||||
Args:
|
||||
canonical_xform: (B, J, 3, 3)
|
||||
global_orient: (B, 3, 3)
|
||||
|
||||
Returns:
|
||||
local_xform: (B, J, 3, 3)
|
||||
"""
|
||||
local_xform = torch.zeros_like(canonical_xform)
|
||||
|
||||
if global_orient is None:
|
||||
global_xform = canonical_xform
|
||||
else:
|
||||
global_xform = torch.matmul(global_orient.unsqueeze(1), canonical_xform)
|
||||
for i in range(global_xform.shape[1]):
|
||||
if i == 0:
|
||||
local_xform[:, i] = global_xform[:, i]
|
||||
else:
|
||||
local_xform[:, i] = torch.bmm(torch.linalg.inv(global_xform[:, self.parents[i]]), global_xform[:, i])
|
||||
|
||||
return local_xform
|
||||
|
||||
def global_to_local(self, global_xform):
|
||||
"""
|
||||
Args:
|
||||
global_xform: (B, J, 3, 3)
|
||||
|
||||
Returns:
|
||||
local_xform: (B, J, 3, 3)
|
||||
"""
|
||||
local_xform = torch.zeros_like(global_xform)
|
||||
|
||||
for i in range(global_xform.shape[1]):
|
||||
if i == 0:
|
||||
local_xform[:, i] = global_xform[:, i]
|
||||
else:
|
||||
local_xform[:, i] = torch.bmm(torch.linalg.inv(global_xform[:, self.parents[i]]), global_xform[:, i])
|
||||
|
||||
return local_xform
|
||||
|
||||
def forward(self, rotations, positions=None):
|
||||
"""
|
||||
Args:
|
||||
rotations (B, J, D)
|
||||
|
||||
Returns:
|
||||
The global position of each joint after FK (B, J, 3)
|
||||
"""
|
||||
# Get the full transform with rotations for skinning
|
||||
b_size = rotations.shape[0]
|
||||
if positions is None:
|
||||
positions = self.positions.repeat(b_size, 1, 1)
|
||||
transforms = self.rotate_global(self.parents, positions, rotations)
|
||||
coordinates = transforms[:, :, :3, 3] / transforms[:, :, 3:, 3]
|
||||
|
||||
return coordinates, transforms
|
||||
@@ -0,0 +1,525 @@
|
||||
import os
|
||||
import yaml
|
||||
|
||||
import holden.BVH as BVH
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from holden.Animation import Animation
|
||||
from holden.Quaternions import Quaternions
|
||||
from human_body_prior.tools.omni_tools import copy2cpu as c2c
|
||||
from rotations import matrix_to_quaternion, matrix_to_rotation_6d, quaternion_to_axis_angle, rotation_6d_to_matrix
|
||||
from utils import align_joints, compute_trajectory, estimate_angular_velocity, estimate_linear_velocity, normalize
|
||||
|
||||
from .base_model import BaseModel
|
||||
from .fk import ForwardKinematicsLayer
|
||||
from .global_motion import GlobalMotionPredictor
|
||||
from .losses import GeodesicLoss
|
||||
from .neural_motion import NeuralMotionField
|
||||
from .prior import GlobalEncoder, LocalEncoder
|
||||
from .skeleton import build_edge_topology
|
||||
|
||||
|
||||
|
||||
class Struct:
|
||||
def __init__(self, **entries):
|
||||
self.__dict__.update(entries)
|
||||
|
||||
class Arguments:
|
||||
def __init__(self, config_path, filename='default.yaml'):
|
||||
with open(os.path.join(config_path, 'mano.yaml'), 'r') as f:
|
||||
smpl = yaml.safe_load(f)
|
||||
|
||||
self.smpl = Struct(**smpl)
|
||||
self.smpl.offsets['right'] = np.array(self.smpl.offsets['right'])
|
||||
self.smpl.offsets['left'] = np.array(self.smpl.offsets['left'])
|
||||
self.smpl.parents = np.array(self.smpl.parents).astype(np.int32)
|
||||
self.smpl.joint_num = 16 # len(self.smpl.joints_to_use)
|
||||
self.smpl.joints_to_use = np.array(self.smpl.joints_to_use)
|
||||
self.smpl.joints_to_use = np.arange(0, 63).reshape((-1, 3))[self.smpl.joints_to_use].reshape(-1)
|
||||
|
||||
self.filename = os.path.splitext(filename)[0]
|
||||
with open(os.path.join(config_path, filename), 'r') as f:
|
||||
config = yaml.safe_load(f)
|
||||
|
||||
for key, value in config.items():
|
||||
if isinstance(value, dict):
|
||||
setattr(self, key, Struct(**value))
|
||||
else:
|
||||
setattr(self, key, value)
|
||||
|
||||
self.json = config
|
||||
|
||||
|
||||
class Architecture(BaseModel):
|
||||
def __init__(self, args, ngpu, batch_num=0):
|
||||
super(Architecture, self).__init__(args)
|
||||
|
||||
self.batch_num = batch_num
|
||||
self.args = args
|
||||
#from pprint import pprint
|
||||
#pprint(vars(self.args))
|
||||
|
||||
self.fk = ForwardKinematicsLayer(args)
|
||||
|
||||
smpl_fname = os.path.join(args.smpl.smpl_body_model, "MANO_RIGHT.pkl")
|
||||
smpl_data = np.load(smpl_fname, allow_pickle=True, encoding='latin1')
|
||||
parents = smpl_data['kintree_table'][0].astype(np.int32)
|
||||
edges = build_edge_topology(parents)
|
||||
|
||||
self.local_encoder = LocalEncoder(args.local_prior, edges).to(self.device)
|
||||
self.global_encoder = GlobalEncoder(args.global_prior).to(self.device)
|
||||
self.field = NeuralMotionField(args.nemf).to(self.device)
|
||||
|
||||
if args.multi_gpu is True:
|
||||
self.local_encoder = nn.DataParallel(self.local_encoder, device_ids=range(ngpu))
|
||||
self.global_encoder = nn.DataParallel(self.global_encoder, device_ids=range(ngpu))
|
||||
self.field = nn.DataParallel(self.field, device_ids=range(ngpu))
|
||||
|
||||
self.models = [self.local_encoder, self.field]
|
||||
if args.data.root_transform:
|
||||
self.models.append(self.global_encoder)
|
||||
|
||||
ordermap = {
|
||||
'x': 0,
|
||||
'y': 1,
|
||||
'z': 2,
|
||||
}
|
||||
self.up_axis = ordermap[self.args.data.up]
|
||||
self.v_axis = [x for x in ordermap.values() if x != self.up_axis]
|
||||
|
||||
print(self.up_axis)
|
||||
print(self.v_axis)
|
||||
|
||||
self.input_data = dict()
|
||||
self.recon_data = dict()
|
||||
print(os.getcwd(), '...', os.path.join(args.dataset_dir, f'mean-{args.data.gender}-{args.data.clip_length}-{args.data.fps}fps.pt'))
|
||||
self.mean = torch.load(os.path.join(args.dataset_dir, f'mean-{args.data.gender}-{args.data.clip_length}-{args.data.fps}fps.pt'), map_location=self.device)
|
||||
print(f'mean: {list(self.mean.keys())}')
|
||||
self.std = torch.load(os.path.join(args.dataset_dir, f'std-{args.data.gender}-{args.data.clip_length}-{args.data.fps}fps.pt'), map_location=self.device)
|
||||
print(f'std: {list(self.std.keys())}')
|
||||
|
||||
#for k, v in self.mean.items():
|
||||
# print(k, v.shape)
|
||||
|
||||
if args.pretrained_gmp:
|
||||
assert args.output_trans is False, 'already output global translation in this model'
|
||||
gmp_args = Arguments('./configs', filename=args.pretrained_gmp)
|
||||
self.gmp = GlobalMotionPredictor(gmp_args, ngpu)
|
||||
self.gmp.load(optimal=True)
|
||||
self.gmp.eval()
|
||||
|
||||
self.criterion_geo = GeodesicLoss().to(self.device)
|
||||
if self.is_train:
|
||||
parameters = list(self.local_encoder.parameters()) + list(self.field.parameters())
|
||||
if args.data.root_transform:
|
||||
parameters += list(self.global_encoder.parameters())
|
||||
|
||||
if args.adam_optimizer:
|
||||
self.optimizer = torch.optim.Adam(parameters, args.learning_rate, weight_decay=args.weight_decay)
|
||||
else:
|
||||
self.optimizer = torch.optim.Adamax(parameters, args.learning_rate, weight_decay=args.weight_decay)
|
||||
self.optimizers = [self.optimizer]
|
||||
|
||||
self.criterion_rec = nn.L1Loss() if args.l1_loss else nn.MSELoss()
|
||||
self.criterion_bce = nn.BCEWithLogitsLoss()
|
||||
|
||||
self.iteration = 0
|
||||
else:
|
||||
self.fps = args.data.fps
|
||||
self.test_index = 0
|
||||
if args.bvh_viz:
|
||||
self.viz_dir = os.path.join(args.save_dir, 'results', 'bvh')
|
||||
else:
|
||||
self.viz_dir = os.path.join(args.save_dir, 'results', 'smpl')
|
||||
|
||||
def set_input(self, input):
|
||||
self.input_data = {k: v.float().to(self.device) for k, v in input.items() if k in [
|
||||
'pos', 'velocity', 'global_xform', 'angular', 'root_orient', 'root_vel', 'trans', 'contacts'
|
||||
]}
|
||||
self.input_data['rotmat'] = rotation_6d_to_matrix(self.input_data['global_xform'])
|
||||
|
||||
def reparameterize(self, mu, logvar):
|
||||
eps = torch.randn_like(mu)
|
||||
std = torch.exp(0.5 * logvar)
|
||||
return mu + eps * std
|
||||
|
||||
def encode_local(self):
|
||||
b_size, t_length = self.input_data['pos'].shape[:2]
|
||||
|
||||
if 'pos' in self.args.data.normalize:
|
||||
pos = normalize(self.input_data['pos'], mean=self.mean['pos'], std=self.std['pos'])
|
||||
else:
|
||||
pos = self.input_data['pos']
|
||||
if 'velocity' in self.args.data.normalize:
|
||||
velocity = normalize(self.input_data['velocity'], mean=self.mean['velocity'], std=self.std['velocity'])
|
||||
else:
|
||||
velocity = self.input_data['velocity']
|
||||
if 'global_xform' in self.args.data.normalize:
|
||||
global_xform = normalize(self.input_data['global_xform'], mean=self.mean['global_xform'], std=self.std['global_xform'])
|
||||
else:
|
||||
global_xform = self.input_data['global_xform']
|
||||
if 'angular' in self.args.data.normalize:
|
||||
angular = normalize(self.input_data['angular'], mean=self.mean['angular'], std=self.std['angular'])
|
||||
else:
|
||||
angular = self.input_data['angular']
|
||||
|
||||
x = torch.cat((pos, velocity, global_xform, angular), dim=-1)
|
||||
x = x.view(b_size, t_length, -1) # (B, T, J x D)
|
||||
x = x.permute(0, 2, 1) # (B, J x D, T)
|
||||
|
||||
mu, logvar = self.local_encoder(x) # (B, D)
|
||||
if self.args.lambda_kl != 0:
|
||||
z = self.reparameterize(mu, logvar)
|
||||
else:
|
||||
z = mu
|
||||
|
||||
return z, mu, logvar
|
||||
|
||||
def encode_global(self):
|
||||
if 'root_orient' in self.args.data.normalize:
|
||||
root_orient = normalize(self.input_data['root_orient'], mean=self.mean['root_orient'], std=self.std['root_orient'])
|
||||
else:
|
||||
root_orient = self.input_data['root_orient']
|
||||
if self.args.global_prior.in_channels == 9:
|
||||
if 'root_vel' in self.args.data.normalize:
|
||||
root_vel = normalize(self.input_data['root_vel'], mean=self.mean['root_vel'], std=self.std['root_vel'])
|
||||
else:
|
||||
root_vel = self.input_data['root_vel']
|
||||
x = torch.cat((root_orient, root_vel), dim=-1) # (B, T, D)
|
||||
elif self.args.global_prior.in_channels == 6:
|
||||
x = root_orient
|
||||
else:
|
||||
raise RuntimeError('unsupported input dimensions for global prior')
|
||||
x = x.permute(0, 2, 1) # (B, D, T)
|
||||
|
||||
mu, logvar = self.global_encoder(x) # (B, D)
|
||||
if self.args.lambda_kl != 0:
|
||||
z = self.reparameterize(mu, logvar)
|
||||
else:
|
||||
z = mu
|
||||
|
||||
return z, mu, logvar
|
||||
|
||||
def forward(self, step=1):
|
||||
z_l, mu_l, logvar_l = self.encode_local() # (B, D)
|
||||
if self.args.data.root_transform:
|
||||
z_g, mu_g, logvar_g = self.encode_global() # (B, D)
|
||||
else:
|
||||
z_g = None
|
||||
|
||||
self.recon_data.clear()
|
||||
self.recon_data = self.decode(z_l, z_g, length=self.args.data.clip_length, step=step)
|
||||
|
||||
self.recon_data['mu_l'] = mu_l
|
||||
self.recon_data['logvar_l'] = logvar_l
|
||||
if self.args.data.root_transform:
|
||||
self.recon_data['mu_g'] = mu_g
|
||||
self.recon_data['logvar_g'] = logvar_g
|
||||
|
||||
def decode(self, z_l, z_g, length, step=1):
|
||||
b_size = z_l.shape[0]
|
||||
n_joints = self.args.smpl.joint_num
|
||||
n_joints = 16
|
||||
|
||||
# generate f(t; z)
|
||||
t = torch.arange(start=0, end=length, step=step).unsqueeze(0) # (1, T)
|
||||
t = t / self.args.data.clip_length * 2 - 1 # map t to [-1, 1]
|
||||
t = t.expand(b_size, -1).unsqueeze(-1).to(self.device) # (B, T, 1)
|
||||
local_motion, global_motion = self.field(t, z_l, z_g) # (B, T, N)
|
||||
|
||||
rot6d_recon = local_motion[:, :, :n_joints * 6].contiguous() # (B, T, J x 6)
|
||||
rot6d_recon = rot6d_recon.view(b_size, -1, n_joints, 6) # (B, T, J, 6)
|
||||
rotmat_recon = rotation_6d_to_matrix(rot6d_recon) # (B, T, J, 3, 3)
|
||||
|
||||
if self.args.data.root_transform:
|
||||
identity = torch.zeros(b_size, t.shape[1], 3, 3).to(self.device)
|
||||
identity[:, :, 0, 0] = 1
|
||||
identity[:, :, 1, 1] = 1
|
||||
identity[:, :, 2, 2] = 1
|
||||
rotmat_recon[:, :, 0] = identity
|
||||
# rotmat_recon[:, :, -2] = rotmat_recon[:, :, -4]
|
||||
# rotmat_recon[:, :, -1] = rotmat_recon[:, :, -3]
|
||||
local_rotmat = self.fk.global_to_local(rotmat_recon.view(-1, n_joints, 3, 3)) # (B x T, J, 3. 3)
|
||||
pos_recon, _ = self.fk(local_rotmat) # (B x T, J, 3)
|
||||
pos_recon = pos_recon.contiguous().view(b_size, -1, n_joints, 3) # (B, T, J, 3)
|
||||
|
||||
output = dict()
|
||||
output['rotmat'] = rotmat_recon
|
||||
output['pos'] = pos_recon
|
||||
|
||||
if self.args.data.root_transform:
|
||||
root_orient = global_motion[:, :, :6] # (B, T, 6)
|
||||
output['root_orient'] = root_orient
|
||||
else:
|
||||
output['root_orient'] = rot6d_recon[:, :, 0] # (B, T, 6)
|
||||
|
||||
dt = 1.0 / self.args.data.fps * step
|
||||
if self.args.output_trans:
|
||||
root_vel = global_motion[:, :, 6:9] # (B, T, 3)
|
||||
root_height = global_motion[:, :, 9] # (B, T)
|
||||
|
||||
output['root_vel'] = root_vel
|
||||
output['root_height'] = root_height
|
||||
|
||||
contacts = global_motion[:, :, 10:] # (B, T, 8)
|
||||
output['contacts'] = contacts
|
||||
|
||||
origin = torch.zeros(b_size, 3).to(self.device)
|
||||
trans = compute_trajectory(root_vel, root_height, origin, dt, up_axis=self.args.data.up)
|
||||
output['trans'] = trans
|
||||
elif self.args.pretrained_gmp:
|
||||
local_rotmat = local_rotmat.view(b_size, -1, n_joints, 3, 3) # (B, T, J, 3, 3)
|
||||
if self.args.data.root_transform:
|
||||
root_orient_rotmat = rotation_6d_to_matrix(output['root_orient']) # (B, T, 3, 3)
|
||||
local_rotmat[:, :, 0] = root_orient_rotmat
|
||||
pos = torch.matmul(root_orient_rotmat.unsqueeze(2), pos_recon.unsqueeze(-1)).squeeze(-1)
|
||||
else:
|
||||
pos = pos_recon
|
||||
|
||||
gmp_data = dict()
|
||||
gmp_data['rot6d'] = matrix_to_rotation_6d(local_rotmat) # (B, T, J, 6)
|
||||
gmp_data['angular'] = estimate_angular_velocity(local_rotmat.clone(), dt=dt) # (B, T, J, 3)
|
||||
|
||||
gmp_data['pos'] = pos
|
||||
gmp_data['velocity'] = estimate_linear_velocity(pos, dt=dt)
|
||||
|
||||
pred_data = self.gmp.predict(gmp_data, dt=dt)
|
||||
|
||||
output |= pred_data # merge two dicts
|
||||
|
||||
return output
|
||||
|
||||
def kl_scheduler(self):
|
||||
"""
|
||||
Cyclical Annealing Schedule
|
||||
"""
|
||||
if (self.epoch_cnt % self.args.annealing_cycles == 0) and (self.iteration % self.batch_num == 0):
|
||||
print('KL annealing restart')
|
||||
return 0.01
|
||||
|
||||
return min(1, 0.01 + self.iteration / (self.args.annealing_warmup * self.batch_num))
|
||||
|
||||
def backward(self, validation=False):
|
||||
loss = 0
|
||||
|
||||
root_orient = rotation_6d_to_matrix(self.recon_data['root_orient'])
|
||||
root_orient_gt = rotation_6d_to_matrix(self.input_data['root_orient'])
|
||||
|
||||
if self.args.geodesic_loss:
|
||||
orient_recon_loss = self.criterion_geo(root_orient.view(-1, 3, 3), root_orient_gt.view(-1, 3, 3))
|
||||
else:
|
||||
orient_recon_loss = self.criterion_rec(root_orient, root_orient_gt)
|
||||
self.loss_recorder.add_scalar('orient_recon_loss', orient_recon_loss, validation=validation)
|
||||
loss += self.args.lambda_orient * orient_recon_loss
|
||||
|
||||
if self.args.geodesic_loss:
|
||||
rotmat_recon_loss = self.criterion_geo(self.recon_data['rotmat'].view(-1, 3, 3), self.input_data['rotmat'].view(-1, 3, 3))
|
||||
else:
|
||||
rotmat_recon_loss = self.criterion_rec(self.recon_data['rotmat'], self.input_data['rotmat'])
|
||||
self.loss_recorder.add_scalar('rotmat_recon_loss', rotmat_recon_loss, validation=validation)
|
||||
loss += self.args.lambda_rotmat * rotmat_recon_loss
|
||||
|
||||
pos_recon_loss = self.criterion_rec(self.recon_data['pos'], self.input_data['pos'])
|
||||
self.loss_recorder.add_scalar('pos_recon_loss', pos_recon_loss, validation=validation)
|
||||
loss += self.args.lambda_pos * pos_recon_loss
|
||||
|
||||
if self.args.output_trans:
|
||||
v_loss = self.criterion_rec(self.recon_data['root_vel'], self.input_data['root_vel'])
|
||||
up_loss = self.criterion_rec(self.recon_data['root_height'], self.input_data['trans'][..., 'xyz'.index(self.args.data.up)])
|
||||
self.loss_recorder.add_scalar('v_loss', v_loss, validation=validation)
|
||||
self.loss_recorder.add_scalar('up_loss', up_loss, validation=validation)
|
||||
loss += self.args.lambda_v * v_loss + self.args.lambda_up * up_loss
|
||||
|
||||
origin = self.input_data['trans'][:, 0]
|
||||
trans = self.recon_data['trans']
|
||||
trans[..., self.v_axis] = trans[..., self.v_axis] + origin[..., self.v_axis].unsqueeze(1)
|
||||
trans_loss = self.criterion_rec(trans, self.input_data['trans'])
|
||||
self.loss_recorder.add_scalar('trans_loss', trans_loss, validation=validation)
|
||||
loss += self.args.lambda_trans * trans_loss
|
||||
|
||||
contacts_loss = self.criterion_bce(self.recon_data['contacts'], self.input_data['contacts'])
|
||||
self.loss_recorder.add_scalar('contacts_loss', contacts_loss, validation=validation)
|
||||
loss += self.args.lambda_contacts * contacts_loss
|
||||
|
||||
if self.args.lambda_kl != 0:
|
||||
mu_l = self.recon_data['mu_l']
|
||||
logvar_l = self.recon_data['logvar_l']
|
||||
local_kl_loss = torch.mean(-0.5 * torch.sum(1 + logvar_l - mu_l.pow(2) - logvar_l.exp(), dim=-1))
|
||||
self.loss_recorder.add_scalar('local_kl_loss', local_kl_loss, validation=validation)
|
||||
|
||||
mu_g = self.recon_data['mu_g']
|
||||
logvar_g = self.recon_data['logvar_g']
|
||||
global_kl_loss = torch.mean(-0.5 * torch.sum(1 + logvar_g - mu_g.pow(2) - logvar_g.exp(), dim=-1))
|
||||
self.loss_recorder.add_scalar('global_kl_loss', global_kl_loss, validation=validation)
|
||||
if not validation:
|
||||
loss += self.args.lambda_kl * self.kl_scheduler() * (local_kl_loss + global_kl_loss)
|
||||
|
||||
self.loss_recorder.add_scalar('total_loss', loss, validation=validation)
|
||||
|
||||
if not validation:
|
||||
loss.backward()
|
||||
|
||||
def optimize_parameters(self):
|
||||
self.optimizer.zero_grad()
|
||||
self.forward()
|
||||
self.backward(validation=False)
|
||||
self.optimizer.step()
|
||||
|
||||
self.iteration += 1
|
||||
|
||||
def report_errors(self):
|
||||
root_orient = rotation_6d_to_matrix(self.recon_data['root_orient']) # (B, T, 3, 3)
|
||||
root_orient_gt = rotation_6d_to_matrix(self.input_data['root_orient']) # (B, T, 3, 3)
|
||||
orientation_error = self.criterion_geo(root_orient.view(-1, 3, 3), root_orient_gt.view(-1, 3, 3))
|
||||
|
||||
local_rotmat = self.fk.global_to_local(self.recon_data['rotmat'].view(-1, self.args.smpl.joint_num, 3, 3)) # (B x T, J, 3, 3)
|
||||
local_rotmat_gt = self.fk.global_to_local(self.input_data['rotmat'].view(-1, self.args.smpl.joint_num, 3, 3)) # (B x T, J, 3, 3)
|
||||
rotation_error = self.criterion_geo(local_rotmat[:, 1:].reshape(-1, 3, 3), local_rotmat_gt[:, 1:].reshape(-1, 3, 3))
|
||||
|
||||
pos = self.recon_data['pos'] # (B, T, J, 3)
|
||||
pos_gt = self.input_data['pos'] # (B, T, J, 3)
|
||||
position_error = torch.linalg.norm((pos - pos_gt), dim=-1).mean()
|
||||
|
||||
if 'trans' in self.recon_data.keys():
|
||||
origin = self.input_data['trans'][:, 0]
|
||||
trans = self.recon_data['trans']
|
||||
trans[..., self.v_axis] = trans[..., self.v_axis] + origin[..., self.v_axis].unsqueeze(1)
|
||||
trans_gt = self.input_data['trans'] # (B, T, 3)
|
||||
translation_error = torch.linalg.norm((trans - trans_gt), dim=-1).mean()
|
||||
else:
|
||||
translation_error = torch.zeros(1).mean()
|
||||
|
||||
return {
|
||||
'rotation': c2c(rotation_error),
|
||||
'position': c2c(position_error),
|
||||
'orientation': c2c(orientation_error),
|
||||
'translation': c2c(translation_error)
|
||||
}
|
||||
|
||||
def verbose(self):
|
||||
res = {}
|
||||
for loss in self.loss_recorder.losses.values():
|
||||
res[loss.name] = {'train': loss.current()[0], 'val': loss.current()[1]}
|
||||
|
||||
return res
|
||||
|
||||
def validate(self):
|
||||
with torch.no_grad():
|
||||
self.forward()
|
||||
self.backward(validation=True)
|
||||
|
||||
def save(self, optimal=False):
|
||||
if optimal:
|
||||
path = os.path.join(self.args.save_dir, 'results', 'model')
|
||||
else:
|
||||
path = os.path.join(self.model_save_dir, f'{self.epoch_cnt:04d}')
|
||||
|
||||
os.makedirs(path, exist_ok=True)
|
||||
local_encoder = self.local_encoder.module if isinstance(self.local_encoder, nn.DataParallel) else self.local_encoder
|
||||
nemf = self.field.module if isinstance(self.field, nn.DataParallel) else self.field
|
||||
torch.save(local_encoder.state_dict(), os.path.join(path, 'local_encoder.pth'))
|
||||
torch.save(nemf.state_dict(), os.path.join(path, 'nemf.pth'))
|
||||
if self.args.data.root_transform:
|
||||
global_encoder = self.global_encoder.module if isinstance(self.global_encoder, nn.DataParallel) else self.global_encoder
|
||||
torch.save(global_encoder.state_dict(), os.path.join(path, 'global_encoder.pth'))
|
||||
torch.save(self.optimizer.state_dict(), os.path.join(path, 'optimizer.pth'))
|
||||
if self.args.scheduler.name:
|
||||
torch.save(self.schedulers[0].state_dict(), os.path.join(path, 'scheduler.pth'))
|
||||
self.loss_recorder.save(path)
|
||||
|
||||
print(f'Save at {path} succeeded')
|
||||
|
||||
def load(self, epoch=None, optimal=False):
|
||||
if optimal:
|
||||
path = os.path.join(self.args.save_dir, 'results', 'model')
|
||||
else:
|
||||
if epoch is None:
|
||||
all = [int(q) for q in os.listdir(self.model_save_dir)]
|
||||
if len(all) == 0:
|
||||
raise RuntimeError(f'Empty loading path {self.model_save_dir}')
|
||||
epoch = sorted(all)[-1]
|
||||
path = os.path.join(self.model_save_dir, f'{epoch:04d}')
|
||||
|
||||
print(f'Loading from {path}')
|
||||
local_encoder = self.local_encoder.module if isinstance(self.local_encoder, nn.DataParallel) else self.local_encoder
|
||||
nemf = self.field.module if isinstance(self.field, nn.DataParallel) else self.field
|
||||
local_encoder.load_state_dict(torch.load(os.path.join(path, 'local_encoder.pth'), map_location=self.device))
|
||||
nemf.load_state_dict(torch.load(os.path.join(path, 'nemf.pth'), map_location=self.device))
|
||||
if self.args.data.root_transform:
|
||||
global_encoder = self.global_encoder.module if isinstance(self.global_encoder, nn.DataParallel) else self.global_encoder
|
||||
global_encoder.load_state_dict(torch.load(os.path.join(path, 'global_encoder.pth'), map_location=self.device))
|
||||
if self.is_train:
|
||||
self.optimizer.load_state_dict(torch.load(os.path.join(path, 'optimizer.pth')))
|
||||
self.loss_recorder.load(path)
|
||||
self.iteration = len(self.loss_recorder.losses['total_loss'].loss_step)
|
||||
print(f'trained iterations: {self.iteration}')
|
||||
if self.args.scheduler.name:
|
||||
self.schedulers[0].load_state_dict(torch.load(os.path.join(path, 'scheduler.pth')))
|
||||
self.epoch_cnt = epoch if not optimal else 0
|
||||
|
||||
print('Load succeeded')
|
||||
|
||||
def super_sampling(self, step=1):
|
||||
with torch.no_grad():
|
||||
self.forward(step)
|
||||
self.fps = self.args.data.fps / step
|
||||
self.compute_test_result()
|
||||
|
||||
def compute_test_result(self):
|
||||
os.makedirs(self.viz_dir, exist_ok=True)
|
||||
|
||||
b_size = self.input_data['pos'].shape[0]
|
||||
n_joints = self.input_data['pos'].shape[2]
|
||||
|
||||
rotmat = self.recon_data['rotmat'] # (B, T, J, 3, 3)
|
||||
local_rotmat = self.fk.global_to_local(rotmat.view(-1, n_joints, 3, 3)) # (B x T, J, 3, 3)
|
||||
local_rotmat = local_rotmat.view(b_size, -1, n_joints, 3, 3) # (B, T, J, 3, 3)
|
||||
rotmat_gt = self.input_data['rotmat'] # (B, T, J, 3, 3)
|
||||
local_rotmat_gt = self.fk.global_to_local(rotmat_gt.view(-1, n_joints, 3, 3)) # (B x T, J, 3, 3)
|
||||
local_rotmat_gt = local_rotmat_gt.view(b_size, -1, n_joints, 3, 3) # (B, T, J, 3, 3)
|
||||
|
||||
if self.args.data.root_transform:
|
||||
root_orient = rotation_6d_to_matrix(self.recon_data['root_orient']) # (B, T, 3, 3)
|
||||
root_orient_gt = rotation_6d_to_matrix(self.input_data['root_orient']) # (B, T, 3, 3)
|
||||
local_rotmat[:, :, 0] = root_orient
|
||||
local_rotmat_gt[:, :, 0] = root_orient_gt
|
||||
rotations = matrix_to_quaternion(local_rotmat) # (B, T, J, 4)
|
||||
rotations_gt = matrix_to_quaternion(local_rotmat_gt) # (B, T, J, 4)
|
||||
|
||||
if 'trans' in self.recon_data.keys():
|
||||
origin = self.input_data['trans'][:, 0]
|
||||
positions = self.recon_data['trans']
|
||||
positions[..., self.v_axis] = positions[..., self.v_axis] + origin[..., self.v_axis].unsqueeze(1)
|
||||
positions_gt = self.input_data['trans'] # (B, T, 3)
|
||||
else:
|
||||
positions = torch.zeros_like(self.input_data['trans']) # (B, T, 3)
|
||||
positions_gt = torch.zeros_like(self.input_data['trans']) # (B, T, 3)
|
||||
|
||||
for i in range(b_size):
|
||||
if self.args.bvh_viz:
|
||||
rotation = align_joints(rotations[i])
|
||||
position = positions[i].unsqueeze(1) # (T, 1, 3)
|
||||
anim = Animation(Quaternions(c2c(rotation)), c2c(position), None, offsets=self.args.smpl.offsets[self.args.data.gender], parents=self.args.smpl.parents)
|
||||
BVH.save(os.path.join(self.viz_dir, f'test_{self.test_index + i:03d}_{int(self.fps)}fps.bvh'), anim, names=self.args.smpl.joint_names, frametime=1 / self.fps)
|
||||
|
||||
rotation_gt = align_joints(rotations_gt[i])
|
||||
position_gt = positions_gt[i].unsqueeze(1) # (T, 1, 3)
|
||||
anim.rotations = Quaternions(c2c(rotation_gt))
|
||||
anim.positions = c2c(position_gt)
|
||||
BVH.save(os.path.join(self.viz_dir, f'test_{self.test_index + i:03d}_gt.bvh'), anim, names=self.args.smpl.joint_names, frametime=1 / self.args.data.fps)
|
||||
else:
|
||||
poses = c2c(quaternion_to_axis_angle(rotations[i])) # (T, J, 3)
|
||||
poses_gt = c2c(quaternion_to_axis_angle(rotations_gt[i])) # (T, J, 3)
|
||||
|
||||
poses = poses.reshape((poses.shape[0], -1)) # (T, J x 3)
|
||||
poses = np.pad(poses, [(0, 0), (0, 93)], mode='constant')
|
||||
poses_gt = poses_gt.reshape((poses_gt.shape[0], -1)) # (T, J x 3)
|
||||
poses_gt = np.pad(poses_gt, [(0, 0), (0, 93)], mode='constant')
|
||||
|
||||
np.savez(os.path.join(self.viz_dir, f'test_{self.test_index + i:03d}_{int(self.fps)}fps.npz'),
|
||||
poses=poses, trans=c2c(positions[i]), betas=np.zeros(10), gender=self.args.data.gender, mocap_framerate=self.fps)
|
||||
np.savez(os.path.join(self.viz_dir, f'test_{self.test_index + i:03d}_gt.npz'),
|
||||
poses=poses_gt, trans=c2c(positions_gt[i]), betas=np.zeros(10), gender=self.args.data.gender, mocap_framerate=self.args.data.fps)
|
||||
|
||||
self.test_index += b_size
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user