src/pipeline/pipeline_steps.py

# ----------------------------------------------------------------------
# IMPORTS
# ----------------------------------------------------------------------
from src.config.constants import BACKGROUND_REMOVAL_MODEL
from src.processing.image_download.image_download import image_download
from src.processing.bounding_box.bounding_box import bounding_box
from src.processing.cropping_padding.cropping_padding import cropping_padding
from src.processing.return_images.return_images import return_images

# ----------------------------------------------------------------------
# CONDITIONAL IMPORTS
# ----------------------------------------------------------------------
if BACKGROUND_REMOVAL_MODEL == 1:
    from src.processing.remove_background.remove_background_RMBG_2_0 import remove_background
elif BACKGROUND_REMOVAL_MODEL == 2:
    from src.processing.remove_background.remove_background_BiRefNet import remove_background
else:
    raise ValueError(f"Invalid BACKGROUND_REMOVAL_MODEL: {BACKGROUND_REMOVAL_MODEL}")

# ----------------------------------------------------------------------
# PIPELINE STEPS
# ----------------------------------------------------------------------
PIPELINE_STEPS = [
    image_download,
    remove_background,
    bounding_box,
    cropping_padding,
    return_images
]