Spaces:
Running
Running
| import os | |
| import cv2 | |
| import gradio as gr | |
| import AnimeGANv3_src | |
| import numpy as np | |
| import logging | |
| # Set up logging | |
| logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') | |
| class AnimeGANv3: | |
| def __init__(self): | |
| # Ensure directories exist | |
| os.makedirs('output', exist_ok=True) | |
| os.makedirs('frames', 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[:, :, ::-1] # Convert back to BGR for OpenCV | |
| def inference(self, video_path, style, if_face=None): | |
| logging.info(f"Starting inference: video={video_path}, style={style}, face_detection={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 input video and extract frames | |
| cap = cv2.VideoCapture(video_path) | |
| if not cap.isOpened(): | |
| raise Exception("Could not open video file") | |
| fps = cap.get(cv2.CAP_PROP_FPS) | |
| frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) | |
| frames = [] | |
| while cap.isOpened(): | |
| ret, frame = cap.read() | |
| if not ret: | |
| break | |
| frames.append(frame) | |
| cap.release() | |
| logging.info(f"Extracted {frame_count} frames at {fps} FPS to process") | |
| # Process each frame and save as PNG with logging | |
| for idx, frame in enumerate(frames): | |
| stylized_frame = self.process_frame(frame, style_code, det_face) | |
| png_filename = f'frames/frame_{idx:04d}.png' | |
| cv2.imwrite(png_filename, stylized_frame) | |
| logging.info(f"Processed and saved frame {idx + 1}/{frame_count} as {png_filename}") | |
| logging.info("All frames processed and saved as PNGs") | |
| # Combine PNGs into video using ffmpeg | |
| save_path = "output/out.mp4" | |
| os.system(f"ffmpeg -framerate {fps} -i frames/frame_%04d.png -c:v libx264 -pix_fmt yuv420p {save_path} -y") | |
| # Check if the video was created | |
| if not os.path.exists(save_path): | |
| raise Exception("Failed to create output video with ffmpeg") | |
| logging.info(f"Video created: {save_path}") | |
| return save_path | |
| except Exception as error: | |
| logging.error(f"Error: {str(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', value='AnimeGANv3_Arcane'), | |
| gr.Radio(choices=["Yes", "No"], label='Extract face', value="No"), | |
| ], | |
| outputs=[ | |
| gr.Video(label="Output Video") | |
| ], | |
| title=title, | |
| description=description, | |
| allow_flagging="never" | |
| ) | |
| # Launch the interface | |
| iface.launch() |