Pipeline Module

This module provides a custom segmentation pipeline for image and video processing.

It extends the functionality of the Hugging Face Transformers library's ImageSegmentationPipeline to support various segmentation models and create detailed segmentation maps with associated metadata.

ModelConfig

ModelConfig(
    name,
    model_type=None,
    max_size=None,
    device=None,
    dataset=None,
    num_workers=8,
    pipe_batch=1,
)

Configuration class for the segmentation model.

ATTRIBUTE	DESCRIPTION
name	The name or path of the pre-trained model to use. TYPE: str
model_type	The type of the model (e.g., 'oneformer', 'mask2former'). TYPE: Optional[str]
max_size	The maximum size for input image resizing. TYPE: Optional[int]
device	The device to use for processing (e.g., 'cpu', 'cuda'). TYPE: Optional[str]

__post_init__

__post_init__()

Post-initialization method to set up the model type if not provided.

Source code in cityseg/config.py

def __post_init__(self):
    """
    Post-initialization method to set up the model type if not provided.
    """
    self.auto_detect_model_type()
    if self.device == "mps" and self.num_workers > 0 or self.num_workers is None:
        logger.warning(
            "MPS is not compatible with multiple workers in pytorch. Setting num_workers to 0."
        )
        self.num_workers = 0

auto_detect_model_type

auto_detect_model_type()

Automatically detect the model type from the model name if not provided.

Source code in cityseg/config.py

def auto_detect_model_type(self):
    """
    Automatically detect the model type from the model name if not provided.
    """

    def auto_model_type(model_name: str) -> str:
        return model_name.split("/")[-1].split("-")[0]

    if self.model_type is None:
        try:
            self.model_type = auto_model_type(self.name)
        except IndexError:
            logger.warning(
                "Unable to auto-detect model type from the model name and none provided."
            )
            return
        logger.info(f"Auto-detected model type: {self.model_type}")
    elif self.model_type != auto_model_type(self.name):
        logger.warning(
            f"Model type does not match auto-detected model type. Using provided model type: {self.model_type}"
        )

SegmentationPipeline

SegmentationPipeline(*args, **kwargs)

Bases: ImageSegmentationPipeline

A custom segmentation pipeline that extends ImageSegmentationPipeline.

This class provides additional functionality for creating and processing segmentation maps, including support for different color palettes and batch processing of images.

ATTRIBUTE	DESCRIPTION
palette	The color palette used for visualization. TYPE: ndarray

Source code in cityseg/pipeline.py

def __init__(self, *args, **kwargs) -> None:
    super().__init__(*args, **kwargs)
    self.palette = self._get_palette()

create_single_segmentation_map

create_single_segmentation_map(annotations, target_size)

Create a single segmentation map from annotations.

PARAMETER	DESCRIPTION
annotations	List of annotation dictionaries. TYPE: List[Dict[str, Any]]
target_size	The target size of the segmentation map. TYPE: tuple

RETURNS	DESCRIPTION
Dict[str, Any]	Dict[str, Any]: A dictionary containing the segmentation map and associated metadata.

Source code in cityseg/pipeline.py

def create_single_segmentation_map(
    self, annotations: List[Dict[str, Any]], target_size: tuple
) -> Dict[str, Any]:
    """
    Create a single segmentation map from annotations.

    Args:
        annotations (List[Dict[str, Any]]): List of annotation dictionaries.
        target_size (tuple): The target size of the segmentation map.

    Returns:
        Dict[str, Any]: A dictionary containing the segmentation map and associated metadata.
    """
    seg_map = np.zeros(target_size, dtype=np.int32)
    for annotation in annotations:
        mask = np.array(annotation["mask"])
        label_id = self.model.config.label2id[annotation["label"]]
        seg_map[mask != 0] = label_id

    return {
        "seg_map": seg_map,
        "label2id": self.model.config.label2id,
        "id2label": self.model.config.id2label,
        "palette": self.palette,
    }

__call__

__call__(images, **kwargs)

Process the input image(s) and create segmentation map(s).

PARAMETER	DESCRIPTION
images	The input image(s) to process. TYPE: Union[Image, List[Image]]
**kwargs	Additional keyword arguments. TYPE: Any DEFAULT: {}

RETURNS	DESCRIPTION
List[Dict[str, Any]]	List[Dict[str, Any]]: A list of dictionaries containing segmentation maps and metadata.

