From 203e571aba6b93b58801857df04e90a785ce365c Mon Sep 17 00:00:00 2001
From: wangzhibo <wangzhibo@shsening.com>
Date: 星期三, 30 七月 2025 22:11:26 +0800
Subject: [PATCH] add face swapper
---
data/auto_media_publisher.sqlite3 | 0
/dev/null | 17 --
auto_media_publisher/modifiers/face_swap/__init__.py | 0
auto_media_publisher/modifiers/face_swap/face_nsfw_predictor.py | 0
auto_media_publisher/modifiers/face_swap/face_enhancer.py | 1
auto_media_publisher/modifiers/face_swap/face_swaper_avpy.py | 166 +++++++++++++++++++++++
auto_media_publisher/modifiers/face_swap/typing.py | 7 +
auto_media_publisher/utils/media_input_output.py | 93 +++++++++++++
auto_media_publisher/modifiers/face_swap/face_analyser.py | 1
auto_media_publisher/modifiers/face_swap/face_swaper_ffmpeg.py | 79 +++++------
auto_media_publisher/modifiers/face_swap/face_swaper_ffmpeg_stream.py | 2
11 files changed, 302 insertions(+), 64 deletions(-)
diff --git a/auto_media_publisher/utils/face_swap/__init__.py b/auto_media_publisher/modifiers/face_swap/__init__.py
similarity index 100%
rename from auto_media_publisher/utils/face_swap/__init__.py
rename to auto_media_publisher/modifiers/face_swap/__init__.py
diff --git a/auto_media_publisher/utils/face_swap/face_analyser.py b/auto_media_publisher/modifiers/face_swap/face_analyser.py
similarity index 98%
rename from auto_media_publisher/utils/face_swap/face_analyser.py
rename to auto_media_publisher/modifiers/face_swap/face_analyser.py
index e927185..1c4538b 100644
--- a/auto_media_publisher/utils/face_swap/face_analyser.py
+++ b/auto_media_publisher/modifiers/face_swap/face_analyser.py
@@ -18,7 +18,6 @@
self._analyser.prepare(ctx_id=0, det_size=(640, 640))
def _select_device(self):
- return 'cpu'
if torch.cuda.is_available():
return 'cuda'
elif torch.backends.mps.is_available():
diff --git a/auto_media_publisher/utils/face_swap/face_enhancer.py b/auto_media_publisher/modifiers/face_swap/face_enhancer.py
similarity index 98%
rename from auto_media_publisher/utils/face_swap/face_enhancer.py
rename to auto_media_publisher/modifiers/face_swap/face_enhancer.py
index d4b6beb..1b600fe 100644
--- a/auto_media_publisher/utils/face_swap/face_enhancer.py
+++ b/auto_media_publisher/modifiers/face_swap/face_enhancer.py
@@ -16,7 +16,6 @@
self._enhancer = GFPGANer(model_path=FACE_ENHANCER_MODEL_PATH, upscale=1, device=self.device)
def _select_device(self):
- return 'cpu'
if torch.cuda.is_available():
return 'cuda'
elif torch.backends.mps.is_available():
diff --git a/auto_media_publisher/utils/face_swap/face_nsfw_predictor.py b/auto_media_publisher/modifiers/face_swap/face_nsfw_predictor.py
similarity index 100%
rename from auto_media_publisher/utils/face_swap/face_nsfw_predictor.py
rename to auto_media_publisher/modifiers/face_swap/face_nsfw_predictor.py
diff --git a/auto_media_publisher/modifiers/face_swap/face_swaper_avpy.py b/auto_media_publisher/modifiers/face_swap/face_swaper_avpy.py
new file mode 100644
index 0000000..52b0f41
--- /dev/null
+++ b/auto_media_publisher/modifiers/face_swap/face_swaper_avpy.py
@@ -0,0 +1,166 @@
+import av
+import torch
+import threading, queue
+from threading import current_thread
+import numpy as np
+import insightface
+import cv2
+import time
+from insightface.app.common import Face
+from gfpgan.utils import GFPGANer
+from .typing import Frame # 浣犲凡鏈夊畾涔�
+from auto_media_publisher.config.conf_base import FACE_ENHANCER_MODEL_PATH,FACE_SWAPPER_MODEL_PATH
+from insightface.app import FaceAnalysis
+
+SIMILAR_FACE_DISTANCE: float = 0.45 # < 0.4 鍚屼竴浜� 0.4 -0.6 澶ф鏄悓涓�浜� >0.6 涓嶆槸鍚屼竴涓�
+
+def _select_device():
+ if torch.cuda.is_available():
+ return 'cuda:0'
+ elif torch.backends.mps.is_available():
+ return 'mps'
+ else:
+ return 'cpu'
+
+NUM_WORKERS = 1
+
+#face_swapper = insightface.model_zoo.get_model(FACE_SWAPPER_MODEL_PATH, providers=['CUDAExecutionProvider'])
+#face_enhancer = GFPGANer(model_path=FACE_ENHANCER_MODEL_PATH, upscale=1, device=_select_device())
+
+frame_q = queue.Queue(maxsize=10)
+processed_q = queue.Queue(maxsize=10)
+
+THREAD_SEMAPHORE =threading.Semaphore(4) # 涓嶉檺鍒朵簡銆�
+
+def enhance_face(thread_face_enhancer,target_face, temp_frame:np.ndarray) -> np.ndarray:
+ start_x, start_y, end_x, end_y = map(int, target_face['bbox'])
+ padding_x = int((end_x - start_x) * 0.5)
+ padding_y = int((end_y - start_y) * 0.5)
+ start_x = max(0, start_x - padding_x)
+ start_y = max(0, start_y - padding_y)
+ end_x = max(0, end_x + padding_x)
+ end_y = max(0, end_y + padding_y)
+ temp_face = temp_frame[start_y:end_y, start_x:end_x]
+ #print(temp_face.size)
+ if temp_face.size:
+ _, _, temp_face = thread_face_enhancer.enhance(
+ temp_face,
+ paste_back=True
+ )
+ temp_frame[start_y:end_y, start_x:end_x] = temp_face
+ return temp_frame
+
+
+def decode_video(video_path):
+ container = av.open(video_path)
+ for idx, frame in enumerate(container.decode(video=0)):
+ img = frame.to_ndarray(format='rgb24')
+ frame_q.put((idx, img))
+ for _ in range(NUM_WORKERS):
+ frame_q.put((None, None))
+
+def process_frame_loop(new_face):
+ start_time = time.time()
+ thread_face_analyser = FaceAnalysis(name='buffalo_l', root='./', providers=['CUDAExecutionProvider'])
+ thread_face_analyser.prepare(ctx_id=3, det_size=(640, 640))
+ #thread_face_enhancer = GFPGANer(model_path=FACE_ENHANCER_MODEL_PATH, upscale=1, device=_select_device())
+ thread_face_enhancer = GFPGANer(model_path=FACE_ENHANCER_MODEL_PATH, upscale=1, device=_select_device())
+ thread_face_swapper = insightface.model_zoo.get_model(FACE_SWAPPER_MODEL_PATH, providers=['CUDAExecutionProvider'])
+
+ while True:
+ idx, frame = frame_q.get()
+ if idx is None:
+ processed_q.put((None, None))
+ break
+ try:
+ many_faces = thread_face_analyser.get(frame)
+ if many_faces:
+ old_face = many_faces[0]
+ frame = thread_face_swapper.get(frame, old_face, new_face, paste_back=True)
+
+ for sinlge_face in many_faces:
+ frame_ = enhance_face(thread_face_enhancer,sinlge_face, frame)
+
+ processed_q.put((idx, frame_))
+ print(f"{current_thread().name} has proceesed frame :{idx}")
+
+ except Exception as e:
+ print(f"[ERROR] Frame {idx} failed: {e}")
+
+ print(f"totol time is : {time.time()- start_time}")
+
+def encode_video(output_path, fps, width, height):
+ out = av.open(output_path, mode='w')
+
+ #stream = out.add_stream('h264_nvenc', rate=fps)
+ stream = out.add_stream('libx264', rate=fps)
+ stream.width = width
+ stream.height = height
+ stream.pix_fmt = 'yuv420p'
+
+ stream.options = {'crf': '18', 'preset': 'veryfast'}
+ #stream.options = {'rc': 'constqp', 'qp': '18','preset': 'p7','profile': 'high'}
+
+ next_index = 0
+ buffer = {}
+ end_count = 0
+
+ while True:
+ idx, frame = processed_q.get()
+ if idx is None:
+ end_count += 1
+ if end_count >= NUM_WORKERS:
+ break
+ continue
+ buffer[idx] = frame
+
+ # 鍙鏈夊綋鍓嶈鍐欏叆鐨勫抚锛屽氨杈撳嚭
+ while next_index in buffer:
+ frame = buffer.pop(next_index)
+ video_frame = av.VideoFrame.from_ndarray(frame, format='rgb24')
+ for packet in stream.encode(video_frame):
+ out.mux(packet)
+ next_index += 1
+
+ # 娓呯┖缂栫爜缂撳瓨
+ for packet in stream.encode():
+ out.mux(packet)
+ out.close()
+# 鍚姩浠诲姟
+video_path = './temp/buhui.mp4'
+output_path = './temp/buhui-wangrong-enhancer.mp4'
+new_face_image_path = './temp/wangrong.jpg'
+face_analyser = FaceAnalysis(name='buffalo_l', root='./', providers=['CUDAExecutionProvider'])
+face_analyser.prepare(ctx_id=0, det_size=(640, 640))
+
+faces = face_analyser.get(cv2.imread(new_face_image_path))
+if faces:
+ new_face = sorted(faces, key=lambda x: x.bbox[0])[0]
+
+ #threading.Thread(target=decode_video, args=(video_path,)).start()
+ #threading.Thread(target=process_frame_loop, args=(new_face,)).start()
+ #threading.Thread(target=encode_video, args=(output_path, 25,576,1080)).start()
+ threads = []
+
+ # 瑙ｇ爜绾跨▼
+ t_decode = threading.Thread(target=decode_video,args=(video_path,))
+ threads.append(t_decode)
+ t_decode.start()
+
+ # 澶氬鐞嗙嚎绋�
+ for i in range(NUM_WORKERS):
+ t = threading.Thread(target=process_frame_loop, args=(new_face,),name=f"process_worker_{i}")
+ threads.append(t)
+ t.start()
+
+ # 缂栫爜绾跨▼
+ t_encode = threading.Thread(target=encode_video,args=(output_path, 25,576,1080))
+ threads.append(t_encode)
+ t_encode.start()
+
+ for t in threads:
+ t.join()
+
+ print("Processing complete.")
+
+
diff --git a/auto_media_publisher/utils/face_swap/face_swaper_ffmpeg.py b/auto_media_publisher/modifiers/face_swap/face_swaper_ffmpeg.py
similarity index 80%
rename from auto_media_publisher/utils/face_swap/face_swaper_ffmpeg.py
rename to auto_media_publisher/modifiers/face_swap/face_swaper_ffmpeg.py
index 7b28497..12f4d11 100644
--- a/auto_media_publisher/utils/face_swap/face_swaper_ffmpeg.py
+++ b/auto_media_publisher/modifiers/face_swap/face_swaper_ffmpeg.py
@@ -15,7 +15,9 @@
from .face_analyser import FaceAnalyser
from auto_media_publisher.utils.logger import get_logger
from auto_media_publisher.config.conf_base import FACE_SWAPPER_MODEL_PATH
-
+from auto_media_publisher.utils.media_input_output import VideoReader,VideoWriter
+from torch.utils.data import DataLoader
+from torchvision import transforms
class FaceSwapWorker:
def __init__(self, enable_enhancer:bool = False,enable_nsfw_predictor:bool=False):
@@ -25,6 +27,7 @@
self.main_analyser = FaceAnalyser()
self.enable_enhancer = enable_enhancer
self.enable_nsfw_predictor = enable_nsfw_predictor
+ self.batch_size = 4
if self.enable_nsfw_predictor:
self.nsfw_predictor = NSFWPredictor()
self._init()
@@ -48,56 +51,46 @@
return 'mps'
else:
return 'cpu'
+
+ def auto_downsample_ratio(self,h, w):
+ """
+ Automatically find a downsample ratio so that the largest side of the resolution be 512px.
+ """
+ return min(512 / max(h, w), 1)
def _decode_video(self):
- # === Step 1: 鐢� ffprobe 鑾峰彇瑙嗛淇℃伅 ===
- probe_cmd = [
- "ffprobe",
- "-v", "error",
- "-select_streams", "v:0",
- "-show_entries", "stream=width,height,r_frame_rate",
- "-of", "json",
- self.input_video_path
- ]
- probe_result = subprocess.run(probe_cmd, capture_output=True, text=True, check=True)
- probe_data = json.loads(probe_result.stdout)
- video_stream = probe_data['streams'][0]
- self.width = int(video_stream['width'])
- self.height = int(video_stream['height'])
- self.fps = eval(video_stream['r_frame_rate']) # e.g. "30/1"
+ transform = transforms.ToTensor()
+ source = VideoReader(self.input_video_path, transform)
+
+ reader = DataLoader(source, batch_size=self.batch_size, pin_memory=True, num_workers=self.num_processors)
- print(f"Video info: {self.width}x{self.height} @ {self.fps}fps")
+ writer = VideoWriter(
+ path=self.output_video_path,
+ frame_rate=30,
+ bit_rate=int(1 * 1000000))
+
+ for src in reader:
- # === Step 2: 鐢� ffmpeg 瑙ｇ爜杈撳嚭 rawvideo 鍒� stdout pipe ===
- decode_cmd = [
- "ffmpeg",
- "-loglevel", "info",
- "-i", self.input_video_path,
- "-f", "rawvideo",
- "-pix_fmt", "rgb24",
- "pipe:1"
- ]
- process = subprocess.Popen(decode_cmd, stdout=subprocess.PIPE, bufsize=10**8)
+ if downsample_ratio is None:
+ downsample_ratio = auto_downsample_ratio(*src.shape[2:])
- frame_size = self.width * self.height * 3
- idx = 0
- while True:
- in_bytes = process.stdout.read(frame_size)
- if not in_bytes:
- break
- frame = np.frombuffer(in_bytes, dtype=np.uint8).reshape((self.height, self.width, 3))
- self.frame_queue.put((idx, frame)) # 鏀逛负 tuple 褰㈠紡
- idx += 1
+ src = src.to(device, dtype, non_blocking=True).unsqueeze(0) # [B, T, C, H, W]
+ fgr, pha, *rec = model(src, *rec, downsample_ratio)
- self.total_frames = idx
- print(f"_decode_video Total frames: {self.total_frames}")
+ if output_foreground is not None:
+ writer_fgr.write(fgr[0])
+ if output_alpha is not None:
+ writer_pha.write(pha[0])
+ if output_composition is not None:
+ if output_type == 'video':
+ com = fgr * pha + bgr * (1 - pha)
+ else:
+ fgr = fgr * pha.gt(0)
+ com = torch.cat([fgr, pha], dim=-3)
+ writer_com.write(com[0])
+
- process.stdout.close()
- process.wait()
- # 閫氱煡澶勭悊绾跨▼缁撴潫
- for _ in range(self.num_processors):
- self.frame_queue.put(None)
def _process_frame(self):
thread_face_analyser = FaceAnalyser()
diff --git a/auto_media_publisher/utils/face_swap/face_swaper_ffmpeg_stream.py b/auto_media_publisher/modifiers/face_swap/face_swaper_ffmpeg_stream.py
similarity index 99%
rename from auto_media_publisher/utils/face_swap/face_swaper_ffmpeg_stream.py
rename to auto_media_publisher/modifiers/face_swap/face_swaper_ffmpeg_stream.py
index 5966fe3..0f2452d 100644
--- a/auto_media_publisher/utils/face_swap/face_swaper_ffmpeg_stream.py
+++ b/auto_media_publisher/modifiers/face_swap/face_swaper_ffmpeg_stream.py
@@ -8,7 +8,6 @@
import torch
import json
import insightface
-from insightface.app.common import Face
from .face_nsfw_predictor import NSFWPredictor
from .face_enhancer import FaceEnhancer
from .face_analyser import FaceAnalyser
@@ -41,7 +40,6 @@
self.stop_event = threading.Event()
def _select_device(self):
- return 'cpu'
if torch.cuda.is_available():
return 'cuda'
elif torch.backends.mps.is_available():
diff --git a/auto_media_publisher/modifiers/face_swap/typing.py b/auto_media_publisher/modifiers/face_swap/typing.py
new file mode 100644
index 0000000..bc8d6fc
--- /dev/null
+++ b/auto_media_publisher/modifiers/face_swap/typing.py
@@ -0,0 +1,7 @@
+
+from typing import Any
+from insightface.app.common import Face
+import numpy
+
+Face = Face
+Frame = numpy.ndarray[Any, Any]
\ No newline at end of file
diff --git a/auto_media_publisher/utils/face_swap/face_swaper_cv2.py b/auto_media_publisher/utils/face_swap/face_swaper_cv2.py
deleted file mode 100644
index 1e2dc41..0000000
--- a/auto_media_publisher/utils/face_swap/face_swaper_cv2.py
+++ /dev/null
@@ -1,138 +0,0 @@
-import cv2
-import numpy as np
-import insightface
-from insightface.app import FaceAnalysis
-from insightface.app.common import Face
-import threading
-from gfpgan.utils import GFPGANer
-from typing import Any
-
-Frame = np.ndarray
-
-
-class FaceProcessor:
- """
- 璋冪敤cv鐨勬ā鍧楀垎瑙p4锛岀粍瑁卪p4锛岃棰戣川閲忓樊
- """
- def __init__(self, face_enhancer_model: str = 'GFPGANv1.4.pth', face_swapper_model: str = 'inswapper_128.onnx', device: str = 'cuda'):
- self.FACE_ENHANCER = None
- self.FACE_SWAPPER = None
- self.THREAD_SEMAPHORE = threading.Semaphore()
- self.THREAD_LOCK = threading.Lock()
-
- # Set device
- self.device = device # 'cuda' for GPU, 'cpu' for CPU
-
- # Load models
- self.face_enhancer_model = face_enhancer_model
- self.face_swapper_model = face_swapper_model
-
- # Initialize face analysis app
- self.app = FaceAnalysis(providers=['CUDAExecutionProvider' if device == 'cuda' else 'CPUExecutionProvider'])
- self.app.prepare(ctx_id=0, det_size=(640, 640))
-
- def get_face_enhancer(self) -> GFPGANer:
- with self.THREAD_LOCK:
- if self.FACE_ENHANCER is None:
- self.FACE_ENHANCER = GFPGANer(
- model_path=self.face_enhancer_model, upscale=1, device=self.device)
- return self.FACE_ENHANCER
-
- def get_face_swapper(self) -> Any:
- with self.THREAD_LOCK:
- if self.FACE_SWAPPER is None:
- self.FACE_SWAPPER = insightface.model_zoo.get_model(
- self.face_swapper_model, providers=['CUDAExecutionProvider' if self.device == 'cuda' else 'CPUExecutionProvider'])
- return self.FACE_SWAPPER
-
- def enhance_face(self, target_face: Face, temp_frame: Frame) -> Frame:
- start_x, start_y, end_x, end_y = map(int, target_face['bbox'])
- padding_x = int((end_x - start_x) * 0.5)
- padding_y = int((end_y - start_y) * 0.5)
- start_x = max(0, start_x - padding_x)
- start_y = max(0, start_y - padding_y)
- end_x = max(0, end_x + padding_x)
- end_y = max(0, end_y + padding_y)
- temp_face = temp_frame[start_y:end_y, start_x:end_x]
- if temp_face.size:
- with self.THREAD_SEMAPHORE:
- _, _, temp_face = self.get_face_enhancer().enhance(
- temp_face,
- paste_back=True
- )
- temp_frame[start_y:end_y, start_x:end_x] = temp_face
- return temp_frame
-
- def swap_face(self, source_face: Face, target_face: Face, temp_frame: Frame) -> Frame:
- return self.get_face_swapper().get(temp_frame, target_face, source_face, paste_back=True)
-
- def process_video(self, src_mp4: str, image_path: str, dst_mp4: str) -> None:
- # Read the target image (source face)
- target_image = cv2.imread(image_path)
- target_faces = self.app.get(target_image)
- if len(target_faces) == 0:
- print("No face detected in the target image.")
- return
- target_face = target_faces[0]
-
- # Open the source video
- cap = cv2.VideoCapture(src_mp4)
- if not cap.isOpened():
- print("Error opening video stream or file")
- return
-
- # Get video properties
- fourcc = cv2.VideoWriter_fourcc(*'mp4v')
- fps = cap.get(cv2.CAP_PROP_FPS)
- frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
- frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-
- # Create VideoWriter object to save the output video
- out = cv2.VideoWriter(dst_mp4, fourcc, fps, (frame_width, frame_height))
-
- # Get the total frame count
- total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-
- # Process each frame of the video
- frame_count = 0
- while cap.isOpened():
- ret, frame = cap.read()
- if not ret:
- break
-
- # Detect faces in the current frame
- faces = self.app.get(frame)
- if len(faces) > 0:
- # Swap face from the target image to the current frame
- swapped_face_frame = self.swap_face(target_face, faces[0], frame)
- self.enhance_face(faces[0], swapped_face_frame)
- out.write(swapped_face_frame)
- else:
- # If no face detected, just write the original frame
- out.write(frame)
-
- # Increment frame count
- frame_count += 1
-
- # Calculate and print progress every 5%
- progress = (frame_count / total_frames) * 100
- if progress >= 5 * (frame_count // (total_frames // 20)):
- print(f"Processing: {int(progress)}% completed")
-
- cap.release()
- out.release()
- print(f"Video processing complete. Output saved to {dst_mp4}")
-
-
-# Usage example:
-
-# Initialize the FaceProcessor (use 'cuda' for GPU)
-face_processor = FaceProcessor(device='cpu')
-
-# Process the video by swapping faces from the target image
-# face_processor.process_video('g1.mp4', 'from-face/girl-1.png', './g1-demo_output.mp4')
-# face_processor.process_video('g1.mp4', 'from-face/girl-2.png', './g2-demo_output.mp4')
-face_processor.process_video('g3.mp4', 'lilizhen.jpg', './lilizhen_output.mp4')
-face_processor.process_video('g3.mp4', 'from-face/girl-1.png', './g1-demo_output.mp4')
-face_processor.process_video('g3.mp4', 'from-face/girl-2.png', './g2-demo_output.mp4')
-
diff --git a/auto_media_publisher/utils/face_swap/typing.py b/auto_media_publisher/utils/face_swap/typing.py
deleted file mode 100644
index 30791ef..0000000
--- a/auto_media_publisher/utils/face_swap/typing.py
+++ /dev/null
@@ -1,17 +0,0 @@
-
-from typing import Any
-from insightface.app.common import Face
-import numpy
-
-Face = Face
-Frame = numpy.ndarray[Any, Any]
-'''
-from typing import Any
-import numpy
-import numpy.typing as npt # Import numpy.typing for NDArray
-
-from insightface.app.common import Face
-
-Face = Face
-Frame = npt.NDArray[numpy.uint8] # Specify the dtype, e.g., uint8 for image frames
-'''
\ No newline at end of file
diff --git a/auto_media_publisher/utils/media_input_output.py b/auto_media_publisher/utils/media_input_output.py
new file mode 100644
index 0000000..4e2e209
--- /dev/null
+++ b/auto_media_publisher/utils/media_input_output.py
@@ -0,0 +1,93 @@
+import av
+import os
+import pims
+import numpy as np
+from torch.utils.data import Dataset
+from torchvision.transforms.functional import to_pil_image
+from PIL import Image
+
+
+class VideoReader(Dataset):
+ def __init__(self, path, transform=None):
+ self.video = pims.PyAVVideoReader(path)
+ self.rate = self.video.frame_rate
+ self.transform = transform
+
+ @property
+ def frame_rate(self):
+ return self.rate
+
+ def __len__(self):
+ return len(self.video)
+
+ def __getitem__(self, idx):
+ frame = self.video[idx]
+ frame = Image.fromarray(np.asarray(frame))
+ if self.transform is not None:
+ frame = self.transform(frame)
+ return frame
+
+
+class VideoWriter:
+ def __init__(self, path, frame_rate, bit_rate=1000000):
+ self.container = av.open(path, mode='w')
+ self.stream = self.container.add_stream('h264', rate=round(frame_rate))
+ self.stream.pix_fmt = 'yuv420p'
+ self.stream.bit_rate = bit_rate
+ self.stream.codec_context.options = {
+ 'crf': '18', # 鎺у埗鐢昏川锛岃秺灏忚秺娓呮櫚锛�18鏄珮璐ㄩ噺锛�
+ 'preset': 'veryfast', # 褰卞搷缂栫爜閫熷害锛宮edium 鏄�1x锛宖ast鏄�2x faster鏄�3x veryfast鏄�4x superfast鏄�6x ultrafast鏄�10x锛屾枃浠朵綋绉ぇ绾︿細澧炲姞(浠edium涓哄熀鍑�) 10%锛�15%锛�20% 30% 50%
+ 'tune': 'film',
+ }
+
+ def write(self, frames):
+ # frames: [T, C, H, W]
+ self.stream.width = frames.size(3)
+ self.stream.height = frames.size(2)
+ if frames.size(1) == 1:
+ frames = frames.repeat(1, 3, 1, 1) # convert grayscale to RGB
+ frames = frames.mul(255).byte().cpu().permute(0, 2, 3, 1).numpy()
+ for t in range(frames.shape[0]):
+ frame = frames[t]
+ frame = av.VideoFrame.from_ndarray(frame, format='rgb24')
+ self.container.mux(self.stream.encode(frame))
+
+ def close(self):
+ self.container.mux(self.stream.encode())
+ self.container.close()
+
+
+class ImageSequenceReader(Dataset):
+ def __init__(self, path, transform=None):
+ self.path = path
+ self.files = sorted(os.listdir(path))
+ self.transform = transform
+
+ def __len__(self):
+ return len(self.files)
+
+ def __getitem__(self, idx):
+ with Image.open(os.path.join(self.path, self.files[idx])) as img:
+ img.load()
+ if self.transform is not None:
+ return self.transform(img)
+ return img
+
+
+class ImageSequenceWriter:
+ def __init__(self, path, extension='jpg'):
+ self.path = path
+ self.extension = extension
+ self.counter = 0
+ os.makedirs(path, exist_ok=True)
+
+ def write(self, frames):
+ # frames: [T, C, H, W]
+ for t in range(frames.shape[0]):
+ to_pil_image(frames[t]).save(os.path.join(
+ self.path, str(self.counter).zfill(4) + '.' + self.extension))
+ self.counter += 1
+
+ def close(self):
+ pass
+
\ No newline at end of file
diff --git a/data/auto_media_publisher.sqlite3 b/data/auto_media_publisher.sqlite3
index 2a3b573..05e7791 100644
--- a/data/auto_media_publisher.sqlite3
+++ b/data/auto_media_publisher.sqlite3
Binary files differ
--
Gitblit v1.9.1