-
Notifications
You must be signed in to change notification settings - Fork 1
/
model.py
33 lines (25 loc) · 852 Bytes
/
model.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
import cv2
import torch
import numpy as np
from vedastd.datasets import build_datasets
from vedastd.datasets.transforms import build_transform
from vedastd.utils.config import Config
from vedastd.models.builder import build_model
def tensor_to_img(t_img: torch.Tensor):
if t_img.ndim == 3:
t_img = t_img.permute(1, 2, 0).cpu().numpy()
else:
t_img = t_img.cpu().numpy()
t_img = (t_img - np.min(t_img)) / (np.max(t_img) - np.min(t_img))
t_img = (t_img * 255).clip(0, 255).astype(np.uint8)
return t_img
def main():
cfg = Config.fromfile('./configs/psenet_resnet50.py')
model = build_model(cfg['model'])
dummy_input = torch.randn((1, 3, 224, 224))
out = model(dummy_input)
for key, value in out.items():
print(key, value.shape)
print('done')
if __name__ == '__main__':
main()