예제 #1
0
    def save_img(self):
        if self.img4save is None:
            return

        lines = [self.img4save]
        if self.learn_from_unlabel:
            lines.append(self.semi_img4save)

        tile_img = None
        for l in lines:
            for i, sect in enumerate(l):
                #l[0] = (l[0] + 1) * 0.5
                if sect is None:
                    l[i] = np.zeros_like(l[0])
                    continue

                if i != 0:
                    sect = onehot2label(sect)
                l[i] = np.transpose(sect, (1, 2, 0))

            l = np.concatenate(l, axis=1)
            if tile_img is None:
                tile_img = l
            else:
                tile_img = np.concatenate((tile_img, l), axis=0)

        out = np.uint8(tile_img * 255)
        out = Image.fromarray(out)

        out_dir_name = self.opt.out_dir + '/out_img'
        os.makedirs(out_dir_name, exist_ok=True)
        out.save(out_dir_name + '/' + str(self.num_saved_img) + '.png')
        self.num_saved_img += 1
예제 #2
0
    def save_img(self):
        out_dir = self.opt.out_dir + '/img/'
        os.makedirs(out_dir, exist_ok=True)
        train = self.img4save
        valid1 = self.evaluate()
        valid2 = self.evaluate()

        tile_img = None
        for seed, real, fake in [train, valid1, valid2]:
            seed = onehot2label(seed, skip_bg=True, dtype='float32')

            tp = (1, 2, 0)
            seed = np.transpose(seed, tp)
            real = np.transpose(real, tp)
            fake = np.transpose(fake, tp)

            real = (real + 1) / 2
            fake = (fake + 1) / 2

            row_img = np.concatenate((seed, real, fake), axis=1)

            if tile_img is None:
                tile_img = row_img
            else:
                tile_img = np.concatenate((tile_img, row_img), axis=0)

        img_array = np.uint8(tile_img * 255)
        img = Image.fromarray(img_array)
        img.save(out_dir + 'tile_img-' + str(self.num_saved_img) + '.png')
        self.num_saved_img += 1
def main():
    out_dir = 'predict_to'
    in_dir = 'predict_from'
    batch_size = 1
    gen_npz = 'pretrained/gen.npz'

    opt = get_options()

    gen = ResNetDeepLab(opt)
    gen.to_gpu(0)
    chainer.serializers.load_npz(gen_npz, gen)
    gen.to_cpu()

    num = 0
    ksize = 5

    os.makedirs(out_dir, exist_ok=True)

    files = glob(in_dir + '/*.*')

    for filename in files:
        print(filename)

        img_array = np.array(Image.open(filename), dtype='float32')
        img_array = img_array.transpose((2, 0, 1)) / 255
        x = chainer.Variable(img_array[np.newaxis, :3, :, :])

        out = gen(x)

        onehot = out.array[0]
        x = x.array[0]

        out = onehot2label(onehot)

        bg_onehot = np.argmax(onehot, axis=0)
        bg_onehot = bg_onehot == 4
        bg_threshold = 0.5

        bg_num = x * bg_onehot
        bg_num = bg_num > bg_threshold
        bg_num = np.sum(bg_num, axis=0)
        bg_num = np.sum(bg_num == 3)

        bg_ratio = bg_num / np.sum(bg_onehot)

        if bg_ratio < 0.6:
            print('bg is black')
            continue

        out = np.transpose(out * 255, (1, 2, 0)).astype('uint8')
        #out = remove_noise(out, ksize=ksize)

        #exist eye ?
        if not is_exist_color(out, [255, 0, 0], threshold_num=32):
            print('not exist eye')
            continue

        #exist face ?
        if not is_exist_color(out, [0, 255, 0], threshold_num=100):
            print('not exist face')
            continue

        #exist hair ?
        if not is_exist_color(out, [0, 0, 255], threshold_num=100):
            print('not exist hair')
            continue

        x = np.transpose(x * 255, (1, 2, 0)).astype('uint8')

        out_img = np.concatenate((x, out), axis=1)
        img = Image.fromarray(out_img)
        path = out_dir + '/' + str(num) + '.png'
        img.save(path)

        num += 1
def main():
    out_predict_dir = 'out'
    device = 0
    gen_npz = 'trained/gen_snapshot_epoch-900.npz'

    opt = get_options()

    opt.spade_ch = 32
    opt.ngf = 64
    opt.ndf = 64

    gen = SPADEGenerator(opt)
    gen.to_gpu(device)
    chainer.serializers.load_npz(gen_npz, gen)

    os.makedirs(out_predict_dir, exist_ok=True)

    out_dir = out_predict_dir + '/predicted'
    os.makedirs(out_dir, exist_ok=True)
    num = 0

    dir_path = 'datasets/resnet-large_hc'
    files = glob(dir_path + '/*.png')
    random.shuffle(files)

    for img_path in files:
        if not os.path.exists(img_path):
            continue

        img = Image.open(img_path)

        if img == None:
            continue
        print(img_path)

        img_array = np.array(img).astype('float32') / 255
        img_array = np.transpose(img_array, (2, 0, 1))

        t_array = img_array[:3, :, :256]
        x_array = img_array[:3, :, 256:512]
        c_array = img_array[:3, :, 512:]

        #to onehot
        x_array = label2onehot(x_array,
                               threshold=0.4,
                               skip_bg=True,
                               dtype='float32').astype('float32')
        c_array = c_array * x_array[2]

        #cast 16bit -> 32bit (cannot use tensor core)
        x = Variable(cuda.to_gpu(x_array[np.newaxis, :, :, :]))
        c = Variable(cuda.to_gpu(c_array[np.newaxis, :, :, :]))

        out = gen([x, c])[0]

        out = cp.asnumpy(out.array[0])
        out = (out + 1) / 2
        x = cp.asnumpy(x.array[0])
        x = onehot2label(x, skip_bg=True, dtype='float32')

        out = np.transpose(out * 255, (1, 2, 0)).astype('uint8')
        x = np.transpose(x * 255, (1, 2, 0)).astype('uint8')

        y = np.transpose(t_array * 255, (1, 2, 0)).astype('uint8')

        out_img = np.concatenate((x, y, out), axis=1)
        img = Image.fromarray(out_img)
        path = out_dir + '/' + str(num) + '.png'
        img.save(path)

        num += 1