search

sjmoran / CURL

Code for the ICPR 2020 paper: "CURL: Neural Curve Layers for Image Enhancement"
205 stars 34 forks source link

getting error on line line 148 in main.py #8

Closed M-Suzy closed 3 years ago

M-Suzy commented 3 years ago

Hi, Could you please help to run your code without errors. Right now I have the following issue: 2021-02-23 21:25:34,587 INFO Loading Adobe5k dataset ... Traceback (most recent call last): File "main.py", line 351, in <module> main() File "main.py", line 148, in main inference_dataset = Dataset(data_dict=inference_data_dict, NameError: name 'Dataset' is not defined

sjmoran commented 3 years ago

Hi, thank you for your interest in our work. I will have updated code in a few days time, including a new pre-trained model.

M-Suzy commented 3 years ago

Yes, I am using the updated commit. Here is the command I use which gives the error python3.8 main.py --inference_img_dirpath=./adobe5k_dpe/ --checkpoint_filepath=./pretrained_models/curl_validpsnr_23.070664790618277_validloss_0.029170328751206398_testpsnr_23.556465119481636_testloss_0.02679653838276863_epoch_24_model.pt

sjmoran commented 3 years ago

Issue now resolved, please see the latest code pushed to the repository today. I have also uploaded a model from epoch 99 that is more accurate (better PSNR, SSIM).

M-Suzy commented 3 years ago

Thank you very much for the fix! By the way, does your code support only GPU, or I can also run on CPU mode?

sjmoran commented 3 years ago

Hi, CPU mode is straightforward. You just need to place the cuda keywords in the code with cpu I believe e.g. net.cuda() with net.cpu() etc. For inference you'll also need to load the model on the CPU, more information is here: https://pytorch.org/tutorials/beginner/saving_loading_models.html. Feel free to raise a PR if you have a CPU/GPU switch implemented for CURL.

M-Suzy commented 3 years ago

Got it, thank you!