# UFV Reference Training

This repository provides a framework for training super-resolution models on YUV image datasets. It includes a dummy model for demonstration and is designed to be easily extended with your own models.

## Getting Started

### 1. Install Requirements

#### pip
    Install Python 3.8+ and the required packages:

    ```bash
    pip install -r requirments.txt
```

#### Conda


### 2. Prepare Your Dataset

- Place your low-resolution and high-resolution YUV images in separate directories.
- Ensure that for each low-res file, a corresponding high-res file exists with the same name in the other directory (E.G., low_res_path/image001.yuv and high_res_path/image001.yuv).
- Update the paths in `config.yaml`:

```yaml
dataset:
  low_sr_dir: "path/to/low_sr"
  high_sr_dir: "path/to/high_sr"
  width: 256
  height: 256
  pix_fmt: "yuv420p"
  # ... other options ...
```

### 3. Configure Training

Edit `config.yaml` to adjust:
- Model scaling factor
- Patch size, batch size
- Scheduler type and parameters
- Whether to use only luminance (Y) channel, upsample UV, etc.

### 4. Run Training

From the project root, run:

```bash
python src/train.py
```

The script will:
- Load the dataset and split into training/validation sets (80/20 random split).
- Train the model for the number of epochs specified in the script (default: 50)
- Save the best model as `best_model.pth`

---

## Using a Custom Model

To use your own model instead of the dummy one:

1. **Create your model class** in `src/model/` (e.g., `my_sr_model.py`). It should inherit from `torch.nn.Module` and accept the following arguments:
   - `in_channels` (int): Number of input channels (1 or 3)
   - `out_channels` (int): Number of output channels (1 or 3)
   - `super_resolution` (int): Scaling factor

2. **Implement the `forward(self, x)` method** for your model.

3. **Import your model in `src/train.py`:**

```python
from model.my_sr_model import MySRModel
```

4. **Replace the model instantiation:**

```python
model = MySRModel(in_channels=in_channels, out_channels=in_channels, super_resolution=settings.scaling)
```

5. **(Optional) Adjust training loop** if your model requires special input/output handling.

---

## Configuration Reference

See `config.yaml` for all available options and their descriptions. Key options include:
- `scaling`: Super-resolution scaling factor
- `dataset.width`, `dataset.height`: Image dimensions
- `dataset.pix_fmt`: YUV pixel format
- `dataset.luminance_only`: Use only Y channel
- `dataset.upsample_uv`: Upsample U/V channels
- `dataset.batch_size`, `dataset.patch_size`: Training batch/patch sizes
- `scheduler`: Learning rate scheduler settings

---

## Notes
- The dummy model is for demonstration only. For real use, implement a proper super-resolution model.
- The DataLoader expects matching filenames in both low and high resolution directories.
- For the range you can use or the image color depth value.  `range: 255.0` for 8-bit images `range: 1023` for 10-nit formats, etc... . Or for normalized data, use `1.0`.

---

## License
MIT