-
Notifications
You must be signed in to change notification settings - Fork 31
/
test.py
72 lines (60 loc) · 2.91 KB
/
test.py
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
# This file is a part of StarNet code.
# https://github.com/nekitmm/starnet
#
# StarNet is a neural network that can remove stars from images leaving only background.
#
# Throughout the code all input and output images are 8 bits per channel tif images.
# This code in original form will not read any images other than these (like jpeg, etc), but you can change that if you like.
#
# Copyright (c) 2018 Nikita Misiura
# http://www.astrobin.com/users/nekitmm/
#
# This code is distributed on an "AS IS" BASIS WITHOUT WARRANTIES OF ANY KIND, express or implied.
# Please review LICENSE file before use.
import tensorflow as tf
import numpy as np
import model
import starnet_utils
from PIL import Image as img
from scipy.misc import toimage
import os
PANELS = 5 # Number of panels in output pictures.
OUT_DIR = './test/' # Output directory for test showcases.
WINDOW_SIZE = 256 # Size of the image fed to net. Do not change until you know what you are doing! Default is 256
# and changing this will force you to train the net anew.
def test(input, numtests = 20):
if not os.path.exists(OUT_DIR):
os.makedirs(OUT_DIR)
# placeholder for tensorflow
X = tf.placeholder(tf.float32, shape = [None, WINDOW_SIZE, WINDOW_SIZE, 3], name = "X")
Y = tf.placeholder(tf.float32, shape = [None, WINDOW_SIZE, WINDOW_SIZE, 3], name = "Y")
# create model
train, avers, outputs = model.model(X, Y)
#initialize variables
init = tf.global_variables_initializer()
# create saver instance to load model
saver = tf.train.Saver()
with tf.Session() as sess:
# initialize all variables and start training
sess.run(init)
# restore state of the model
print("Restoring current state of the model...")
saver.restore(sess, "./model.ckpt")
print("Done!")
print("Opening test image...")
test = np.array(img.open(input), dtype = np.float32)
print("Done!")
head = np.array(img.open("./test_head.tif"), dtype = np.float32)
head /= 255
for i in range(numtests):
print("Test %d..." % i)
all = np.zeros((1, WINDOW_SIZE * 2, 3))
all = np.concatenate((head, all), axis = 0)
for d in range(PANELS):
X_test = starnet_utils.get_test_samples(test, 1)
output = sess.run(outputs, feed_dict = {X:X_test})
im = (np.concatenate((X_test[0], output[0]), axis = 1) + 1 ) / 2
all = np.concatenate((all, im), axis = 0)
all = np.concatenate((all, np.zeros((1, WINDOW_SIZE * 2, 3))), axis = 0)
toimage(all * 255, cmin = 0, cmax = 255).save(OUT_DIR + '/test_' + str(i) + '.tif')
print("Done!")