post_install_open_terminal.py

from typing import (
    Tuple,
    List
)
from deepomatic.workflows.v2.core import (
    task,
    TaskGroup,
    Region
)
from deepomatic.workflows.v2.cv import (
    detection,
    tagging
)
from deepomatic.workflows.v2.entries import (
    Image
)

from deepomatic.workflows.v2.cv.utils import create_region_from_prediction


class TG(TaskGroup):

    def context_validation(self, image_input: Image) -> bool:

        # will set `post_install_open_terminal`
        # Init
        post_install_open_terminal = False

        context_preds = tagging("context_validator", image_input)
        if self.wo_metadata.get('parameter_meter_location') == "MDU":
            accepted_tags_for_mdu = ['mdu_open_terminal', 'open_terminal', 'close_terminal_pull']
            ctx_pred = context_preds.get_predictions_for_concept(accepted_tags_for_mdu, threshold=0.1)
            post_install_open_terminal = (ctx_pred is not None)
        else:
            accepted_tags = 'open_terminal'
            ctx_pred = context_preds.get_predictions_for_concept(accepted_tags, threshold=0.1)
            post_install_open_terminal = (ctx_pred is not None)

        return post_install_open_terminal

    def main(self, image_input: Image) -> None:

        # If we are outdoor, we will need some detections
        self.items_preds = detection("post_install_open_terminal_items_detector", image_input)

    @task(name="post_install_open_terminal_orange_fuse_or_shroud")
    def post_install_open_terminal_orange_fuse_or_shroud(self, image_input: Image) -> Tuple[bool, List[Region]]:
        # Don't execute if we are in MDU
        self.assert_wo_metadata_is_not("parameter_meter_location", "MDU")

        # Init
        task_value = False
        task_regions = []

        # Check if we have a fuse or a shroud
        fuse_or_shroud_pred = self.items_preds.get_top_prediction_for_concept(
            ['cutout_orange_fuse-1', 'clear_plastic_shroud-1'], threshold=0.3)

        if fuse_or_shroud_pred:
            task_value = True
            task_regions.extend(create_region_from_prediction('image_input', fuse_or_shroud_pred))
            # Store the fuse or shroud pred in the workspace as it will be used in the next task
            self.workspace['fuse_or_shroud_pred'] = fuse_or_shroud_pred

        return task_value, task_regions

    @task(name="post_install_open_terminal_rare_cut_out")
    def post_install_open_terminal_rare_cut_out(self, image_input: Image) -> bool:
        # Don't execute if we are in MDU
        self.assert_wo_metadata_is_not("parameter_meter_location", "MDU")

        # Init
        task_value = False

        # Check if we have a rare cutout
        if 'fuse_or_shroud_pred' in self.workspace and self.workspace['fuse_or_shroud_pred'].get('concept') == 'clear_plastic_shroud-1':
            task_value = True
        return task_value