failing to compute hair_mask with CDGnet

[op10ds]

I am trying to pre-process custom data, by following your indications. For simplicity, i am using the monocular data provided by you in /monocular/person_0_image. I get an error when executing the command:

python preprocess_custom_data/calc_masks.py --scene_path ./implicit-hair-data/data/SCENE_TYPE/CASE/ --MODNET_ckpt path_to_modnet --CDGNET_ckpt path_to_cdgnet

calc_masks.py", line 159, in main
for key, nkey in zip(state_dict_old.keys(), state_dict.keys()):
RuntimeError: OrderedDict mutated during iteration

I tried to find a workaround by rewriting the lines from 156 to 164 like this (following the discussion here: https://github.com/pytorch/pytorch/issues/40859

current_model_dict = model.state_dict()
loaded_state_dict = torch.load(args.CDGNET_ckpt, map_location='cpu')
new_state_dict={k:v if v.size()==current_model_dict[k].size()  else  current_model_dict[k] for k,v in zip(current_model_dict.keys(), loaded_state_dict.values())}
model.load_state_dict(new_state_dict, strict=False)

and then the images are generated, but they are all black. So what am i missing?
By contrast, mask images via MODnet are generated succesfully.

[Kenshi Takayama]

I encountered the same issue, and was able to fix it by modifying the code as

     state_dict = model.state_dict().copy()
     state_dict_old = torch.load(args.CDGNET_ckpt, map_location='cpu')
 
-    for key, nkey in zip(state_dict_old.keys(), state_dict.keys()):
+    state_dict_keys = list(state_dict.keys())
+    for key, nkey in zip(state_dict_old.keys(), state_dict_keys):
         if key != nkey:
             # remove the 'module.' in the 'key'
             state_dict[key[7:]] = deepcopy(state_dict_old[key])

and then downloading a different version of LIP_epoch_149.pth from this Google Drive link, posted in tjpulkl/CDGNet#5.

Hope this helps!

[op10ds]

Hi Kenshi,
yeah the checkpoint file you linked works well. The size is bigger than the one indicated by the authors, so that may be corrupted.
Thanks a lot.