Method process

BiRefNetModule/wrapper.py:106–206

Process a single video or directory of images.

(self, input_path, alpha_output_dir=None, dilate_radius=0, on_frame_complete=None)

            torch.cuda.ipc_collect()

    def process(self, input_path, alpha_output_dir=None, dilate_radius=0, on_frame_complete=None):
        """
        Process a single video or directory of images.
        """
        input_path = Path(input_path)
        file_name = input_path.stem
        is_video = input_path.suffix.lower() in [".mp4", ".mkv", ".gif", ".mov", ".avi"]

        def get_frames():
            """Yields tuples of (image_numpy_array, output_file_name)"""
            if is_video:
                cap = cv2.VideoCapture(str(input_path))
                count = 0
                while True:
                    success, img = cap.read()
                    if not success:
                        break
                    yield img, f"{file_name}_alpha_{count:05d}.png"
                    count += 1
                cap.release()
            else:
                image_files = sorted(
                    [
                        f
                        for f in input_path.iterdir()
                        if f.is_file() and f.suffix.lower() in [".jpg", ".png", ".jpeg", ".exr"]
                    ]
                )
                if not image_files:
                    logging.warning(f"No images found in {input_path}")
                    return

                # Setup EXR support once if needed
                if "OPENCV_IO_ENABLE_OPENEXR" not in os.environ:
                    os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"

                for img_path in image_files:
                    img = cv2.imread(str(img_path), cv2.IMREAD_UNCHANGED)
                    if img is None:
                        continue
                    # Keep original filename for image sequences
                    yield img, f"alphaSeq_{img_path.stem}.png"

        count = 0
        for image, out_name in get_frames():
            # Ensure correct conversion to RGB regardless of input format (EXR/PNG/JPG)
            if len(image.shape) == 2:
                image_rgb = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
            elif image.shape[2] == 4:
                image_rgb = cv2.cvtColor(image, cv2.COLOR_BGRA2RGB)
            else:
                image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

            # EXR images load as float32. PIL expects uint8. Normalize if necessary.
            if image_rgb.dtype != np.uint8:
                image_rgb = cv2.normalize(image_rgb, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)

            pil_image = Image.fromarray(image_rgb)