PyCS/pycs/interfaces/MediaFile.py

from typing import List
from typing import Optional
from typing import Union

from pycs.database.File import File
from pycs.database.Result import Result
from pycs.frontend.notifications.NotificationList import NotificationList
from pycs.interfaces.MediaBoundingBox import MediaBoundingBox
from pycs.interfaces.MediaImageLabel import MediaImageLabel
from pycs.interfaces.MediaLabel import MediaLabel


class MediaFile:
    """
    contains various attributes of a saved media file
    """

    def __init__(self, file: File, notifications: NotificationList):
        self.__file = file
        self.__notifications = notifications

        self.type = file.type
        self.size = file.size
        self.frames = file.frames
        self.fps = file.fps

        self.path = file.absolute_path

    def set_collection(self, reference: Optional[str]):
        """
        set this file's collection

        :param reference: use None to remove this file's collection
        """
        self.__file.set_collection_by_reference(reference)
        self.__notifications.add(self.__notifications.notifications.edit_file, self.__file)

    def set_image_label(self, label: Union[int, MediaLabel], frame: int = None):
        """
        create a labeled-image result

        :param label: label identifier
        :param frame: frame index (only set for videos)
        """
        if label is not None and isinstance(label, MediaLabel):
            label = label.identifier

        if frame is not None:
            data = {'frame': frame}
        else:
            data = None

        created = self.__file.create_result('pipeline', 'labeled-image', label, data)
        self.__notifications.add(self.__notifications.notifications.create_result, created)

    def add_bounding_box(self, x: float, y: float, w: float, h: float,
                         label: Union[int, MediaLabel] = None, frame: int = None):
        """
        create a bounding-box result

        :param x: relative x coordinate [0, 1]
        :param y: relative y coordinate [0, 1]
        :param w: relative width [0, 1]
        :param h: relative height [0, 1]
        :param label: label
        :param frame: frame index (only set for videos)
        """
        result = {
            'x': x,
            'y': y,
            'w': w,
            'h': h
        }
        if frame is not None:
            result['frame'] = frame

        if label is not None and isinstance(label, MediaLabel):
            label = label.identifier

        created = self.__file.create_result('pipeline', 'bounding-box', label, result)
        self.__notifications.add(self.__notifications.notifications.create_result, created)

    def remove_predictions(self):
        """
        remove and return all predictions added from pipelines
        """
        removed = self.__file.remove_results(origin='pipeline')
        for result in removed:
            self.__notifications.add(self.__notifications.notifications.remove_result, result)

    def __get_results(self, origin: str) -> List[Union[MediaImageLabel, MediaBoundingBox]]:

        def result_to_media(result: Result) -> Union[MediaImageLabel, MediaBoundingBox]:
            cls = MediaImageLabel if result.type == 'labeled-image' else MediaBoundingBox
            return cls(result)

        return [result_to_media(r) for r in self.__file.results.filter_by(origin=origin).all()]

    def results(self) -> List[Union[MediaImageLabel, MediaBoundingBox]]:
        """
        receive results added by users

        :return: list of results
        """
        return self.__get_results('user')

    def predictions(self) -> List[Union[MediaImageLabel, MediaBoundingBox]]:
        """
        receive results added by pipelines

        :return: list of predictions
        """
        return self.__get_results('pipeline')

    def serialize(self) -> dict:
        """
        serialize all object properties to a dict

        :return: dict
        """
        return {
            'type': self.type,
            'size': self.size,
            'frames': self.frames,
            'fps': self.fps,
            'path': self.path,
            'filename': self.__file.filename,
            'results': list(map(lambda r: r.serialize(), self.results())),
            'predictions': list(map(lambda r: r.serialize(), self.predictions())),
        }