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.min_layer_length = config["minLayerLength"]
        self.max_layer_length = config["maxLayerLength"]
        self.resize_width = config["resizeWidth"]
        self.footage_path = config["inputPath"]
        self.config = config
        print("LayerFactory constructed")
        self.data = data
        if data is not None:
            self.extract_layers(data)

    def extract_layers(self, data, mask_arr):
        """Bundle given contours together into Layer Objects"""

        frame_number = min(data)
        contours = data[frame_number]
        masks = mask_arr[frame_number]

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

        self.old_layer_i_ds = []

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

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

        # self.joinLayers()
        return self.layers

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

        found_layer_i_ds = set()
        for i, layer in enumerate(self.layers):
            if frame_number - layer.last_frame > self.ttolerance:
                continue

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

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

        found_layer_i_ds = sorted(list(found_layer_i_ds))
        if len(found_layer_i_ds) == 0:
            self.layers.append(Layer(frame_number, (x, y, w, h), mask, self.config))
        if len(found_layer_i_ds) > 1:
            self.merge_layers(found_layer_i_ds)

    def merge_layers(self, found_layer_i_ds):
        layers = self.get_layers_by_id(found_layer_i_ds)
        merged_layers = 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):
                    merged_layers.add(layer.startFrame + i, contour, mask)

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

        self.layers.append(merged_layers)

    def join_layers(self):
        self.layers.sort(key=lambda c: c.startFrame)
        min_frame = self.get_min_start(self.layers)
        max_frame = self.get_max_end(self.layers)

        for i in range(min_frame, max_frame):
            p_l, indexes = self.get_possible_layers(i)
            if len(p_l) <= 1:
                continue
            merge = set()
            inner_max = self.get_max_end(p_l)
            for x in range(self.get_min_start(p_l), inner_max):
                for lc, l in enumerate(p_l):
                    if l.startFrame < x or l.lastFrame > x:
                        continue
                    for lc2, l2 in enumerate(p_l):
                        if lc2 == lc:
                            continue
                        for cnt in l.bounds[x - l.startFrame]:
                            for cnt2 in l2.bounds[x - l2.startFrame]:
                                if self.contours_overlay(cnt, cnt2):
                                    merge.add(indexes[lc])
                                    merge.add(indexes[lc2])
            merge = list(merge)
            if len(merge) > 1:
                self.merge_layers(merge)
            i = inner_max

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

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

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

    def contours_overlay(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 get_layers_by_id(self, found_layer_i_ds):
        layers = []
        for layer_id in found_layer_i_ds:
            layers.append(self.layers[layer_id])

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