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:
- `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`.
