floatlazer / semantic_slam

Real time semantic slam in ROS with a hand held RGB-D camera
GNU General Public License v3.0
612 stars 177 forks source link

How can I use another PSPNet model with this? #59

Open jolags313 opened 3 months ago

jolags313 commented 3 months ago

Hello, is there any way to use another PSPNet model with this? I found a newer and more accurate one, but when I change the parameter in semantic_cloud.yaml to point to it, I get the following error

Traceback (most recent call last):
  File "/home/conlab/ssv2_ws/src/semantic_slam/semantic_cloud/src/semantic_cloud.py", line 295, in <module>
    main(sys.argv)
  File "/home/conlab/ssv2_ws/src/semantic_slam/semantic_cloud/src/semantic_cloud.py", line 288, in main
    seg_cnn = SemanticCloud(gen_pcl = True)
  File "/home/conlab/ssv2_ws/src/semantic_slam/semantic_cloud/src/semantic_cloud.py", line 124, in __init__
    self.model.load_state_dict(convert_state_dict(state['model_state'])) # Remove 'module' from dictionary keys
KeyError: 'model_state'

Then, when I try to replace line 111 model_name ='pspnet' with model_name = {"model_state":"pspnet"}, I get another error

Traceback (most recent call last):
  File "/home/conlab/ssv2_ws/src/semantic_slam/semantic_cloud/src/semantic_cloud.py", line 296, in <module>
    main(sys.argv)
  File "/home/conlab/ssv2_ws/src/semantic_slam/semantic_cloud/src/semantic_cloud.py", line 289, in main
    seg_cnn = SemanticCloud(gen_pcl = True)
  File "/home/conlab/ssv2_ws/src/semantic_slam/semantic_cloud/src/semantic_cloud.py", line 123, in __init__
    self.model = get_model(model_name, self.n_classes, version = 'ade20k')
  File "/home/conlab/ssv2_ws/src/semantic_slam/semantic_cloud/include/ptsemseg/models/__init__.py", line 44, in get_model
    model = model(n_classes=n_classes)
TypeError: 'NoneType' object is not callable

Is there any way I can fix this? The model is the pspnet101 one pretrained on ADE20k I got from the Google Drive link here https://github.com/hszhao/semseg?tab=readme-ov-file