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

Refine show.

This commit is contained in:
Colin 2024-08-18 16:38:13 +08:00
parent b860d794a6
commit 950055c210
1 changed files with 2 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
tools/show.py
View File

@ -11,6 +11,8 @@ import os
def DumpTensorToImage(tensor, name, forceSquare=True, scale=1.0, AutoContrast=True, GridValue=0): def DumpTensorToImage(tensor, name, forceSquare=True, scale=1.0, AutoContrast=True, GridValue=0):
if len(tensor.shape) != 2 and len(tensor.shape) != 1 and len(tensor.shape) != 3: if len(tensor.shape) != 2 and len(tensor.shape) != 1 and len(tensor.shape) != 3:
raise ("Error input dims") raise ("Error input dims")
if ("." not in name) or (name.split(".")[-1] not in {"jpg", "png", "bmp"}):
raise ("Error input name")
if len(tensor.shape) == 3: if len(tensor.shape) == 3:
channel = tensor.shape[0] channel = tensor.shape[0]