├── LICENSE ├── README.md ├── data ├── x_0_0.png └── y_0_0.png ├── datagen.py ├── figs └── overview.png ├── input ├── demo1.png └── demo2.png ├── model1.py ├── model2.py ├── model3.py ├── model4.py ├── model5.py ├── output ├── demo1.png └── demo2.png ├── predict.py ├── predict_block.py ├── pytorch_model1.py ├── pytorch_model2.py ├── pytorch_predict.py ├── pytorch_predict_block.py └── weight └── .gitkeep /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2018 HEPESU 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # LineNormalizer 2 |

3 | 4 |

5 | 6 | ## Overview 7 | Normalize rasterization line-drawings to uniform width using deep learning with model from [Smart Inker](http://hi.cs.waseda.ac.jp/~esimo/en/research/inking/). 8 | 9 | This model can serve as line-drawings preprocessor for [LineRelifer](https://github.com/hepesu/LineRelifer/). Line-drawings can be normlized to an intermediate representation and then be used as training data or input for it. 10 | Also by using this method, we can achieve uniform line width during scaling up or down the rasterization line-drawings, which is a feature of vector line-drawings. The train data is generated by code, so you can get model for any width easily. 11 | 12 | ## Dependencies 13 | * Keras2 (Tensorflow1 backend) 14 | * Pytorch 15 | * OpenCV3 16 | * CairoSVG 17 | 18 | ## Usage 19 | 1. Set up directories. 20 | 21 | 2. Download the model from release and put it in the same folder with code. 22 | 23 | 3. Run `predict.py` for prediction. Run `model{NUM}.py` for train. 24 | 25 | Files with name starts with `pytorch` are Pytorch version. 26 | 27 | ## Models 28 | Models are licensed under a CC-BY-NC-SA 4.0 international license. 29 | * [LineNormalizer Release Page](https://github.com/hepesu/LineNormalizer/releases) 30 | * [BaiduPan](https://pan.baidu.com/s/1tooIPfJX9LVA1aY5spNWiw) [Code: bnaw] 31 | 32 | #### Keras 33 | * model_180913 34 | 35 | #### Pytorch 36 | * model_200801 37 | 38 | From **Project HAT** by Hepesu With :heart: 39 | -------------------------------------------------------------------------------- /data/x_0_0.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hepesu/LineNormalizer/47e41d3973c255185b696e62dd46a850a84e1323/data/x_0_0.png -------------------------------------------------------------------------------- /data/y_0_0.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hepesu/LineNormalizer/47e41d3973c255185b696e62dd46a850a84e1323/data/y_0_0.png -------------------------------------------------------------------------------- /datagen.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Modified from https://github.com/byungsook/vectornet/blob/master/data_line.py 3 | 4 | Semantic Segmentation for Line Drawing Vectorization Using Neural Networks 5 | Tensorflow implementation of Semantic Segmentation for Line Drawing Vectorization Using Neural Networks. 6 | 7 | Byungsoo Kim1, Oliver Wang2, Cengiz ?ztireli1, Markus Gross1 8 | 9 | 1ETH Zurich, 2Adobe Research 10 | 11 | Computer Graphics Forum (Proceedings of Eurographics 2018) 12 | 13 | ''' 14 | 15 | import numpy as np 16 | import cv2 17 | from PIL import Image 18 | import cairosvg 19 | import io 20 | 21 | SEED = 1 22 | WIDTH, HEIGHT = 128, 128 23 | 24 | MIN_STROKE_WIDTH, MAX_STROKE_WIDTH = 0.2, 2 25 | MAX_STROKE_COLOR = 50 26 | 27 | NORM_STROKE_WIDTH = 0.5 28 | 29 | MAX_NUM_STROKES = 10 30 | 31 | SVG_START_TEMPLATE = """ 32 | """ 45 | 46 | 47 | def draw_line(id, w, h, rng): 48 | stroke_color = rng.randint(MAX_STROKE_COLOR) 49 | stroke_width = rng.rand() * (MAX_STROKE_WIDTH - MIN_STROKE_WIDTH) + MIN_STROKE_WIDTH 50 | x = rng.randint(w, size=2) 51 | y = rng.randint(h, size=2) 52 | 53 | return SVG_LINE_TEMPLATE.format( 54 | id=id, 55 | x1=x[0], y1=y[0], 56 | x2=x[1], y2=y[1], 57 | r=stroke_color, g=stroke_color, b=stroke_color, 58 | sw=stroke_width 59 | ) 60 | 61 | 62 | def draw_cubic_bezier_curve(id, w, h, rng): 63 | stroke_color = rng.randint(MAX_STROKE_COLOR) 64 | stroke_width = rng.rand() * (MAX_STROKE_WIDTH - MIN_STROKE_WIDTH) + MIN_STROKE_WIDTH 65 | x = rng.randint(w, size=4) 66 | y = rng.randint(h, size=4) 67 | 68 | return SVG_CUBIC_BEZIER_TEMPLATE.format( 69 | id=id, 70 | sx=x[0], sy=y[0], 71 | cx1=x[1], cy1=y[1], 72 | cx2=x[2], cy2=y[2], 73 | tx=x[3], ty=y[3], 74 | r=stroke_color, g=stroke_color, b=stroke_color, 75 | sw=stroke_width 76 | ) 77 | 78 | 79 | def draw_rect(id, w, h, rng): 80 | stroke_color = rng.randint(MAX_STROKE_COLOR) 81 | stroke_width = rng.rand() * (MAX_STROKE_WIDTH - MIN_STROKE_WIDTH) + MIN_STROKE_WIDTH 82 | x = rng.randint(w) 83 | y = rng.randint(h) 84 | w = rng.randint(low=w // 4, high=w // 2) 85 | h = rng.randint(low=h // 4, high=h // 2) 86 | 87 | return SVG_RECT_TEMPLATE.format( 88 | id=id, 89 | x=x, y=y, 90 | w=w, h=h, 91 | r=stroke_color, g=stroke_color, b=stroke_color, 92 | sw=stroke_width 93 | ) 94 | 95 | 96 | def draw_ellipse(id, w, h, rng): 97 | stroke_color = rng.randint(MAX_STROKE_COLOR) 98 | stroke_width = rng.rand() * (MAX_STROKE_WIDTH - MIN_STROKE_WIDTH) + MIN_STROKE_WIDTH 99 | x = rng.randint(w) 100 | y = rng.randint(h) 101 | rx = rng.randint(low=w // 4, high=w // 2) 102 | ry = rng.randint(low=h // 4, high=h // 2) 103 | 104 | return SVG_ELLIPSE_TEMPLATE.format( 105 | id=id, 106 | x=x, y=y, 107 | rx=rx, ry=ry, 108 | r=stroke_color, g=stroke_color, b=stroke_color, 109 | sw=stroke_width 110 | ) 111 | 112 | 113 | def draw_path(id, w, h, rng): 114 | path_selector = { 115 | 0: draw_line, 116 | 1: draw_cubic_bezier_curve, 117 | 2: draw_rect, 118 | 3: draw_ellipse 119 | } 120 | 121 | stroke_type = rng.randint(len(path_selector)) 122 | 123 | return path_selector[stroke_type](id, w, h, rng) 124 | 125 | 126 | def gen_data(rng, batch_size): 127 | x = [] 128 | y = [] 129 | 130 | norm_stroke_width_txt = """stroke-width="{sw}" _stroke-width""".format(sw=NORM_STROKE_WIDTH) 131 | for file_id in range(batch_size): 132 | while True: 133 | svg = SVG_START_TEMPLATE.format( 134 | w=WIDTH, 135 | h=HEIGHT, 136 | rot=rng.randint(0, 180) 137 | ) 138 | svgpre = SVG_START_TEMPLATE 139 | 140 | for i in range(rng.randint(MAX_NUM_STROKES) + 1): 141 | path = draw_path( 142 | id=i, 143 | w=WIDTH, 144 | h=HEIGHT, 145 | rng=rng 146 | ) 147 | svg += path + '\n' 148 | svgpre += path + '\n' 149 | 150 | svg += SVG_END_TEMPLATE 151 | 152 | x_png = cairosvg.svg2png(bytestring=svg.encode('utf-8')) 153 | x_img = Image.open(io.BytesIO(x_png)) 154 | x_arr = np.array(x_img, np.float) 155 | 156 | # with open('data/s.svg', 'w') as f: 157 | # f.write(svg.replace('stroke-width', norm_stroke_width_txt)) 158 | 159 | y_png = cairosvg.svg2png(bytestring=svg.replace('stroke-width', norm_stroke_width_txt).encode('utf-8')) 160 | y_img = Image.open(io.BytesIO(y_png)) 161 | y_arr = np.array(y_img, np.float) 162 | 163 | if np.mean(x_arr) < 200 or np.mean(x_arr) > 245: 164 | continue 165 | else: 166 | x.append(np.reshape(x_arr[:, :, 0], (HEIGHT, WIDTH, 1))) 167 | y.append(np.reshape(y_arr[:, :, 0], (HEIGHT, WIDTH, 1))) 168 | break 169 | 170 | return np.array(x).astype(np.float32) / 255.0, np.array(y).astype(np.float32) / 255.0 171 | 172 | 173 | def test(): 174 | rnd = np.random.RandomState(SEED) 175 | 176 | for i in range(5): 177 | x_data, y_data = gen_data(rnd, 4) 178 | for j in range(4): 179 | cv2.imwrite('data/x_%d_%d.png' % (i, j), x_data[j] * 255) 180 | cv2.imwrite('data/y_%d_%d.png' % (i, j), y_data[j] * 255) 181 | 182 | 183 | if __name__ == "__main__": 184 | test() 185 | -------------------------------------------------------------------------------- /figs/overview.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hepesu/LineNormalizer/47e41d3973c255185b696e62dd46a850a84e1323/figs/overview.png -------------------------------------------------------------------------------- /input/demo1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hepesu/LineNormalizer/47e41d3973c255185b696e62dd46a850a84e1323/input/demo1.png -------------------------------------------------------------------------------- /input/demo2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hepesu/LineNormalizer/47e41d3973c255185b696e62dd46a850a84e1323/input/demo2.png -------------------------------------------------------------------------------- /model1.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import cv2 3 | from keras.preprocessing import image 4 | from keras.models import Model 5 | from keras.layers import Input, Conv2D, Deconv2D, Activation, BatchNormalization, add 6 | from keras.callbacks import ModelCheckpoint 7 | 8 | from datagen import gen_data 9 | 10 | SEED = 1 11 | 12 | EPOCHS = 40 13 | BATCH_SIZE = 4 14 | LOAD_WEIGHTS = False 15 | 16 | IMG_HEIGHT, IMG_WIDTH = 128, 128 17 | 18 | inputs = Input((None, None, 1)) 19 | 20 | x = Conv2D(64, 9, padding='same')(inputs) 21 | x = BatchNormalization()(x) 22 | x = Activation('relu')(x) 23 | 24 | x = Conv2D(64, 3, padding='same')(x) 25 | x = BatchNormalization()(x) 26 | x = Activation('relu')(x) 27 | 28 | x = Conv2D(64, 3, padding='same')(x) 29 | x = BatchNormalization()(x) 30 | x = Activation('relu')(x) 31 | 32 | x = Conv2D(64, 3, padding='same')(x) 33 | x = BatchNormalization()(x) 34 | x = Activation('relu')(x) 35 | 36 | x = Conv2D(64, 3, padding='same')(x) 37 | x = BatchNormalization()(x) 38 | x = Activation('relu')(x) 39 | 40 | x = Conv2D(64, 3, padding='same')(x) 41 | x = BatchNormalization()(x) 42 | x = Activation('relu')(x) 43 | 44 | x = Conv2D(64, 3, padding='same')(x) 45 | x = BatchNormalization()(x) 46 | x = Activation('relu')(x) 47 | 48 | x = Conv2D(64, 3, padding='same')(x) 49 | x = BatchNormalization()(x) 50 | x = Activation('relu')(x) 51 | 52 | outputs = Conv2D(1, 3, padding='same', activation='sigmoid')(x) 53 | 54 | model = Model(inputs=inputs, outputs=outputs) 55 | model.summary() 56 | 57 | if LOAD_WEIGHTS: 58 | model.load_weights('model1.h5') 59 | 60 | model.compile(loss='MSE', optimizer='Adam') 61 | 62 | checkpointer = ModelCheckpoint(filepath='model1.h5', verbose=1) 63 | 64 | 65 | def _train_generator(): 66 | rnd = np.random.RandomState(SEED) 67 | while True: 68 | yield gen_data(rnd, BATCH_SIZE) 69 | 70 | 71 | def _val_generator(): 72 | rnd = np.random.RandomState(SEED + 1) 73 | while True: 74 | yield gen_data(rnd, BATCH_SIZE) 75 | 76 | 77 | train_generator = _train_generator() 78 | val_generator = _val_generator() 79 | 80 | history = model.fit_generator( 81 | train_generator, 82 | steps_per_epoch=512 // BATCH_SIZE, 83 | epochs=EPOCHS, 84 | validation_data=val_generator, 85 | validation_steps=32 // BATCH_SIZE, 86 | callbacks=[checkpointer] 87 | ) 88 | 89 | model.save('model1_final.h5') 90 | -------------------------------------------------------------------------------- /model2.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import cv2 3 | from keras.preprocessing import image 4 | from keras.models import Model 5 | from keras.layers import Input, Conv2D, Deconv2D, Activation, BatchNormalization, add 6 | from keras.callbacks import ModelCheckpoint 7 | 8 | from datagen import gen_data 9 | 10 | SEED = 1 11 | 12 | EPOCHS = 40 13 | BATCH_SIZE = 4 14 | LOAD_WEIGHTS = False 15 | 16 | IMG_HEIGHT, IMG_WIDTH = 128, 128 17 | 18 | inputs = Input((None, None, 1)) 19 | 20 | x = Conv2D(64, 9, padding='same')(inputs) 21 | x = BatchNormalization()(x) 22 | x = Activation('relu')(x) 23 | 24 | x = Conv2D(64, 3, padding='same')(x) 25 | x = BatchNormalization()(x) 26 | x = Activation('relu')(x) 27 | 28 | x = Conv2D(64, 3, padding='same')(x) 29 | x = BatchNormalization()(x) 30 | x = Activation('relu')(x) 31 | 32 | x = Conv2D(64, 3, padding='same')(x) 33 | x = BatchNormalization()(x) 34 | x = Activation('relu')(x) 35 | 36 | outputs = Conv2D(1, 3, padding='same', activation='sigmoid')(x) 37 | 38 | model = Model(inputs=inputs, outputs=outputs) 39 | model.summary() 40 | 41 | if LOAD_WEIGHTS: 42 | model.load_weights('model2.h5') 43 | 44 | model.compile(loss='MSE', optimizer='Adam') 45 | 46 | checkpointer = ModelCheckpoint(filepath='model2.h5', verbose=1) 47 | 48 | 49 | def _train_generator(): 50 | rnd = np.random.RandomState(SEED) 51 | while True: 52 | yield gen_data(rnd, BATCH_SIZE) 53 | 54 | 55 | def _val_generator(): 56 | rnd = np.random.RandomState(SEED + 1) 57 | while True: 58 | yield gen_data(rnd, BATCH_SIZE) 59 | 60 | 61 | train_generator = _train_generator() 62 | val_generator = _val_generator() 63 | 64 | history = model.fit_generator( 65 | train_generator, 66 | steps_per_epoch=512 // BATCH_SIZE, 67 | epochs=EPOCHS, 68 | validation_data=val_generator, 69 | validation_steps=32 // BATCH_SIZE, 70 | callbacks=[checkpointer] 71 | ) 72 | 73 | model.save('model2_final.h5') 74 | -------------------------------------------------------------------------------- /model3.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import cv2 3 | from keras.preprocessing import image 4 | from keras.models import Model 5 | from keras.layers import Input, Conv2D, Deconv2D, Activation, BatchNormalization, add 6 | from keras.callbacks import ModelCheckpoint 7 | 8 | from datagen import gen_data 9 | 10 | SEED = 1 11 | 12 | EPOCHS = 40 13 | BATCH_SIZE = 4 14 | LOAD_WEIGHTS = False 15 | 16 | IMG_HEIGHT, IMG_WIDTH = 128, 128 17 | 18 | inputs = Input((None, None, 1)) 19 | 20 | x = Conv2D(32, 9, padding='same')(inputs) 21 | x = BatchNormalization()(x) 22 | x = Activation('relu')(x) 23 | 24 | x = Conv2D(32, 3, padding='same')(x) 25 | x = BatchNormalization()(x) 26 | x = Activation('relu')(x) 27 | 28 | x = Conv2D(32, 3, padding='same')(x) 29 | x = BatchNormalization()(x) 30 | x = Activation('relu')(x) 31 | 32 | x = Conv2D(32, 3, padding='same')(x) 33 | x = BatchNormalization()(x) 34 | x = Activation('relu')(x) 35 | 36 | outputs = Conv2D(1, 3, padding='same', activation='sigmoid')(x) 37 | 38 | model = Model(inputs=inputs, outputs=outputs) 39 | model.summary() 40 | 41 | if LOAD_WEIGHTS: 42 | model.load_weights('model3.h5') 43 | 44 | model.compile(loss='MSE', optimizer='Adam') 45 | 46 | checkpointer = ModelCheckpoint(filepath='model3.h5', verbose=1) 47 | 48 | 49 | def _train_generator(): 50 | rnd = np.random.RandomState(SEED) 51 | while True: 52 | yield gen_data(rnd, BATCH_SIZE) 53 | 54 | 55 | def _val_generator(): 56 | rnd = np.random.RandomState(SEED + 1) 57 | while True: 58 | yield gen_data(rnd, BATCH_SIZE) 59 | 60 | 61 | train_generator = _train_generator() 62 | val_generator = _val_generator() 63 | 64 | history = model.fit_generator( 65 | train_generator, 66 | steps_per_epoch=512 // BATCH_SIZE, 67 | epochs=EPOCHS, 68 | validation_data=val_generator, 69 | validation_steps=32 // BATCH_SIZE, 70 | callbacks=[checkpointer] 71 | ) 72 | 73 | model.save('model3_final.h5') 74 | -------------------------------------------------------------------------------- /model4.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import cv2 3 | from keras.preprocessing import image 4 | from keras.models import Model 5 | from keras.layers import Input, Conv2D, Deconv2D, Activation, BatchNormalization, Add 6 | from keras.callbacks import ModelCheckpoint 7 | 8 | from datagen import gen_data 9 | 10 | SEED = 1 11 | 12 | EPOCHS = 40 13 | BATCH_SIZE = 4 14 | LOAD_WEIGHTS = False 15 | 16 | 17 | def resnet_block(layer_input, filters=64, f_size=3): 18 | y = Conv2D(filters, (f_size, f_size), padding='same')(layer_input) 19 | y = BatchNormalization()(y) 20 | y = Activation('relu')(y) 21 | 22 | y = Conv2D(filters, (f_size, f_size), padding='same')(y) 23 | y = BatchNormalization()(y) 24 | 25 | y = Add()([layer_input, y]) 26 | 27 | return Activation('relu')(y) 28 | 29 | 30 | IMG_HEIGHT, IMG_WIDTH = 128, 128 31 | 32 | inputs = Input((None, None, 1)) 33 | 34 | x = Conv2D(64, 9, padding='same')(inputs) 35 | x = BatchNormalization()(x) 36 | x = Activation('relu')(x) 37 | 38 | x = resnet_block(x) 39 | x = resnet_block(x) 40 | x = resnet_block(x) 41 | 42 | outputs = Conv2D(1, 3, padding='same', activation='sigmoid')(x) 43 | 44 | model = Model(inputs=inputs, outputs=outputs) 45 | model.summary() 46 | 47 | if LOAD_WEIGHTS: 48 | model.load_weights('model4.h5') 49 | 50 | model.compile(loss='MSE', optimizer='Adam') 51 | 52 | checkpointer = ModelCheckpoint(filepath='model4.h5', verbose=1) 53 | 54 | 55 | def _train_generator(): 56 | rnd = np.random.RandomState(SEED) 57 | while True: 58 | yield gen_data(rnd, BATCH_SIZE) 59 | 60 | 61 | def _val_generator(): 62 | rnd = np.random.RandomState(SEED + 1) 63 | while True: 64 | yield gen_data(rnd, BATCH_SIZE) 65 | 66 | 67 | train_generator = _train_generator() 68 | val_generator = _val_generator() 69 | 70 | history = model.fit_generator( 71 | train_generator, 72 | steps_per_epoch=512 // BATCH_SIZE, 73 | epochs=EPOCHS, 74 | validation_data=val_generator, 75 | validation_steps=32 // BATCH_SIZE, 76 | callbacks=[checkpointer] 77 | ) 78 | 79 | model.save('model4_final.h5') 80 | -------------------------------------------------------------------------------- /model5.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import cv2 3 | from keras.preprocessing import image 4 | from keras.models import Model 5 | from keras.layers import Input, Conv2D, Deconv2D, Activation, BatchNormalization, Add 6 | from keras.callbacks import ModelCheckpoint 7 | 8 | from datagen import gen_data 9 | 10 | SEED = 1 11 | 12 | EPOCHS = 40 13 | BATCH_SIZE = 4 14 | LOAD_WEIGHTS = False 15 | 16 | 17 | def resnet_block(layer_input, filters=32, f_size=3): 18 | y = Conv2D(filters, (f_size, f_size), padding='same')(layer_input) 19 | y = BatchNormalization()(y) 20 | y = Activation('relu')(y) 21 | 22 | y = Conv2D(filters, (f_size, f_size), padding='same')(y) 23 | y = BatchNormalization()(y) 24 | 25 | y = Add()([layer_input, y]) 26 | 27 | return Activation('relu')(y) 28 | 29 | 30 | IMG_HEIGHT, IMG_WIDTH = 128, 128 31 | 32 | inputs = Input((None, None, 1)) 33 | 34 | x = Conv2D(32, 9, padding='same')(inputs) 35 | x = BatchNormalization()(x) 36 | x = Activation('relu')(x) 37 | 38 | x = resnet_block(x) 39 | x = resnet_block(x) 40 | x = resnet_block(x) 41 | x = resnet_block(x) 42 | x = resnet_block(x) 43 | x = resnet_block(x) 44 | x = resnet_block(x) 45 | 46 | outputs = Conv2D(1, 3, padding='same', activation='sigmoid')(x) 47 | 48 | model = Model(inputs=inputs, outputs=outputs) 49 | model.summary() 50 | 51 | if LOAD_WEIGHTS: 52 | model.load_weights('model5.h5') 53 | 54 | model.compile(loss='MSE', optimizer='Adam') 55 | 56 | checkpointer = ModelCheckpoint(filepath='model5.h5', verbose=1) 57 | 58 | 59 | def _train_generator(): 60 | rnd = np.random.RandomState(SEED) 61 | while True: 62 | yield gen_data(rnd, BATCH_SIZE) 63 | 64 | 65 | def _val_generator(): 66 | rnd = np.random.RandomState(SEED + 1) 67 | while True: 68 | yield gen_data(rnd, BATCH_SIZE) 69 | 70 | 71 | train_generator = _train_generator() 72 | val_generator = _val_generator() 73 | 74 | history = model.fit_generator( 75 | train_generator, 76 | steps_per_epoch=512 // BATCH_SIZE, 77 | epochs=EPOCHS, 78 | validation_data=val_generator, 79 | validation_steps=32 // BATCH_SIZE, 80 | callbacks=[checkpointer] 81 | ) 82 | 83 | model.save('model5_final.h5') 84 | -------------------------------------------------------------------------------- /output/demo1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hepesu/LineNormalizer/47e41d3973c255185b696e62dd46a850a84e1323/output/demo1.png -------------------------------------------------------------------------------- /output/demo2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hepesu/LineNormalizer/47e41d3973c255185b696e62dd46a850a84e1323/output/demo2.png -------------------------------------------------------------------------------- /predict.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | # # Try running on CPU 4 | os.environ['CUDA_VISIBLE_DEVICES'] = '-1' 5 | 6 | import numpy as np 7 | import cv2 8 | from keras.models import load_model 9 | 10 | MODEL_NAME = 'model1' 11 | model = load_model('./%s.h5' % MODEL_NAME) 12 | 13 | for root, dirs, files in os.walk('./input', topdown=False): 14 | for name in files: 15 | print(os.path.join(root, name)) 16 | 17 | im = cv2.imread(os.path.join(root, name), cv2.IMREAD_GRAYSCALE) 18 | 19 | im_predict = im.reshape((1, im.shape[0], im.shape[1], 1)) 20 | im_predict = im_predict.astype(np.float32) / 255 21 | 22 | result = model.predict(im_predict) 23 | 24 | im_res = result.reshape((result.shape[1], result.shape[2])) 25 | im_res = im_res * 255 26 | 27 | cv2.imwrite(os.path.join('./output', name), im_res) 28 | -------------------------------------------------------------------------------- /predict_block.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | import cv2 4 | from keras import backend as K 5 | from keras.models import load_model 6 | 7 | MODEL_NAME = 'model1' 8 | SIZE = 128 9 | BATCH_SIZE = 4 10 | PAD_SIZE = 8 11 | 12 | K.set_learning_phase(0) 13 | 14 | 15 | def to_crop(im, size): 16 | height, width = im.shape[:2] 17 | 18 | pad_height = size * int(np.ceil(height / float(size))) - height 19 | pad_width = size * int(np.ceil(width / float(size))) - width 20 | 21 | im_pad = cv2.copyMakeBorder(im, 0, pad_height, 0, pad_width, cv2.BORDER_REFLECT) 22 | 23 | im_crops = [] 24 | for i in range(0, height, size): 25 | for j in range(0, width, size): 26 | im_crop = im_pad[i:i + size, j:j + size] 27 | im_crops.append(cv2.copyMakeBorder(im_crop, PAD_SIZE, PAD_SIZE, PAD_SIZE, PAD_SIZE, cv2.BORDER_REFLECT)) 28 | 29 | return np.array(im_crops) 30 | 31 | 32 | def from_crop(im, im_crops, size): 33 | height, width = im.shape[:2] 34 | 35 | im_pad = np.zeros(( 36 | size * int(np.ceil(height / float(size))), 37 | size * int(np.ceil(width / float(size))), 38 | 1 39 | )) 40 | 41 | idx = 0 42 | for i in range(0, height, size): 43 | for j in range(0, width, size): 44 | im_pad[i:i + size, j:j + size] = im_crops[idx][PAD_SIZE:-PAD_SIZE, PAD_SIZE:-PAD_SIZE] 45 | idx += 1 46 | 47 | return im_pad[:height, :width, :] 48 | 49 | 50 | def preprocess(x): 51 | return np.reshape(x, (1, x.shape[0], x.shape[1], 1)).astype(np.float32) / 255 52 | 53 | 54 | def main(): 55 | model = load_model('./%s.h5' % MODEL_NAME) 56 | 57 | for root, dirs, files in os.walk('./input', topdown=False): 58 | for name in files: 59 | print(os.path.join(root, name)) 60 | 61 | im = cv2.imread(os.path.join(root, name), cv2.IMREAD_GRAYSCALE) 62 | 63 | im_crops = to_crop(im, SIZE) 64 | 65 | im_crops_res = [] 66 | for c in range(0, im_crops.shape[0], BATCH_SIZE): 67 | batch = np.concatenate([preprocess(im_crop) for im_crop in im_crops[c:c + BATCH_SIZE]], 0) 68 | 69 | res = model.predict_on_batch(batch) 70 | 71 | for r in res: 72 | im_crops_res.append(r * 255) 73 | 74 | im_res = from_crop(im, im_crops_res, SIZE) 75 | cv2.imwrite(os.path.join('./output', name), im_res) 76 | 77 | 78 | if __name__ == "__main__": 79 | main() 80 | -------------------------------------------------------------------------------- /pytorch_model1.py: -------------------------------------------------------------------------------- 1 | import os 2 | import random 3 | import numpy as np 4 | 5 | import torch 6 | import torch.nn as nn 7 | import torch.nn.functional as F 8 | 9 | from datagen import gen_data 10 | 11 | SIZE = 128 12 | ITERATIONS = 5100 13 | BATCH_SIZE = 8 14 | SEED = 1 15 | DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 16 | 17 | 18 | class Net(nn.Module): 19 | def __init__(self): 20 | super(Net, self).__init__() 21 | self.conv = nn.Sequential( 22 | nn.Conv2d(1, 64, kernel_size=9, stride=1, padding=4, bias=False), 23 | nn.BatchNorm2d(64), 24 | nn.ReLU(), 25 | 26 | nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=False), 27 | nn.BatchNorm2d(64), 28 | nn.ReLU(), 29 | 30 | nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=False), 31 | nn.BatchNorm2d(64), 32 | nn.ReLU(), 33 | 34 | nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=False), 35 | nn.BatchNorm2d(64), 36 | nn.ReLU(), 37 | 38 | nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=False), 39 | nn.BatchNorm2d(64), 40 | nn.ReLU(), 41 | 42 | nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=False), 43 | nn.BatchNorm2d(64), 44 | nn.ReLU(), 45 | 46 | nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=False), 47 | nn.BatchNorm2d(64), 48 | nn.ReLU(), 49 | 50 | nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=False), 51 | nn.BatchNorm2d(64), 52 | nn.ReLU(), 53 | 54 | nn.Conv2d(64, 1, kernel_size=3, stride=1, padding=1, bias=False), 55 | nn.Sigmoid() 56 | ) 57 | 58 | def forward(self, x): 59 | return self.conv(x) 60 | 61 | 62 | def data_generator(): 63 | rnd = np.random.RandomState(SEED) 64 | while True: 65 | raw, norm = gen_data(rnd, BATCH_SIZE) 66 | yield torch.from_numpy(raw).permute(0, 3, 1, 2), \ 67 | torch.from_numpy(norm).permute(0, 3, 1, 2) 68 | 69 | 70 | def main(): 71 | random.seed(SEED) 72 | np.random.seed(SEED) 73 | torch.manual_seed(SEED) 74 | 75 | # ---------- 76 | # Model and Optimizer 77 | # ---------- 78 | model = Net().to(DEVICE) 79 | optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) 80 | scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 1000, gamma=0.5, last_epoch=-1) 81 | 82 | # ---------- 83 | # Data 84 | # ---------- 85 | dataIter = data_generator() 86 | 87 | # ---------- 88 | # Training 89 | # ---------- 90 | for iteration in range(1, ITERATIONS + 1): 91 | rawData, normData = next(dataIter) 92 | 93 | rawData = rawData.to(DEVICE) 94 | normData = normData.to(DEVICE) 95 | 96 | optimizer.zero_grad() 97 | 98 | loss = F.mse_loss(model(rawData), normData) 99 | 100 | loss.backward() 101 | optimizer.step() 102 | 103 | print("[Iteration %d] [loss: %f]" % (iteration, loss.item())) 104 | 105 | if iteration % 100 == 0: 106 | torch.save(model.state_dict(), os.path.join('weight', '{}.pth'.format(iteration))) 107 | 108 | if iteration < 4000: 109 | scheduler.step() 110 | 111 | 112 | if __name__ == "__main__": 113 | main() 114 | -------------------------------------------------------------------------------- /pytorch_model2.py: -------------------------------------------------------------------------------- 1 | import os 2 | import random 3 | import numpy as np 4 | 5 | import torch 6 | import torch.nn as nn 7 | import torch.nn.functional as F 8 | 9 | from datagen import gen_data 10 | 11 | SIZE = 128 12 | ITERATIONS = 5100 13 | BATCH_SIZE = 8 14 | SEED = 1 15 | DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 16 | 17 | 18 | class Net(nn.Module): 19 | def __init__(self): 20 | super(Net, self).__init__() 21 | self.conv = nn.Sequential( 22 | nn.Conv2d(1, 64, kernel_size=9, stride=1, padding=4, bias=False), 23 | nn.BatchNorm2d(64), 24 | nn.ReLU(), 25 | 26 | nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=False), 27 | nn.BatchNorm2d(64), 28 | nn.ReLU(), 29 | 30 | nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=False), 31 | nn.BatchNorm2d(64), 32 | nn.ReLU(), 33 | 34 | nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=False), 35 | nn.BatchNorm2d(64), 36 | nn.ReLU(), 37 | 38 | nn.Conv2d(64, 1, kernel_size=3, stride=1, padding=1, bias=False), 39 | nn.Sigmoid() 40 | ) 41 | 42 | def forward(self, x): 43 | return self.conv(x) 44 | 45 | 46 | def data_generator(): 47 | rnd = np.random.RandomState(SEED) 48 | while True: 49 | raw, norm = gen_data(rnd, BATCH_SIZE) 50 | yield torch.from_numpy(raw).permute(0, 3, 1, 2), \ 51 | torch.from_numpy(norm).permute(0, 3, 1, 2) 52 | 53 | 54 | def main(): 55 | random.seed(SEED) 56 | np.random.seed(SEED) 57 | torch.manual_seed(SEED) 58 | 59 | # ---------- 60 | # Model and Optimizer 61 | # ---------- 62 | model = Net().to(DEVICE) 63 | optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) 64 | scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 1000, gamma=0.5, last_epoch=-1) 65 | 66 | # ---------- 67 | # Data 68 | # ---------- 69 | dataIter = data_generator() 70 | 71 | # ---------- 72 | # Training 73 | # ---------- 74 | for iteration in range(1, ITERATIONS + 1): 75 | rawData, normData = next(dataIter) 76 | 77 | rawData = rawData.to(DEVICE) 78 | normData = normData.to(DEVICE) 79 | 80 | optimizer.zero_grad() 81 | 82 | loss = F.mse_loss(model(rawData), normData) 83 | 84 | loss.backward() 85 | optimizer.step() 86 | 87 | print("[Iteration %d] [loss: %f]" % (iteration, loss.item())) 88 | 89 | if iteration % 100 == 0: 90 | torch.save(model.state_dict(), os.path.join('weight', '{}.pth'.format(iteration))) 91 | 92 | if iteration < 4000: 93 | scheduler.step() 94 | 95 | 96 | if __name__ == "__main__": 97 | main() 98 | -------------------------------------------------------------------------------- /pytorch_predict.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | import cv2 4 | 5 | import torch 6 | from torchvision import transforms 7 | 8 | MODEL_NAME = 'model1' 9 | DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 10 | 11 | 12 | def main(): 13 | # Provided model are entirely saved. if testing, use following code 14 | # replace NUM to test model name 15 | 16 | # from pytorch_model1 import Net 17 | # model = Net().to(DEVICE) 18 | # model.load_state_dict(torch.load('./NUM.pth', map_location=DEVICE)) 19 | # torch.save(model, './%s.pth' % MODEL_NAME) 20 | 21 | model = torch.load('./%s.pth' % MODEL_NAME, map_location=DEVICE) 22 | 23 | model.eval() 24 | 25 | preprocess = transforms.Compose([ 26 | transforms.ToTensor(), 27 | ]) 28 | 29 | for root, dirs, files in os.walk('./input', topdown=False): 30 | for name in files: 31 | print(os.path.join(root, name)) 32 | 33 | im = cv2.imread(os.path.join(root, name), cv2.IMREAD_GRAYSCALE) 34 | 35 | res = model(preprocess(im).unsqueeze(0).to(DEVICE)) 36 | 37 | im_res = (res.squeeze(0).permute(1, 2, 0).detach().cpu().numpy()) * 255 38 | 39 | cv2.imwrite(os.path.join('./output', name), im_res) 40 | 41 | 42 | if __name__ == "__main__": 43 | main() 44 | -------------------------------------------------------------------------------- /pytorch_predict_block.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | import cv2 4 | 5 | import torch 6 | from torchvision import transforms 7 | 8 | MODEL_NAME = 'model1' 9 | SIZE = 128 10 | BATCH_SIZE = 4 11 | PAD_SIZE = 8 12 | DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 13 | 14 | 15 | def to_crop(im, size): 16 | height, width = im.shape[:2] 17 | 18 | pad_height = size * int(np.ceil(height / float(size))) - height 19 | pad_width = size * int(np.ceil(width / float(size))) - width 20 | 21 | im_pad = cv2.copyMakeBorder(im, 0, pad_height, 0, pad_width, cv2.BORDER_REFLECT) 22 | 23 | im_crops = [] 24 | for i in range(0, height, size): 25 | for j in range(0, width, size): 26 | im_crop = im_pad[i:i + size, j:j + size] 27 | im_crops.append(cv2.copyMakeBorder(im_crop, PAD_SIZE, PAD_SIZE, PAD_SIZE, PAD_SIZE, cv2.BORDER_REFLECT)) 28 | 29 | return np.array(im_crops) 30 | 31 | 32 | def from_crop(im, im_crops, size): 33 | height, width = im.shape[:2] 34 | 35 | im_pad = np.zeros(( 36 | size * int(np.ceil(height / float(size))), 37 | size * int(np.ceil(width / float(size))), 38 | 1 39 | )) 40 | 41 | idx = 0 42 | for i in range(0, height, size): 43 | for j in range(0, width, size): 44 | im_pad[i:i + size, j:j + size] = im_crops[idx][PAD_SIZE:-PAD_SIZE, PAD_SIZE:-PAD_SIZE] 45 | idx += 1 46 | 47 | return im_pad[:height, :width, :] 48 | 49 | 50 | def main(): 51 | model = torch.load('./%s.pth' % MODEL_NAME, map_location=DEVICE) 52 | model.eval() 53 | 54 | preprocess = transforms.Compose([ 55 | transforms.ToTensor() 56 | ]) 57 | 58 | for root, dirs, files in os.walk('./input', topdown=False): 59 | for name in files: 60 | print(os.path.join(root, name)) 61 | 62 | im = cv2.imread(os.path.join(root, name), cv2.IMREAD_GRAYSCALE) 63 | 64 | im_crops = to_crop(im, SIZE) 65 | 66 | im_crops_res = [] 67 | for c in range(0, im_crops.shape[0], BATCH_SIZE): 68 | batch = torch.cat([preprocess(im_crop).unsqueeze(0) for im_crop in im_crops[c:c + BATCH_SIZE]], 0) 69 | 70 | batch = batch.to(DEVICE) 71 | 72 | res = model(batch).permute(0, 2, 3, 1).detach().cpu().numpy() 73 | 74 | for r in res: 75 | im_crops_res.append(r * 255) 76 | 77 | im_res = from_crop(im, im_crops_res, SIZE) 78 | cv2.imwrite(os.path.join('./output', name), im_res) 79 | 80 | 81 | if __name__ == "__main__": 82 | main() 83 | -------------------------------------------------------------------------------- /weight/.gitkeep: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hepesu/LineNormalizer/47e41d3973c255185b696e62dd46a850a84e1323/weight/.gitkeep --------------------------------------------------------------------------------