0.0.2
This commit is contained in:
华成炀
+147
-24
@@ -1,7 +1,6 @@
|
||||
import os
|
||||
import torch
|
||||
import numpy as np
|
||||
import time
|
||||
|
||||
from diffusers.models.attention_processor import XFormersAttnProcessor
|
||||
from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
|
||||
@@ -18,7 +17,12 @@ from .StableAnimator.animation.modules.pose_net import PoseNet
|
||||
from .StableAnimator.animation.modules.unet import UNetSpatioTemporalConditionModel
|
||||
from .StableAnimator.animation.pipelines.inference_pipeline_animation import InferenceAnimationPipeline
|
||||
|
||||
from .utils.image_utils import tensor_to_pil, tensor_to_np, np_to_tensor, load_images_from_folder, save_frames_as_mp4
|
||||
from .utils.image_utils import tensor_to_pil, tensor_to_np, np_to_tensor, load_images_from_folder
|
||||
|
||||
|
||||
# from .StableAnimator.DWPose.dwpose_utils.dwpose_detector import dwpose_detector_aligned
|
||||
from .utils.dwpose.dwpose_detector import DWposeDetectorAligned # 这里重写StableAnimator.DWPose.dwpose_utils.dwpose_detector.dwpose_detector_aligned 函数，以适应ComfyUI的模型加载机制
|
||||
from .utils.dwpose.skeleton_extraction import draw_pose
|
||||
|
||||
GLOBAL_CATEGORY = "HJH_StableAnimatorNode🪅"
|
||||
|
||||
@@ -189,7 +193,8 @@ class StableAnimatorNode:
|
||||
"required": {
|
||||
"stable_animator_models":("STABLEANIMATORMODELS",),
|
||||
"reference_image": ("IMAGE",),
|
||||
"pose_images_dir": ("STRING",{"default":"",}), # 输入为图像序列文件夹路径
|
||||
# "pose_images_dir": ("STRING",{"default":"",}), # 输入为图像序列文件夹路径
|
||||
"pose_images": ("IMAGE",),
|
||||
"format":(["512x512","576x1024"],{}),
|
||||
"tile_size":("INT",{"default":16,"min":16,"max":64,"step":16}),
|
||||
"frames_overlap":("INT",{"default":4,"min":4,"max":64,"step":4}),
|
||||
@@ -231,8 +236,8 @@ class StableAnimatorNode:
|
||||
FUNCTION = "generate"
|
||||
CATEGORY = GLOBAL_CATEGORY
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("video_path",)
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("video_frames",)
|
||||
# OUTPUT_IS_LIST = (True,)
|
||||
|
||||
def generate(self, **kwargs):
|
||||
@@ -244,8 +249,13 @@ class StableAnimatorNode:
|
||||
models = kwargs["stable_animator_models"]
|
||||
reference_image = tensor_to_pil(kwargs["reference_image"])
|
||||
reference_image_np = tensor_to_np(kwargs["reference_image"])
|
||||
pose_images_dir = kwargs["pose_images_dir"]
|
||||
pose_images = load_images_from_folder(pose_images_dir, width=width, height=height)
|
||||
# pose_images_dir = kwargs["pose_images_dir"]
|
||||
# pose_images = load_images_from_folder(pose_images_dir, width=width, height=height)
|
||||
tensor_pose_images = kwargs["pose_images"]
|
||||
pose_images = []
|
||||
for pose_image in tensor_pose_images:
|
||||
pose_images.append(tensor_to_pil(pose_image.unsqueeze(0)).resize((width,height)))
|
||||
|
||||
seed = kwargs["seed"] if kwargs["seed"] != -1 else torch.randint(0, 2**32-1, (1,)).item()
|
||||
generator = torch.Generator(device=self.device).manual_seed(seed)
|
||||
num_frames = len(pose_images)
|
||||
@@ -293,7 +303,7 @@ class StableAnimatorNode:
|
||||
tile_overlap=kwargs["frames_overlap"],
|
||||
decode_chunk_size=kwargs["decode_chunk_size"],
|
||||
motion_bucket_id=127.,
|
||||
fps=7,
|
||||
fps=kwargs["fps"],
|
||||
min_guidance_scale=kwargs["guidance_scale"],
|
||||
max_guidance_scale=kwargs["guidance_scale"],
|
||||
noise_aug_strength=kwargs["noise_aug_strength"],
|
||||
@@ -305,28 +315,141 @@ class StableAnimatorNode:
|
||||
|
||||
# 转换为张量输出
|
||||
output_tensors = []
|
||||
np_frames=np.array(video_frames)
|
||||
print("********************")
|
||||
print(np_frames.shape)
|
||||
print(np_frames[0].shape)
|
||||
# for frame in video_frames:
|
||||
# np_frames = np.array(frame)
|
||||
# output_tensors.append(np_to_tensor(np.array(frame)))
|
||||
for frame in video_frames:
|
||||
output_tensors.append(np_to_tensor(np.array(frame)))
|
||||
|
||||
output_root = folder_paths.get_output_directory()
|
||||
output_dir = os.path.join(output_root, "stable_animator_output")
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
output_file = os.path.join(output_dir, f"{time.time()}.mp4")
|
||||
save_frames_as_mp4(np_frames, output_file, fps=kwargs["fps"])
|
||||
return (output_file,)
|
||||
# return (output_tensors,)
|
||||
return torch.cat(output_tensors, dim=0),
|
||||
|
||||
|
||||
class StableAnimatorDWPoseDetectorAlignedModels:
|
||||
"""
|
||||
StableAnimator加载DWPose模型,
|
||||
"""
|
||||
def __init__(self):
|
||||
pass
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
"""定义输入参数"""
|
||||
return {
|
||||
"required": {
|
||||
},
|
||||
}
|
||||
|
||||
FUNCTION = "load"
|
||||
CATEGORY = GLOBAL_CATEGORY
|
||||
RETURN_TYPES = ("DWPOSEDETECTORALIGNED",)
|
||||
RETURN_NAMES = ("dwpose_detector_aligned",)
|
||||
|
||||
def load(self, ):
|
||||
"""执行模型加载"""
|
||||
return DWposeDetectorAligned(),
|
||||
|
||||
class StableAnimatorSkeletonNode:
|
||||
"""
|
||||
StableAnimator生成视频POSE骨架
|
||||
"""
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
"""定义输入参数"""
|
||||
return {
|
||||
"required": {
|
||||
"dwpose_detector_aligned":("DWPOSEDETECTORALIGNED",),
|
||||
"reference_image":("IMAGE",),
|
||||
"video_frames":("IMAGE",),
|
||||
},
|
||||
}
|
||||
|
||||
FUNCTION = "extraction"
|
||||
CATEGORY = GLOBAL_CATEGORY
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("pose_frames",)
|
||||
|
||||
def extraction(self, **kwargs):
|
||||
'''
|
||||
以下代码参考 StableAnimator.DWPose.dwpose_utils.dwpose_detector.get_video_pose, 以适应ComfyUI的参数输入
|
||||
'''
|
||||
dwpose_detector_aligned = kwargs["dwpose_detector_aligned"]
|
||||
ref_image = tensor_to_np(kwargs["reference_image"])
|
||||
height, width, _ = ref_image.shape
|
||||
ref_pose = dwpose_detector_aligned(ref_image)
|
||||
ref_keypoint_id = [0, 1, 2, 5, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17]
|
||||
ref_keypoint_id = [i for i in ref_keypoint_id \
|
||||
if len(ref_pose['bodies']['subset']) > 0 and ref_pose['bodies']['subset'][0][i] >= .0]
|
||||
ref_body = ref_pose['bodies']['candidate'][ref_keypoint_id]
|
||||
|
||||
video_frames = kwargs["video_frames"]
|
||||
|
||||
detected_poses = []
|
||||
for video_frame in video_frames:
|
||||
frame = tensor_to_np(video_frame.unsqueeze(0))
|
||||
pose = dwpose_detector_aligned(frame)
|
||||
detected_poses.append(pose)
|
||||
|
||||
detected_bodies = np.stack([p['bodies']['candidate'] for p in detected_poses if p['bodies']['candidate'].shape[0] == 18])[:,ref_keypoint_id]
|
||||
ay, by = np.polyfit(detected_bodies[:, :, 1].flatten(), np.tile(ref_body[:, 1], len(detected_bodies)), 1)
|
||||
fh = height
|
||||
fw = width
|
||||
ax = ay / (fh / fw / height * width)
|
||||
bx = np.mean(np.tile(ref_body[:, 0], len(detected_bodies)) - detected_bodies[:, :, 0].flatten() * ax)
|
||||
a = np.array([ax, ay])
|
||||
b = np.array([bx, by])
|
||||
|
||||
output_pose = []
|
||||
for detected_pose in detected_poses:
|
||||
detected_pose['bodies']['candidate'] = detected_pose['bodies']['candidate'] * a + b
|
||||
detected_pose['faces'] = detected_pose['faces'] * a + b
|
||||
detected_pose['hands'] = detected_pose['hands'] * a + b
|
||||
im = draw_pose(detected_pose, height, width)
|
||||
output_pose.append(np_to_tensor(np.array(im)))
|
||||
|
||||
output = torch.cat(output_pose, dim=0)
|
||||
return output,
|
||||
|
||||
class StableAnimatorLoadFramesFromFolderNode:
|
||||
"""
|
||||
StableAnimator从文件夹加载帧
|
||||
"""
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
"""定义输入参数"""
|
||||
return {
|
||||
"required": {
|
||||
"pose_images_folder": ("STRING",{"default":"",}), # 输入为图像序列文件夹路径
|
||||
},
|
||||
}
|
||||
|
||||
FUNCTION = "load"
|
||||
CATEGORY = GLOBAL_CATEGORY
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("frames",)
|
||||
|
||||
def load(self, **kwargs):
|
||||
frames = []
|
||||
folder = kwargs["pose_images_folder"]
|
||||
frames = load_images_from_folder(folder,)
|
||||
|
||||
return frames,
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"StableAnimatorNode": StableAnimatorNode,
|
||||
"StableAnimatorModels": StableAnimatorModels
|
||||
"StableAnimatorModels": StableAnimatorModels,
|
||||
"StableAnimatorSkeletonNode": StableAnimatorSkeletonNode,
|
||||
"StableAnimatorDWPoseDetectorAlignedModels": StableAnimatorDWPoseDetectorAlignedModels,
|
||||
"StableAnimatorLoadFramesFromFolderNode": StableAnimatorLoadFramesFromFolderNode,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"StableAnimatorNode": "HJH-StableAnimator - Video Generation",
|
||||
"StableAnimatorModels":"HJH-StableAnimator - Load Models"
|
||||
"StableAnimatorModels":"HJH-StableAnimator - Load Models",
|
||||
"StableAnimatorSkeletonNode": "HJH-StableAnimator - Pose Extraction",
|
||||
"StableAnimatorDWPoseDetectorAlignedModels": "HJH-StableAnimator - Load DWPose Models",
|
||||
"StableAnimatorLoadFramesFromFolderNode": "HJH-StableAnimator - Load Pose Frames From Folder",
|
||||
}
|
||||
@@ -0,0 +1,58 @@
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from .wholebody import Wholebody
|
||||
|
||||
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
|
||||
# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
class DWposeDetectorAligned:
|
||||
def __init__(self,):
|
||||
self.pose_estimation = Wholebody()
|
||||
|
||||
|
||||
def release_memory(self):
|
||||
if hasattr(self, 'pose_estimation'):
|
||||
del self.pose_estimation
|
||||
import gc; gc.collect()
|
||||
|
||||
def __call__(self, oriImg):
|
||||
oriImg = oriImg.copy()
|
||||
H, W, C = oriImg.shape
|
||||
with torch.no_grad():
|
||||
candidate, score = self.pose_estimation(oriImg)
|
||||
nums, _, locs = candidate.shape
|
||||
candidate[..., 0] /= float(W)
|
||||
candidate[..., 1] /= float(H)
|
||||
body = candidate[:, :18].copy()
|
||||
body = body.reshape(nums * 18, locs)
|
||||
subset = score[:, :18].copy()
|
||||
for i in range(len(subset)):
|
||||
for j in range(len(subset[i])):
|
||||
if subset[i][j] > 0.3:
|
||||
subset[i][j] = int(18 * i + j)
|
||||
else:
|
||||
subset[i][j] = -1
|
||||
|
||||
# un_visible = subset < 0.3
|
||||
# candidate[un_visible] = -1
|
||||
|
||||
# foot = candidate[:, 18:24]
|
||||
|
||||
faces = candidate[:, 24:92]
|
||||
|
||||
hands = candidate[:, 92:113]
|
||||
hands = np.vstack([hands, candidate[:, 113:]])
|
||||
|
||||
faces_score = score[:, 24:92]
|
||||
hands_score = np.vstack([score[:, 92:113], score[:, 113:]])
|
||||
|
||||
bodies = dict(candidate=body, subset=subset, score=score[:, :18])
|
||||
pose = dict(bodies=bodies, hands=hands, hands_score=hands_score, faces=faces, faces_score=faces_score)
|
||||
|
||||
return pose
|
||||
|
||||
|
||||
# dwpose_detector_aligned = DWposeDetectorAligned(device=device)
|
||||
@@ -0,0 +1,133 @@
|
||||
import math
|
||||
import matplotlib
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
eps = 0.01
|
||||
|
||||
def alpha_blend_color(color, alpha):
|
||||
"""blend color according to point conf
|
||||
"""
|
||||
return [int(c * alpha) for c in color]
|
||||
|
||||
def draw_bodypose(canvas, candidate, subset, score):
|
||||
H, W, C = canvas.shape
|
||||
candidate = np.array(candidate)
|
||||
subset = np.array(subset)
|
||||
|
||||
stickwidth = 4
|
||||
|
||||
limbSeq = [[2, 3], [2, 6], [3, 4], [4, 5], [6, 7], [7, 8], [2, 9], [9, 10], \
|
||||
[10, 11], [2, 12], [12, 13], [13, 14], [2, 1], [1, 15], [15, 17], \
|
||||
[1, 16], [16, 18], [3, 17], [6, 18]]
|
||||
|
||||
colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0], \
|
||||
[0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255], \
|
||||
[170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85]]
|
||||
|
||||
for i in range(17):
|
||||
for n in range(len(subset)):
|
||||
index = subset[n][np.array(limbSeq[i]) - 1]
|
||||
conf = score[n][np.array(limbSeq[i]) - 1]
|
||||
if conf[0] < 0.3 or conf[1] < 0.3:
|
||||
continue
|
||||
Y = candidate[index.astype(int), 0] * float(W)
|
||||
X = candidate[index.astype(int), 1] * float(H)
|
||||
mX = np.mean(X)
|
||||
mY = np.mean(Y)
|
||||
length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5
|
||||
angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
|
||||
polygon = cv2.ellipse2Poly((int(mY), int(mX)), (int(length / 2), stickwidth), int(angle), 0, 360, 1)
|
||||
cv2.fillConvexPoly(canvas, polygon, alpha_blend_color(colors[i], conf[0] * conf[1]))
|
||||
|
||||
canvas = (canvas * 0.6).astype(np.uint8)
|
||||
|
||||
for i in range(18):
|
||||
for n in range(len(subset)):
|
||||
index = int(subset[n][i])
|
||||
if index == -1:
|
||||
continue
|
||||
x, y = candidate[index][0:2]
|
||||
conf = score[n][i]
|
||||
x = int(x * W)
|
||||
y = int(y * H)
|
||||
cv2.circle(canvas, (int(x), int(y)), 4, alpha_blend_color(colors[i], conf), thickness=-1)
|
||||
|
||||
return canvas
|
||||
|
||||
def draw_handpose(canvas, all_hand_peaks, all_hand_scores):
|
||||
H, W, C = canvas.shape
|
||||
|
||||
edges = [[0, 1], [1, 2], [2, 3], [3, 4], [0, 5], [5, 6], [6, 7], [7, 8], [0, 9], [9, 10], \
|
||||
[10, 11], [11, 12], [0, 13], [13, 14], [14, 15], [15, 16], [0, 17], [17, 18], [18, 19], [19, 20]]
|
||||
|
||||
for peaks, scores in zip(all_hand_peaks, all_hand_scores):
|
||||
|
||||
for ie, e in enumerate(edges):
|
||||
x1, y1 = peaks[e[0]]
|
||||
x2, y2 = peaks[e[1]]
|
||||
x1 = int(x1 * W)
|
||||
y1 = int(y1 * H)
|
||||
x2 = int(x2 * W)
|
||||
y2 = int(y2 * H)
|
||||
score = int(scores[e[0]] * scores[e[1]] * 255)
|
||||
if x1 > eps and y1 > eps and x2 > eps and y2 > eps:
|
||||
cv2.line(canvas, (x1, y1), (x2, y2),
|
||||
matplotlib.colors.hsv_to_rgb([ie / float(len(edges)), 1.0, 1.0]) * score, thickness=2)
|
||||
|
||||
for i, keyponit in enumerate(peaks):
|
||||
x, y = keyponit
|
||||
x = int(x * W)
|
||||
y = int(y * H)
|
||||
score = int(scores[i] * 255)
|
||||
if x > eps and y > eps:
|
||||
cv2.circle(canvas, (x, y), 4, (0, 0, score), thickness=-1)
|
||||
return canvas
|
||||
|
||||
def draw_facepose(canvas, all_lmks, all_scores):
|
||||
H, W, C = canvas.shape
|
||||
for lmks, scores in zip(all_lmks, all_scores):
|
||||
for lmk, score in zip(lmks, scores):
|
||||
x, y = lmk
|
||||
x = int(x * W)
|
||||
y = int(y * H)
|
||||
conf = int(score * 255)
|
||||
if x > eps and y > eps:
|
||||
cv2.circle(canvas, (x, y), 3, (conf, conf, conf), thickness=-1)
|
||||
return canvas
|
||||
|
||||
def draw_pose(pose, H, W, ref_w=2160):
|
||||
"""vis dwpose outputs
|
||||
|
||||
Args:
|
||||
pose (List): DWposeDetector outputs in dwpose_detector.py
|
||||
H (int): height
|
||||
W (int): width
|
||||
ref_w (int, optional) Defaults to 2160.
|
||||
|
||||
Returns:
|
||||
np.ndarray: image pixel value in RGB mode
|
||||
"""
|
||||
bodies = pose['bodies']
|
||||
faces = pose['faces']
|
||||
hands = pose['hands']
|
||||
candidate = bodies['candidate']
|
||||
subset = bodies['subset']
|
||||
|
||||
sz = min(H, W)
|
||||
sr = (ref_w / sz) if sz != ref_w else 1
|
||||
|
||||
########################################## create zero canvas ##################################################
|
||||
canvas = np.zeros(shape=(int(H*sr), int(W*sr), 3), dtype=np.uint8)
|
||||
|
||||
########################################### draw body pose #####################################################
|
||||
canvas = draw_bodypose(canvas, candidate, subset, score=bodies['score'])
|
||||
|
||||
########################################### draw hand pose #####################################################
|
||||
canvas = draw_handpose(canvas, hands, pose['hands_score'])
|
||||
|
||||
########################################### draw face pose #####################################################
|
||||
canvas = draw_facepose(canvas, faces, pose['faces_score'])
|
||||
|
||||
# return cv2.cvtColor(cv2.resize(canvas, (W, H)), cv2.COLOR_BGR2RGB).transpose(2, 0, 1)
|
||||
return cv2.cvtColor(cv2.resize(canvas, (W, H)), cv2.COLOR_BGR2RGB)
|
||||
@@ -0,0 +1,54 @@
|
||||
import numpy as np
|
||||
|
||||
import onnxruntime as ort
|
||||
from ...StableAnimator.DWPose.dwpose_utils.onnxdet import inference_detector
|
||||
from ...StableAnimator.DWPose.dwpose_utils.onnxpose import inference_pose
|
||||
|
||||
import folder_paths
|
||||
|
||||
class Wholebody:
|
||||
def __init__(self):
|
||||
|
||||
root_dir = folder_paths.get_folder_paths("stable_animator")[0]
|
||||
|
||||
device = 'cuda:0'
|
||||
# device = 'cuda'
|
||||
providers = ['CPUExecutionProvider'
|
||||
] if device == 'cpu' else ['CUDAExecutionProvider']
|
||||
onnx_det = f'{root_dir}/DWPose/yolox_l.onnx'
|
||||
onnx_pose = f'{root_dir}/DWPose/dw-ll_ucoco_384.onnx'
|
||||
|
||||
self.session_det = ort.InferenceSession(path_or_bytes=onnx_det, providers=providers)
|
||||
self.session_pose = ort.InferenceSession(path_or_bytes=onnx_pose, providers=providers)
|
||||
|
||||
def __call__(self, oriImg):
|
||||
det_result = inference_detector(self.session_det, oriImg)
|
||||
keypoints, scores = inference_pose(self.session_pose, det_result, oriImg)
|
||||
|
||||
keypoints_info = np.concatenate(
|
||||
(keypoints, scores[..., None]), axis=-1)
|
||||
# compute neck joint
|
||||
neck = np.mean(keypoints_info[:, [5, 6]], axis=1)
|
||||
# neck score when visualizing pred
|
||||
neck[:, 2:4] = np.logical_and(
|
||||
keypoints_info[:, 5, 2:4] > 0.3,
|
||||
keypoints_info[:, 6, 2:4] > 0.3).astype(int)
|
||||
new_keypoints_info = np.insert(
|
||||
keypoints_info, 17, neck, axis=1)
|
||||
mmpose_idx = [
|
||||
17, 6, 8, 10, 7, 9, 12, 14, 16, 13, 15, 2, 1, 4, 3
|
||||
]
|
||||
openpose_idx = [
|
||||
1, 2, 3, 4, 6, 7, 8, 9, 10, 12, 13, 14, 15, 16, 17
|
||||
]
|
||||
new_keypoints_info[:, openpose_idx] = \
|
||||
new_keypoints_info[:, mmpose_idx]
|
||||
keypoints_info = new_keypoints_info
|
||||
|
||||
keypoints, scores = keypoints_info[
|
||||
..., :2], keypoints_info[..., 2]
|
||||
|
||||
return keypoints, scores
|
||||
|
||||
|
||||
|
||||
+29
-14
@@ -6,6 +6,7 @@ import cv2
|
||||
|
||||
from PIL import ImageOps, Image
|
||||
|
||||
|
||||
def pil_to_tensor(image_pil):
|
||||
i = ImageOps.exif_transpose(image_pil)
|
||||
image = i.convert("RGB")
|
||||
@@ -30,25 +31,39 @@ def np_to_tensor(image_np):
|
||||
image_tensor = torch.from_numpy(image)[None,]
|
||||
return image_tensor
|
||||
|
||||
def load_images_from_folder(folder, width, height):
|
||||
# def load_images_from_folder(folder, width, height):
|
||||
# """从目录读取图片"""
|
||||
# images = []
|
||||
# files = os.listdir(folder)
|
||||
# png_files = [f for f in files if f.endswith('.png')]
|
||||
# png_files.sort(key=lambda x: int(x.split('_')[1].split('.')[0]))
|
||||
# for filename in png_files:
|
||||
# img = Image.open(os.path.join(folder, filename)).convert('RGB')
|
||||
# img = img.resize((width, height))
|
||||
# images.append(img)
|
||||
|
||||
# return images
|
||||
|
||||
def load_images_from_folder(folder, ):
|
||||
"""从目录读取图片"""
|
||||
images = []
|
||||
files = os.listdir(folder)
|
||||
png_files = [f for f in files if f.endswith('.png')]
|
||||
png_files.sort(key=lambda x: int(x.split('_')[1].split('.')[0]))
|
||||
# png_files.sort(key=lambda x: int(x.split('_')[1].split('.')[0]))
|
||||
for filename in png_files:
|
||||
img = Image.open(os.path.join(folder, filename)).convert('RGB')
|
||||
img = img.resize((width, height))
|
||||
images.append(img)
|
||||
# img = img.resize((width, height))
|
||||
images.append(pil_to_tensor(img))
|
||||
|
||||
return images
|
||||
return torch.cat(images, dim=0)
|
||||
|
||||
def save_frames_as_mp4(frames, output_mp4_path, fps):
|
||||
print("Starting saving the frames as mp4")
|
||||
height, width, _ = frames[0].shape
|
||||
fourcc = cv2.VideoWriter_fourcc(*'mp4v') # 'H264' for better quality
|
||||
out = cv2.VideoWriter(output_mp4_path, fourcc, float(fps), (width, height))
|
||||
for frame in frames:
|
||||
frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
|
||||
out.write(frame_bgr)
|
||||
out.release()
|
||||
# def save_frames_as_mp4(frames, output_mp4_path, fps):
|
||||
# print("Starting saving the frames as mp4")
|
||||
# height, width, _ = frames[0].shape
|
||||
# # fourcc = cv2.VideoWriter_fourcc(*'mp4v') # 'H264' for better quality
|
||||
# fourcc = cv2.VideoWriter_fourcc(*'avc1')
|
||||
# out = cv2.VideoWriter(output_mp4_path, fourcc, float(fps), (width, height))
|
||||
# for frame in frames:
|
||||
# frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
|
||||
# out.write(frame_bgr)
|
||||
# out.release()
|
||||
Reference in New Issue
Block a user