Video-Summary/Application/LayerFactory.py

import os
from multiprocessing.pool import ThreadPool

import numpy as np

from Application.Config import Config
from Application.Exporter import Exporter
from Application.Layer import Layer
from Application.VideoReader import VideoReader


class LayerFactory:
    def __init__(self, config, data=None):
        self.data = {}
        self.layers = []
        self.tolerance = config["tolerance"]
        self.ttolerance = config["ttolerance"]
        self.minLayerLength = config["minLayerLength"]
        self.maxLayerLength = config["maxLayerLength"]
        self.resizeWidth = config["resizeWidth"]
        self.footagePath = config["inputPath"]
        self.config = config
        print("LayerFactory constructed")
        self.data = data
        if data is not None:
            self.extractLayers(data)

    def extractLayers(self, data, maskArr):
        """Bundle given contours together into Layer Objects"""

        frameNumber = min(data)
        contours = data[frameNumber]
        masks = maskArr[frameNumber]

        for contour, mask in zip(contours, masks):
            mask = np.unpackbits(mask, axis=0)
            self.layers.append(Layer(frameNumber, contour, mask, self.config))

        self.oldLayerIDs = []

        with ThreadPool(os.cpu_count()) as pool:
            for frameNumber in sorted(data.keys()):
                contours = data[frameNumber]
                masks = maskArr[frameNumber]
                masks = [np.unpackbits(mask, axis=0) for mask, contours in zip(masks, contours)]
                if frameNumber % 100 == 0:
                    print(
                        f" {int(round(frameNumber/max(data.keys()), 2)*100)}% done with Layer extraction {len(self.layers)} Layers",
                        end="\r",
                    )

                tmp = [[frameNumber, contour, mask] for contour, mask in zip(contours, masks)]
                # pool.map(self.getLayers, tmp)
                for x in tmp:
                    self.getLayers(x)

        # self.joinLayers()
        return self.layers

    def getLayers(self, data):
        frameNumber = data[0]
        bounds = data[1]
        mask = data[2]
        (x, y, w, h) = bounds
        tol = self.tolerance

        foundLayerIDs = set()
        for i, layer in enumerate(self.layers):
            if frameNumber - layer.lastFrame > self.ttolerance:
                continue

            lastXframes = min(40, len(layer))
            lastBounds = [bound for bounds in layer.bounds[-lastXframes:] for bound in bounds]

            for j, bounds in enumerate(sorted(lastBounds, reverse=True)):
                if bounds is None:
                    break
                (x2, y2, w2, h2) = bounds
                if self.contoursOverlay((x - tol, y + h + tol), (x + w + tol, y - tol), (x2, y2 + h2), (x2 + w2, y2)):
                    layer.add(frameNumber, (x, y, w, h), mask)
                    foundLayerIDs.add(i)
                    break

        foundLayerIDs = sorted(list(foundLayerIDs))
        if len(foundLayerIDs) == 0:
            self.layers.append(Layer(frameNumber, (x, y, w, h), mask, self.config))
        if len(foundLayerIDs) > 1:
            self.mergeLayers(foundLayerIDs)

    def mergeLayers(self, foundLayerIDs):
        layers = self.getLayersByID(foundLayerIDs)
        mergedLayers = layers[0]
        for layer in layers[1:]:
            for i, (contours, masks) in enumerate(zip(layer.bounds, layer.masks)):
                for contour, mask in zip(contours, masks):
                    mergedLayers.add(layer.startFrame + i, contour, mask)

        for i, id in enumerate(foundLayerIDs):
            del self.layers[id - i]

        self.layers.append(mergedLayers)

    def joinLayers(self):
        self.layers.sort(key=lambda c: c.startFrame)
        minFrame = self.getMinStart(self.layers)
        maxFrame = self.getMaxEnd(self.layers)

        for i in range(minFrame, maxFrame):
            pL, indexes = self.getPossibleLayers(i)
            if len(pL) <= 1:
                continue
            merge = set()
            innerMax = self.getMaxEnd(pL)
            for x in range(self.getMinStart(pL), innerMax):
                for lc, l in enumerate(pL):
                    if l.startFrame < x or l.lastFrame > x:
                        continue
                    for lc2, l2 in enumerate(pL):
                        if lc2 == lc:
                            continue
                        for cnt in l.bounds[x - l.startFrame]:
                            for cnt2 in l2.bounds[x - l2.startFrame]:
                                if self.contoursOverlay(cnt, cnt2):
                                    merge.add(indexes[lc])
                                    merge.add(indexes[lc2])
            merge = list(merge)
            if len(merge) > 1:
                self.mergeLayers(merge)
            i = innerMax

    def getPossibleLayers(self, t):
        ret = []
        ii = []
        for i, layer in enumerate(self.layers):
            if layer.startFrame <= t and layer.lastFrame <= t:
                ret.append(layer)
                ii.append(i)
        return (ret, ii)

    def getMinStart(self, layers):
        minFrame = layers[0].startFrame
        for l in layers:
            if l.startFrame < minFrame:
                minFrame = l.startFrame
        return minFrame

    def getMaxEnd(self, layers):
        maxFrame = layers[0].lastFrame
        for l in layers:
            if l.lastFrame < maxFrame:
                maxFrame = l.lastFrame
        return maxFrame

    def contoursOverlay(self, l1, r1, l2, r2):
        if l1[0] >= r2[0] or l2[0] >= r1[0]:
            return False
        if l1[1] <= r2[1] or l2[1] <= r1[1]:
            return False
        return True

    def getLayersByID(self, foundLayerIDs):
        layers = []
        for layerID in foundLayerIDs:
            layers.append(self.layers[layerID])

        layers.sort(key=lambda c: c.startFrame)
        return layers