PyCS/models/fixed_model/Pipeline.py

from json import dump, load
from os import path
from time import sleep

from pycs.interfaces.MediaFile import MediaFile
from pycs.interfaces.MediaStorage import MediaStorage
from pycs.interfaces.Pipeline import Pipeline as Interface


class Pipeline(Interface):
    def __init__(self, root_folder, distribution):
        print('fmv1 init')
        self.root_folder = root_folder

    def close(self):
        print('fmv1 close')

    def execute(self, storage: MediaStorage, file: MediaFile):
        print('fmv1 execute')

        data_file = path.join(self.root_folder, 'data.json')
        if path.exists(data_file):
            with open(data_file, 'r') as f:
                result = load(f)
        else:
            result = {}

        if file.path in result:
            for r in result[file.path]:
                if r['type'] == 'MediaBoundingBox':
                    file.add_bounding_box(r['x'], r['y'], r['w'], r['h'], r['label'], r['frame'])
                if r['type'] == 'MediaImageLabel':
                    file.set_image_label(r['label'], r['frame'])

    def fit(self, storage: MediaStorage):
        print('fmv1 fit')

        for i in range(10):
            yield i / 10
            sleep(1)

        result = {}
        for f in storage.files().iter():
            result[f.path] = list(map(lambda r: dict(r.__dict__, **{'type': type(r).__name__}),
                                      f.results()))

        data_file = path.join(self.root_folder, 'data.json')
        with open(data_file, 'w') as file:
            dump(result, file, indent=4)