-
Notifications
You must be signed in to change notification settings - Fork 10
/
base.py
81 lines (60 loc) · 1.96 KB
/
base.py
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
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
import os
import shutil
import torch
import torch.nn as nn
from torchvision.datasets.utils import download_and_extract_archive
MODELS_PATH = ".models"
class BaseModel(nn.Module):
"""Base Super-Resolution module"""
def load_pretrained(self, weights_url: str, weights_path: str) -> None:
"""Download pretrained weights and load as state dict
Parameters
----------
weights_url : str
Base URL to pretrained weights.
weights_path : str
Path to save pretrained weights.
Returns
-------
None
"""
base_file = os.path.basename(weights_path)
if not os.path.exists(os.path.join(MODELS_PATH, base_file)):
self.download(weights_url, weights_path)
self.load_state_dict(torch.load(os.path.join(MODELS_PATH, base_file)))
@staticmethod
def download(url: str, weights_path: str) -> None:
"""Download pretrained weights
Parameters
----------
weights_path : str
Path to save pretrained weights.
Returns
-------
None
"""
base_file = os.path.basename(weights_path)
if not os.path.exists(MODELS_PATH):
os.mkdir(MODELS_PATH)
download_and_extract_archive(url, MODELS_PATH, remove_finished=True)
shutil.copyfile(weights_path, os.path.join(MODELS_PATH, base_file))
shutil.rmtree(os.path.dirname(weights_path))
@torch.no_grad()
def enhance(self, x: torch.Tensor) -> torch.Tensor:
"""Super-resolve x and cast as image
Parameters
----------
x : torch.Tensor
Input Low-Resolution image as tensor
Returns
-------
torch.Tensor
Super-Resolved image as tensor
"""
x = self.forward(x)
x *= 255.0
x = x.clamp(0, 255)
x = x.to(torch.uint8)
x = x.squeeze()
x = x.transpose((1, 2, 0))
return x