liav
/
lut_reproduce
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

update

Browse Source
main
Vladimir Protsenko 8 months ago
parent 3b2c9ee8e3
commit 3885617c8f
5 changed files with 39 additions and 43 deletions
Show Stats Download Patch File Download Diff File

3
readme.md
Unescape Escape View File

@ -10,5 +10,4 @@ python validate.py --val_datasets Set5,Set14,B100,Urban100,Manga109 --model_path
``` ```
Requierements: Requierements:
- [shedulefree](https://github.com/facebookresearch/schedule_free) - [shedulefree](https://github.com/facebookresearch/schedule_free)
- einops

44
src/scripts/image_demo.py
Unescape Escape View File

@ -9,25 +9,23 @@ import cv2
from PIL import Image from PIL import Image
from datetime import datetime from datetime import datetime
import argparse import argparse
class DemoOptions(): class ImageDemoOptions():
def __init__(self): def __init__(self):
self.parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) self.parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
self.parser.add_argument('--net_model_path', '-n', type=str, default=None, help="Net model path folder") self.parser.add_argument('--net_model_path', '-n', type=str, default="../../models/last_transfered_net.pth", help="Net model path folder")
self.parser.add_argument('--lut_model_path', '-l', type=str, default=None, help="Lut model path folder") self.parser.add_argument('--lut_model_path', '-l', type=str, default="../../models/last_transfered_lut.pth", help="Lut model path folder")
self.parser.add_argument('--project_path', '-q', type=str, default="../../", help="Project path.") self.parser.add_argument('--hr_image_path', '-a', type=str, default="../../data/Set14/HR/monarch.png", help="HR image path")
self.parser.add_argument('--batch_size', '-b', type=int, default=2**10, help="Size of the batch for the input domain values.") self.parser.add_argument('--lr_image_path', '-b', type=str, default="../../data/Set14/LR/X4/monarch.png", help="LR image path")
self.parser.add_argument('--project_path', type=str, default="../../", help="Project path.")
self.parser.add_argument('--batch_size', type=int, default=2**10, help="Size of the batch for the input domain values.")
def parse_args(self): def parse_args(self):
args = self.parser.parse_args() args = self.parser.parse_args()
args.project_path = Path(args.project_path).resolve() args.project_path = Path(args.project_path).resolve()
if args.net_model_path is None: args.hr_image_path = Path(args.hr_image_path).resolve()
args.project_path / "models" / "last_transfered_net.pth" args.lr_image_path = Path(args.lr_image_path).resolve()
else: args.net_model_path = Path(args.net_model_path).resolve()
args.net_model_path = Path(args.net_model_path).resolve() args.lut_model_path = Path(args.lut_model_path).resolve()
if args.lut_model_path is None:
args.project_path / "models" / "last_transfered_lut.pth"
else:
args.lut_model_path = Path(args.lut_model_path).resolve()
return args return args
def print_options(self, opt): def print_options(self, opt):
@ -43,21 +41,10 @@ class DemoOptions():
print(message) print(message)
print() print()
config_inst = DemoOptions() config_inst = ImageDemoOptions()
config = config_inst.parse_args() config = config_inst.parse_args()
start_script_time = datetime.now() start_script_time = datetime.now()
# net_model = LoadCheckpoint("/wd/luts/models/rcnet_centered/checkpoints/RCNetCentered_10000.pth")
# lut_model = LoadCheckpoint("/wd/luts/models/rcnet_centered/checkpoints/RCLutCentered_0.pth")
# net_model = LoadCheckpoint("/wd/luts/models/rcnet_rot90_7x7/checkpoints/RCNetRot90_7x7_10000.pth")
# lut_model = LoadCheckpoint("/wd/luts/models/rcnet_rot90_7x7/checkpoints/RCLutRot90_7x7_0.pth")
# net_model = LoadCheckpoint("/wd/luts/models/rcnet_rot90_3x3/checkpoints/RCNetRot90_3x3_10000.pth")
# lut_model = LoadCheckpoint("/wd/luts/models/rcnet_rot90_3x3/checkpoints/RCLutRot90_3x3_0.pth")
# net_model = LoadCheckpoint("/wd/luts/models/rcnet_x1/checkpoints/RCNetx1_46000.pth")
# lut_model = LoadCheckpoint("/wd/luts/models/rcnet_x1/checkpoints/RCLutx1_0.pth")
net_model = LoadCheckpoint(config.net_model_path).cuda() net_model = LoadCheckpoint(config.net_model_path).cuda()
lut_model = LoadCheckpoint(config.lut_model_path).cuda() lut_model = LoadCheckpoint(config.lut_model_path).cuda()
@ -65,11 +52,8 @@ lut_model = LoadCheckpoint(config.lut_model_path).cuda()
print(net_model) print(net_model)
print(lut_model) print(lut_model)
lr_image = cv2.imread(str(config.project_path / "data" / "Set14/LR/X4/lenna.png"))[:,:,::-1].copy() lr_image = cv2.imread(str(config.project_path / "data" / "Set14/LR/X4/monarch.png"))[:,::-1,::-1].copy()
image_gt = cv2.imread(str(config.project_path / "data" / "Set14/HR/lenna.png"))[:,:,::-1].copy() image_gt = cv2.imread(str(config.project_path / "data" / "Set14/HR/monarch.png"))[:,::-1,::-1].copy()
# lr_image = cv2.imread(str(project_path / "data" / "Synthetic/LR/X4/linear.png"))[:,:,::-1].copy()
# image_gt = cv2.imread(str(project_path / "data" / "Synthetic/HR/linear.png"))[:,:,::-1].copy()
input_image = torch.tensor(lr_image).type(torch.float32).permute(2,0,1)[None,...].cuda() input_image = torch.tensor(lr_image).type(torch.float32).permute(2,0,1)[None,...].cuda()

