DSEBImageInpainting/utils.py at master · cvlab-tohoku/DSEBImageInpainting · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
#!/usr/bin/env python3

import numpy as np
import os
import sys
import time
import shutil
import matplotlib as mpl
# mpl.use("Agg")
import matplotlib.pyplot as plt
from scipy.misc import imresize


def cropCenter(img, cropx, cropy):
    sh = img.shape
    x = sh[0]
    y = sh[1]
    startx = x//2-(cropx//2)
    starty = y//2-(cropy//2)
    return img[startx:startx+cropx, starty:starty+cropy, :]


def createBatch(images, path):
    root = path
    resizeSize = 64

    imgList = os.listdir(root)
    numImg = len(imgList)

    dataX = np.empty((len(images), 64, 64, 3))
    dataY = np.empty((len(images), 64, 64, 3))

    areaX = 32  # num of occluder pixels (x-axis)
    areaY = 32  # num of occluder pixels (y-axis)
    spX = 16
    spY = 16
    n = 0
    for k in images:
        img = plt.imread(root + imgList[k])
        new = cropCenter(img, 96, 96)
        new = imresize(new, (resizeSize, resizeSize), interp='bicubic')
        new = new.astype(float)
        newImg = np.copy(new)
        dataY[n] = newImg
        for i in range(areaX):
            for j in range(areaY):
                new[spX+i, spY+j, :] = 0.
        newImg = np.copy(new)
        dataX[n] = newImg
        n = n + 1
    return dataX, dataY


def createBatchSpec(path):
    root = path
    resizeSize = 64

    imgList = os.listdir(root)
    numImg = len(imgList)
    images = np.arange(numImg)

    dataX = np.empty((len(images), 64, 64, 3))
    dataY = np.empty((len(images), 64, 64, 3))

    areaX = 32  # num of occluder pixels (x-axis)
    areaY = 32  # num of occluder pixels (y-axis)
    spX = 16
    spY = 16
    n = 0
    for k in images:
        img = plt.imread(root + imgList[k])
        new = img.astype(float)
        newImg = np.copy(new)
        dataY[n] = newImg
        for i in range(areaX):
            for j in range(areaY):
                new[spX+i, spY+j, :] = 0.
        newImg = np.copy(new)
        dataX[n] = newImg
        n = n + 1
    return dataX, dataY, numImg


def saveImgs(xs, ys, trueY, save, colWidth=10):
    nImgs = xs.shape[0]
    assert(nImgs == ys.shape[0])

    if not os.path.exists(save):
        os.makedirs(save)

    fnames = []
    for i in range(nImgs):
        xy = np.clip(np.squeeze(np.concatenate([trueY[i], xs[i], ys[i]], axis=1)), 0., 255.)
        fname = "{}/{:04d}.jpg".format(save, i)
        plt.imsave(fname, xy/255.)
        fnames.append(fname)

    os.system('montage -geometry +0+0 -tile {}x {} {}.png'.format(
        colWidth, ' '.join(fnames), save))
    shutil.rmtree(save)