SamsungLabs / NeuralHaircut

Neural Haircut: Prior-Guided Strand-Based Hair Reconstruction. ICCV 2023
Other
516 stars 48 forks source link

failing to compute hair_mask with CDGnet #13

Closed op10ds closed 11 months ago

op10ds commented 1 year ago

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.

kenshi84 commented 12 months ago

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 https://github.com/tjpulkl/CDGNet/issues/5.

Hope this helps!

op10ds commented 11 months ago

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.