10
src/scripts/train.py
Unescape Escape View File

@ -201,5 +201,15 @@ if __name__ == "__main__":
config.current_iter = i config.current_iter = i
valid_steps(model=model, datasets=test_datasets, config=config, log_prefix=f"Iter {i}") valid_steps(model=model, datasets=test_datasets, config=config, log_prefix=f"Iter {i}")
model_path = (Path(config.checkpoint_dir) / f"{model.__class__.__name__}_{i}.pth").resolve()
SaveCheckpoint(model=model, path=model_path)
print("Saved to ", model_path)
# check if it is network or lut
if hasattr(model, 'get_lut_model'):
Path(config.models_dir / f"last_trained_net.pth").symlink_to(model_path)
else:
Path(config.models_dir / f"last_trained_lut.pth").symlink_to(model_path)
total_script_time = datetime.now() - script_start_time total_script_time = datetime.now() - script_start_time
config.logger.info(f"Completed after {total_script_time}") config.logger.info(f"Completed after {total_script_time}")

19
src/scripts/transfer_to_lut.py
Unescape Escape View File

@ -26,7 +26,8 @@ class TransferToLutOptions():
def parse_args(self): def parse_args(self):
args = self.parser.parse_args() args = self.parser.parse_args()
args.model_path = Path(args.model_path) args.model_path = Path(args.model_path)
args.checkpoint_dir = Path(args.model_path).absolute().parent args.models_dir = Path(args.model_path).resolve().parent.parent.parent
args.checkpoint_dir = Path(args.model_path).resolve().parent
return args return args
def print_options(self, opt): def print_options(self, opt):
@ -65,14 +66,14 @@ if __name__ == "__main__":
lut_path = Path(config.checkpoint_dir) / f"{lut_model.__class__.__name__}_0.pth" lut_path = Path(config.checkpoint_dir) / f"{lut_model.__class__.__name__}_0.pth"
models.SaveCheckpoint(model=lut_model, path=lut_path) models.SaveCheckpoint(model=lut_model, path=lut_path)
lut_model_size = np.sum([x.nelement()*x.element_size() for x in lut_model.parameters()]) lut_model_size = np.sum([x.nelement()*x.element_size() for x in lut_model.parameters()])
print()
print(datetime.now()-start_time)
print("Saved to", lut_path, f"{lut_model_size/(2**20):.3f} MB") print("Saved to", lut_path, f"{lut_model_size/(2**20):.3f} MB")
models.SaveCheckpoint(model=model, path=Path(config.model_path).absolute().parent.parent.parent / f"last_transfered_net.pth") Path(config.models_dir / f"last_transfered_net.pth").symlink_to(config.model_path.resolve())
models.SaveCheckpoint(model=lut_model, path=Path(config.model_path).absolute().parent.parent.parent / f"last_transfered_lut.pth") Path(config.models_dir / f"last_transfered_lut.pth").symlink_to(lut_path.resolve())
print("Updated", Path(config.model_path).absolute().parent.parent.parent / f"last_transfered_net.pth") print("Updated link", config.models_dir / f"last_transfered_net.pth")
print("Updated", Path(config.model_path).absolute().parent.parent.parent / f"last_transfered_lut.pth") print("Updated link", config.models_dir / f"last_transfered_lut.pth")
print()
print("Completed after", datetime.now()-start_time)

6
src/scripts/validate.py
Unescape Escape View File

@ -26,7 +26,7 @@ class ValOptions():
self.parser.add_argument('--model_path', type=str, help="Model path.") self.parser.add_argument('--model_path', type=str, help="Model path.")
self.parser.add_argument('--datasets_dir', type=str, default="../../data/", help="Path to datasets.") self.parser.add_argument('--datasets_dir', type=str, default="../../data/", help="Path to datasets.")
self.parser.add_argument('--val_datasets', type=str, default='Set5,Set14', help="Names of validation datasets.") self.parser.add_argument('--val_datasets', type=str, default='Set5,Set14', help="Names of validation datasets.")
self.parser.add_argument('--save_val_predictions', action='store_true', default=False, help='Save model predictions to exp_dir/val/dataset_name') self.parser.add_argument('--save_predictions', action='store_true', default=True, help='Save model predictions to exp_dir/val/dataset_name')
def parse_args(self): def parse_args(self):
args = self.parser.parse_args() args = self.parser.parse_args()
@ -65,6 +65,7 @@ class ValOptions():
# TODO with unified save/load function any model file of net or lut can be tested with the same script. # TODO with unified save/load function any model file of net or lut can be tested with the same script.
if __name__ == "__main__": if __name__ == "__main__":
script_start_time = datetime.now()
config_inst = ValOptions() config_inst = ValOptions()
config = config_inst.parse_args() config = config_inst.parse_args()
@ -83,4 +84,5 @@ if __name__ == "__main__":
valid_steps(model=model, datasets=test_datasets, config=config, log_prefix=f"Model {config.model_name}") valid_steps(model=model, datasets=test_datasets, config=config, log_prefix=f"Model {config.model_name}")
config.logger.info("Complete") total_script_time = datetime.now() - script_start_time
config.logger.info(f"Completed after {total_script_time}")
Write Preview
Loading…
Cancel
Save