Source code in cityseg/pipeline.py

def __call__(
    self, images: Union[Image, List[Image]], **kwargs: Any
) -> List[Dict[str, Any]]:
    """
    Process the input image(s) and create segmentation map(s).

    Args:
        images (Union[Image, List[Image]]): The input image(s) to process.
        **kwargs: Additional keyword arguments.

    Returns:
        List[Dict[str, Any]]: A list of dictionaries containing segmentation maps and metadata.
    """
    # logger.debug("Pass image(s) up to HF pipeline...")
    result = super().__call__(images, subtask="semantic", **kwargs)
    # logger.debug("Received result from HF pipeline")
    if self._is_single_image_result(result):
        return [
            self.create_single_segmentation_map(
                result, result[0]["mask"].size[::-1]
            )
        ]
    else:
        return [
            self.create_single_segmentation_map(
                img_result, img_result[0]["mask"].size[::-1]
            )
            for img_result in result
        ]

create_segmentation_pipeline

create_segmentation_pipeline(config, **kwargs)

Create and return a SegmentationPipeline instance based on the specified model.

This function initializes the appropriate model and image processor based on the model name, and creates a SegmentationPipeline instance with these components.

PARAMETER	DESCRIPTION
config	TYPE: ModelConfig
**kwargs	Additional keyword arguments to pass to the SegmentationPipeline constructor. TYPE: Any DEFAULT: {}

RETURNS	DESCRIPTION
SegmentationPipeline	An instance of the SegmentationPipeline class. TYPE: SegmentationPipeline

Source code in cityseg/pipeline.py

@logger.catch
def create_segmentation_pipeline(
    config: ModelConfig, **kwargs: Any
) -> SegmentationPipeline:
    """
    Create and return a SegmentationPipeline instance based on the specified model.

    This function initializes the appropriate model and image processor based on the
    model name, and creates a SegmentationPipeline instance with these components.

    Args:
        config:
        **kwargs: Additional keyword arguments to pass to the SegmentationPipeline constructor.

    Returns:
        SegmentationPipeline: An instance of the SegmentationPipeline class.
    """
    model_name = config.name
    model_type = config.model_type
    device = config.device
    dataset = config.dataset

    if device is None:
        device = (
            "cuda"
            if torch.cuda.is_available()
            else "mps"
            if torch.backends.mps.is_available()
            else "cpu"
        )

    # Initialize the appropriate model and image processor based on the model name
    if "oneformer" == model_type:
        warnings.warn(
            "OneFormer models are experimental and may not be fully supported"
        )
        try:
            model = OneFormerForUniversalSegmentation.from_pretrained(model_name)
            image_processor = AutoProcessor.from_pretrained(model_name)
        except ValueError as e:
            logger.error(f"Error loading model: {e}")

    elif "mask2former" == model_type:
        model = Mask2FormerForUniversalSegmentation.from_pretrained(model_name)
        image_processor = AutoImageProcessor.from_pretrained(model_name)

    elif "maskformer" == model_type:
        model = MaskFormerForInstanceSegmentation.from_pretrained(model_name)
        image_processor = AutoImageProcessor.from_pretrained(model_name)

    elif "beit" == model_type:
        if device != "cpu":
            logger.warning(
                "Beit models are not supported on GPU and will be loaded on CPU"
            )
        device = "cpu"
        model = BeitForSemanticSegmentation.from_pretrained(model_name)
        image_processor = AutoImageProcessor.from_pretrained(model_name)

    elif "segformer" == model_type:
        model = SegformerForSemanticSegmentation.from_pretrained(model_name)
        image_processor = AutoImageProcessor.from_pretrained(model_name)

        if dataset == "sidewalk-semantic":
            logging.debug("Loading Sidewalk Semantic dataset label mappings...")
            with open("SemanticSidewalk_id2label.json") as f:
                id2label = json.load(f)
            model.config.id2label = id2label
    else:
        model = AutoModelForSemanticSegmentation.from_pretrained(model_name)
        image_processor = AutoImageProcessor.from_pretrained(model_name)

    return SegmentationPipeline(
        model=model,
        image_processor=image_processor,
        device=device,
        subtask="semantic",
        num_workers=config.num_workers,
        **kwargs,
    )