Update app.py

app.py

@@ -3,11 +3,16 @@ 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 the output directory exists
+        # 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."""
@@ -16,7 +21,7 @@ class AnimeGANv3:
         return output[:, :, ::-1]  # Convert back to BGR for OpenCV
 
     def inference(self, video_path, style, if_face=None):
-        print(video_path, style, if_face)
+        logging.info(f"Starting inference: video={video_path}, style={style}, face_detection={if_face}")
         try:
             # Map style names to codes
             style_codes = {
@@ -32,38 +37,45 @@ class AnimeGANv3:
             style_code = style_codes.get(style, "U")
             det_face = if_face == "Yes"
 
-            # Open the input video
+            # Open the input video and extract frames
             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)
-            width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-            height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
             frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-            print(f"Processing {frame_count} frames at {fps} FPS, {width}x{height}")
-
-            # Set up video writer
-            save_path = "output/out.mp4"
-            fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # Use 'mp4v' codec for MP4
-            out = cv2.VideoWriter(save_path, fourcc, fps, (width, height))
+            frames = []
 
-            # Process frames
-            while True:
+            while cap.isOpened():
                 ret, frame = cap.read()
                 if not ret:
                     break
-                processed_frame = self.process_frame(frame, style_code, det_face)
-                out.write(processed_frame)
+                frames.append(frame)
 
-            # Release resources
             cap.release()
-            out.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:
-            print('Error:', error)
+            logging.error(f"Error: {str(error)}")
             return None
 
 # Create an instance of the AnimeGANv3 class
@@ -89,8 +101,8 @@ iface = gr.Interface(
             'AnimeGANv3_PortraitSketch',
             'AnimeGANv3_JP_face v1.0',
             'AnimeGANv3_Kpop v2.0',
-        ], label='AnimeGANv3 Style', value='AnimeGANv3_Arcane'),  # Changed 'default' to 'value'
-        gr.Radio(choices=["Yes", "No"], label='Extract face', value="No"),  # Also updated here
+        ], label='AnimeGANv3 Style', value='AnimeGANv3_Arcane'),
+        gr.Radio(choices=["Yes", "No"], label='Extract face', value="No"),
     ],
     outputs=[
         gr.Video(label="Output Video")