Witllm/tools/show.py

33 lines
1.0 KiB
Python
Raw Normal View History

2023-12-21 19:52:19 +08:00
import plotly_express as px
import torch
import torch.nn.functional as F
import torchvision.transforms.functional as Vision
import cv2
2023-12-21 21:20:49 +08:00
import math
import numpy as np
2023-12-21 19:52:19 +08:00
2023-12-21 21:20:49 +08:00
def DumpTensorToImage(tensor, name, autoPad=True, scale=1.0, auto2d=True):
if len(tensor.shape) != 2 and len(tensor.shape) != 1:
2023-12-21 19:52:19 +08:00
raise ("Error input dims")
2023-12-21 21:20:49 +08:00
2023-12-21 19:52:19 +08:00
tensor = tensor.float()
maxv = torch.max(tensor)
minv = torch.min(tensor)
tensor = (((tensor - minv) / (maxv - minv)) * 256).byte().cpu()
img = tensor.numpy()
srp = img.shape
2023-12-21 21:20:49 +08:00
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
2023-12-21 19:52:19 +08:00
if autoPad and (max(srp) / min(srp) > 16):
2023-12-21 21:20:49 +08:00
img = cv2.resize(img, [max(srp), max(srp)])
srp = img.shape
2023-12-21 19:52:19 +08:00
if scale != 1.0:
img = cv2.resize(img, [int(srp[0] * scale), int(srp[1] * scale)])
2023-12-21 21:20:49 +08:00
srp = img.shape
2023-12-21 19:52:19 +08:00
cv2.imwrite(name, img)