53 lines
1.7 KiB
Python
53 lines
1.7 KiB
Python
import plotly_express as px
|
|
import torch
|
|
import torch.nn.functional as F
|
|
import torchvision.transforms.functional as Vision
|
|
import cv2
|
|
|
|
|
|
def DumpTensorToImage(tensor, name, autoPad=True, scale=1.0):
|
|
if len(tensor.shape) != 2:
|
|
raise ("Error input dims")
|
|
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
|
|
if autoPad and (max(srp) / min(srp) > 16):
|
|
img = cv2.resize(img,[max(srp),max(srp)])
|
|
srp = img.shape
|
|
if scale != 1.0:
|
|
img = cv2.resize(img, [int(srp[0] * scale), int(srp[1] * scale)])
|
|
srp = img.shape
|
|
cv2.imwrite(name, img)
|
|
|
|
|
|
# def DumpTensorToImage(tensor, name, autoPad=True, scale=1.0):
|
|
# if len(tensor.shape) != 2:
|
|
# raise ("Error input dims")
|
|
# tensor = tensor.float()
|
|
# maxv = torch.max(tensor)
|
|
# minv = torch.min(tensor)
|
|
# tensor = (((tensor - minv) / (maxv - minv)) * 256).byte().cpu()
|
|
# srp = tensor.shape
|
|
# if autoPad and (max(srp) / min(srp) > 16):
|
|
# if srp[0] == min(srp):
|
|
# tensor = F.pad(tensor, [max(srp) - min(srp), 0], "replicate")
|
|
# else:
|
|
# tensor = F.pad(tensor, [0, max(srp) - min(srp)], "replicate")
|
|
# srp = tensor.shape
|
|
|
|
# tensor = tensor.unsqueeze(0)
|
|
# if scale != 1.0:
|
|
# tensor = Vision.resize(tensor, [int(srp[0] * scale), int(srp[1] * scale)])
|
|
# tensor = tensor.view([int(srp[0] * scale), int(srp[1] * scale)])
|
|
# srp = tensor.shape
|
|
|
|
# w = 1024 if max(srp) > 1024 else max(srp)
|
|
# scale = max(srp) / w
|
|
# # img = px.imshow(tensor)
|
|
# # img.write_image(name)
|
|
# cv2.imwrite(name, tensor.numpy())
|
|
# cv2.CreateMat(name, tensor.numpy())
|