Generating density maps after training doesnt work

[JavierUrenaPhDProjects]

So I've been testing this model for a few days and I am still not able to replicate the results. The density maps overlapped in the pictures for qualitative results doesn't turn out as good as in the paper, nor it doesn't work on all the images.

I trained the model with the example_training config, obtaining multiple 'ckpt' files, which I chose the one I think fitted the best (I don't really have a way to decide the best checkpoint since all of them have similar results with the validation set). Then I run the example_localisation_training with said checkpoint and use the generated ckpt (again, one of the generated checkpoints) to run the example_localisation_vis config, expecting good qualitative results.

Have to say, when I run the example_test with the checkpoint from the training, the results are quite similar to the one on the paper:
'test/loss': 0.3367506265640259,
'test_DDP_MAE': 15.221793174743652,
'test_DDP_RMSE': 53.706050872802734
In the paper the results for the test set are MAE = 14.23 and RMSE = 43.83.

But for some reason using that checkpoint for the visualization generate bad results, compared to the ones on the paper. Even more, not all the images generated have a density/heatmap overlapped. Only half of them generate a density map on them, the rest are in grey-scale. I know this is wrong because when I used the "example weights" provided for FSC-133 (counting.ckpt and localisation.ckpt) all the images turned out well, so even if I trained the model wrong, I don't understand why all the images aren't overlapped with a density map.

This is the result of one of the pictures with the example weights provided:

This is the result for the same image I got after training the model by myself:

And this is one example of many of a picture that did not generate a density map:

So maybe im following the steps wrong? should I consider a different kind of configuration?