remove background remover for now

handler.py

@@ -10,7 +10,6 @@ from omegaconf import OmegaConf
 from transformers import CLIPVisionModelWithProjection
 import cv2
 import os
-from backgroundremover.bg import remove as remove_bg
 from src.models.pose_guider import PoseGuider
 from src.models.unet_2d_condition import UNet2DConditionModel
 from src.models.unet_3d import UNet3DConditionModel
@@ -133,55 +132,6 @@ class EndpointHandler():
 
         return os.path.join(os.getcwd(), output_path)
 
-
-    def remove_bg_from_image(self, image_data):
-        model_name = "u2net"  # Choose your preferred model: "u2net", "u2net_human_seg", "u2netp"
-        processed_image_data = remove_bg(
-            image_data,
-            model_name=model_name,
-            alpha_matting=True,
-            alpha_matting_foreground_threshold=240,
-            alpha_matting_background_threshold=10,
-            alpha_matting_erode_structure_size=10,
-            alpha_matting_base_size=1000
-        )
-        return processed_image_data
-
-    def _remove_background(self, input_path, output_path):
-        cap = cv2.VideoCapture(input_path)
-        if not cap.isOpened():
-            raise IOError(f"Error opening video file {input_path}")
-
-        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-        fps = int(cap.get(cv2.CAP_PROP_FPS))
-
-        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-        out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
-
-        frame_count = 0
-        while cap.isOpened():
-            ret, frame = cap.read()
-            if not ret:
-                break
-
-            frame_count += 1
-            pil_frame = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
-            frame_data = BytesIO()
-            pil_frame.save(frame_data, format="PNG")
-            frame_data = frame_data.getvalue()
-            processed_frame_data = self.remove_bg_from_image(frame_data)
-            processed_pil_frame = Image.open(BytesIO(processed_frame_data))
-            processed_frame = cv2.cvtColor(np.array(processed_pil_frame), cv2.COLOR_RGB2BGR)
-
-            out.write(processed_frame)
-
-        cap.release()
-        out.release()
-
-        if frame_count == 0:
-            raise IOError(f"No frames processed. Error with video file {input_path}")
-
     def __call__(self, data: Any) -> Dict[str, str]:
         inputs = data.get("inputs", {})
         ref_image_base64 = inputs.get("ref_image", "")
@@ -225,21 +175,10 @@ class EndpointHandler():
         cropped_face = self._crop_face(ref_image, save_path=cropped_face_path)
 
         # Perform face swapping
-        print(f"Starting face swap with cropped face: {cropped_face_path} and animation: {animation_path}")
         final_video_path = self._swap_face(cropped_face, animation_path)
-        print(f"Face swap completed. Final video path: {final_video_path}")
-
-        # Ensure the output file exists before trying to open it
-        if not os.path.exists(final_video_path):
-            raise FileNotFoundError(f"Expected output file not found: {final_video_path}")
-
-        # Remove the background from the final video
-        bg_removed_video_path = os.path.join(save_dir, "bg_removed_output.mp4")
-        self._remove_background(final_video_path, bg_removed_video_path)
-        print(f"Background removal completed. Output saved to: {bg_removed_video_path}")
 
         # Encode the final video in base64
-        with open(bg_removed_video_path, "rb") as video_file:
+        with open(final_video_path, "rb") as video_file:
             video_base64 = base64.b64encode(video_file.read()).decode("utf-8")
 
         torch.cuda.empty_cache()

requirements.txt

@@ -32,8 +32,4 @@ omegaconf==2.2.3
 
 # Face swap related dependencies
 facenet-pytorch==2.5.2
-dlib==19.22.0
-
-
-# Background removal
-backgroundremover
+dlib==19.22.0