Update app.py

app.py

@@ -10,20 +10,18 @@ logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(
 
 class AnimeGANv3:
     def __init__(self):
-        # Ensure directories exist
         os.makedirs('output', exist_ok=True)
         os.makedirs('frames', exist_ok=True)
+        logging.info(f"Available ONNX Runtime providers: {ort.get_available_providers()}")
 
     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
+        return output[:, :, ::-1]
 
     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",
@@ -37,40 +35,49 @@ class AnimeGANv3:
             style_code = style_codes.get(style, "U")
             det_face = if_face == "Yes"
 
-            # Open the input video and extract frames
+            # Open video
             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}")
+            # Process in batches
+            batch_size = 50  # Adjust based on testing (e.g., 50 frames per batch)
+            save_path = "output/out.mp4"
+            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+            width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+            height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+            out = None  # Video writer initialized later
 
-            logging.info("All frames processed and saved as PNGs")
+            frame_idx = 0
+            while cap.isOpened():
+                batch_frames = []
+                for _ in range(batch_size):
+                    ret, frame = cap.read()
+                    if not ret:
+                        break
+                    batch_frames.append(frame)
+                    frame_idx += 1
+
+                if not batch_frames:
+                    break
 
-            # 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")
+                # Process batch
+                for idx, frame in enumerate(batch_frames):
+                    stylized_frame = self.process_frame(frame, style_code, det_face)
+                    if out is None:  # Initialize writer on first frame
+                        out = cv2.VideoWriter(save_path, fourcc, fps, (width, height))
+                    out.write(stylized_frame)
+                    logging.info(f"Processed frame {frame_idx - len(batch_frames) + idx + 1}/{frame_count}")
 
-            # Check if the video was created
-            if not os.path.exists(save_path):
-                raise Exception("Failed to create output video with ffmpeg")
+            cap.release()
+            if out:
+                out.release()
+            else:
+                raise Exception("No frames processed")
 
             logging.info(f"Video created: {save_path}")
             return save_path
@@ -78,10 +85,10 @@ class AnimeGANv3:
             logging.error(f"Error: {str(error)}")
             return None
 
-# Create an instance of the AnimeGANv3 class
+# Create an instance
 anime_gan = AnimeGANv3()
 
-# Define the Gradio interface
+# 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>
@@ -112,5 +119,4 @@ iface = gr.Interface(
     allow_flagging="never"
 )
 
-# Launch the interface
 iface.launch()