PyCS/models/fixed_model/Pipeline.py

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

from pycs.interfaces.AnnotatedMediaFile import AnnotatedMediaFile
from pycs.interfaces.MediaFile import MediaFile
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, file: MediaFile) -> List[dict]:
        print('fmv1 execute')
        sleep(0.01)

        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:
            return result[file.path]
        else:
            return []

    def fit(self, files: List[AnnotatedMediaFile]):
        print('fmv1 fit')
        sleep(5)

        result = {}
        for f in files:
            result[f.path] = f.results

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