197 lines
6.5 KiB
Python
197 lines
6.5 KiB
Python
# !/usr/bin/env python
|
|
# -*- coding: UTF-8 -*-
|
|
from ultralytics.nn.tasks import DetectionModel
|
|
from safetensors.torch import load_file
|
|
from PIL import Image, ImageOps
|
|
import numpy as np
|
|
import cv2
|
|
import os
|
|
import torch
|
|
from comfy.utils import common_upscale,ProgressBar
|
|
import folder_paths
|
|
|
|
cur_path = os.path.dirname(os.path.abspath(__file__))
|
|
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
|
|
|
|
POINT_SIZE = 4 # Size of the query point in the preview video
|
|
|
|
|
|
def convert_safetensor_to_pt():
|
|
tensor_dict = load_file(os.path.join(cur_path,"OmniParser/weights/icon_detect/model.safetensors"))
|
|
|
|
model = DetectionModel(os.path.join(cur_path,'OmniParser/weights/icon_detect/model.yaml'))
|
|
model.load_state_dict(tensor_dict)
|
|
torch.save({'model': model}, os.path.join(cur_path,'OmniParser/weights/icon_detect/best.pt'))
|
|
|
|
|
|
def get_config(platform="pc"):
|
|
#platform = 'pc'
|
|
if platform == 'pc':
|
|
draw_bbox_config = {
|
|
'text_scale': 0.8,
|
|
'text_thickness': 2,
|
|
'text_padding': 2,
|
|
'thickness': 2,
|
|
}
|
|
elif platform == 'web':
|
|
draw_bbox_config = {
|
|
'text_scale': 0.8,
|
|
'text_thickness': 2,
|
|
'text_padding': 3,
|
|
'thickness': 3,
|
|
}
|
|
elif platform == 'mobile':
|
|
draw_bbox_config = {
|
|
'text_scale': 0.8,
|
|
'text_thickness': 2,
|
|
'text_padding': 3,
|
|
'thickness': 3,
|
|
}
|
|
else:
|
|
raise "error"
|
|
return draw_bbox_config
|
|
|
|
def spilit_tensor2list(img_tensor):#[B,H,W,C], C=3,B>=1
|
|
video_list = []
|
|
if isinstance(img_tensor, list):
|
|
if isinstance(img_tensor[0], torch.Tensor):
|
|
video_list = img_tensor
|
|
elif isinstance(img_tensor, torch.Tensor):
|
|
b, _, _, _ = img_tensor.size()
|
|
if b == 1:
|
|
img = [b]
|
|
while img is not []:
|
|
video_list += img
|
|
else:
|
|
video_list = torch.chunk(img_tensor, chunks=b)
|
|
return video_list
|
|
|
|
def tensor2imglist(image,np_out=True):# pil first
|
|
B, _, _, _ = image.size()
|
|
if B == 1:
|
|
if np_out:
|
|
list_out = [tensor2pil(image)]
|
|
else:
|
|
list_out = [tensor2cv(image.squeeze())]
|
|
else:
|
|
image_list = torch.chunk(image, chunks=B)
|
|
if np_out:
|
|
list_out = [tensor2pil(i) for i in image_list]
|
|
else:
|
|
list_out = [tensor2cv(i.squeeze()) for i in image_list]
|
|
return list_out,B
|
|
|
|
|
|
def tensor2pil(tensor):
|
|
image_np = tensor.squeeze().mul(255).clamp(0, 255).byte().numpy()
|
|
image = Image.fromarray(image_np, mode='RGB')
|
|
return image
|
|
|
|
def pil2narry(img):
|
|
narry = torch.from_numpy(np.array(img).astype(np.float32) / 255.0).unsqueeze(0)
|
|
return narry
|
|
|
|
def tensor_upscale2pil(img_tensor, width, height):
|
|
samples = img_tensor.movedim(-1, 1)
|
|
img = common_upscale(samples, width, height, "nearest-exact", "center")
|
|
samples = img.movedim(1, -1)
|
|
img_pil = tensor2pil(samples)
|
|
return img_pil
|
|
|
|
def tensor_upscale(img_tensor, width, height): #torch tensor
|
|
samples = img_tensor.movedim(-1, 1)
|
|
img = common_upscale(samples, width, height, "nearest-exact", "center")
|
|
samples = img.movedim(1, -1)
|
|
return samples
|
|
|
|
def convert_from_cv2_to_image(img: np.ndarray) -> Image:
|
|
# return Image.fromarray(img)
|
|
return Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
|
|
|
|
|
|
def convert_from_image_to_cv2(img: Image) -> np.ndarray:
|
|
# return np.asarray(img)
|
|
return cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
|
|
|
|
def tensor2cv(tensor_image):
|
|
if tensor_image.is_cuda:
|
|
tensor_image = tensor_image.detach().cpu()
|
|
tensor_image=tensor_image.numpy()
|
|
#反归一化
|
|
maxValue=tensor_image.max()
|
|
tensor_image=tensor_image*255/maxValue
|
|
img_cv2=np.uint8(tensor_image)#32 to uint8
|
|
img_cv2=cv2.cvtColor(img_cv2,cv2.COLOR_RGB2BGR)
|
|
return img_cv2
|
|
|
|
def cvargb2tensor(img):
|
|
assert type(img) == np.ndarray, 'the img type is {}, but ndarry expected'.format(type(img))
|
|
img = torch.from_numpy(img.transpose((2, 0, 1)))
|
|
return img.float().div(255).unsqueeze(0) # 255也可以改为256
|
|
|
|
def cv2tensor(img):
|
|
assert type(img) == np.ndarray, 'the img type is {}, but ndarry expected'.format(type(img))
|
|
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
|
img = torch.from_numpy(img.transpose((2, 0, 1)))
|
|
return img.float().div(255).unsqueeze(0) # 255也可以改为256
|
|
|
|
def images_generator(img_list: list,):
|
|
#get img size
|
|
sizes = {}
|
|
for image_ in img_list:
|
|
if isinstance(image_,Image.Image):
|
|
count = sizes.get(image_.size, 0)
|
|
sizes[image_.size] = count + 1
|
|
elif isinstance(image_,np.ndarray):
|
|
count = sizes.get(image_.shape[:2][::-1], 0)
|
|
sizes[image_.shape[:2][::-1]] = count + 1
|
|
else:
|
|
raise "unsupport image list,must be pil or cv2!!!"
|
|
size = max(sizes.items(), key=lambda x: x[1])[0]
|
|
yield size[0], size[1]
|
|
|
|
# any to tensor
|
|
def load_image(img_in):
|
|
if isinstance(img_in, Image.Image):
|
|
img_in=img_in.convert("RGB")
|
|
i = np.array(img_in, dtype=np.float32)
|
|
i = torch.from_numpy(i).div_(255)
|
|
if i.shape[0] != size[1] or i.shape[1] != size[0]:
|
|
i = torch.from_numpy(i).movedim(-1, 0).unsqueeze(0)
|
|
i = common_upscale(i, size[0], size[1], "lanczos", "center")
|
|
i = i.squeeze(0).movedim(0, -1).numpy()
|
|
return i
|
|
elif isinstance(img_in,np.ndarray):
|
|
i=cv2.cvtColor(img_in,cv2.COLOR_BGR2RGB).astype(np.float32)
|
|
i = torch.from_numpy(i).div_(255)
|
|
print(i.shape)
|
|
return i
|
|
else:
|
|
raise "unsupport image list,must be pil,cv2 or tensor!!!"
|
|
|
|
total_images = len(img_list)
|
|
processed_images = 0
|
|
pbar = ProgressBar(total_images)
|
|
images = map(load_image, img_list)
|
|
try:
|
|
prev_image = next(images)
|
|
while True:
|
|
next_image = next(images)
|
|
yield prev_image
|
|
processed_images += 1
|
|
pbar.update_absolute(processed_images, total_images)
|
|
prev_image = next_image
|
|
except StopIteration:
|
|
pass
|
|
if prev_image is not None:
|
|
yield prev_image
|
|
|
|
def load_list_images(img_list: list,):
|
|
gen = images_generator(img_list)
|
|
(width, height) = next(gen)
|
|
images = torch.from_numpy(np.fromiter(gen, np.dtype((np.float32, (height, width, 3)))))
|
|
if len(images) == 0:
|
|
raise FileNotFoundError(f"No images could be loaded .")
|
|
return images
|
|
|