This is an implementation of PoseCNN for 6D pose estimation on PROPSP dataset

We tested the codes on

PyTorch version: 2.3.1
CUDA version: 12.1
Ubuntu 22.04
GeForce RTX 4070 and 4090

The project requires the following Python libraries and versions:

You can install the required dependencies using the requirements.txt file:

pip install -r requirements.txt

To use this project, you need to download the required dataset and extract it to the root path of the project.

1. Download the Dataset:

   • Download the dataset from the following link:
     PROPS-Pose-Dataset

2. Place the Dataset:

   • Move the downloaded file PROPS-Pose-Dataset.tar.gz to the root directory of the project.

3. Extract the Dataset:

   • Use the following command to extract the dataset:

     tar -xvzf PROPS-Pose-Dataset.tar.gz

   • This will create a folder named PROPS-Pose-Dataset in the root directory.

4. Verify the Dataset Structure:

   • Ensure the folder structure matches the following:

     PROPS-Pose-Dataset/
         ├── train/
         │   ├── rgb/
         │   ├── depth/
         │   ├── mask_visib/
         │   ├── train_gt.json
         │   ├── train_gt_info.json
         ├── val/
         │   ├── rgb/
         │   ├── depth/
         │   ├── mask_visib/
         │   ├── val_gt.json
         │   ├── val_gt_info.json
         ├── model/
             ├── 1_master_chef_can/
             ├── ...
     

5. Set Dataset Path in Code:

   • The project will automatically locate the dataset in PROPS-Pose-Dataset under the root path during execution. Ensure this directory exists before running the code.

To train the model, run the train.py script:

python train.py

To visualize the results, follow these steps to set up and run the inference.py script:

1. Download Pretrained Weights:

   • Download the pretrained model weights from the following link:
     PoseCNN Pretrained Weights

2. Place the Weights:

   • Save the downloaded weights file (e.g., posecnn_weights.pth) to your desired directory.

3. Set the Weights Path in Code:

   • Open the inference.py script and locate the following line:

     posecnn_model.load_state_dict(torch.load(os.path.join("your weight here")))

   • Replace "your weight here" with the path to your weights file. For example:

     posecnn_model.load_state_dict(torch.load(os.path.join("models/posecnn_weights.pth")))

4. Run the Inference Script:

   • Execute the script to visualize predictions:

     python inference.py