這是論文SG-Bot:透過場景圖上從粗到細的機器人想像力進行物件重新排列 (ICRA 2024),arxiv的最小實現。
conda env create - f environment . yml
cd extension
python setup . py install
另請安裝 Pytorch。我們使用 Pytorch 1.12.1 和 CUDA 11.6 對其進行測試。
請參閱此頁面下載論文中使用的數據和更多資訊。
我們設定了兩個形狀自動編碼器,稱為AtlasNet
和AtlastNet2
。 AtlasNet
在規範座標下使用完整形狀進行訓練,而AtlasNet2
在相機框架下進行訓練,它提供了目標場景圖的先驗形狀以指導想像。我們也提供在這裡下載的訓練模式:訓練有素的 AtlasNet 和訓練有素的 AtlasNet2。
用於生成形狀
AtlasNet
。需要調整--batchSize
, --nepoch
以使訓練最優。 cd AtlasNet
python training/train_AE_AtlasNet.py
推理點雲[可選]:運行AtlasNet/inference/run_AE_AtlasNet.py
。結果將會將產生的點儲存在AtlasNet/log/atlasnet_separate_cultery/network
下。
取得用於訓練 Graph-to-3D 的點特徵:執行AtlasNet/inference/create_features_gt.py
,特徵儲存在objs_features_gt_atlasnet_separate_cultery.json
中。 json 檔案中的鍵是物件的名稱,例如“cup_1”,值是潛在特徵(128 維)。
用於產生形狀先驗
AtlasNet2
。這些檔案可以從這裡下載:partial_pcs。您也可以修改檔案路徑並執行AtlasNet2/auxiliary/generate_partial_pc_for_object.py
。最終輸出會作為 pickle 檔案儲存在AtlasNet2/partial_pc_data
下。AtlasNet2/auxiliary/generate_partial_pc_for_object.py
中的函式generate_train_sample
將AtlasNet2/partial_pc_data
分割為train(90%)和test(10%)。檔案名稱儲存為AtlasNet2/partial_pc_data_splits.json
AtlasNet2
:過程與AtlasNet
相同。我們基於 Graph-to-3D(GCN-VAE 架構)建構了場景產生器。與原始的 Graph-to-3D 不同,我們利用形狀感知場景圖來將生成的形狀與初始場景中觀察到的形狀對齊。我們在這裡提供經過訓練的模型:trained graph_to_3d。
如果你想重新訓練網絡, --batchSize
, --nepoch
, --exp
需要設定適當的數字。
cd graphto3d
python scripts/train_vaegan.py
更多詳細資訊可以在原始存儲庫中找到。
有兩種模式── robot
和oracle
。 robot
模式支援機械手臂根據想像操縱物體。該模式需要一個抓取姿勢預測網絡,我們使用Contact-GraspNet。這需要下載tensorflow。
pip install tensorflow - estimator == 2.7 .0 tensorflow - gpu == 2.7 .0
檢查點可以從原始儲存庫或此處下載。下載檢查點後,將它們移至./contact_graspnet
。
oracle
模式不需要代理,而是直接將物件擺成相對姿勢。要讓腳本正常運作,可以修改裡面的變數mode
,然後執行:
python sgbot_pybullet.py
論文中的結果是在oracle
模式下的。我們直接使用預先定義的場景圖作為目標。
我們提供了錄製的rosbag來展示性能。要進行此試驗,需要從此處下載 MaskRCNN 檢查點。需要安裝額外的要求。