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(origin='pipeline',
            result_type='labeled-image', label=label, data=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,
                         origin:str = None, origin_user: str = None) -> Result:
        """
        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)
        :param origin: Either pipeline or user
        :param origin_user: Username of the user that provided the bounding box
        :return: Created Result
        """
        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

        if origin is None:
            origin = 'pipeline'
        created = self.__file.create_result(origin=origin, origin_user=origin_user,
            result_type='bounding-box', label=label, data=result)
        self.__notifications.add(self.__notifications.notifications.create_result, created)

        return 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 remove_result(self, result_id):
        """
        Removes the result with the given id.

        :param result_id: id of the result to delete
        """

        removed = self.__file.remove_result(id=result_id)
        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())),
        }