Book a Demo!
CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutPoliciesSign UpSign In
iperov
GitHub Repository: iperov/deepfacelab
 Path: blob/master/core/imagelib/equalize_and_stack_square.py
628 views
1
import numpy as np

2
import cv2

3


4
def equalize_and_stack_square (images, axis=1):

5
    max_c = max ([ 1 if len(image.shape) == 2 else image.shape[2]  for image in images ] )

6


7
    target_wh = 99999

8
    for i,image in enumerate(images):

9
        if len(image.shape) == 2:

10
            h,w = image.shape

11
            c = 1

12
        else:

13
            h,w,c = image.shape

14


15
        if h < target_wh:

16
            target_wh = h

17


18
        if w < target_wh:

19
            target_wh = w

20


21
    for i,image in enumerate(images):

22
        if len(image.shape) == 2:

23
            h,w = image.shape

24
            c = 1

25
        else:

26
            h,w,c = image.shape

27


28
        if c < max_c:

29
            if c == 1:

30
                if len(image.shape) == 2:

31
                    image = np.expand_dims ( image, -1 )

32
                image = np.concatenate ( (image,)*max_c, -1 )

33
            elif c == 2: #GA

34
                image = np.expand_dims ( image[...,0], -1 )

35
                image = np.concatenate ( (image,)*max_c, -1 )

36
            else:

37
                image = np.concatenate ( (image, np.ones((h,w,max_c - c))), -1 )

38


39
        if h != target_wh or w != target_wh:

40
            image = cv2.resize ( image, (target_wh, target_wh) )

41
            h,w,c = image.shape

42


43
        images[i] = image

44


45
    return np.concatenate ( images, axis = 1 )

46