flandre/src/utils/RfMat.py

import cv2
import numpy as np
import cupy as cp

from utils.RfFile import RfFrame, RfSequenceMeta
from utils.RfMeta import RfFrameMeta


def bypass(f):
    def wrapper(self, *args, **kwargs):
        if 'bypass' not in kwargs:
            return f(self, *args, **kwargs)
        if kwargs['bypass']:
            return self
        else:
            del kwargs['bypass']
            return f(self, *args, **kwargs)

    return wrapper


class RfMat:
    def __init__(self,
                 data: cp.ndarray,
                 frame_meta: RfFrameMeta = None,
                 seq_meta: RfSequenceMeta = None,
                 ):
        self.m = data
        self.cv = False
        self.frame_meta = frame_meta
        self.seq_meta = seq_meta
        if isinstance(data, np.ndarray):
            self.device = 'cpu'
        elif isinstance(data, cp.ndarray):
            self.device = 'gpu'
        else:
            raise NotImplementedError

    def call(self, f, *args, **kwargs):
        return self.copy(f(self.m, *args, **kwargs))

    def apply(self, f, *args, **kwargs):
        f(self.m, *args, **kwargs)
        return self

    def copy(self, data=None):
        if data is None:
            return RfMat(self.m.copy(), self.frame_meta, self.seq_meta)
        return RfMat(data, self.frame_meta, self.seq_meta)

    @property
    def p(self):
        if self.device == 'cpu':
            return np
        return cp

    def __bytes__(self):
        return self.m.tobytes()

    def init_cv(self):
        cv2.namedWindow('image')
        self.cv = True

    def norm(self):
        m = self.m.astype(self.p.float32)
        m -= m.min()
        mmax = m.max()
        if mmax == 0:
            return self.copy(self.p.zeros_like(m))
        m /= mmax
        return self.copy(m)

    def grey(self):
        m = self.norm().m
        return self.copy((m * 255).astype(self.p.uint8))

    def cpu(self):
        if self.device == 'cpu':
            return self
        return self.copy(self.m.get())

    def crop(self, t_start: int, t_end: int):
        return self.copy(self.m[:, t_start:t_end])

    def watermark(self, watermark=None):
        assert self.m.dtype == np.uint8
        canvas = np.zeros(self.m.shape, dtype=np.uint8)
        ccp = self.copy()

        line1 = ''
        line2 = ''
        if watermark is not None:
            line1 = watermark
        else:
            if self.frame_meta is not None:
                line1 = self.frame_meta.name
            if self.seq_meta is not None:
                line2 = self.seq_meta.name

        cv2.putText(canvas, line1, (0, 60), cv2.FONT_HERSHEY_PLAIN, 4, (255,), 8)
        cv2.putText(canvas, line1, (0, 60), cv2.FONT_HERSHEY_PLAIN, 4, (128,), 4)

        cv2.putText(canvas, line2, (0, 120), cv2.FONT_HERSHEY_PLAIN, 4, (255,), 8)
        cv2.putText(canvas, line2, (0, 120), cv2.FONT_HERSHEY_PLAIN, 4, (128,), 4)

        if canvas.shape.__len__() == 2:
            ccp.m[canvas == 255] = 255
            ccp.m[canvas == 128] = 0
        elif canvas.shape[2] == 3:
            ccp.m[canvas[:, :] == 255] = 255
            ccp.m[canvas[:, :] == 128] = 0
        else:
            raise NotImplementedError()
        return ccp

    def show(self, shape=None, watermark=None):
        if not self.cv:
            self.init_cv()
        cv2.imshow('image', self
                   .grey()
                   .cpu()
                   .resize(shape, bypass=shape is None)
                   .watermark(watermark).m
                   )
        return cv2.waitKey(0)

    def info(self):
        print(f'shape: {self.m.shape},device: {self.device}')
        print(self.frame_meta)
        print(self.seq_meta)
        return self

    @bypass
    def resize(self, shape):
        if self.device == 'cpu':
            return self.copy(cv2.resize(self.m, shape))
        raise NotImplementedError()

    def rotate90(self):
        return self.copy(self.p.rot90(self.m, k=3))


if __name__ == '__main__':
    cp.zeros((1, 2, 3)) + 1