colin
/
Witllm
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add auto2d.

This commit is contained in:
Colin 2023-12-21 21:20:49 +08:00
parent bfc3fb6706
commit 539392c843
3 changed files with 20 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

BIN
test.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 217 B

19
tools/show.py
View File

@ -3,21 +3,30 @@ import torch
import torch.nn.functional as F import torch.nn.functional as F
import torchvision.transforms.functional as Vision import torchvision.transforms.functional as Vision
import cv2 import cv2
import math
import numpy as np
def DumpTensorToImage(tensor, name, autoPad=True, scale=1.0): def DumpTensorToImage(tensor, name, autoPad=True, scale=1.0, auto2d=True):
if len(tensor.shape) != 2: if len(tensor.shape) != 2 and len(tensor.shape) != 1:
raise ("Error input dims") raise ("Error input dims")
tensor = tensor.float() tensor = tensor.float()
maxv = torch.max(tensor) maxv = torch.max(tensor)
minv = torch.min(tensor) minv = torch.min(tensor)
tensor = (((tensor - minv) / (maxv - minv)) * 256).byte().cpu() tensor = (((tensor - minv) / (maxv - minv)) * 256).byte().cpu()
img = tensor.numpy() img = tensor.numpy()
srp = img.shape srp = img.shape
if auto2d and len(srp) == 1:
ceiled = math.ceil((srp[0]) ** 0.5)
img = cv2.copyMakeBorder(img, 0, ceiled * ceiled - srp[0], 0, 0, 0)
img = img.reshape((ceiled, ceiled))
srp = img.shape
if autoPad and (max(srp) / min(srp) > 16): if autoPad and (max(srp) / min(srp) > 16):
img = cv2.resize(img,[max(srp),max(srp)]) img = cv2.resize(img, [max(srp), max(srp)])
srp = img.shape srp = img.shape
if scale != 1.0: if scale != 1.0:
img = cv2.resize(img, [int(srp[0] * scale), int(srp[1] * scale)]) img = cv2.resize(img, [int(srp[0] * scale), int(srp[1] * scale)])
srp = img.shape srp = img.shape
cv2.imwrite(name, img) cv2.imwrite(name, img)

8
tools/test.py
View File

@ -2,5 +2,9 @@ import show
import torch import torch
radata = torch.randn(8192, 128) # radata = torch.randn(8192, 128)
show.DumpTensorToImage(radata, "test.png", autoPad=True,scale=0.2) # show.DumpTensorToImage(radata, "test.png", autoPad=True,scale=0.2)
radata = torch.randn(127)
show.DumpTensorToImage(radata, "test.png")