Get the preview working
This commit is contained in:
parent
eeae9e46f4
commit
9c66fd9712
12
roop/capturer.py
Normal file
12
roop/capturer.py
Normal file
@ -0,0 +1,12 @@
|
|||||||
|
import cv2
|
||||||
|
|
||||||
|
|
||||||
|
def get_video_frame(video_path: str, frame_number: int = 1):
|
||||||
|
capture = cv2.VideoCapture(video_path)
|
||||||
|
frame_total = capture.get(cv2.CAP_PROP_FRAME_COUNT)
|
||||||
|
capture.set(cv2.CAP_PROP_POS_FRAMES, min(frame_total, frame_number - 1))
|
||||||
|
has_frame, frame = capture.read()
|
||||||
|
capture.release()
|
||||||
|
if has_frame:
|
||||||
|
return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||||||
|
return None
|
34
roop/ui.py
34
roop/ui.py
@ -6,6 +6,9 @@ from typing import Callable, Any, Tuple
|
|||||||
import cv2
|
import cv2
|
||||||
from PIL import Image, ImageTk, ImageOps
|
from PIL import Image, ImageTk, ImageOps
|
||||||
import roop.globals
|
import roop.globals
|
||||||
|
from roop.analyser import get_one_face
|
||||||
|
from roop.capturer import get_video_frame
|
||||||
|
from roop.swapper import process_faces
|
||||||
from roop.utilities import is_image, is_video
|
from roop.utilities import is_image, is_video
|
||||||
|
|
||||||
PRIMARY_COLOR = '#2d3436'
|
PRIMARY_COLOR = '#2d3436'
|
||||||
@ -81,7 +84,7 @@ def create_root(start: Callable, destroy: Callable) -> tk.Tk:
|
|||||||
|
|
||||||
|
|
||||||
def create_preview(parent) -> tk.Toplevel:
|
def create_preview(parent) -> tk.Toplevel:
|
||||||
global preview_label
|
global preview_label, preview_scale
|
||||||
|
|
||||||
preview = tk.Toplevel(parent)
|
preview = tk.Toplevel(parent)
|
||||||
preview.withdraw()
|
preview.withdraw()
|
||||||
@ -94,9 +97,8 @@ def create_preview(parent) -> tk.Toplevel:
|
|||||||
preview_label = tk.Label(preview, bg=PRIMARY_COLOR)
|
preview_label = tk.Label(preview, bg=PRIMARY_COLOR)
|
||||||
preview_label.pack(fill='both', expand=True)
|
preview_label.pack(fill='both', expand=True)
|
||||||
|
|
||||||
frame_value = tk.IntVar()
|
preview_scale = tk.Scale(preview, orient='horizontal', command=lambda frame_value: update_preview(int(frame_value)))
|
||||||
frame_slider = tk.Scale(preview, orient='horizontal', variable=frame_value)
|
preview_scale.pack(fill='x')
|
||||||
frame_slider.pack(fill='x')
|
|
||||||
|
|
||||||
return preview
|
return preview
|
||||||
|
|
||||||
@ -183,7 +185,7 @@ def select_target_path():
|
|||||||
|
|
||||||
def select_output_path(start):
|
def select_output_path(start):
|
||||||
output_path = filedialog.asksaveasfilename(title='Save to output file', initialfile='output.mp4')
|
output_path = filedialog.asksaveasfilename(title='Save to output file', initialfile='output.mp4')
|
||||||
if output_path and os.path.isfile(output_path):
|
if output_path:
|
||||||
roop.globals.output_path = output_path
|
roop.globals.output_path = output_path
|
||||||
start()
|
start()
|
||||||
|
|
||||||
@ -195,14 +197,14 @@ def render_image_preview(image_path: str, dimensions: Tuple[int, int] = None) ->
|
|||||||
return ImageTk.PhotoImage(image)
|
return ImageTk.PhotoImage(image)
|
||||||
|
|
||||||
|
|
||||||
def render_video_preview(video_path: str, dimensions: Tuple[int, int] = None, frame: int = 1) -> ImageTk.PhotoImage:
|
def render_video_preview(video_path: str, dimensions: Tuple[int, int] = None, frame_number: int = 1) -> ImageTk.PhotoImage:
|
||||||
capture = cv2.VideoCapture(video_path)
|
capture = cv2.VideoCapture(video_path)
|
||||||
if frame:
|
if frame_number:
|
||||||
capture.set(cv2.CAP_PROP_POS_FRAMES, frame)
|
capture.set(cv2.CAP_PROP_POS_FRAMES, frame_number)
|
||||||
has_frame, frame = capture.read()
|
has_frame, frame = capture.read()
|
||||||
if has_frame:
|
if has_frame:
|
||||||
if dimensions:
|
|
||||||
image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
|
image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
|
||||||
|
if dimensions:
|
||||||
image = ImageOps.fit(image, dimensions, Image.LANCZOS)
|
image = ImageOps.fit(image, dimensions, Image.LANCZOS)
|
||||||
return ImageTk.PhotoImage(image)
|
return ImageTk.PhotoImage(image)
|
||||||
capture.release()
|
capture.release()
|
||||||
@ -213,4 +215,18 @@ def toggle_preview() -> None:
|
|||||||
if PREVIEW.state() == 'normal':
|
if PREVIEW.state() == 'normal':
|
||||||
PREVIEW.withdraw()
|
PREVIEW.withdraw()
|
||||||
else:
|
else:
|
||||||
|
update_preview(1)
|
||||||
PREVIEW.deiconify()
|
PREVIEW.deiconify()
|
||||||
|
|
||||||
|
|
||||||
|
def update_preview(frame_number: int) -> None:
|
||||||
|
if roop.globals.source_path and roop.globals.target_path and frame_number:
|
||||||
|
video_frame = process_faces(
|
||||||
|
get_one_face(cv2.imread(roop.globals.source_path)),
|
||||||
|
get_video_frame(roop.globals.target_path, frame_number)
|
||||||
|
)
|
||||||
|
img = Image.fromarray(video_frame)
|
||||||
|
img = ImageOps.contain(img, (PREVIEW_WIDTH, PREVIEW_HEIGHT), Image.LANCZOS)
|
||||||
|
img = ImageTk.PhotoImage(img)
|
||||||
|
preview_label.configure(image=img)
|
||||||
|
preview_label.image = img
|
||||||
|
@ -35,7 +35,7 @@ def extract_frames(target_path: str) -> None:
|
|||||||
|
|
||||||
|
|
||||||
def create_video(target_path: str, fps: int) -> None:
|
def create_video(target_path: str, fps: int) -> None:
|
||||||
run_ffmpeg(['-i', get_temp_directory_path(target_path) + os.sep + '%04d.png', '-framerate', str(fps), '-c:v', 'libx264', '-crf', roop.globals.video_quality, '-pix_fmt', 'yuv420p', '-y', get_temp_file_path(target_path)])
|
run_ffmpeg(['-i', get_temp_directory_path(target_path) + os.sep + '%04d.png', '-framerate', str(fps), '-c:v', 'libx264', '-crf', str(roop.globals.video_quality), '-pix_fmt', 'yuv420p', '-y', get_temp_file_path(target_path)])
|
||||||
|
|
||||||
|
|
||||||
def restore_audio(target_path: str, output_path: str) -> None:
|
def restore_audio(target_path: str, output_path: str) -> None:
|
||||||
@ -67,8 +67,9 @@ def move_temp(target_path: str, output_path: str) -> None:
|
|||||||
|
|
||||||
|
|
||||||
def clean_temp(target_path: str) -> None:
|
def clean_temp(target_path: str) -> None:
|
||||||
if not roop.globals.keep_frames:
|
temp_directory_path = get_temp_directory_path(target_path)
|
||||||
shutil.rmtree(get_temp_directory_path(target_path))
|
if not roop.globals.keep_frames and os.path.isdir(temp_directory_path):
|
||||||
|
shutil.rmtree(temp_directory_path)
|
||||||
|
|
||||||
|
|
||||||
def has_image_extention(image_path: str) -> bool:
|
def has_image_extention(image_path: str) -> bool:
|
||||||
|
Loading…
Reference in New Issue
Block a user