from pycs.pipeline.Job import Job


class Fit(Job):
    def __init__(self, project, data: dict):
        super().__init__(project, data)

        self.result = []
        for identifier in data['predictionResults']:
            if data['predictionResults'][identifier]['origin'] != 'user':
                continue

            copy = data['predictionResults'][identifier].copy()
            del copy['id']
            del copy['origin']

            # TODO remove and add in webui endpoint
            if 'x' not in copy:
                copy['type'] = 'labeled-image'
            elif 'label' in copy:
                copy['type'] = 'labeled-bounding-box'
            else:
                copy['type'] = 'bounding-box'

            self.result.append(copy)