settings.py 1.19 KB
Newer Older
valentini's avatar
valentini committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
import yaml

class Settings:
    def __init__(self, d, config):
        self.low_sr_dir = d['low_sr_dir']
        self.high_sr_dir = d['high_sr_dir']
        self.width = d['width']
        self.height = d['height']
        self.pix_fmt = d.get('pix_fmt', 'yuv420p')
        self.range = d.get('model_range', 1)
        self.luminance_only = d.get('luminance_only', False)
        self.clone_luminance_as_rgb = d.get('clone_luminance_as_rgb', False)
        self.upsample_uv = d.get('upsample_uv', False)
        self.uv_interpolation = d.get('uv_interpolation', 'nearest')
        self.scaling = config.get('scaling', 2)
        self.batch_size = d.get('batch_size', 4)
        self.patch_size = d.get('patch_size', None)
        # Scheduler settings
        scheduler_cfg = config.get('scheduler', {})
        self.use_plateau_scheduler = scheduler_cfg.get('use_plateau_scheduler', False)
        self.lr_step_size = scheduler_cfg.get('lr_step_size', 10)
        self.lr_gamma = scheduler_cfg.get('lr_gamma', 0.5)

def load_settings(config_path):
    with open(config_path, 'r') as f:
        config = yaml.safe_load(f)
    dataset_cfg = config['dataset']
    settings = Settings(dataset_cfg, config)
    return settings