Spaces:
Running
Running
import os | |
import cv2 | |
import gradio as gr | |
import AnimeGANv3_src | |
import numpy as np | |
from moviepy.editor import VideoFileClip, ImageSequenceClip | |
class AnimeGANv3: | |
def __init__(self): | |
# Ensure the output directory exists | |
os.makedirs('output', exist_ok=True) | |
def process_frame(self, frame, style_code, det_face): | |
"""Process a single frame with AnimeGANv3.""" | |
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) | |
output = AnimeGANv3_src.Convert(frame_rgb, style_code, det_face) | |
return output | |
def inference(self, video_path, style, if_face=None): | |
print(video_path, style, if_face) | |
try: | |
# Map style names to codes | |
style_codes = { | |
"AnimeGANv3_Arcane": "A", | |
"AnimeGANv3_Trump v1.0": "T", | |
"AnimeGANv3_Shinkai": "S", | |
"AnimeGANv3_PortraitSketch": "P", | |
"AnimeGANv3_Hayao": "H", | |
"AnimeGANv3_Disney v1.0": "D", | |
"AnimeGANv3_JP_face v1.0": "J", | |
"AnimeGANv3_Kpop v2.0": "K", | |
} | |
style_code = style_codes.get(style, "U") | |
det_face = if_face == "Yes" | |
# Open the video | |
cap = cv2.VideoCapture(video_path) | |
if not cap.isOpened(): | |
raise Exception("Could not open video file") | |
# Get video properties | |
fps = cap.get(cv2.CAP_PROP_FPS) | |
frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) | |
print(f"Processing {frame_count} frames at {fps} FPS") | |
# Extract and process frames | |
frames = [] | |
while True: | |
ret, frame = cap.read() | |
if not ret: | |
break | |
processed_frame = self.process_frame(frame, style_code, det_face) | |
frames.append(processed_frame) | |
cap.release() | |
# Convert frames to video | |
save_path = "output/out.mp4" | |
clip = ImageSequenceClip(frames, fps=fps) | |
clip.write_videofile(save_path, codec="libx264", audio=False) | |
return save_path | |
except Exception as error: | |
print('Error:', error) | |
return None | |
# Create an instance of the AnimeGANv3 class | |
anime_gan = AnimeGANv3() | |
# Define the Gradio interface | |
title = "AnimeGANv3: Video to Anime Converter" | |
description = r"""Upload a video to convert it into anime style using AnimeGANv3.<br> | |
Select a style and choose whether to optimize for faces.<br> | |
<a href='https://github.com/TachibanaYoshino/AnimeGANv3' target='_blank'><b>AnimeGANv3 GitHub</b></a> | | |
<a href='https://www.patreon.com/Asher_Chan' target='_blank'><b>Patreon</b></a>""" | |
iface = gr.Interface( | |
fn=anime_gan.inference, | |
inputs=[ | |
gr.Video(label="Input Video"), | |
gr.Dropdown(choices=[ | |
'AnimeGANv3_Hayao', | |
'AnimeGANv3_Shinkai', | |
'AnimeGANv3_Arcane', | |
'AnimeGANv3_Trump v1.0', | |
'AnimeGANv3_Disney v1.0', | |
'AnimeGANv3_PortraitSketch', | |
'AnimeGANv3_JP_face v1.0', | |
'AnimeGANv3_Kpop v2.0', | |
], label='AnimeGANv3 Style', default='AnimeGANv3_Arcane'), | |
gr.Radio(choices=["Yes", "No"], label='Extract face', default="No"), | |
], | |
outputs=[ | |
gr.Video(label="Output Video") | |
], | |
title=title, | |
description=description, | |
allow_flagging="never" | |
) | |
# Launch the interface | |
iface.launch() |