File size: 1,304 Bytes
7c4d825
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
# ----------------------------------------------------------------------
# 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
]