main.py 536 Bytes
Newer Older
1
2
3
4
import os 
import torch
from PIL import Image
import numpy as np
boomb0om's avatar
boomb0om committed
5
from RealESRGAN import RealESRGAN
6
7
8
9
10


def main() -> int:
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = RealESRGAN(device, scale=4)
boomb0om's avatar
boomb0om committed
11
    model.load_weights('weights/RealESRGAN_x4.pth', download=True)
12
13
14
15
16
17
18
19
    for i, image in enumerate(os.listdir("inputs")):
        image = Image.open(f"inputs/{image}").convert('RGB')
        sr_image = model.predict(image)
        sr_image.save(f'results/{i}.png')


if __name__ == '__main__':
    main()