├── README.md ├── example.ipynb ├── expts └── Allx768xDnCNN │ └── logs │ └── BF32:BLK5:BN0,0:M8:LArelu:LCbefore_decoder:SC0:DSdrop:USbilinear:BD1:N0.00:P1 │ ├── events.out.tfevents.1608988492.viplab-SYS-7049GP-TRT │ └── events.out.tfevents.1608988866.viplab-SYS-7049GP-TRT ├── image └── Structure.png ├── requirements.txt ├── samples ├── 001-P17.png ├── 001-P18.png ├── 001-P19.png ├── 001-P20.png ├── 002-P16.png ├── 002-P17.png ├── 002-P18.png ├── 003-P16.png ├── 003-P17.png ├── 021.tif ├── 022.tif ├── 023.tif ├── 024.tif ├── 025.tif ├── 027.tif ├── 028.tif ├── 029.tif ├── 030.tif ├── 031.tif ├── 032.tif ├── book1_page19.png ├── book2_page2.png ├── book2_page5.png ├── book2_page6.png ├── moc_test_6.png └── moc_test_7.png ├── src ├── .ipynb_checkpoints │ ├── CLLayers-checkpoint.py │ ├── CLLosses-checkpoint.py │ ├── CLUtils-checkpoint.py │ ├── Untitled-checkpoint.ipynb │ └── __init__-checkpoint.py ├── CLLayers.py ├── CLLosses.py ├── CLUtils.py ├── Untitled.ipynb ├── __init__.py ├── __pycache__ │ ├── CLLayers.cpython-37.pyc │ ├── CLLosses.cpython-37.pyc │ ├── CLUtils.cpython-37.pyc │ └── __init__.cpython-37.pyc └── robust_loss │ ├── .ipynb_checkpoints │ ├── adaptive-checkpoint.py │ ├── distribution-checkpoint.py │ ├── example-checkpoint.ipynb │ ├── fit_partition_spline-checkpoint.py │ ├── general-checkpoint.py │ ├── util-checkpoint.py │ ├── vae-checkpoint.py │ └── wavelet-checkpoint.py │ ├── README.md │ ├── __pycache__ │ ├── adaptive.cpython-37.pyc │ ├── cubic_spline.cpython-37.pyc │ ├── distribution.cpython-37.pyc │ ├── general.cpython-37.pyc │ ├── util.cpython-37.pyc │ └── wavelet.cpython-37.pyc │ ├── adaptive.py │ ├── adaptive_test.py │ ├── cubic_spline.py │ ├── cubic_spline_test.py │ ├── data │ ├── partition_spline.npz │ ├── wavelet_golden.mat │ └── wavelet_vis_golden.png │ ├── distribution.py │ ├── distribution_test.py │ ├── example.ipynb │ ├── fit_partition_spline.py │ ├── fit_partition_spline_test.py │ ├── general.py │ ├── general_test.py │ ├── partition_spline.npz │ ├── requirements.txt │ ├── run.sh │ ├── util.py │ ├── util_test.py │ ├── vae.py │ ├── wavelet.py │ └── wavelet_test.py ├── thinplate ├── __init__.py ├── __pycache__ │ ├── __init__.cpython-37.pyc │ └── numpy.cpython-37.pyc ├── numpy.py ├── pytorch.py └── tests │ ├── __init__.py │ ├── test_tps_numpy.py │ └── test_tps_pytorch.py └── trainLineCounterV2.py /README.md: -------------------------------------------------------------------------------- 1 | # LineCounter: Learning Handwritten Text Line Segmentation by Counting 2 | 3 |
4 | 5 |
6 | 7 | *** 8 | 9 | This is the official repo for LineCounter (ICIP 2021). For details of LineCounter, please refer to 10 | 11 | ``` 12 | @INPROCEEDINGS{9506664, 13 | author={Li, Deng and Wu, Yue and Zhou, Yicong}, 14 | booktitle={2021 IEEE International Conference on Image Processing (ICIP)}, 15 | title={Linecounter: Learning Handwritten Text Line Segmentation By Counting}, 16 | year={2021}, 17 | volume={}, 18 | number={}, 19 | pages={929-933}, 20 | doi={10.1109/ICIP42928.2021.9506664}} 21 | ``` 22 | 23 | *** 24 | 25 | # Overview 26 | 27 | 28 | 29 | 30 | # Dependency 31 | 32 | SauvolaNet is written in the TensorFlow. 33 | 34 | - TensorFlow-GPU: 1.15.0 35 | 36 | Other versions might also work but are not tested. 37 | 38 | 39 | # Demo 40 | 41 | Download the repo and create the virtual environment by following commands 42 | 43 | ``` 44 | conda create --name LineCounter --file requirements.txt 45 | ``` 46 | 47 | Download [trained-model](https://drive.google.com/file/d/1fMUkyg67QLLzyDMkU1vgnsgDIKb9SPF9/view?usp=sharing) 48 | 49 | *** 50 | Put trained-model in **LineCounter/expts/Allx768xDnCNN/models/BF32:BLK5:BN0,0:M8:LArelu:LCbefore_decoder:SC0:DSdrop:USbilinear:BD1:N0.00:P1/** 51 | *** 52 | 53 | Then play with the provided ipython notebook. 54 | 55 | Alternatively, one may play with the inference code using this [google colab link](https://colab.research.google.com/drive/1v-h7eSNhxfCTqQZC_IPGEp_s-sfA6dxn?usp=sharing). 56 | 57 | # Datasets 58 | We do not own the copyright of the dataset used in this repo. 59 | 60 | Below is a summary table of the datasets used in this work along with a link from which they can be downloaded: 61 | 62 | 63 | | Dataset | URL | 64 | | ------------ | ------- | 65 | | ICDAR-HCS2013 | https://users.iit.demokritos.gr/~nstam/ICDAR2013HandSegmCont/ | 66 | | HIT-MW | http://space.hit.edu.cn/article/2019/03/11/10660 (Chinese) | 67 | | VML-AHTE | https://www.cs.bgu.ac.il/~berat/ | 68 | 69 | # Concat 70 | 71 | For any paper related questions, please feel free to contact leedengsh@gmail.com. 72 | -------------------------------------------------------------------------------- /expts/Allx768xDnCNN/logs/BF32:BLK5:BN0,0:M8:LArelu:LCbefore_decoder:SC0:DSdrop:USbilinear:BD1:N0.00:P1/events.out.tfevents.1608988492.viplab-SYS-7049GP-TRT: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/expts/Allx768xDnCNN/logs/BF32:BLK5:BN0,0:M8:LArelu:LCbefore_decoder:SC0:DSdrop:USbilinear:BD1:N0.00:P1/events.out.tfevents.1608988492.viplab-SYS-7049GP-TRT -------------------------------------------------------------------------------- /expts/Allx768xDnCNN/logs/BF32:BLK5:BN0,0:M8:LArelu:LCbefore_decoder:SC0:DSdrop:USbilinear:BD1:N0.00:P1/events.out.tfevents.1608988866.viplab-SYS-7049GP-TRT: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/expts/Allx768xDnCNN/logs/BF32:BLK5:BN0,0:M8:LArelu:LCbefore_decoder:SC0:DSdrop:USbilinear:BD1:N0.00:P1/events.out.tfevents.1608988866.viplab-SYS-7049GP-TRT -------------------------------------------------------------------------------- /image/Structure.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/image/Structure.png -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | keras-gpu=2.2.4 2 | numba=0.50.1 3 | numpy=1.16.4 4 | opencv=4.1.1 5 | parse=1.12.1 6 | pillow=6.1.0 7 | scikit-image=0.15.0 8 | scikit-learn=0.21.3 9 | tensorflow-gpu=1.15.0 10 | python=3.7.4 11 | -------------------------------------------------------------------------------- /samples/001-P17.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/samples/001-P17.png -------------------------------------------------------------------------------- /samples/001-P18.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/samples/001-P18.png -------------------------------------------------------------------------------- /samples/001-P19.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/samples/001-P19.png -------------------------------------------------------------------------------- /samples/001-P20.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/samples/001-P20.png -------------------------------------------------------------------------------- /samples/002-P16.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/samples/002-P16.png -------------------------------------------------------------------------------- /samples/002-P17.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/samples/002-P17.png -------------------------------------------------------------------------------- /samples/002-P18.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/samples/002-P18.png -------------------------------------------------------------------------------- /samples/003-P16.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/samples/003-P16.png -------------------------------------------------------------------------------- /samples/003-P17.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/samples/003-P17.png -------------------------------------------------------------------------------- /samples/021.tif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/samples/021.tif -------------------------------------------------------------------------------- /samples/022.tif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/samples/022.tif -------------------------------------------------------------------------------- /samples/023.tif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/samples/023.tif -------------------------------------------------------------------------------- /samples/024.tif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/samples/024.tif -------------------------------------------------------------------------------- /samples/025.tif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/samples/025.tif -------------------------------------------------------------------------------- /samples/027.tif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/samples/027.tif -------------------------------------------------------------------------------- /samples/028.tif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/samples/028.tif -------------------------------------------------------------------------------- /samples/029.tif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/samples/029.tif -------------------------------------------------------------------------------- /samples/030.tif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/samples/030.tif -------------------------------------------------------------------------------- /samples/031.tif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/samples/031.tif -------------------------------------------------------------------------------- /samples/032.tif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/samples/032.tif -------------------------------------------------------------------------------- /samples/book1_page19.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/samples/book1_page19.png -------------------------------------------------------------------------------- /samples/book2_page2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/samples/book2_page2.png -------------------------------------------------------------------------------- /samples/book2_page5.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/samples/book2_page5.png -------------------------------------------------------------------------------- /samples/book2_page6.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/samples/book2_page6.png -------------------------------------------------------------------------------- /samples/moc_test_6.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/samples/moc_test_6.png -------------------------------------------------------------------------------- /samples/moc_test_7.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/samples/moc_test_7.png -------------------------------------------------------------------------------- /src/.ipynb_checkpoints/CLLosses-checkpoint.py: -------------------------------------------------------------------------------- 1 | from keras.metrics import mse, mae 2 | from scipy.optimize import linear_sum_assignment 3 | from keras import backend as K 4 | import numpy as np 5 | import tensorflow as tf 6 | 7 | 8 | import os 9 | import sys 10 | 11 | def lineClf_np( y_true, y_pred, th=1 ) : 12 | #y_true = assign_label_np( y_true, y_pred ) 13 | zz = np.round(y_pred) 14 | yy = np.round(y_true) 15 | return np.mean( ( zz == yy )[yy >=th] ).astype('float32') 16 | 17 | def acc( y_true, y_pred, th=1 ) : 18 | #y_true = assign_label( y_true, y_pred ) 19 | #y_true = K.stop_gradient( y_true ) 20 | y_pred_int = K.round( y_pred ) 21 | 22 | y_true_int = K.round( y_true ) 23 | mask = K.cast( y_true>=th, 'float32' ) 24 | matched = K.cast( K.abs(y_pred_int-y_true_int) < .1, 'float32' ) * mask 25 | return K.sum( matched, axis=(1,2,3) ) / K.sum( mask, axis=(1,2,3) ) 26 | 27 | def MatchScore( y_true, y_pred, th=1 ) : 28 | 29 | y_pred_int = K.round( y_pred ) 30 | y_true_int = K.round( y_true ) 31 | mask = K.cast( y_true>=th, 'float32' ) 32 | matched = K.cast( K.abs(y_pred_int-y_true_int) ==0, 'float32' ) * mask 33 | return K.sum( matched, axis=(1,2,3) ) / K.sum( mask, axis=(1,2,3) ) 34 | def acc0( y_true, y_pred, th=0 ) : 35 | y_pred_int = K.round( y_pred ) 36 | y_true_int = K.round( y_true ) 37 | mask = K.cast( K.abs(y_true-th)<.1, 'float32' ) 38 | matched = K.cast( K.abs(y_pred_int-y_true_int) < .1, 'float32' ) * mask 39 | return K.sum( matched, axis=(1,2,3) ) / K.sum( mask, axis=(1,2,3) ) 40 | 41 | def assign_label( y_true, y_pred ) : 42 | y_modi = tf.py_func( assign_label_np, [ y_true, y_pred ], 'float32', stateful=False ) 43 | y_modi.set_shape( K.int_shape( y_true ) ) 44 | return y_modi 45 | 46 | def assign_label_np( y_true, y_pred ) : 47 | mask = ( y_true > 0 ) 48 | y_pred = np.round( y_pred ).astype('int') 49 | y_true = np.round( y_true ).astype('int') 50 | y_modi = [] 51 | for idx in range( len(y_pred) ): 52 | p = y_pred[idx] 53 | t = y_true[idx] 54 | #print "-" * 50 55 | #print "idx=", idx 56 | true_line_labels = filter( lambda v : v>0, np.unique(t) ) 57 | pred_line_labels = filter( lambda v : v>0, np.unique(p) ) 58 | print ("true_line_labels", true_line_labels) 59 | print ("pred_line_labels", pred_line_labels) 60 | mat = [] 61 | for tl in true_line_labels : 62 | mm = t==tl 63 | row = [] 64 | for pl in pred_line_labels : 65 | vv = -np.sum(p[mm] == pl) 66 | row.append(vv) 67 | mat.append( row ) 68 | mat = np.row_stack(mat) 69 | row_ind, col_ind = linear_sum_assignment( mat ) 70 | true_ind = [ true_line_labels[k] for k in row_ind ] 71 | pred_ind = [ pred_line_labels[k] for k in col_ind ] 72 | for tl, pl in zip( true_ind, pred_ind ) : 73 | t[ t==tl ] = pltf.compat.v1.disable_eager_execution() 74 | #print "assign", tl, "to", pltf.compat.v1.disable_eager_execution() 75 | y_modi.append( np.expand_dims(t,axis=0)) 76 | return np.concatenate(y_modi).astype('float32') 77 | 78 | def seg( y_true, y_pred, th=1 ) : 79 | loss = K.maximum( K.square( y_true - y_pred ), K.abs( y_true - y_pred ) ) 80 | mask = K.cast( y_true>=th, 'float32' ) 81 | return K.sum( mask * loss,axis=(1,2,3)) / K.sum( mask, axis=(1,2,3)) 82 | 83 | 84 | 85 | 86 | 87 | 88 | def prepare_group_loss_numpy(y_true, y_pred) : 89 | """Implement group loss in Sec3.3 of https://arxiv.org/pdf/1611.05424.pdf 90 | NOTE: the 2nd term of the loss has been simplified to the closest mu 91 | """ 92 | within, between = [], [] 93 | for a_true, a_pred in zip(y_true, y_pred) : 94 | N = int(a_true.max()) 95 | 96 | within_true = np.zeros_like(a_true) 97 | #between_true = np.ones_like(a_true) * N 98 | between_true = np.zeros_like(a_true) 99 | masks = [] 100 | mu_list = [] 101 | for line_idx in range(1, N+1) : 102 | #mask = np.abs(a_true - line_idx) < 0.1 103 | mask = (a_true==line_idx) 104 | vals = a_pred[mask] 105 | mu = vals.mean() 106 | within_true[mask] = mu 107 | mu_list.append(mu) 108 | masks.append(mask) 109 | mu_arr = np.array(mu_list) 110 | 111 | 112 | for mask, mu in zip(masks, mu_arr): 113 | #indices = np.argsort(np.abs(mu_arr - mu)) 114 | ind_mu = np.where(mu_arr==mu) 115 | ind_mu = int(ind_mu[0]) 116 | closest_mu = mu_arr[np.minimum(ind_mu+1,len(mu_arr)-1)] 117 | 118 | between_true[mask] = closest_mu 119 | # update output 120 | within.append(within_true) 121 | between.append(between_true) 122 | 123 | return np.stack(within, axis=0), np.stack(between, axis=0) 124 | 125 | def grouping_loss(y_true, y_pred, sigma=0.025) : 126 | y_within, y_between = tf.py_func(prepare_group_loss_numpy, 127 | [y_true,y_pred], 128 | [tf.float32, tf.float32]) 129 | y_within.set_shape(K.int_shape(y_true)) 130 | y_between.set_shape(K.int_shape(y_true)) 131 | y_within = K.stop_gradient(y_within) 132 | y_between = K.stop_gradient(y_between) 133 | mask = K.cast(y_true >=1, 'float32') 134 | diff = (y_within-y_pred)**2 + tf.exp(-(y_between-y_pred)**2/(2.*sigma)) 135 | #diff = tf.exp(-(y_between-y_pred)**2/(2.*sigma)) 136 | loss = K.sum(diff * mask) / (K.sum(mask)) 137 | 138 | origin = K.maximum( K.square( y_true - y_pred ), K.abs( y_true - y_pred ) ) 139 | 140 | origin_loss = K.sum( mask * origin,axis=(1,2,3)) / K.sum( mask, axis=(1,2,3)) 141 | return origin_loss 142 | def MSE_loss(y_true, y_pred, sigma=0.025) : 143 | y_within, y_between = tf.py_func(prepare_group_loss_numpy, 144 | [y_true,y_pred], 145 | [tf.float32, tf.float32]) 146 | y_within.set_shape(K.int_shape(y_true)) 147 | y_between.set_shape(K.int_shape(y_true)) 148 | y_within = K.stop_gradient(y_within) 149 | y_between = K.stop_gradient(y_between) 150 | mask = K.cast(y_true >=1, 'float32') 151 | diff = (y_within-y_pred)**2 + tf.exp(-(y_between-y_pred)**2/(2.*sigma)) 152 | #diff = tf.exp(-(y_between-y_pred)**2/(2.*sigma)) 153 | loss = K.sum(diff * mask) / (K.sum(mask)) 154 | 155 | 156 | origin = K.square( y_true - y_pred ) 157 | 158 | origin_loss = K.sum( mask * origin,axis=(1,2,3)) / K.sum( mask, axis=(1,2,3)) 159 | return origin_loss 160 | 161 | def MAE_loss(y_true, y_pred, sigma=0.025) : 162 | y_within, y_between = tf.py_func(prepare_group_loss_numpy, 163 | [y_true,y_pred], 164 | [tf.float32, tf.float32]) 165 | y_within.set_shape(K.int_shape(y_true)) 166 | y_between.set_shape(K.int_shape(y_true)) 167 | y_within = K.stop_gradient(y_within) 168 | y_between = K.stop_gradient(y_between) 169 | mask = K.cast(y_true >=1, 'float32') 170 | diff = (y_within-y_pred)**2 + tf.exp(-(y_between-y_pred)**2/(2.*sigma)) 171 | #diff = tf.exp(-(y_between-y_pred)**2/(2.*sigma)) 172 | loss = K.sum(diff * mask) / (K.sum(mask)) 173 | origin = K.abs( y_true - y_pred ) 174 | origin_loss = K.sum( mask * origin,axis=(1,2,3)) / K.sum( mask, axis=(1,2,3)) 175 | return origin_loss 176 | 177 | 178 | 179 | 180 | 181 | 182 | class RobustAdaptativeLoss(object): 183 | def __init__(self): 184 | z = np.array([[0]]) 185 | self.v_alpha = K.zeros(shape=(1088, 768, 1)) 186 | self.v_scale = K.zeros(shape=(1088, 768, 1)) 187 | 188 | def loss(self, y_true, y_pred, **kwargs): 189 | mask = K.cast(y_true >=1, 'float32') 190 | x = y_true - y_pred*mask 191 | #origin_loss = K.sum( mask * x,axis=(1,2,3)) / K.sum( mask, axis=(1,2,3)) 192 | 193 | #x = K.reshape(x, shape=(-1, 1)) 194 | lossfun = robust_loss.adaptive.AdaptiveImageLossFunction((1088, 768, 1), float_dtype='float32',color_space='RGB',representation='PIXEL') 195 | alpha = lossfun.alpha() 196 | scale = lossfun.scale() 197 | #loss, alpha, scale = robust_loss.adaptive.AdaptiveLossFunction(num_channels=1,float_dtype="float32") 198 | a = K.update(self.v_alpha, alpha) 199 | s = K.update(self.v_scale, scale) 200 | # The alpha update must be part of the graph but it should 201 | # not influence the result. 202 | 203 | #mask = K.cast(y_true >=1, 'float32') 204 | origin = lossfun(x) 205 | #origin_loss = K.sum( origin,axis=(1,2,3)) / K.sum( mask, axis=(1,2,3)) 206 | #origin_loss = K.sum( mask * origin,axis=(1,2,3)) / K.sum( mask, axis=(1,2,3)) 207 | return origin + 0 * a + 0 * s 208 | 209 | def alpha(self, y_true, y_pred): 210 | return self.v_alpha 211 | def scale(self, y_true, y_pred): 212 | return self.v_scale 213 | 214 | #lossfun = robust_loss.adaptive.AdaptiveImageLossFunction((1088, 768, 1), float_dtype='float32',color_space='RGB') 215 | def Robustloss(y_true, y_pred): 216 | 217 | mask = K.cast(y_true >=1, 'float32') 218 | #lossfun = robust_loss.adaptive.AdaptiveImageLossFunction((1088, 768, 1), float_dtype='float32',color_space='RGB') 219 | x = y_true-y_pred 220 | #Cauchy 221 | #origin = K.log(0.5*x*x + 1) 222 | # Welsch 223 | #origin = 1-K.exp(-0.5*x*x) 224 | #Charbonnier 225 | origin = K.sqrt(x*x+1) - 1 226 | #Geman 227 | #origin = -2*(1/(((x*x)/4)+1) - 1) 228 | #origin = lossfun(x) 229 | 230 | origin_loss = K.sum( mask * origin,axis=(0,1,2)) / K.sum( mask, axis=(1,2,3)) 231 | return origin_loss 232 | 233 | def seg_2( y_true, y_pred, th=1 ) : 234 | loss = K.cast( y_true == y_pred, 'float32' ) 235 | mask = K.cast( y_true>=th, 'float32' ) 236 | return K.sum( mask * loss,axis=(1,2,3)) / K.sum( mask, axis=(1,2,3)) 237 | 238 | def IOU_calc(y_true, y_pred): 239 | smooth = 1.0 240 | y_true_f = K.flatten(y_true) 241 | y_pred_f = K.flatten(y_pred) 242 | intersection = K.sum(y_true_f * y_pred_f) 243 | return 2*(intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) + smooth) 244 | 245 | -------------------------------------------------------------------------------- /src/.ipynb_checkpoints/Untitled-checkpoint.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "code", 5 | "execution_count": 1, 6 | "metadata": {}, 7 | "outputs": [], 8 | "source": [ 9 | "import sys\n", 10 | "sys.path.append('robust_loss')\n", 11 | "\n", 12 | "import robust_loss.adaptive " 13 | ] 14 | }, 15 | { 16 | "cell_type": "code", 17 | "execution_count": 7, 18 | "metadata": {}, 19 | "outputs": [], 20 | "source": [ 21 | "import numpy as np\n", 22 | "import tensorflow as tf\n", 23 | "from tensorflow import keras\n", 24 | "from tensorflow.keras.layers import *\n", 25 | "from tensorflow.keras.models import Model\n", 26 | "from tensorflow.keras import backend as K" 27 | ] 28 | }, 29 | { 30 | "cell_type": "code", 31 | "execution_count": 9, 32 | "metadata": {}, 33 | "outputs": [ 34 | { 35 | "ename": "TypeError", 36 | "evalue": "cannot unpack non-iterable AdaptiveLossFunction object", 37 | "output_type": "error", 38 | "traceback": [ 39 | "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", 40 | "\u001b[0;31mTypeError\u001b[0m Traceback (most recent call last)", 41 | "\u001b[0;32m\u001b[0m in \u001b[0;36m\u001b[0;34m\u001b[0m\n\u001b[1;32m 27\u001b[0m \u001b[0;32mreturn\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 28\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 29\u001b[0;31m \u001b[0mmodel\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmake_model\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 30\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msummary\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 31\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n", 42 | "\u001b[0;32m\u001b[0m in \u001b[0;36mmake_model\u001b[0;34m()\u001b[0m\n\u001b[1;32m 24\u001b[0m \u001b[0mmodel\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mModel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minp\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mout\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 25\u001b[0m \u001b[0mloss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mRobustAdaptativeLoss\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 26\u001b[0;31m \u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcompile\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'adam'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mloss\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmetrics\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0malpha\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 27\u001b[0m \u001b[0;32mreturn\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 28\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n", 43 | "\u001b[0;32m~/anaconda3/envs/hyperLine2/lib/python3.7/site-packages/tensorflow_core/python/training/tracking/base.py\u001b[0m in \u001b[0;36m_method_wrapper\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 455\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_self_setattr_tracking\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mFalse\u001b[0m \u001b[0;31m# pylint: disable=protected-access\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 456\u001b[0m \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 457\u001b[0;31m \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmethod\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 458\u001b[0m \u001b[0;32mfinally\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 459\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_self_setattr_tracking\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mprevious_value\u001b[0m \u001b[0;31m# pylint: disable=protected-access\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", 44 | "\u001b[0;32m~/anaconda3/envs/hyperLine2/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/training.py\u001b[0m in \u001b[0;36mcompile\u001b[0;34m(self, optimizer, loss, metrics, loss_weights, sample_weight_mode, weighted_metrics, target_tensors, distribute, **kwargs)\u001b[0m\n\u001b[1;32m 371\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 372\u001b[0m \u001b[0;31m# Creates the model loss and weighted metrics sub-graphs.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 373\u001b[0;31m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_compile_weights_loss_and_weighted_metrics\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 374\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 375\u001b[0m \u001b[0;31m# Functions for train, test and predict will\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", 45 | "\u001b[0;32m~/anaconda3/envs/hyperLine2/lib/python3.7/site-packages/tensorflow_core/python/training/tracking/base.py\u001b[0m in \u001b[0;36m_method_wrapper\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 455\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_self_setattr_tracking\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mFalse\u001b[0m \u001b[0;31m# pylint: disable=protected-access\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 456\u001b[0m \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 457\u001b[0;31m \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmethod\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 458\u001b[0m \u001b[0;32mfinally\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 459\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_self_setattr_tracking\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mprevious_value\u001b[0m \u001b[0;31m# pylint: disable=protected-access\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", 46 | "\u001b[0;32m~/anaconda3/envs/hyperLine2/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/training.py\u001b[0m in \u001b[0;36m_compile_weights_loss_and_weighted_metrics\u001b[0;34m(self, sample_weights)\u001b[0m\n\u001b[1;32m 1650\u001b[0m \u001b[0;31m# loss_weight_2 * output_2_loss_fn(...) +\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 1651\u001b[0m \u001b[0;31m# layer losses.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1652\u001b[0;31m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtotal_loss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_prepare_total_loss\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmasks\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 1653\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 1654\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0m_prepare_skip_target_masks\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", 47 | "\u001b[0;32m~/anaconda3/envs/hyperLine2/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/training.py\u001b[0m in \u001b[0;36m_prepare_total_loss\u001b[0;34m(self, masks)\u001b[0m\n\u001b[1;32m 1710\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 1711\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mhasattr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mloss_fn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'reduction'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1712\u001b[0;31m \u001b[0mper_sample_losses\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mloss_fn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcall\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0my_true\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my_pred\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 1713\u001b[0m weighted_losses = losses_utils.compute_weighted_loss(\n\u001b[1;32m 1714\u001b[0m \u001b[0mper_sample_losses\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", 48 | "\u001b[0;32m~/anaconda3/envs/hyperLine2/lib/python3.7/site-packages/tensorflow_core/python/keras/losses.py\u001b[0m in \u001b[0;36mcall\u001b[0;34m(self, y_true, y_pred)\u001b[0m\n\u001b[1;32m 214\u001b[0m \u001b[0mLoss\u001b[0m \u001b[0mvalues\u001b[0m \u001b[0mper\u001b[0m \u001b[0msample\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 215\u001b[0m \"\"\"\n\u001b[0;32m--> 216\u001b[0;31m \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0my_true\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my_pred\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_fn_kwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 217\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 218\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mget_config\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", 49 | "\u001b[0;32m\u001b[0m in \u001b[0;36mloss\u001b[0;34m(self, y_true, y_pred, **kwargs)\u001b[0m\n\u001b[1;32m 10\u001b[0m \u001b[0;31m#x = K.reshape(x, shape=(-1, 1))\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 11\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 12\u001b[0;31m \u001b[0mloss\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0malpha\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mscale\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mrobust_loss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0madaptive\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mAdaptiveLossFunction\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnum_channels\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mfloat_dtype\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m\"float32\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 13\u001b[0m \u001b[0mop\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mK\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mupdate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mv_alpha\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0malpha\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 14\u001b[0m \u001b[0;31m# The alpha update must be part of the graph but it should\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", 50 | "\u001b[0;31mTypeError\u001b[0m: cannot unpack non-iterable AdaptiveLossFunction object" 51 | ] 52 | } 53 | ], 54 | "source": [ 55 | "\n", 56 | "\n", 57 | "class RobustAdaptativeLoss(object):\n", 58 | " \n", 59 | " def __init__(self):\n", 60 | " \n", 61 | " z = np.array([[0]])\n", 62 | " self.v_alpha = K.variable(z)\n", 63 | "\n", 64 | " def loss(self, y_true, y_pred, **kwargs):\n", 65 | " x = y_true - y_pred\n", 66 | " #x = K.reshape(x, shape=(-1, 1))\n", 67 | " \n", 68 | " loss = robust_loss.adaptive.AdaptiveLossFunction(num_channels=1,float_dtype=\"float32\")\n", 69 | " alpha = loss.aplha()\n", 70 | " scale = loss.scale()\n", 71 | " op = K.update(self.v_alpha, alpha)\n", 72 | " # The alpha update must be part of the graph but it should\n", 73 | " # not influence the result.\n", 74 | " return loss + 0 * op\n", 75 | "\n", 76 | " def alpha(self, y_true, y_pred):\n", 77 | " return self.v_alpha\n", 78 | "\n", 79 | "def make_model():\n", 80 | " inp = Input(shape=(3,))\n", 81 | " out = Dense(1, use_bias=False)(inp)\n", 82 | " model = Model(inp, out)\n", 83 | " loss = RobustAdaptativeLoss()\n", 84 | " model.compile('adam', loss.loss, metrics=[loss.alpha])\n", 85 | " return model\n", 86 | "\n", 87 | "model = make_model()\n", 88 | "model.summary()\n", 89 | "\n", 90 | "\n", 91 | "\n", 92 | "import numpy as np\n", 93 | "\n", 94 | "FACTORS = np.array([0.5, 2.0, 5.0])\n", 95 | "def target_fn(x):\n", 96 | " return np.dot(x, FACTORS.T)\n", 97 | "\n", 98 | "N_SAMPLES=100\n", 99 | "X = np.random.rand(N_SAMPLES, 3)\n", 100 | "Y = np.apply_along_axis(target_fn, 1, X)\n", 101 | "\n", 102 | "\n", 103 | "history = model.fit(X, Y, epochs=2, verbose=True)\n", 104 | "print('final loss:', history.history['loss'][-1])" 105 | ] 106 | }, 107 | { 108 | "cell_type": "code", 109 | "execution_count": null, 110 | "metadata": {}, 111 | "outputs": [], 112 | "source": [] 113 | } 114 | ], 115 | "metadata": { 116 | "kernelspec": { 117 | "display_name": "hyperLine2", 118 | "language": "python", 119 | "name": "hyperline2" 120 | }, 121 | "language_info": { 122 | "codemirror_mode": { 123 | "name": "ipython", 124 | "version": 3 125 | }, 126 | "file_extension": ".py", 127 | "mimetype": "text/x-python", 128 | "name": "python", 129 | "nbconvert_exporter": "python", 130 | "pygments_lexer": "ipython3", 131 | "version": "3.7.4" 132 | } 133 | }, 134 | "nbformat": 4, 135 | "nbformat_minor": 4 136 | } 137 | -------------------------------------------------------------------------------- /src/.ipynb_checkpoints/__init__-checkpoint.py: -------------------------------------------------------------------------------- 1 | from . import CLLayers 2 | from . import CLLosses 3 | from . import CLUtils -------------------------------------------------------------------------------- /src/CLLosses.py: -------------------------------------------------------------------------------- 1 | from keras.metrics import mse, mae 2 | from scipy.optimize import linear_sum_assignment 3 | from keras import backend as K 4 | import numpy as np 5 | import tensorflow as tf 6 | 7 | 8 | import os 9 | import sys 10 | 11 | def lineClf_np( y_true, y_pred, th=1 ) : 12 | #y_true = assign_label_np( y_true, y_pred ) 13 | zz = np.round(y_pred) 14 | yy = np.round(y_true) 15 | return np.mean( ( zz == yy )[yy >=th] ).astype('float32') 16 | 17 | def acc( y_true, y_pred, th=1 ) : 18 | #y_true = assign_label( y_true, y_pred ) 19 | #y_true = K.stop_gradient( y_true ) 20 | y_pred_int = K.round( y_pred ) 21 | 22 | y_true_int = K.round( y_true ) 23 | mask = K.cast( y_true>=th, 'float32' ) 24 | matched = K.cast( K.abs(y_pred_int-y_true_int) < .1, 'float32' ) * mask 25 | return K.sum( matched, axis=(1,2,3) ) / K.sum( mask, axis=(1,2,3) ) 26 | 27 | def MatchScore( y_true, y_pred, th=1 ) : 28 | 29 | y_pred_int = K.round( y_pred ) 30 | y_true_int = K.round( y_true ) 31 | mask = K.cast( y_true>=th, 'float32' ) 32 | matched = K.cast( K.abs(y_pred_int-y_true_int) ==0, 'float32' ) * mask 33 | return K.sum( matched, axis=(1,2,3) ) / K.sum( mask, axis=(1,2,3) ) 34 | def acc0( y_true, y_pred, th=0 ) : 35 | y_pred_int = K.round( y_pred ) 36 | y_true_int = K.round( y_true ) 37 | mask = K.cast( K.abs(y_true-th)<.1, 'float32' ) 38 | matched = K.cast( K.abs(y_pred_int-y_true_int) < .1, 'float32' ) * mask 39 | return K.sum( matched, axis=(1,2,3) ) / K.sum( mask, axis=(1,2,3) ) 40 | 41 | def assign_label( y_true, y_pred ) : 42 | y_modi = tf.py_func( assign_label_np, [ y_true, y_pred ], 'float32', stateful=False ) 43 | y_modi.set_shape( K.int_shape( y_true ) ) 44 | return y_modi 45 | 46 | def assign_label_np( y_true, y_pred ) : 47 | mask = ( y_true > 0 ) 48 | y_pred = np.round( y_pred ).astype('int') 49 | y_true = np.round( y_true ).astype('int') 50 | y_modi = [] 51 | for idx in range( len(y_pred) ): 52 | p = y_pred[idx] 53 | t = y_true[idx] 54 | #print "-" * 50 55 | #print "idx=", idx 56 | true_line_labels = filter( lambda v : v>0, np.unique(t) ) 57 | pred_line_labels = filter( lambda v : v>0, np.unique(p) ) 58 | print ("true_line_labels", true_line_labels) 59 | print ("pred_line_labels", pred_line_labels) 60 | mat = [] 61 | for tl in true_line_labels : 62 | mm = t==tl 63 | row = [] 64 | for pl in pred_line_labels : 65 | vv = -np.sum(p[mm] == pl) 66 | row.append(vv) 67 | mat.append( row ) 68 | mat = np.row_stack(mat) 69 | row_ind, col_ind = linear_sum_assignment( mat ) 70 | true_ind = [ true_line_labels[k] for k in row_ind ] 71 | pred_ind = [ pred_line_labels[k] for k in col_ind ] 72 | for tl, pl in zip( true_ind, pred_ind ) : 73 | t[ t==tl ] = pltf.compat.v1.disable_eager_execution() 74 | #print "assign", tl, "to", pltf.compat.v1.disable_eager_execution() 75 | y_modi.append( np.expand_dims(t,axis=0)) 76 | return np.concatenate(y_modi).astype('float32') 77 | 78 | def seg( y_true, y_pred, th=1 ) : 79 | loss = K.maximum( K.square( y_true - y_pred ), K.abs( y_true - y_pred ) ) 80 | mask = K.cast( y_true>=th, 'float32' ) 81 | return K.sum( mask * loss,axis=(1,2,3)) / K.sum( mask, axis=(1,2,3)) 82 | 83 | 84 | 85 | 86 | 87 | 88 | def prepare_group_loss_numpy(y_true, y_pred) : 89 | """Implement group loss in Sec3.3 of https://arxiv.org/pdf/1611.05424.pdf 90 | NOTE: the 2nd term of the loss has been simplified to the closest mu 91 | """ 92 | within, between = [], [] 93 | for a_true, a_pred in zip(y_true, y_pred) : 94 | N = int(a_true.max()) 95 | 96 | within_true = np.zeros_like(a_true) 97 | #between_true = np.ones_like(a_true) * N 98 | between_true = np.zeros_like(a_true) 99 | masks = [] 100 | mu_list = [] 101 | for line_idx in range(1, N+1) : 102 | #mask = np.abs(a_true - line_idx) < 0.1 103 | mask = (a_true==line_idx) 104 | vals = a_pred[mask] 105 | mu = vals.mean() 106 | within_true[mask] = mu 107 | mu_list.append(mu) 108 | masks.append(mask) 109 | mu_arr = np.array(mu_list) 110 | 111 | 112 | for mask, mu in zip(masks, mu_arr): 113 | #indices = np.argsort(np.abs(mu_arr - mu)) 114 | ind_mu = np.where(mu_arr==mu) 115 | ind_mu = int(ind_mu[0]) 116 | closest_mu = mu_arr[np.minimum(ind_mu+1,len(mu_arr)-1)] 117 | 118 | between_true[mask] = closest_mu 119 | # update output 120 | within.append(within_true) 121 | between.append(between_true) 122 | 123 | return np.stack(within, axis=0), np.stack(between, axis=0) 124 | 125 | def grouping_loss(y_true, y_pred, sigma=0.025) : 126 | y_within, y_between = tf.py_func(prepare_group_loss_numpy, 127 | [y_true,y_pred], 128 | [tf.float32, tf.float32]) 129 | y_within.set_shape(K.int_shape(y_true)) 130 | y_between.set_shape(K.int_shape(y_true)) 131 | y_within = K.stop_gradient(y_within) 132 | y_between = K.stop_gradient(y_between) 133 | mask = K.cast(y_true >=1, 'float32') 134 | diff = (y_within-y_pred)**2 + tf.exp(-(y_between-y_pred)**2/(2.*sigma)) 135 | #diff = tf.exp(-(y_between-y_pred)**2/(2.*sigma)) 136 | loss = K.sum(diff * mask) / (K.sum(mask)) 137 | 138 | origin = K.maximum( K.square( y_true - y_pred ), K.abs( y_true - y_pred ) ) 139 | 140 | origin_loss = K.sum( mask * origin,axis=(1,2,3)) / K.sum( mask, axis=(1,2,3)) 141 | return origin_loss 142 | def MSE_loss(y_true, y_pred, sigma=0.025) : 143 | y_within, y_between = tf.py_func(prepare_group_loss_numpy, 144 | [y_true,y_pred], 145 | [tf.float32, tf.float32]) 146 | y_within.set_shape(K.int_shape(y_true)) 147 | y_between.set_shape(K.int_shape(y_true)) 148 | y_within = K.stop_gradient(y_within) 149 | y_between = K.stop_gradient(y_between) 150 | mask = K.cast(y_true >=1, 'float32') 151 | diff = (y_within-y_pred)**2 + tf.exp(-(y_between-y_pred)**2/(2.*sigma)) 152 | #diff = tf.exp(-(y_between-y_pred)**2/(2.*sigma)) 153 | loss = K.sum(diff * mask) / (K.sum(mask)) 154 | 155 | 156 | origin = K.square( y_true - y_pred ) 157 | 158 | origin_loss = K.sum( mask * origin,axis=(1,2,3)) / K.sum( mask, axis=(1,2,3)) 159 | return origin_loss 160 | 161 | def MAE_loss(y_true, y_pred, sigma=0.025) : 162 | y_within, y_between = tf.py_func(prepare_group_loss_numpy, 163 | [y_true,y_pred], 164 | [tf.float32, tf.float32]) 165 | y_within.set_shape(K.int_shape(y_true)) 166 | y_between.set_shape(K.int_shape(y_true)) 167 | y_within = K.stop_gradient(y_within) 168 | y_between = K.stop_gradient(y_between) 169 | mask = K.cast(y_true >=1, 'float32') 170 | diff = (y_within-y_pred)**2 + tf.exp(-(y_between-y_pred)**2/(2.*sigma)) 171 | #diff = tf.exp(-(y_between-y_pred)**2/(2.*sigma)) 172 | loss = K.sum(diff * mask) / (K.sum(mask)) 173 | origin = K.abs( y_true - y_pred ) 174 | origin_loss = K.sum( mask * origin,axis=(1,2,3)) / K.sum( mask, axis=(1,2,3)) 175 | return origin_loss 176 | 177 | 178 | 179 | 180 | 181 | 182 | class RobustAdaptativeLoss(object): 183 | def __init__(self): 184 | z = np.array([[0]]) 185 | self.v_alpha = K.zeros(shape=(1088, 768, 1)) 186 | self.v_scale = K.zeros(shape=(1088, 768, 1)) 187 | 188 | def loss(self, y_true, y_pred, **kwargs): 189 | mask = K.cast(y_true >=1, 'float32') 190 | x = y_true - y_pred*mask 191 | #origin_loss = K.sum( mask * x,axis=(1,2,3)) / K.sum( mask, axis=(1,2,3)) 192 | 193 | #x = K.reshape(x, shape=(-1, 1)) 194 | lossfun = robust_loss.adaptive.AdaptiveImageLossFunction((1088, 768, 1), float_dtype='float32',color_space='RGB',representation='PIXEL') 195 | alpha = lossfun.alpha() 196 | scale = lossfun.scale() 197 | #loss, alpha, scale = robust_loss.adaptive.AdaptiveLossFunction(num_channels=1,float_dtype="float32") 198 | a = K.update(self.v_alpha, alpha) 199 | s = K.update(self.v_scale, scale) 200 | # The alpha update must be part of the graph but it should 201 | # not influence the result. 202 | 203 | #mask = K.cast(y_true >=1, 'float32') 204 | origin = lossfun(x) 205 | #origin_loss = K.sum( origin,axis=(1,2,3)) / K.sum( mask, axis=(1,2,3)) 206 | #origin_loss = K.sum( mask * origin,axis=(1,2,3)) / K.sum( mask, axis=(1,2,3)) 207 | return origin + 0 * a + 0 * s 208 | 209 | def alpha(self, y_true, y_pred): 210 | return self.v_alpha 211 | def scale(self, y_true, y_pred): 212 | return self.v_scale 213 | 214 | #lossfun = robust_loss.adaptive.AdaptiveImageLossFunction((1088, 768, 1), float_dtype='float32',color_space='RGB') 215 | def Robustloss(y_true, y_pred): 216 | 217 | mask = K.cast(y_true >=1, 'float32') 218 | #lossfun = robust_loss.adaptive.AdaptiveImageLossFunction((1088, 768, 1), float_dtype='float32',color_space='RGB') 219 | x = y_true-y_pred 220 | #Cauchy 221 | #origin = K.log(0.5*x*x + 1) 222 | # Welsch 223 | #origin = 1-K.exp(-0.5*x*x) 224 | #Charbonnier 225 | origin = K.sqrt(x*x+1) - 1 226 | #Geman 227 | #origin = -2*(1/(((x*x)/4)+1) - 1) 228 | #origin = lossfun(x) 229 | 230 | origin_loss = K.sum( mask * origin,axis=(0,1,2)) / K.sum( mask, axis=(1,2,3)) 231 | return origin_loss 232 | 233 | def seg_2( y_true, y_pred, th=1 ) : 234 | loss = K.cast( y_true == y_pred, 'float32' ) 235 | mask = K.cast( y_true>=th, 'float32' ) 236 | return K.sum( mask * loss,axis=(1,2,3)) / K.sum( mask, axis=(1,2,3)) 237 | 238 | def IOU_calc(y_true, y_pred): 239 | smooth = 1.0 240 | y_true_f = K.flatten(y_true) 241 | y_pred_f = K.flatten(y_pred) 242 | intersection = K.sum(y_true_f * y_pred_f) 243 | return 2*(intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) + smooth) 244 | 245 | -------------------------------------------------------------------------------- /src/Untitled.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "code", 5 | "execution_count": 1, 6 | "metadata": {}, 7 | "outputs": [], 8 | "source": [ 9 | "import sys\n", 10 | "sys.path.append('robust_loss')\n", 11 | "\n", 12 | "import robust_loss.adaptive " 13 | ] 14 | }, 15 | { 16 | "cell_type": "code", 17 | "execution_count": 7, 18 | "metadata": {}, 19 | "outputs": [], 20 | "source": [ 21 | "import numpy as np\n", 22 | "import tensorflow as tf\n", 23 | "from tensorflow import keras\n", 24 | "from tensorflow.keras.layers import *\n", 25 | "from tensorflow.keras.models import Model\n", 26 | "from tensorflow.keras import backend as K" 27 | ] 28 | }, 29 | { 30 | "cell_type": "code", 31 | "execution_count": 15, 32 | "metadata": {}, 33 | "outputs": [ 34 | { 35 | "name": "stdout", 36 | "output_type": "stream", 37 | "text": [ 38 | "Model: \"model_7\"\n", 39 | "_________________________________________________________________\n", 40 | "Layer (type) Output Shape Param # \n", 41 | "=================================================================\n", 42 | "input_8 (InputLayer) [(None, 3)] 0 \n", 43 | "_________________________________________________________________\n", 44 | "dense_7 (Dense) (None, 1) 3 \n", 45 | "=================================================================\n", 46 | "Total params: 3\n", 47 | "Trainable params: 3\n", 48 | "Non-trainable params: 0\n", 49 | "_________________________________________________________________\n", 50 | "Train on 100 samples\n", 51 | "Epoch 1/2\n" 52 | ] 53 | }, 54 | { 55 | "ename": "FailedPreconditionError", 56 | "evalue": "Error while reading resource variable loss_7/dense_7_loss/LatentAlpha from Container: localhost. This could mean that the variable was uninitialized. Not found: Resource localhost/loss_7/dense_7_loss/LatentAlpha/N10tensorflow3VarE does not exist.\n\t [[{{node loss_7/dense_7_loss/Sigmoid/ReadVariableOp}}]]", 57 | "output_type": "error", 58 | "traceback": [ 59 | "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", 60 | "\u001b[0;31mFailedPreconditionError\u001b[0m Traceback (most recent call last)", 61 | "\u001b[0;32m\u001b[0m in \u001b[0;36m\u001b[0;34m\u001b[0m\n\u001b[1;32m 45\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 46\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 47\u001b[0;31m \u001b[0mhistory\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mX\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mY\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mepochs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mverbose\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 48\u001b[0m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'final loss:'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhistory\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mhistory\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'loss'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", 62 | "\u001b[0;32m~/anaconda3/envs/hyperLine2/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/training.py\u001b[0m in \u001b[0;36mfit\u001b[0;34m(self, x, y, batch_size, epochs, verbose, callbacks, validation_split, validation_data, shuffle, class_weight, sample_weight, initial_epoch, steps_per_epoch, validation_steps, validation_freq, max_queue_size, workers, use_multiprocessing, **kwargs)\u001b[0m\n\u001b[1;32m 725\u001b[0m \u001b[0mmax_queue_size\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mmax_queue_size\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 726\u001b[0m \u001b[0mworkers\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mworkers\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 727\u001b[0;31m use_multiprocessing=use_multiprocessing)\n\u001b[0m\u001b[1;32m 728\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 729\u001b[0m def evaluate(self,\n", 63 | "\u001b[0;32m~/anaconda3/envs/hyperLine2/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/training_arrays.py\u001b[0m in \u001b[0;36mfit\u001b[0;34m(self, model, x, y, batch_size, epochs, verbose, callbacks, validation_split, validation_data, shuffle, class_weight, sample_weight, initial_epoch, steps_per_epoch, validation_steps, validation_freq, **kwargs)\u001b[0m\n\u001b[1;32m 673\u001b[0m \u001b[0mvalidation_steps\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mvalidation_steps\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 674\u001b[0m \u001b[0mvalidation_freq\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mvalidation_freq\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 675\u001b[0;31m steps_name='steps_per_epoch')\n\u001b[0m\u001b[1;32m 676\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 677\u001b[0m def evaluate(self,\n", 64 | "\u001b[0;32m~/anaconda3/envs/hyperLine2/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/training_arrays.py\u001b[0m in \u001b[0;36mmodel_iteration\u001b[0;34m(model, inputs, targets, sample_weights, batch_size, epochs, verbose, callbacks, val_inputs, val_targets, val_sample_weights, shuffle, initial_epoch, steps_per_epoch, validation_steps, validation_freq, mode, validation_in_fit, prepared_feed_values_from_dataset, steps_name, **kwargs)\u001b[0m\n\u001b[1;32m 392\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 393\u001b[0m \u001b[0;31m# Get outputs.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 394\u001b[0;31m \u001b[0mbatch_outs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mf\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mins_batch\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 395\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbatch_outs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlist\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 396\u001b[0m \u001b[0mbatch_outs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mbatch_outs\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", 65 | "\u001b[0;32m~/anaconda3/envs/hyperLine2/lib/python3.7/site-packages/tensorflow_core/python/keras/backend.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, inputs)\u001b[0m\n\u001b[1;32m 3474\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 3475\u001b[0m fetched = self._callable_fn(*array_vals,\n\u001b[0;32m-> 3476\u001b[0;31m run_metadata=self.run_metadata)\n\u001b[0m\u001b[1;32m 3477\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call_fetch_callbacks\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfetched\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_fetches\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 3478\u001b[0m output_structure = nest.pack_sequence_as(\n", 66 | "\u001b[0;32m~/anaconda3/envs/hyperLine2/lib/python3.7/site-packages/tensorflow_core/python/client/session.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 1470\u001b[0m ret = tf_session.TF_SessionRunCallable(self._session._session,\n\u001b[1;32m 1471\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_handle\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1472\u001b[0;31m run_metadata_ptr)\n\u001b[0m\u001b[1;32m 1473\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mrun_metadata\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 1474\u001b[0m \u001b[0mproto_data\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtf_session\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mTF_GetBuffer\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mrun_metadata_ptr\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", 67 | "\u001b[0;31mFailedPreconditionError\u001b[0m: Error while reading resource variable loss_7/dense_7_loss/LatentAlpha from Container: localhost. This could mean that the variable was uninitialized. Not found: Resource localhost/loss_7/dense_7_loss/LatentAlpha/N10tensorflow3VarE does not exist.\n\t [[{{node loss_7/dense_7_loss/Sigmoid/ReadVariableOp}}]]" 68 | ] 69 | } 70 | ], 71 | "source": [ 72 | "\n", 73 | "\n", 74 | "class RobustAdaptativeLoss(object):\n", 75 | " \n", 76 | " def __init__(self):\n", 77 | " \n", 78 | " z = np.array([[0]])\n", 79 | " self.v_alpha = K.variable(z)\n", 80 | "\n", 81 | " def loss(self, y_true, y_pred, **kwargs):\n", 82 | " x = y_true - y_pred\n", 83 | " #x = K.reshape(x, shape=(-1, 1))\n", 84 | " \n", 85 | " loss = robust_loss.adaptive.AdaptiveLossFunction(num_channels=1,float_dtype=\"float32\")\n", 86 | " alpha = loss.alpha()\n", 87 | " scale = loss.scale()\n", 88 | " op = K.update(self.v_alpha, alpha)\n", 89 | " # The alpha update must be part of the graph but it should\n", 90 | " # not influence the result.\n", 91 | " return loss(x) + 0 * op\n", 92 | "\n", 93 | " def alpha(self, y_true, y_pred):\n", 94 | " return self.v_alpha\n", 95 | "\n", 96 | "def make_model():\n", 97 | " inp = Input(shape=(3,))\n", 98 | " out = Dense(1, use_bias=False)(inp)\n", 99 | " model = Model(inp, out)\n", 100 | " loss = RobustAdaptativeLoss()\n", 101 | " model.compile('adam', loss.loss, metrics=[loss.alpha])\n", 102 | " return model\n", 103 | "\n", 104 | "model = make_model()\n", 105 | "model.summary()\n", 106 | "\n", 107 | "\n", 108 | "\n", 109 | "import numpy as np\n", 110 | "\n", 111 | "FACTORS = np.array([0.5, 2.0, 5.0])\n", 112 | "def target_fn(x):\n", 113 | " return np.dot(x, FACTORS.T)\n", 114 | "\n", 115 | "N_SAMPLES=100\n", 116 | "X = np.random.rand(N_SAMPLES, 3)\n", 117 | "Y = np.apply_along_axis(target_fn, 1, X)\n", 118 | "\n", 119 | "\n", 120 | "history = model.fit(X, Y, epochs=2, verbose=True)\n", 121 | "print('final loss:', history.history['loss'][-1])" 122 | ] 123 | }, 124 | { 125 | "cell_type": "code", 126 | "execution_count": null, 127 | "metadata": {}, 128 | "outputs": [], 129 | "source": [] 130 | } 131 | ], 132 | "metadata": { 133 | "kernelspec": { 134 | "display_name": "hyperLine2", 135 | "language": "python", 136 | "name": "hyperline2" 137 | }, 138 | "language_info": { 139 | "codemirror_mode": { 140 | "name": "ipython", 141 | "version": 3 142 | }, 143 | "file_extension": ".py", 144 | "mimetype": "text/x-python", 145 | "name": "python", 146 | "nbconvert_exporter": "python", 147 | "pygments_lexer": "ipython3", 148 | "version": "3.7.4" 149 | } 150 | }, 151 | "nbformat": 4, 152 | "nbformat_minor": 4 153 | } 154 | -------------------------------------------------------------------------------- /src/__init__.py: -------------------------------------------------------------------------------- 1 | from . import CLLayers 2 | from . import CLLosses 3 | from . import CLUtils -------------------------------------------------------------------------------- /src/__pycache__/CLLayers.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/src/__pycache__/CLLayers.cpython-37.pyc -------------------------------------------------------------------------------- /src/__pycache__/CLLosses.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/src/__pycache__/CLLosses.cpython-37.pyc -------------------------------------------------------------------------------- /src/__pycache__/CLUtils.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/src/__pycache__/CLUtils.cpython-37.pyc -------------------------------------------------------------------------------- /src/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/src/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /src/robust_loss/.ipynb_checkpoints/distribution-checkpoint.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | # Copyright 2020 The Google Research Authors. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | r"""Implements the distribution corresponding to the loss function. 17 | 18 | This library implements the parts of Section 2 of "A General and Adaptive Robust 19 | Loss Function", Jonathan T. Barron, https://arxiv.org/abs/1701.03077, that are 20 | required for evaluating the negative log-likelihood (NLL) of the distribution 21 | and for sampling from the distribution. 22 | """ 23 | 24 | import numbers 25 | 26 | import mpmath 27 | import numpy as np 28 | import tensorflow.compat.v2 as tf 29 | import tensorflow_probability as tfp 30 | from robust_loss import cubic_spline 31 | from robust_loss import general 32 | from robust_loss import util 33 | 34 | 35 | def analytical_base_partition_function(numer, denom): 36 | r"""Accurately approximate the partition function Z(numer / denom). 37 | 38 | This uses the analytical formulation of the true partition function Z(alpha), 39 | as described in the paper (the math after Equation 18), where alpha is a 40 | positive rational value numer/denom. This is expensive to compute and not 41 | differentiable, so it's not implemented in TensorFlow and is only used for 42 | unit tests. 43 | 44 | Args: 45 | numer: the numerator of alpha, an integer >= 0. 46 | denom: the denominator of alpha, an integer > 0. 47 | 48 | Returns: 49 | Z(numer / denom), a double-precision float, accurate to around 9 digits 50 | of precision. 51 | 52 | Raises: 53 | ValueError: If `numer` is not a non-negative integer or if `denom` is not 54 | a positive integer. 55 | """ 56 | if not isinstance(numer, numbers.Integral): 57 | raise ValueError('Expected `numer` of type int, but is of type {}'.format( 58 | type(numer))) 59 | if not isinstance(denom, numbers.Integral): 60 | raise ValueError('Expected `denom` of type int, but is of type {}'.format( 61 | type(denom))) 62 | if not numer >= 0: 63 | raise ValueError('Expected `numer` >= 0, but is = {}'.format(numer)) 64 | if not denom > 0: 65 | raise ValueError('Expected `denom` > 0, but is = {}'.format(denom)) 66 | 67 | alpha = numer / denom 68 | 69 | # The Meijer-G formulation of the partition function has singularities at 70 | # alpha = 0 and alpha = 2, but at those special cases the partition function 71 | # has simple closed forms which we special-case here. 72 | if alpha == 0: 73 | return np.pi * np.sqrt(2) 74 | if alpha == 2: 75 | return np.sqrt(2 * np.pi) 76 | 77 | # Z(n/d) as described in the paper. 78 | a_p = (np.arange(1, numer, dtype=np.float64) / numer).tolist() 79 | b_q = ((np.arange(-0.5, numer - 0.5, dtype=np.float64)) / 80 | numer).tolist() + (np.arange(1, 2 * denom, dtype=np.float64) / 81 | (2 * denom)).tolist() 82 | z = (1. / numer - 1. / (2 * denom))**(2 * denom) 83 | mult = np.exp(np.abs(2 * denom / numer - 1.)) * np.sqrt( 84 | np.abs(2 * denom / numer - 1.)) * (2 * np.pi)**(1 - denom) 85 | return mult * np.float64(mpmath.meijerg([[], a_p], [b_q, []], z)) 86 | 87 | 88 | def partition_spline_curve(alpha): 89 | """Applies a curve to alpha >= 0 to compress its range before interpolation. 90 | 91 | This is a weird hand-crafted function designed to take in alpha values and 92 | curve them to occupy a short finite range that works well when using spline 93 | interpolation to model the partition function Z(alpha). Because Z(alpha) 94 | is only varied in [0, 4] and is especially interesting around alpha=2, this 95 | curve is roughly linear in [0, 4] with a slope of ~1 at alpha=0 and alpha=4 96 | but a slope of ~10 at alpha=2. When alpha > 4 the curve becomes logarithmic. 97 | Some (input, output) pairs for this function are: 98 | [(0, 0), (1, ~1.2), (2, 4), (3, ~6.8), (4, 8), (8, ~8.8), (400000, ~12)] 99 | This function is continuously differentiable. 100 | 101 | Args: 102 | alpha: A numpy array or TF tensor (float32 or float64) with values >= 0. 103 | 104 | Returns: 105 | An array/tensor of curved values >= 0 with the same type as `alpha`, to be 106 | used as input x-coordinates for spline interpolation. 107 | """ 108 | c = lambda z: tf.cast(z, alpha.dtype) 109 | assert_ops = [tf.Assert(tf.reduce_all(alpha >= 0.), [alpha])] 110 | with tf.control_dependencies(assert_ops): 111 | x = tf.where(alpha < 4, (c(2.25) * alpha - c(4.5)) / 112 | (tf.abs(alpha - c(2)) + c(0.25)) + alpha + c(2), 113 | c(5) / c(18) * util.log_safe(c(4) * alpha - c(15)) + c(8)) 114 | return x 115 | 116 | 117 | def inv_partition_spline_curve(x): 118 | """The inverse of partition_spline_curve().""" 119 | c = lambda z: tf.cast(z, x.dtype) 120 | assert_ops = [tf.Assert(tf.reduce_all(x >= 0.), [x])] 121 | with tf.control_dependencies(assert_ops): 122 | alpha = tf.where( 123 | x < 8, 124 | c(0.5) * x + tf.where( 125 | x <= 4, 126 | c(1.25) - tf.sqrt(c(1.5625) - x + c(.25) * tf.square(x)), 127 | c(-1.25) + tf.sqrt(c(9.5625) - c(3) * x + c(.25) * tf.square(x))), 128 | c(3.75) + c(0.25) * util.exp_safe(x * c(3.6) - c(28.8))) 129 | return alpha 130 | 131 | 132 | class Distribution(object): 133 | """A wrapper class around the distribution.""" 134 | 135 | def __init__(self): 136 | 137 | 138 | """Initialize the distribution. 139 | 140 | Load the values, tangents, and x-coordinate scaling of a spline that 141 | approximates the partition function. The spline was produced by running 142 | the script in fit_partition_spline.py. 143 | """ 144 | 145 | #with util.get_resource_as_file( 146 | #'robust_loss/data/partition_spline.npz') as spline_file: 147 | with np.load('partition_spline.npz', allow_pickle=False) as f: 148 | self._spline_x_scale = f['x_scale'] 149 | self._spline_values = f['values'] 150 | self._spline_tangents = f['tangents'] 151 | 152 | def log_base_partition_function(self, alpha): 153 | r"""Approximate the distribution's log-partition function with a 1D spline. 154 | 155 | Because the partition function (Z(\alpha) in the paper) of the distribution 156 | is difficult to model analytically, we approximate it with a (transformed) 157 | cubic hermite spline: Each alpha is pushed through a nonlinearity before 158 | being used to interpolate into a spline, which allows us to use a relatively 159 | small spline to accurately model the log partition function over the range 160 | of all non-negative input values. 161 | 162 | Args: 163 | alpha: A tensor or scalar of single or double precision floats containing 164 | the set of alphas for which we would like an approximate log partition 165 | function. Must be non-negative, as the partition function is undefined 166 | when alpha < 0. 167 | 168 | Returns: 169 | An approximation of log(Z(alpha)) accurate to within 1e-6 170 | """ 171 | float_dtype = alpha.dtype 172 | 173 | # The partition function is undefined when `alpha`< 0. 174 | assert_ops = [tf.Assert(tf.reduce_all(alpha >= 0.), [alpha])] 175 | with tf.control_dependencies(assert_ops): 176 | # Transform `alpha` to the form expected by the spline. 177 | x = partition_spline_curve(alpha) 178 | # Interpolate into the spline. 179 | return cubic_spline.interpolate1d( 180 | x * tf.cast(self._spline_x_scale, float_dtype), 181 | tf.cast(self._spline_values, float_dtype), 182 | tf.cast(self._spline_tangents, float_dtype)) 183 | 184 | def nllfun(self, x, alpha, scale): 185 | r"""Implements the negative log-likelihood (NLL). 186 | 187 | Specifically, we implement -log(p(x | 0, \alpha, c) of Equation 16 in the 188 | paper as nllfun(x, alpha, shape). 189 | 190 | Args: 191 | x: The residual for which the NLL is being computed. x can have any shape, 192 | and alpha and scale will be broadcasted to match x's shape if necessary. 193 | Must be a tensorflow tensor or numpy array of floats. 194 | alpha: The shape parameter of the NLL (\alpha in the paper), where more 195 | negative values cause outliers to "cost" more and inliers to "cost" 196 | less. Alpha can be any non-negative value, but the gradient of the NLL 197 | with respect to alpha has singularities at 0 and 2 so you may want to 198 | limit usage to (0, 2) during gradient descent. Must be a tensorflow 199 | tensor or numpy array of floats. Varying alpha in that range allows for 200 | smooth interpolation between a Cauchy distribution (alpha = 0) and a 201 | Normal distribution (alpha = 2) similar to a Student's T distribution. 202 | scale: The scale parameter of the loss. When |x| < scale, the NLL is like 203 | that of a (possibly unnormalized) normal distribution, and when |x| > 204 | scale the NLL takes on a different shape according to alpha. Must be a 205 | tensorflow tensor or numpy array of floats. 206 | 207 | Returns: 208 | The NLLs for each element of x, in the same shape as x. This is returned 209 | as a TensorFlow graph node of floats with the same precision as x. 210 | """ 211 | # `scale` and `alpha` must have the same type as `x`. 212 | tf.debugging.assert_type(scale, x.dtype) 213 | tf.debugging.assert_type(alpha, x.dtype) 214 | assert_ops = [ 215 | # `scale` must be > 0. 216 | tf.Assert(tf.reduce_all(scale > 0.), [scale]), 217 | # `alpha` must be >= 0. 218 | tf.Assert(tf.reduce_all(alpha >= 0.), [alpha]), 219 | ] 220 | with tf.control_dependencies(assert_ops): 221 | loss = general.lossfun(x, alpha, scale, approximate=False) 222 | log_partition = ( 223 | tf.math.log(scale) + self.log_base_partition_function(alpha)) 224 | nll = loss + log_partition 225 | return nll 226 | 227 | def draw_samples(self, alpha, scale): 228 | r"""Draw samples from the robust distribution. 229 | 230 | This function implements Algorithm 1 the paper. This code is written to 231 | allow for sampling from a set of different distributions, each parametrized 232 | by its own alpha and scale values, as opposed to the more standard approach 233 | of drawing N samples from the same distribution. This is done by repeatedly 234 | performing N instances of rejection sampling for each of the N distributions 235 | until at least one proposal for each of the N distributions has been 236 | accepted. All samples assume a zero mean --- to get non-zero mean samples, 237 | just add each mean to each sample. 238 | 239 | Args: 240 | alpha: A TF tensor/scalar or numpy array/scalar of floats where each 241 | element is the shape parameter of that element's distribution. 242 | scale: A TF tensor/scalar or numpy array/scalar of floats where each 243 | element is the scale parameter of that element's distribution. Must be 244 | the same shape as `alpha`. 245 | 246 | Returns: 247 | A TF tensor with the same shape and precision as `alpha` and `scale` where 248 | each element is a sample drawn from the zero-mean distribution specified 249 | for that element by `alpha` and `scale`. 250 | """ 251 | # `scale` must have the same type as `alpha`. 252 | float_dtype = alpha.dtype 253 | tf.debugging.assert_type(scale, float_dtype) 254 | assert_ops = [ 255 | # `scale` must be > 0. 256 | tf.Assert(tf.reduce_all(scale > 0.), [scale]), 257 | # `alpha` must be >= 0. 258 | tf.Assert(tf.reduce_all(alpha >= 0.), [alpha]), 259 | # `alpha` and `scale` must have the same shape. 260 | tf.Assert( 261 | tf.reduce_all(tf.equal(tf.shape(alpha), tf.shape(scale))), 262 | [tf.shape(alpha), tf.shape(scale)]), 263 | ] 264 | 265 | with tf.control_dependencies(assert_ops): 266 | shape = tf.shape(alpha) 267 | 268 | # The distributions we will need for rejection sampling. The sqrt(2) 269 | # scaling of the Cauchy distribution corrects for our differing 270 | # conventions for standardization. 271 | cauchy = tfp.distributions.Cauchy(loc=0., scale=tf.sqrt(2.)) 272 | uniform = tfp.distributions.Uniform(low=0., high=1.) 273 | 274 | def while_cond(_, accepted): 275 | """Terminate the loop only when all samples have been accepted.""" 276 | return ~tf.reduce_all(accepted) 277 | 278 | def while_body(samples, accepted): 279 | """Generate N proposal samples, and then perform rejection sampling.""" 280 | # Draw N samples from a Cauchy, our proposal distribution. 281 | cauchy_sample = tf.cast(cauchy.sample(shape), float_dtype) 282 | 283 | # Compute the likelihood of each sample under its target distribution. 284 | nll = self.nllfun(cauchy_sample, alpha, tf.cast(1, float_dtype)) 285 | # Bound the NLL. We don't use the approximate loss as it may cause 286 | # unpredictable behavior in the context of sampling. 287 | nll_bound = general.lossfun( 288 | cauchy_sample, 289 | tf.cast(0, float_dtype), 290 | tf.cast(1, float_dtype), 291 | approximate=False) + self.log_base_partition_function(alpha) 292 | 293 | # Draw N samples from a uniform distribution, and use each uniform 294 | # sample to decide whether or not to accept each proposal sample. 295 | uniform_sample = tf.cast(uniform.sample(shape), float_dtype) 296 | accept = uniform_sample <= tf.math.exp(nll_bound - nll) 297 | 298 | # If a sample is accepted, replace its element in `samples` with the 299 | # proposal sample, and set its bit in `accepted` to True. 300 | samples = tf.where(accept, cauchy_sample, samples) 301 | accepted = accept | accepted 302 | return (samples, accepted) 303 | 304 | # Initialize the loop. The first item does not matter as it will get 305 | # overwritten, the second item must be all False. 306 | while_loop_vars = (tf.zeros(shape, 307 | float_dtype), tf.zeros(shape, dtype=bool)) 308 | 309 | # Perform rejection sampling until all N samples have been accepted. 310 | terminal_state = tf.while_loop( 311 | cond=while_cond, body=while_body, loop_vars=while_loop_vars) 312 | 313 | # Because our distribution is a location-scale family, we sample from 314 | # p(x | 0, \alpha, 1) and then scale each sample by `scale`. 315 | samples = tf.multiply(terminal_state[0], scale) 316 | 317 | return samples 318 | -------------------------------------------------------------------------------- /src/robust_loss/.ipynb_checkpoints/fit_partition_spline-checkpoint.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | # Copyright 2020 The Google Research Authors. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | # Lint as: python3 17 | r"""Approximate the distribution's partition function with a spline. 18 | 19 | This script generates values for the distribution's partition function and then 20 | fits a cubic hermite spline to those values, which is then stored to disk. 21 | To run this script, assuming you're in this directory, run: 22 | python -m robust_loss.fit_partition_spline_test 23 | This script will likely never have to be run again, and is provided here for 24 | completeness and reproducibility, or in case someone decides to modify 25 | distribution.partition_spline_curve() in the future in case they find a better 26 | curve. If the user wants a more accurate spline approximation, this can be 27 | obtained by modifying the `x_max`, `x_scale`, and `redundancy` parameters in the 28 | code below, but this should only be done with care. 29 | """ 30 | 31 | from absl import app 32 | import numpy as np 33 | import tensorflow.compat.v2 as tf 34 | from robust_loss import cubic_spline 35 | from robust_loss import distribution 36 | from robust_loss import general 37 | 38 | tf.enable_v2_behavior() 39 | 40 | 41 | def numerical_base_partition_function(alpha): 42 | """Numerically approximate the partition function Z(alpha).""" 43 | # Generate values `num_samples` values in [-x_max, x_max], with more samples 44 | # near the origin as `power` is set to larger values. 45 | num_samples = 2**24 + 1 # We want an odd value so that 0 gets sampled. 46 | x_max = 10**10 47 | power = 6 48 | t = t = tf.linspace( 49 | tf.constant(-1, tf.float64), tf.constant(1, tf.float64), num_samples) 50 | t = tf.sign(t) * tf.abs(t)**power 51 | x = t * x_max 52 | 53 | # Compute losses for the values, then exponentiate the negative losses and 54 | # integrate with the trapezoid rule to get the partition function. 55 | losses = general.lossfun(x, alpha, np.float64(1)) 56 | y = tf.math.exp(-losses) 57 | partition = tf.reduce_sum((y[1:] + y[:-1]) * (x[1:] - x[:-1])) / 2. 58 | return partition 59 | 60 | 61 | def main(argv): 62 | if len(argv) > 1: 63 | raise app.UsageError('Too many command-line arguments.') 64 | 65 | # Parameters governing how the x coordinate of the spline will be laid out. 66 | # We will construct a spline with knots at 67 | # [0 : 1 / x_scale : x_max], 68 | # by fitting it to values sampled at 69 | # [0 : 1 / (x_scale * redundancy) : x_max] 70 | x_max = 12 71 | x_scale = 1024 72 | redundancy = 4 # Must be >= 2 for the spline to be useful. 73 | 74 | spline_spacing = 1. / (x_scale * redundancy) 75 | x_knots = np.arange( 76 | 0, x_max + spline_spacing, spline_spacing, dtype=np.float64) 77 | table = [] 78 | # We iterate over knots, and for each knot recover the alpha value 79 | # corresponding to that knot with inv_partition_spline_curve(), and then 80 | # with that alpha we accurately approximate its partition function using 81 | # numerical_base_partition_function(). 82 | for x_knot in x_knots: 83 | alpha = distribution.inv_partition_spline_curve(x_knot).numpy() 84 | partition = numerical_base_partition_function(alpha).numpy() 85 | table.append((x_knot, alpha, partition)) 86 | print(table[-1]) 87 | 88 | table = np.array(table) 89 | x = table[:, 0] 90 | alpha = table[:, 1] 91 | y_gt = np.log(table[:, 2]) 92 | 93 | # We grab the values from the true log-partition table that correpond to 94 | # knots, by looking for where x * x_scale is an integer. 95 | mask = np.abs(np.round(x * x_scale) - (x * x_scale)) <= 1e-8 96 | values = y_gt[mask] 97 | 98 | # Initialize `tangents` using a central differencing scheme. 99 | values_pad = np.concatenate([[values[0] - values[1] + values[0]], values, 100 | [values[-1] - values[-2] + values[-1]]], 0) 101 | tangents = (values_pad[2:] - values_pad[:-2]) / 2. 102 | 103 | # Construct the spline's value and tangent TF variables, constraining the last 104 | # knot to have a fixed value Z(infinity) and a tangent of zero. 105 | n = len(values) 106 | tangents = tf.Variable(tangents, tf.float64) 107 | values = tf.Variable(values, tf.float64) 108 | 109 | # Fit the spline. 110 | num_iters = 10001 111 | 112 | optimizer = tf.keras.optimizers.SGD(learning_rate=1e-9, momentum=0.99) 113 | 114 | trace = [] 115 | for ii in range(num_iters): 116 | with tf.GradientTape() as tape: 117 | tape.watch([values, tangents]) 118 | # Fix the endpoint to be a known constant with a zero tangent. 119 | i_values = tf.where( 120 | np.arange(n) == (n - 1), 121 | tf.ones_like(values) * 0.70526025442689566, values) 122 | i_tangents = tf.where( 123 | np.arange(n) == (n - 1), tf.zeros_like(tangents), tangents) 124 | i_y = cubic_spline.interpolate1d(x * x_scale, i_values, i_tangents) 125 | # We minimize the maximum residual, which makes for a very ugly 126 | # optimization problem but works well in practice. 127 | i_loss = tf.reduce_max(tf.abs(i_y - y_gt)) 128 | grads = tape.gradient(i_loss, [values, tangents]) 129 | optimizer.apply_gradients(zip(grads, [values, tangents])) 130 | trace.append(i_loss.numpy()) 131 | if (ii % 200) == 0: 132 | print('{:5d}: {:e}'.format(ii, trace[-1])) 133 | 134 | mask = alpha <= 4 135 | max_error_a4 = np.max(np.abs(i_y[mask] - y_gt[mask])) 136 | max_error = np.max(np.abs(i_y - y_gt)) 137 | print('Max Error (a <= 4): {:e}'.format(max_error_a4)) 138 | print('Max Error: {:e}'.format(max_error)) 139 | 140 | # Just a sanity-check on the error. 141 | assert max_error_a4 <= 5e-7 142 | assert max_error <= 5e-7 143 | 144 | # Save the spline to disk. 145 | np.savez( 146 | './data/partition_spline.npz', 147 | x_scale=x_scale, 148 | values=i_values.numpy(), 149 | tangents=i_tangents.numpy()) 150 | 151 | 152 | if __name__ == '__main__': 153 | app.run(main) 154 | -------------------------------------------------------------------------------- /src/robust_loss/.ipynb_checkpoints/general-checkpoint.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | # Copyright 2020 The Google Research Authors. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | r"""Implements the general form of the loss. 17 | 18 | This is the simplest way of using this loss. No parameters will be tuned 19 | automatically, it's just a simple function that takes in parameters (likely 20 | hand-tuned ones) and return a loss. For an adaptive loss, look at adaptive.py 21 | or distribution.py. 22 | """ 23 | 24 | import numpy as np 25 | import tensorflow.compat.v2 as tf 26 | from robust_loss import util 27 | 28 | 29 | def lossfun(x, alpha, scale, approximate=False, epsilon=1e-6): 30 | r"""Implements the general form of the loss. 31 | 32 | This implements the rho(x, \alpha, c) function described in "A General and 33 | Adaptive Robust Loss Function", Jonathan T. Barron, 34 | https://arxiv.org/abs/1701.03077. 35 | 36 | Args: 37 | x: The residual for which the loss is being computed. x can have any shape, 38 | and alpha and scale will be broadcasted to match x's shape if necessary. 39 | Must be a tensorflow tensor or numpy array of floats. 40 | alpha: The shape parameter of the loss (\alpha in the paper), where more 41 | negative values produce a loss with more robust behavior (outliers "cost" 42 | less), and more positive values produce a loss with less robust behavior 43 | (outliers are penalized more heavily). Alpha can be any value in 44 | [-infinity, infinity], but the gradient of the loss with respect to alpha 45 | is 0 at -infinity, infinity, 0, and 2. Must be a tensorflow tensor or 46 | numpy array of floats with the same precision as `x`. Varying alpha allows 47 | for smooth interpolation between a number of discrete robust losses: 48 | alpha=-Infinity: Welsch/Leclerc Loss. 49 | alpha=-2: Geman-McClure loss. 50 | alpha=0: Cauchy/Lortentzian loss. 51 | alpha=1: Charbonnier/pseudo-Huber loss. 52 | alpha=2: L2 loss. 53 | scale: The scale parameter of the loss. When |x| < scale, the loss is an 54 | L2-like quadratic bowl, and when |x| > scale the loss function takes on a 55 | different shape according to alpha. Must be a tensorflow tensor or numpy 56 | array of single-precision floats. 57 | approximate: a bool, where if True, this function returns an approximate and 58 | faster form of the loss, as described in the appendix of the paper. This 59 | approximation holds well everywhere except as x and alpha approach zero. 60 | epsilon: A float that determines how inaccurate the "approximate" version of 61 | the loss will be. Larger values are less accurate but more numerically 62 | stable. Must be great than single-precision machine epsilon. 63 | 64 | Returns: 65 | The losses for each element of x, in the same shape as x. This is returned 66 | as a TensorFlow graph node of single precision floats. 67 | """ 68 | # `scale` and `alpha` must have the same type as `x`. 69 | float_dtype = x.dtype 70 | tf.debugging.assert_type(scale, float_dtype) 71 | tf.debugging.assert_type(alpha, float_dtype) 72 | # `scale` must be > 0. 73 | assert_ops = [tf.Assert(tf.reduce_all(tf.greater(scale, 0.)), [scale])] 74 | with tf.control_dependencies(assert_ops): 75 | # Broadcast `alpha` and `scale` to have the same shape as `x`. 76 | alpha = tf.broadcast_to(alpha, tf.shape(x)) 77 | scale = tf.broadcast_to(scale, tf.shape(x)) 78 | 79 | if approximate: 80 | # `epsilon` must be greater than single-precision machine epsilon. 81 | assert epsilon > np.finfo(np.float32).eps 82 | # Compute an approximate form of the loss which is faster, but innacurate 83 | # when x and alpha are near zero. 84 | b = tf.abs(alpha - tf.cast(2., float_dtype)) + epsilon 85 | d = tf.where( 86 | tf.greater_equal(alpha, 0.), alpha + epsilon, alpha - epsilon) 87 | loss = (b / d) * (tf.pow(tf.square(x / scale) / b + 1., 0.5 * d) - 1.) 88 | else: 89 | # Compute the exact loss. 90 | 91 | # This will be used repeatedly. 92 | squared_scaled_x = tf.square(x / scale) 93 | 94 | # The loss when alpha == 2. 95 | loss_two = 0.5 * squared_scaled_x 96 | # The loss when alpha == 0. 97 | loss_zero = util.log1p_safe(0.5 * squared_scaled_x) 98 | # The loss when alpha == -infinity. 99 | loss_neginf = -tf.math.expm1(-0.5 * squared_scaled_x) 100 | # The loss when alpha == +infinity. 101 | loss_posinf = util.expm1_safe(0.5 * squared_scaled_x) 102 | 103 | # The loss when not in one of the above special cases. 104 | machine_epsilon = tf.cast(np.finfo(np.float32).eps, float_dtype) 105 | # Clamp |2-alpha| to be >= machine epsilon so that it's safe to divide by. 106 | beta_safe = tf.maximum(machine_epsilon, tf.abs(alpha - 2.)) 107 | # Clamp |alpha| to be >= machine epsilon so that it's safe to divide by. 108 | alpha_safe = tf.where( 109 | tf.greater_equal(alpha, 0.), tf.ones_like(alpha), 110 | -tf.ones_like(alpha)) * tf.maximum(machine_epsilon, tf.abs(alpha)) 111 | loss_otherwise = (beta_safe / alpha_safe) * ( 112 | tf.pow(squared_scaled_x / beta_safe + 1., 0.5 * alpha) - 1.) 113 | 114 | # Select which of the cases of the loss to return. 115 | loss = tf.where( 116 | tf.equal(alpha, -tf.cast(float('inf'), float_dtype)), loss_neginf, 117 | tf.where( 118 | tf.equal(alpha, 0.), loss_zero, 119 | tf.where( 120 | tf.equal(alpha, 2.), loss_two, 121 | tf.where( 122 | tf.equal(alpha, tf.cast(float('inf'), float_dtype)), 123 | loss_posinf, loss_otherwise)))) 124 | 125 | return loss 126 | -------------------------------------------------------------------------------- /src/robust_loss/.ipynb_checkpoints/util-checkpoint.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | # Copyright 2020 The Google Research Authors. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | """Helper functions.""" 17 | 18 | import numpy as np 19 | import tensorflow.compat.v2 as tf 20 | 21 | 22 | 23 | def log_safe(x): 24 | """The same as tf.math.log(x), but clamps the input to prevent NaNs.""" 25 | return tf.math.log(tf.minimum(x, tf.cast(3e37, x.dtype))) 26 | 27 | 28 | def log1p_safe(x): 29 | """The same as tf.math.log1p(x), but clamps the input to prevent NaNs.""" 30 | return tf.math.log1p(tf.minimum(x, tf.cast(3e37, x.dtype))) 31 | 32 | 33 | def exp_safe(x): 34 | """The same as tf.math.exp(x), but clamps the input to prevent NaNs.""" 35 | return tf.math.exp(tf.minimum(x, tf.cast(87.5, x.dtype))) 36 | 37 | 38 | def expm1_safe(x): 39 | """The same as tf.math.expm1(x), but clamps the input to prevent NaNs.""" 40 | return tf.math.expm1(tf.minimum(x, tf.cast(87.5, x.dtype))) 41 | 42 | 43 | def inv_softplus(y): 44 | """The inverse of tf.nn.softplus().""" 45 | return tf.where(y > 87.5, y, tf.math.log(tf.math.expm1(y))) 46 | 47 | 48 | def logit(y): 49 | """The inverse of tf.nn.sigmoid().""" 50 | return -tf.math.log(1. / y - 1.) 51 | 52 | 53 | def affine_sigmoid(real, lo=0, hi=1): 54 | """Maps reals to (lo, hi), where 0 maps to (lo+hi)/2.""" 55 | if not lo < hi: 56 | raise ValueError('`lo` (%g) must be < `hi` (%g)' % (lo, hi)) 57 | alpha = tf.sigmoid(real) * (hi - lo) + lo 58 | return alpha 59 | 60 | 61 | def inv_affine_sigmoid(alpha, lo=0, hi=1): 62 | """The inverse of affine_sigmoid(., lo, hi).""" 63 | if not lo < hi: 64 | raise ValueError('`lo` (%g) must be < `hi` (%g)' % (lo, hi)) 65 | real = logit((alpha - lo) / (hi - lo)) 66 | return real 67 | 68 | 69 | def affine_softplus(real, lo=0, ref=1): 70 | """Maps real numbers to (lo, infinity), where 0 maps to ref.""" 71 | if not lo < ref: 72 | raise ValueError('`lo` (%g) must be < `ref` (%g)' % (lo, ref)) 73 | shift = inv_softplus(tf.cast(1., real.dtype)) 74 | scale = (ref - lo) * tf.nn.softplus(real + shift) + lo 75 | return scale 76 | 77 | 78 | def inv_affine_softplus(scale, lo=0, ref=1): 79 | """The inverse of affine_softplus(., lo, ref).""" 80 | if not lo < ref: 81 | raise ValueError('`lo` (%g) must be < `ref` (%g)' % (lo, ref)) 82 | shift = inv_softplus(tf.cast(1., scale.dtype)) 83 | real = inv_softplus((scale - lo) / (ref - lo)) - shift 84 | return real 85 | 86 | 87 | def students_t_nll(x, df, scale): 88 | """The NLL of a Generalized Student's T distribution (w/o including TFP).""" 89 | return 0.5 * ((df + 1.) * tf.math.log1p( 90 | (x / scale)**2. / df) + tf.math.log(df)) + tf.math.log( 91 | tf.abs(scale)) + tf.math.lgamma( 92 | 0.5 * df) - tf.math.lgamma(0.5 * df + 0.5) + 0.5 * np.log(np.pi) 93 | 94 | 95 | # A constant scale that makes tf.image.rgb_to_yuv() volume preserving. 96 | _VOLUME_PRESERVING_YUV_SCALE = 1.580227820074 97 | 98 | 99 | def rgb_to_syuv(rgb): 100 | """A volume preserving version of tf.image.rgb_to_yuv(). 101 | 102 | By "volume preserving" we mean that rgb_to_syuv() is in the "special linear 103 | group", or equivalently, that the Jacobian determinant of the transformation 104 | is 1. 105 | 106 | Args: 107 | rgb: A tensor whose last dimension corresponds to RGB channels and is of 108 | size 3. 109 | 110 | Returns: 111 | A scaled YUV version of the input tensor, such that this transformation is 112 | volume-preserving. 113 | """ 114 | return _VOLUME_PRESERVING_YUV_SCALE * tf.image.rgb_to_yuv(rgb) 115 | 116 | 117 | def syuv_to_rgb(yuv): 118 | """A volume preserving version of tf.image.yuv_to_rgb(). 119 | 120 | By "volume preserving" we mean that rgb_to_syuv() is in the "special linear 121 | group", or equivalently, that the Jacobian determinant of the transformation 122 | is 1. 123 | 124 | Args: 125 | yuv: A tensor whose last dimension corresponds to scaled YUV channels and is 126 | of size 3 (ie, the output of rgb_to_syuv()). 127 | 128 | Returns: 129 | An RGB version of the input tensor, such that this transformation is 130 | volume-preserving. 131 | """ 132 | return tf.image.yuv_to_rgb(yuv / _VOLUME_PRESERVING_YUV_SCALE) 133 | 134 | 135 | def image_dct(image): 136 | """Does a type-II DCT (aka "The DCT") on axes 1 and 2 of a rank-3 tensor.""" 137 | dct_y = tf.transpose( 138 | a=tf.signal.dct(image, type=2, norm='ortho'), perm=[0, 2, 1]) 139 | dct_x = tf.transpose( 140 | a=tf.signal.dct(dct_y, type=2, norm='ortho'), perm=[0, 2, 1]) 141 | return dct_x 142 | 143 | 144 | def image_idct(dct_x): 145 | """Inverts image_dct(), by performing a type-III DCT.""" 146 | dct_y = tf.signal.idct( 147 | tf.transpose(dct_x, perm=[0, 2, 1]), type=2, norm='ortho') 148 | image = tf.signal.idct( 149 | tf.transpose(dct_y, perm=[0, 2, 1]), type=2, norm='ortho') 150 | return image 151 | 152 | 153 | def compute_jacobian(f, x): 154 | """Computes the Jacobian of function `f` with respect to input `x`.""" 155 | x = tf.convert_to_tensor(x) 156 | with tf.GradientTape(persistent=True) as tape: 157 | tape.watch(x) 158 | vec = lambda x: tf.reshape(x, [-1]) 159 | jacobian = tf.stack( 160 | [vec(tape.gradient(vec(f(x))[d], x)) for d in range(tf.size(x))]) 161 | return jacobian 162 | 163 | 164 | def get_resource_as_file(path): 165 | """A uniform interface for internal/open-source files.""" 166 | 167 | class NullContextManager(object): 168 | 169 | def __init__(self, dummy_resource=None): 170 | self.dummy_resource = dummy_resource 171 | 172 | def __enter__(self): 173 | return self.dummy_resource 174 | 175 | def __exit__(self, *args): 176 | pass 177 | 178 | return NullContextManager('./' + path) 179 | 180 | 181 | def get_resource_filename(path): 182 | """A uniform interface for internal/open-source filenames.""" 183 | return './' + path 184 | -------------------------------------------------------------------------------- /src/robust_loss/README.md: -------------------------------------------------------------------------------- 1 | # A General and Adaptive Robust Loss Function 2 | 3 | This directory contains Tensorflow 2 reference code for the paper 4 | [A General and Adaptive Robust Loss Function](https://arxiv.org/abs/1701.03077), 5 | Jonathan T. Barron CVPR, 2019 6 | 7 | To use this code, include `general.py` or `adaptive.py` and call the loss 8 | function. `general.py` implements the "general" form of the loss, which assumes 9 | you are prepared to set and tune hyperparameters yourself, and `adaptive.py` 10 | implements the "adaptive" form of the loss, which tries to adapt the 11 | hyperparameters automatically and also includes support for imposing losses in 12 | different image representations. The probability distribution underneath the 13 | adaptive loss is implemented in `distribution.py`. 14 | 15 | The VAE experiment from the paper can be reproduced by running `vae.py`. See 16 | `example.ipynb` for a simple toy example of how this loss can be used. 17 | 18 | This code repository is shared with all of Google Research, so it's not very 19 | useful for reporting or tracking bugs. If you have any issues using this code, 20 | please do not open an issue, and instead just email jonbarron@gmail.com. 21 | 22 | If you use this code, please cite it: 23 | ``` 24 | @article{BarronCVPR2019, 25 | Author = {Jonathan T. Barron}, 26 | Title = {A General and Adaptive Robust Loss Function}, 27 | Journal = {CVPR}, 28 | Year = {2019} 29 | } 30 | ``` 31 | -------------------------------------------------------------------------------- /src/robust_loss/__pycache__/adaptive.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/src/robust_loss/__pycache__/adaptive.cpython-37.pyc -------------------------------------------------------------------------------- /src/robust_loss/__pycache__/cubic_spline.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/src/robust_loss/__pycache__/cubic_spline.cpython-37.pyc -------------------------------------------------------------------------------- /src/robust_loss/__pycache__/distribution.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/src/robust_loss/__pycache__/distribution.cpython-37.pyc -------------------------------------------------------------------------------- /src/robust_loss/__pycache__/general.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/src/robust_loss/__pycache__/general.cpython-37.pyc -------------------------------------------------------------------------------- /src/robust_loss/__pycache__/util.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/src/robust_loss/__pycache__/util.cpython-37.pyc -------------------------------------------------------------------------------- /src/robust_loss/__pycache__/wavelet.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/src/robust_loss/__pycache__/wavelet.cpython-37.pyc -------------------------------------------------------------------------------- /src/robust_loss/cubic_spline.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | # Copyright 2020 The Google Research Authors. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | """Implements 1D cubic Hermite spline interpolation.""" 17 | 18 | import tensorflow.compat.v2 as tf 19 | 20 | 21 | def interpolate1d(x, values, tangents): 22 | r"""Perform cubic hermite spline interpolation on a 1D spline. 23 | 24 | The x coordinates of the spline knots are at [0 : 1 : len(values)-1]. 25 | Queries outside of the range of the spline are computed using linear 26 | extrapolation. See https://en.wikipedia.org/wiki/Cubic_Hermite_spline 27 | for details, where "x" corresponds to `x`, "p" corresponds to `values`, and 28 | "m" corresponds to `tangents`. 29 | 30 | Args: 31 | x: A tensor of any size of single or double precision floats containing the 32 | set of values to be used for interpolation into the spline. 33 | values: A vector of single or double precision floats containing the value 34 | of each knot of the spline being interpolated into. Must be the same 35 | length as `tangents` and the same type as `x`. 36 | tangents: A vector of single or double precision floats containing the 37 | tangent (derivative) of each knot of the spline being interpolated into. 38 | Must be the same length as `values` and the same type as `x`. 39 | 40 | Returns: 41 | The result of interpolating along the spline defined by `values`, and 42 | `tangents`, using `x` as the query values. Will be the same length and type 43 | as `x`. 44 | """ 45 | # `values` and `tangents` must have the same type as `x`. 46 | tf.debugging.assert_type(values, x.dtype) 47 | tf.debugging.assert_type(tangents, x.dtype) 48 | float_dtype = x.dtype 49 | assert_ops = [ 50 | # `values` must be a vector. 51 | tf.Assert(tf.equal(tf.rank(values), 1), [tf.shape(values)]), 52 | # `tangents` must be a vector. 53 | tf.Assert(tf.equal(tf.rank(tangents), 1), [tf.shape(values)]), 54 | # `values` and `tangents` must have the same length. 55 | tf.Assert( 56 | tf.equal(tf.shape(values)[0], 57 | tf.shape(tangents)[0]), 58 | [tf.shape(values)[0], tf.shape(tangents)[0]]), 59 | ] 60 | with tf.control_dependencies(assert_ops): 61 | # Find the indices of the knots below and above each x. 62 | x_lo = tf.cast( 63 | tf.floor( 64 | tf.clip_by_value(x, 0., 65 | tf.cast(tf.shape(values)[0] - 2, float_dtype))), 66 | tf.int32) 67 | x_hi = x_lo + 1 68 | 69 | # Compute the relative distance between each `x` and the knot below it. 70 | t = x - tf.cast(x_lo, float_dtype) 71 | 72 | # Compute the cubic hermite expansion of `t`. 73 | t_sq = tf.square(t) 74 | t_cu = t * t_sq 75 | h01 = -2. * t_cu + 3. * t_sq 76 | h00 = 1. - h01 77 | h11 = t_cu - t_sq 78 | h10 = h11 - t_sq + t 79 | 80 | # Linearly extrapolate above and below the extents of the spline for all 81 | # values. 82 | value_before = tangents[0] * t + values[0] 83 | value_after = tangents[-1] * (t - 1.) + values[-1] 84 | 85 | # Cubically interpolate between the knots below and above each query point. 86 | neighbor_values_lo = tf.gather(values, x_lo) 87 | neighbor_values_hi = tf.gather(values, x_hi) 88 | neighbor_tangents_lo = tf.gather(tangents, x_lo) 89 | neighbor_tangents_hi = tf.gather(tangents, x_hi) 90 | value_mid = ( 91 | neighbor_values_lo * h00 + neighbor_values_hi * h01 + 92 | neighbor_tangents_lo * h10 + neighbor_tangents_hi * h11) 93 | 94 | # Return the interpolated or extrapolated values for each query point, 95 | # depending on whether or not the query lies within the span of the spline. 96 | return tf.where(t < 0., value_before, 97 | tf.where(t > 1., value_after, value_mid)) 98 | -------------------------------------------------------------------------------- /src/robust_loss/cubic_spline_test.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | # Copyright 2020 The Google Research Authors. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | """Tests for cubic_spline.py.""" 17 | 18 | import numpy as np 19 | import tensorflow.compat.v2 as tf 20 | 21 | from robust_loss import cubic_spline 22 | 23 | tf.enable_v2_behavior() 24 | 25 | 26 | class CubicSplineTest(tf.test.TestCase): 27 | 28 | def setUp(self): 29 | super(CubicSplineTest, self).setUp() 30 | np.random.seed(0) 31 | 32 | def _interpolate1d(self, x, values, tangents): 33 | """Compute interpolate1d(x, values, tangents) and its derivative. 34 | 35 | This is just a helper function around cubic_spline.interpolate1d() that does 36 | the necessary work to get derivatives and handle TensorFlow sessions. 37 | 38 | Args: 39 | x: A np.array of values to interpolate with. 40 | values: A np.array of knot values for the spline. 41 | tangents: A np.array of knot tangents for the spline. 42 | 43 | Returns: 44 | A tuple containing: 45 | (An np.array of interpolated values, 46 | A np.array of derivatives of interpolated values wrt `x`) 47 | 48 | Typical usage example: 49 | y, dy_dx = self._interpolate1d(x, values, tangents) 50 | """ 51 | x = tf.convert_to_tensor(x) 52 | with tf.GradientTape() as tape: 53 | tape.watch(x) 54 | y = cubic_spline.interpolate1d(x, values, tangents) 55 | dy_dx = tape.gradient(y, x) 56 | return y, dy_dx 57 | 58 | def _interpolation_preserves_dtype(self, float_dtype): 59 | """Check that interpolating at a knot produces the value at that knot.""" 60 | n = 16 61 | x = float_dtype(np.random.normal(size=n)) 62 | values = float_dtype(np.random.normal(size=n)) 63 | tangents = float_dtype(np.random.normal(size=n)) 64 | y = cubic_spline.interpolate1d(x, values, tangents) 65 | self.assertDTypeEqual(y, float_dtype) 66 | 67 | def testInterpolationPreservesDtypeSingle(self): 68 | self._interpolation_preserves_dtype(np.float32) 69 | 70 | def testInterpolationPreservesDtypeDouble(self): 71 | self._interpolation_preserves_dtype(np.float64) 72 | 73 | def _interpolation_reproduces_values_at_knots(self, float_dtype): 74 | """Check that interpolating at a knot produces the value at that knot.""" 75 | n = 32768 76 | x = np.arange(n, dtype=float_dtype) 77 | values = float_dtype(np.random.normal(size=n)) 78 | tangents = float_dtype(np.random.normal(size=n)) 79 | y = cubic_spline.interpolate1d(x, values, tangents) 80 | self.assertAllClose(y, values) 81 | 82 | def testInterpolationReproducesValuesAtKnotsSingle(self): 83 | self._interpolation_reproduces_values_at_knots(np.float32) 84 | 85 | def testInterpolationReproducesValuesAtKnotsDouble(self): 86 | self._interpolation_reproduces_values_at_knots(np.float64) 87 | 88 | def _interpolation_reproduces_tangents_at_knots(self, float_dtype): 89 | """Check that the derivative at a knot produces the tangent at that knot.""" 90 | n = 32768 91 | x = np.arange(n, dtype=float_dtype) 92 | values = float_dtype(np.random.normal(size=n)) 93 | tangents = float_dtype(np.random.normal(size=n)) 94 | _, dy_dx = self._interpolate1d(x, values, tangents) 95 | self.assertAllClose(dy_dx, tangents) 96 | 97 | def testInterpolationReproducesTangentsAtKnotsSingle(self): 98 | self._interpolation_reproduces_tangents_at_knots(np.float32) 99 | 100 | def testInterpolationReproducesTangentsAtKnotsDouble(self): 101 | self._interpolation_reproduces_tangents_at_knots(np.float64) 102 | 103 | def _zero_tangent_midpoint_values_and_derivatives_are_correct( 104 | self, float_dtype): 105 | """Check that splines with zero tangents behave correctly at midpoints. 106 | 107 | Make a spline whose tangents are all zeros, and then verify that 108 | midpoints between each pair of knots have the mean value of their adjacent 109 | knots, and have a derivative that is 1.5x the difference between their 110 | adjacent knots. 111 | 112 | Args: 113 | float_dtype: the dtype of the floats to be tested. 114 | """ 115 | # Make a spline with random values and all-zero tangents. 116 | n = 32768 117 | values = float_dtype(np.random.normal(size=n)) 118 | tangents = np.zeros_like(values) 119 | 120 | # Query n-1 points placed exactly in between each pair of knots. 121 | x = float_dtype(np.arange(n - 1)) + float_dtype(0.5) 122 | 123 | # Get the interpolated values and derivatives. 124 | y, dy_dx = self._interpolate1d(x, values, tangents) 125 | 126 | # Check that the interpolated values of all queries lies at the midpoint of 127 | # its surrounding knot values. 128 | y_true = (values[0:-1] + values[1:]) / 2. 129 | self.assertAllClose(y, y_true) 130 | 131 | # Check that the derivative of all interpolated values is (fun fact!) 1.5x 132 | # the numerical difference between adjacent knot values. 133 | dy_dx_true = 1.5 * (values[1:] - values[0:-1]) 134 | self.assertAllClose(dy_dx, dy_dx_true) 135 | 136 | def testZeroTangentMidpointValuesAndDerivativesAreCorrectSingle(self): 137 | self._zero_tangent_midpoint_values_and_derivatives_are_correct(np.float32) 138 | 139 | def testZeroTangentMidpointValuesAndDerivativesAreCorrectDouble(self): 140 | self._zero_tangent_midpoint_values_and_derivatives_are_correct(np.float64) 141 | 142 | def _zero_tangent_intermediate_values_and_derivatives_do_not_overshoot( 143 | self, float_dtype): 144 | """Check that splines with zero tangents behave correctly between knots. 145 | 146 | Make a spline whose tangents are all zeros, and then verify that points 147 | between each knot lie in between the knot values, and have derivatives 148 | are between 0 and 1.5x the numerical difference between knot values 149 | (mathematically, 1.5x is the max derivative if the tangents are zero). 150 | 151 | Args: 152 | float_dtype: the dtype of the floats to be tested. 153 | """ 154 | 155 | # Make a spline with all-zero tangents and random values. 156 | n = 32768 157 | values = float_dtype(np.random.normal(size=n)) 158 | tangents = np.zeros_like(values) 159 | 160 | # Query n-1 points placed somewhere randomly in between all adjacent knots. 161 | x = np.arange( 162 | n - 1, dtype=float_dtype) + float_dtype(np.random.uniform(size=n - 1)) 163 | 164 | # Get the interpolated values and derivatives. 165 | y, dy_dx = self._interpolate1d(x, values, tangents) 166 | 167 | # Check that the interpolated values of all queries lies between its 168 | # surrounding knot values. 169 | self.assertTrue( 170 | np.all(((values[0:-1] <= y) & (y <= values[1:])) 171 | | ((values[0:-1] >= y) & (y >= values[1:])))) 172 | 173 | # Check that all derivatives of interpolated values are between 0 and 1.5x 174 | # the numerical difference between adjacent knot values. 175 | max_dy_dx = (1.5 + 1e-3) * (values[1:] - values[0:-1]) 176 | self.assertTrue( 177 | np.all(((0 <= dy_dx) & (dy_dx <= max_dy_dx)) 178 | | ((0 >= dy_dx) & (dy_dx >= max_dy_dx)))) 179 | 180 | def testZeroTangentIntermediateValuesAndDerivativesDoNotOvershootSingle(self): 181 | self._zero_tangent_intermediate_values_and_derivatives_do_not_overshoot( 182 | np.float32) 183 | 184 | def testZeroTangentIntermediateValuesAndDerivativesDoNotOvershootDouble(self): 185 | self._zero_tangent_intermediate_values_and_derivatives_do_not_overshoot( 186 | np.float64) 187 | 188 | def _linear_ramps_reproduce_correctly(self, float_dtype): 189 | """Check that interpolating a ramp reproduces a ramp. 190 | 191 | Generate linear ramps, render them into splines, and then interpolate and 192 | extrapolate the splines and verify that they reproduce the ramp. 193 | 194 | Args: 195 | float_dtype: the dtype of the floats to be tested. 196 | """ 197 | n = 256 198 | # Generate queries inside and outside the support of the spline. 199 | x = float_dtype((np.random.uniform(size=1024) * 2 - 0.5) * (n - 1)) 200 | idx = np.arange(n, dtype=float_dtype) 201 | for _ in range(8): 202 | slope = np.random.normal() 203 | bias = np.random.normal() 204 | values = slope * idx + bias 205 | tangents = np.ones_like(values) * slope 206 | y = cubic_spline.interpolate1d(x, values, tangents) 207 | y_true = slope * x + bias 208 | self.assertAllClose(y, y_true) 209 | 210 | def testLinearRampsReproduceCorrectlySingle(self): 211 | self._linear_ramps_reproduce_correctly(np.float32) 212 | 213 | def testLinearRampsReproduceCorrectlyDouble(self): 214 | self._linear_ramps_reproduce_correctly(np.float64) 215 | 216 | def _extrapolation_is_linear(self, float_dtype): 217 | """Check that extrapolation is linear with respect to the endpoint knots. 218 | 219 | Generate random splines and query them outside of the support of the 220 | spline, and veify that extrapolation is linear with respect to the 221 | endpoint knots. 222 | 223 | Args: 224 | float_dtype: the dtype of the floats to be tested. 225 | """ 226 | n = 256 227 | # Generate queries above and below the support of the spline. 228 | x_below = float_dtype(-(np.random.uniform(size=1024)) * (n - 1)) 229 | x_above = float_dtype((np.random.uniform(size=1024) + 1.) * (n - 1)) 230 | for _ in range(8): 231 | values = float_dtype(np.random.normal(size=n)) 232 | tangents = float_dtype(np.random.normal(size=n)) 233 | 234 | # Query the spline below its support and check that it's a linear ramp 235 | # with the slope and bias of the beginning of the spline. 236 | y_below = cubic_spline.interpolate1d(x_below, values, tangents) 237 | y_below_true = tangents[0] * x_below + values[0] 238 | self.assertAllClose(y_below, y_below_true) 239 | 240 | # Query the spline above its support and check that it's a linear ramp 241 | # with the slope and bias of the end of the spline. 242 | y_above = cubic_spline.interpolate1d(x_above, values, tangents) 243 | y_above_true = tangents[-1] * (x_above - (n - 1)) + values[-1] 244 | self.assertAllClose(y_above, y_above_true) 245 | 246 | def testExtrapolationIsLinearSingle(self): 247 | self._extrapolation_is_linear(np.float32) 248 | 249 | def testExtrapolationIsLinearDouble(self): 250 | self._extrapolation_is_linear(np.float64) 251 | 252 | 253 | if __name__ == '__main__': 254 | tf.test.main() 255 | -------------------------------------------------------------------------------- /src/robust_loss/data/partition_spline.npz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/src/robust_loss/data/partition_spline.npz -------------------------------------------------------------------------------- /src/robust_loss/data/wavelet_golden.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/src/robust_loss/data/wavelet_golden.mat -------------------------------------------------------------------------------- /src/robust_loss/data/wavelet_vis_golden.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/src/robust_loss/data/wavelet_vis_golden.png -------------------------------------------------------------------------------- /src/robust_loss/distribution.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | # Copyright 2020 The Google Research Authors. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | r"""Implements the distribution corresponding to the loss function. 17 | 18 | This library implements the parts of Section 2 of "A General and Adaptive Robust 19 | Loss Function", Jonathan T. Barron, https://arxiv.org/abs/1701.03077, that are 20 | required for evaluating the negative log-likelihood (NLL) of the distribution 21 | and for sampling from the distribution. 22 | """ 23 | 24 | import numbers 25 | 26 | import mpmath 27 | import numpy as np 28 | import tensorflow.compat.v2 as tf 29 | import tensorflow_probability as tfp 30 | from robust_loss import cubic_spline 31 | from robust_loss import general 32 | from robust_loss import util 33 | 34 | 35 | def analytical_base_partition_function(numer, denom): 36 | r"""Accurately approximate the partition function Z(numer / denom). 37 | 38 | This uses the analytical formulation of the true partition function Z(alpha), 39 | as described in the paper (the math after Equation 18), where alpha is a 40 | positive rational value numer/denom. This is expensive to compute and not 41 | differentiable, so it's not implemented in TensorFlow and is only used for 42 | unit tests. 43 | 44 | Args: 45 | numer: the numerator of alpha, an integer >= 0. 46 | denom: the denominator of alpha, an integer > 0. 47 | 48 | Returns: 49 | Z(numer / denom), a double-precision float, accurate to around 9 digits 50 | of precision. 51 | 52 | Raises: 53 | ValueError: If `numer` is not a non-negative integer or if `denom` is not 54 | a positive integer. 55 | """ 56 | if not isinstance(numer, numbers.Integral): 57 | raise ValueError('Expected `numer` of type int, but is of type {}'.format( 58 | type(numer))) 59 | if not isinstance(denom, numbers.Integral): 60 | raise ValueError('Expected `denom` of type int, but is of type {}'.format( 61 | type(denom))) 62 | if not numer >= 0: 63 | raise ValueError('Expected `numer` >= 0, but is = {}'.format(numer)) 64 | if not denom > 0: 65 | raise ValueError('Expected `denom` > 0, but is = {}'.format(denom)) 66 | 67 | alpha = numer / denom 68 | 69 | # The Meijer-G formulation of the partition function has singularities at 70 | # alpha = 0 and alpha = 2, but at those special cases the partition function 71 | # has simple closed forms which we special-case here. 72 | if alpha == 0: 73 | return np.pi * np.sqrt(2) 74 | if alpha == 2: 75 | return np.sqrt(2 * np.pi) 76 | 77 | # Z(n/d) as described in the paper. 78 | a_p = (np.arange(1, numer, dtype=np.float64) / numer).tolist() 79 | b_q = ((np.arange(-0.5, numer - 0.5, dtype=np.float64)) / 80 | numer).tolist() + (np.arange(1, 2 * denom, dtype=np.float64) / 81 | (2 * denom)).tolist() 82 | z = (1. / numer - 1. / (2 * denom))**(2 * denom) 83 | mult = np.exp(np.abs(2 * denom / numer - 1.)) * np.sqrt( 84 | np.abs(2 * denom / numer - 1.)) * (2 * np.pi)**(1 - denom) 85 | return mult * np.float64(mpmath.meijerg([[], a_p], [b_q, []], z)) 86 | 87 | 88 | def partition_spline_curve(alpha): 89 | """Applies a curve to alpha >= 0 to compress its range before interpolation. 90 | 91 | This is a weird hand-crafted function designed to take in alpha values and 92 | curve them to occupy a short finite range that works well when using spline 93 | interpolation to model the partition function Z(alpha). Because Z(alpha) 94 | is only varied in [0, 4] and is especially interesting around alpha=2, this 95 | curve is roughly linear in [0, 4] with a slope of ~1 at alpha=0 and alpha=4 96 | but a slope of ~10 at alpha=2. When alpha > 4 the curve becomes logarithmic. 97 | Some (input, output) pairs for this function are: 98 | [(0, 0), (1, ~1.2), (2, 4), (3, ~6.8), (4, 8), (8, ~8.8), (400000, ~12)] 99 | This function is continuously differentiable. 100 | 101 | Args: 102 | alpha: A numpy array or TF tensor (float32 or float64) with values >= 0. 103 | 104 | Returns: 105 | An array/tensor of curved values >= 0 with the same type as `alpha`, to be 106 | used as input x-coordinates for spline interpolation. 107 | """ 108 | c = lambda z: tf.cast(z, alpha.dtype) 109 | assert_ops = [tf.Assert(tf.reduce_all(alpha >= 0.), [alpha])] 110 | with tf.control_dependencies(assert_ops): 111 | x = tf.where(alpha < 4, (c(2.25) * alpha - c(4.5)) / 112 | (tf.abs(alpha - c(2)) + c(0.25)) + alpha + c(2), 113 | c(5) / c(18) * util.log_safe(c(4) * alpha - c(15)) + c(8)) 114 | return x 115 | 116 | 117 | def inv_partition_spline_curve(x): 118 | """The inverse of partition_spline_curve().""" 119 | c = lambda z: tf.cast(z, x.dtype) 120 | assert_ops = [tf.Assert(tf.reduce_all(x >= 0.), [x])] 121 | with tf.control_dependencies(assert_ops): 122 | alpha = tf.where( 123 | x < 8, 124 | c(0.5) * x + tf.where( 125 | x <= 4, 126 | c(1.25) - tf.sqrt(c(1.5625) - x + c(.25) * tf.square(x)), 127 | c(-1.25) + tf.sqrt(c(9.5625) - c(3) * x + c(.25) * tf.square(x))), 128 | c(3.75) + c(0.25) * util.exp_safe(x * c(3.6) - c(28.8))) 129 | return alpha 130 | 131 | 132 | class Distribution(object): 133 | """A wrapper class around the distribution.""" 134 | 135 | def __init__(self): 136 | 137 | 138 | """Initialize the distribution. 139 | 140 | Load the values, tangents, and x-coordinate scaling of a spline that 141 | approximates the partition function. The spline was produced by running 142 | the script in fit_partition_spline.py. 143 | """ 144 | 145 | #with util.get_resource_as_file( 146 | #'robust_loss/data/partition_spline.npz') as spline_file: 147 | with np.load('partition_spline.npz', allow_pickle=False) as f: 148 | self._spline_x_scale = f['x_scale'] 149 | self._spline_values = f['values'] 150 | self._spline_tangents = f['tangents'] 151 | 152 | def log_base_partition_function(self, alpha): 153 | r"""Approximate the distribution's log-partition function with a 1D spline. 154 | 155 | Because the partition function (Z(\alpha) in the paper) of the distribution 156 | is difficult to model analytically, we approximate it with a (transformed) 157 | cubic hermite spline: Each alpha is pushed through a nonlinearity before 158 | being used to interpolate into a spline, which allows us to use a relatively 159 | small spline to accurately model the log partition function over the range 160 | of all non-negative input values. 161 | 162 | Args: 163 | alpha: A tensor or scalar of single or double precision floats containing 164 | the set of alphas for which we would like an approximate log partition 165 | function. Must be non-negative, as the partition function is undefined 166 | when alpha < 0. 167 | 168 | Returns: 169 | An approximation of log(Z(alpha)) accurate to within 1e-6 170 | """ 171 | float_dtype = alpha.dtype 172 | 173 | # The partition function is undefined when `alpha`< 0. 174 | assert_ops = [tf.Assert(tf.reduce_all(alpha >= 0.), [alpha])] 175 | with tf.control_dependencies(assert_ops): 176 | # Transform `alpha` to the form expected by the spline. 177 | x = partition_spline_curve(alpha) 178 | # Interpolate into the spline. 179 | return cubic_spline.interpolate1d( 180 | x * tf.cast(self._spline_x_scale, float_dtype), 181 | tf.cast(self._spline_values, float_dtype), 182 | tf.cast(self._spline_tangents, float_dtype)) 183 | 184 | def nllfun(self, x, alpha, scale): 185 | r"""Implements the negative log-likelihood (NLL). 186 | 187 | Specifically, we implement -log(p(x | 0, \alpha, c) of Equation 16 in the 188 | paper as nllfun(x, alpha, shape). 189 | 190 | Args: 191 | x: The residual for which the NLL is being computed. x can have any shape, 192 | and alpha and scale will be broadcasted to match x's shape if necessary. 193 | Must be a tensorflow tensor or numpy array of floats. 194 | alpha: The shape parameter of the NLL (\alpha in the paper), where more 195 | negative values cause outliers to "cost" more and inliers to "cost" 196 | less. Alpha can be any non-negative value, but the gradient of the NLL 197 | with respect to alpha has singularities at 0 and 2 so you may want to 198 | limit usage to (0, 2) during gradient descent. Must be a tensorflow 199 | tensor or numpy array of floats. Varying alpha in that range allows for 200 | smooth interpolation between a Cauchy distribution (alpha = 0) and a 201 | Normal distribution (alpha = 2) similar to a Student's T distribution. 202 | scale: The scale parameter of the loss. When |x| < scale, the NLL is like 203 | that of a (possibly unnormalized) normal distribution, and when |x| > 204 | scale the NLL takes on a different shape according to alpha. Must be a 205 | tensorflow tensor or numpy array of floats. 206 | 207 | Returns: 208 | The NLLs for each element of x, in the same shape as x. This is returned 209 | as a TensorFlow graph node of floats with the same precision as x. 210 | """ 211 | # `scale` and `alpha` must have the same type as `x`. 212 | tf.debugging.assert_type(scale, x.dtype) 213 | tf.debugging.assert_type(alpha, x.dtype) 214 | assert_ops = [ 215 | # `scale` must be > 0. 216 | tf.Assert(tf.reduce_all(scale > 0.), [scale]), 217 | # `alpha` must be >= 0. 218 | tf.Assert(tf.reduce_all(alpha >= 0.), [alpha]), 219 | ] 220 | with tf.control_dependencies(assert_ops): 221 | loss = general.lossfun(x, alpha, scale, approximate=False) 222 | log_partition = ( 223 | tf.math.log(scale) + self.log_base_partition_function(alpha)) 224 | nll = loss + log_partition 225 | return nll 226 | 227 | def draw_samples(self, alpha, scale): 228 | r"""Draw samples from the robust distribution. 229 | 230 | This function implements Algorithm 1 the paper. This code is written to 231 | allow for sampling from a set of different distributions, each parametrized 232 | by its own alpha and scale values, as opposed to the more standard approach 233 | of drawing N samples from the same distribution. This is done by repeatedly 234 | performing N instances of rejection sampling for each of the N distributions 235 | until at least one proposal for each of the N distributions has been 236 | accepted. All samples assume a zero mean --- to get non-zero mean samples, 237 | just add each mean to each sample. 238 | 239 | Args: 240 | alpha: A TF tensor/scalar or numpy array/scalar of floats where each 241 | element is the shape parameter of that element's distribution. 242 | scale: A TF tensor/scalar or numpy array/scalar of floats where each 243 | element is the scale parameter of that element's distribution. Must be 244 | the same shape as `alpha`. 245 | 246 | Returns: 247 | A TF tensor with the same shape and precision as `alpha` and `scale` where 248 | each element is a sample drawn from the zero-mean distribution specified 249 | for that element by `alpha` and `scale`. 250 | """ 251 | # `scale` must have the same type as `alpha`. 252 | float_dtype = alpha.dtype 253 | tf.debugging.assert_type(scale, float_dtype) 254 | assert_ops = [ 255 | # `scale` must be > 0. 256 | tf.Assert(tf.reduce_all(scale > 0.), [scale]), 257 | # `alpha` must be >= 0. 258 | tf.Assert(tf.reduce_all(alpha >= 0.), [alpha]), 259 | # `alpha` and `scale` must have the same shape. 260 | tf.Assert( 261 | tf.reduce_all(tf.equal(tf.shape(alpha), tf.shape(scale))), 262 | [tf.shape(alpha), tf.shape(scale)]), 263 | ] 264 | 265 | with tf.control_dependencies(assert_ops): 266 | shape = tf.shape(alpha) 267 | 268 | # The distributions we will need for rejection sampling. The sqrt(2) 269 | # scaling of the Cauchy distribution corrects for our differing 270 | # conventions for standardization. 271 | cauchy = tfp.distributions.Cauchy(loc=0., scale=tf.sqrt(2.)) 272 | uniform = tfp.distributions.Uniform(low=0., high=1.) 273 | 274 | def while_cond(_, accepted): 275 | """Terminate the loop only when all samples have been accepted.""" 276 | return ~tf.reduce_all(accepted) 277 | 278 | def while_body(samples, accepted): 279 | """Generate N proposal samples, and then perform rejection sampling.""" 280 | # Draw N samples from a Cauchy, our proposal distribution. 281 | cauchy_sample = tf.cast(cauchy.sample(shape), float_dtype) 282 | 283 | # Compute the likelihood of each sample under its target distribution. 284 | nll = self.nllfun(cauchy_sample, alpha, tf.cast(1, float_dtype)) 285 | # Bound the NLL. We don't use the approximate loss as it may cause 286 | # unpredictable behavior in the context of sampling. 287 | nll_bound = general.lossfun( 288 | cauchy_sample, 289 | tf.cast(0, float_dtype), 290 | tf.cast(1, float_dtype), 291 | approximate=False) + self.log_base_partition_function(alpha) 292 | 293 | # Draw N samples from a uniform distribution, and use each uniform 294 | # sample to decide whether or not to accept each proposal sample. 295 | uniform_sample = tf.cast(uniform.sample(shape), float_dtype) 296 | accept = uniform_sample <= tf.math.exp(nll_bound - nll) 297 | 298 | # If a sample is accepted, replace its element in `samples` with the 299 | # proposal sample, and set its bit in `accepted` to True. 300 | samples = tf.where(accept, cauchy_sample, samples) 301 | accepted = accept | accepted 302 | return (samples, accepted) 303 | 304 | # Initialize the loop. The first item does not matter as it will get 305 | # overwritten, the second item must be all False. 306 | while_loop_vars = (tf.zeros(shape, 307 | float_dtype), tf.zeros(shape, dtype=bool)) 308 | 309 | # Perform rejection sampling until all N samples have been accepted. 310 | terminal_state = tf.while_loop( 311 | cond=while_cond, body=while_body, loop_vars=while_loop_vars) 312 | 313 | # Because our distribution is a location-scale family, we sample from 314 | # p(x | 0, \alpha, 1) and then scale each sample by `scale`. 315 | samples = tf.multiply(terminal_state[0], scale) 316 | 317 | return samples 318 | -------------------------------------------------------------------------------- /src/robust_loss/distribution_test.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | # Copyright 2020 The Google Research Authors. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | """Tests for distribution.py.""" 17 | 18 | from absl.testing import parameterized 19 | import numpy as np 20 | import scipy.stats 21 | import tensorflow.compat.v2 as tf 22 | from robust_loss import distribution 23 | 24 | tf.enable_v2_behavior() 25 | 26 | 27 | class DistributionTest(parameterized.TestCase, tf.test.TestCase): 28 | 29 | def setUp(self): 30 | self._distribution = distribution.Distribution() 31 | super(DistributionTest, self).setUp() 32 | np.random.seed(0) 33 | 34 | def testSplineCurveIsC1Smooth(self): 35 | """Tests that partition_spline_curve() and its derivative are continuous.""" 36 | x1 = np.linspace(0., 8., 10000, dtype=np.float64) 37 | x2 = x1 + 1e-7 38 | 39 | x1 = tf.convert_to_tensor(x1) 40 | x2 = tf.convert_to_tensor(x2) 41 | 42 | with tf.GradientTape(persistent=True) as tape: 43 | tape.watch(x1) 44 | tape.watch(x2) 45 | y1 = distribution.partition_spline_curve(x1) 46 | y2 = distribution.partition_spline_curve(x2) 47 | dy1 = tape.gradient(tf.reduce_sum(y1), x1) 48 | dy2 = tape.gradient(tf.reduce_sum(y2), x2) 49 | self.assertAllClose(y1, y2) 50 | self.assertAllClose(dy1, dy2) 51 | 52 | def testAnalyaticalPartitionIsCorrect(self): 53 | """Tests _analytical_base_partition_function against some golden data.""" 54 | # Here we enumerate a set of positive rational numbers n/d alongside 55 | # numerically approximated values of Z(n / d) up to 10 digits of precision, 56 | # stored as (n, d, Z(n/d)). This was generated with an external mathematica 57 | # script. 58 | ground_truth_rational_partitions = ( 59 | (1, 7, 4.080330073), (1, 6, 4.038544331), (1, 5, 3.984791180), 60 | (1, 4, 3.912448576), (1, 3, 3.808203509), (2, 5, 3.735479786), 61 | (3, 7, 3.706553276), (1, 2, 3.638993131), (3, 5, 3.553489270), 62 | (2, 3, 3.501024540), (3, 4, 3.439385624), (4, 5, 3.404121259), 63 | (1, 1, 3.272306973), (6, 5, 3.149249092), (5, 4, 3.119044506), 64 | (4, 3, 3.068687433), (7, 5, 3.028084866), (3, 2, 2.965924889), 65 | (8, 5, 2.901059987), (5, 3, 2.855391798), (7, 4, 2.794052016), 66 | (7, 3, 2.260434598), (5, 2, 2.218882601), (8, 3, 2.190349858), 67 | (3, 1, 2.153202857), (4, 1, 2.101960916), (7, 2, 2.121140098), 68 | (5, 1, 2.080000512), (9, 2, 2.089161164), (6, 1, 2.067751267), 69 | (7, 1, 2.059929623), (8, 1, 2.054500222), (10, 3, 2.129863884), 70 | (11, 3, 2.113763384), (13, 3, 2.092928254), (14, 3, 2.085788350), 71 | (16, 3, 2.075212740), (11, 2, 2.073116001), (17, 3, 2.071185791), 72 | (13, 2, 2.063452243), (15, 2, 2.056990258)) # pyformat: disable 73 | for numer, denom, z_true in ground_truth_rational_partitions: 74 | z = distribution.analytical_base_partition_function(numer, denom) 75 | self.assertAllClose(z, z_true, atol=1e-9, rtol=1e-9) 76 | 77 | def testSplineCurveInverseIsCorrect(self): 78 | """Tests that the inverse curve is indeed the inverse of the curve.""" 79 | x_knot = np.arange(0, 16, 0.01, dtype=np.float64) 80 | alpha = distribution.inv_partition_spline_curve(x_knot) 81 | x_recon = distribution.partition_spline_curve(alpha) 82 | self.assertAllClose(x_recon, x_knot) 83 | 84 | @parameterized.named_parameters(('Single', np.float32), 85 | ('Double', np.float64)) 86 | def testLogPartitionInfinityIsAccurate(self, float_dtype): 87 | """Tests that the partition function is accurate at infinity.""" 88 | alpha = float_dtype(float('inf')) 89 | log_z_true = np.float64(0.70526025442) # From mathematica. 90 | log_z = self._distribution.log_base_partition_function(alpha) 91 | self.assertAllClose(log_z, log_z_true, atol=1e-7, rtol=1e-7) 92 | 93 | @parameterized.named_parameters(('Single', np.float32), 94 | ('Double', np.float64)) 95 | def testLogPartitionFractionsAreAccurate(self, float_dtype): 96 | """Test that the partition function is correct for [0/11, ... 22/11].""" 97 | numers = range(0, 23) 98 | denom = 11 99 | log_zs_true = [ 100 | np.log(distribution.analytical_base_partition_function(n, denom)) 101 | for n in numers 102 | ] 103 | log_zs = self._distribution.log_base_partition_function( 104 | float_dtype(np.array(numers)) / float_dtype(denom)) 105 | self.assertAllClose(log_zs, log_zs_true, atol=1e-7, rtol=1e-7) 106 | 107 | @parameterized.named_parameters(('Single', np.float32), 108 | ('Double', np.float64)) 109 | def testAlphaZeroSamplesMatchACauchyDistribution(self, float_dtype): 110 | """Tests that samples when alpha=0 match a Cauchy distribution.""" 111 | num_samples = 16384 112 | scale = float_dtype(1.7) 113 | samples = self._distribution.draw_samples( 114 | np.zeros(num_samples, dtype=float_dtype), 115 | scale * np.ones(num_samples, dtype=float_dtype)) 116 | # Perform the Kolmogorov-Smirnov test against a Cauchy distribution. 117 | ks_statistic = scipy.stats.kstest(samples, 'cauchy', 118 | (0., scale * np.sqrt(2.))).statistic 119 | self.assertLess(ks_statistic, 0.02) 120 | 121 | @parameterized.named_parameters(('Single', np.float32), 122 | ('Double', np.float64)) 123 | def testAlphaTwoSamplesMatchANormalDistribution(self, float_dtype): 124 | """Tests that samples when alpha=2 match a normal distribution.""" 125 | num_samples = 16384 126 | scale = float_dtype(1.7) 127 | samples = self._distribution.draw_samples( 128 | 2. * np.ones(num_samples, dtype=float_dtype), 129 | scale * np.ones(num_samples, dtype=float_dtype)) 130 | # Perform the Kolmogorov-Smirnov test against a normal distribution. 131 | ks_statistic = scipy.stats.kstest(samples, 'norm', (0., scale)).statistic 132 | self.assertLess(ks_statistic, 0.01) 133 | 134 | @parameterized.named_parameters(('Single', np.float32), 135 | ('Double', np.float64)) 136 | def testAlphaZeroNllsMatchACauchyDistribution(self, float_dtype): 137 | """Tests that NLLs when alpha=0 match a Cauchy distribution.""" 138 | x = np.linspace(-10., 10, 1000, dtype=float_dtype) 139 | scale = float_dtype(1.7) 140 | nll = self._distribution.nllfun(x, float_dtype(0.), scale) 141 | nll_true = -scipy.stats.cauchy(0., scale * np.sqrt(2.)).logpdf(x) 142 | self.assertAllClose(nll, nll_true) 143 | 144 | @parameterized.named_parameters(('Single', np.float32), 145 | ('Double', np.float64)) 146 | def testAlphaTwoNllsMatchANormalDistribution(self, float_dtype): 147 | """Tests that NLLs when alpha=2 match a normal distribution.""" 148 | x = np.linspace(-10., 10, 1000, dtype=float_dtype) 149 | scale = float_dtype(1.7) 150 | nll = self._distribution.nllfun(x, float_dtype(2.), scale) 151 | nll_true = -scipy.stats.norm(0., scale).logpdf(x) 152 | self.assertAllClose(nll, nll_true) 153 | 154 | @parameterized.named_parameters(('Single', np.float32), 155 | ('Double', np.float64)) 156 | def testPdfIntegratesToOne(self, float_dtype): 157 | """Tests that the PDF integrates to 1 for different alphas.""" 158 | alphas = np.exp(np.linspace(-4., 8., 8, dtype=float_dtype)) 159 | scale = float_dtype(1.7) 160 | x = np.arange(-128., 128., 1 / 256., dtype=float_dtype) * scale 161 | for alpha in alphas: 162 | nll = self._distribution.nllfun(x, alpha, scale) 163 | pdf_sum = np.sum(np.exp(-nll)) * (x[1] - x[0]) 164 | self.assertAllClose(pdf_sum, 1., atol=0.005, rtol=0.005) 165 | 166 | @parameterized.named_parameters(('Single', np.float32), 167 | ('Double', np.float64)) 168 | def testNllfunPreservesDtype(self, float_dtype): 169 | """Checks that the loss's output has the same precision as its input.""" 170 | n = 16 171 | x = float_dtype(np.random.normal(size=n)) 172 | alpha = float_dtype(np.exp(np.random.normal(size=n))) 173 | scale = float_dtype(np.exp(np.random.normal(size=n))) 174 | y = self._distribution.nllfun(x, alpha, scale) 175 | self.assertDTypeEqual(y, float_dtype) 176 | 177 | @parameterized.named_parameters(('Single', np.float32), 178 | ('Double', np.float64)) 179 | def testSamplingPreservesDtype(self, float_dtype): 180 | """Checks that sampling's output has the same precision as its input.""" 181 | n = 16 182 | alpha = float_dtype(np.exp(np.random.normal(size=n))) 183 | scale = float_dtype(np.exp(np.random.normal(size=n))) 184 | y = self._distribution.draw_samples(alpha, scale) 185 | self.assertDTypeEqual(y, float_dtype) 186 | 187 | 188 | if __name__ == '__main__': 189 | tf.test.main() 190 | -------------------------------------------------------------------------------- /src/robust_loss/fit_partition_spline.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | # Copyright 2020 The Google Research Authors. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | # Lint as: python3 17 | r"""Approximate the distribution's partition function with a spline. 18 | 19 | This script generates values for the distribution's partition function and then 20 | fits a cubic hermite spline to those values, which is then stored to disk. 21 | To run this script, assuming you're in this directory, run: 22 | python -m robust_loss.fit_partition_spline_test 23 | This script will likely never have to be run again, and is provided here for 24 | completeness and reproducibility, or in case someone decides to modify 25 | distribution.partition_spline_curve() in the future in case they find a better 26 | curve. If the user wants a more accurate spline approximation, this can be 27 | obtained by modifying the `x_max`, `x_scale`, and `redundancy` parameters in the 28 | code below, but this should only be done with care. 29 | """ 30 | 31 | from absl import app 32 | import numpy as np 33 | import tensorflow.compat.v2 as tf 34 | from robust_loss import cubic_spline 35 | from robust_loss import distribution 36 | from robust_loss import general 37 | 38 | tf.enable_v2_behavior() 39 | 40 | 41 | def numerical_base_partition_function(alpha): 42 | """Numerically approximate the partition function Z(alpha).""" 43 | # Generate values `num_samples` values in [-x_max, x_max], with more samples 44 | # near the origin as `power` is set to larger values. 45 | num_samples = 2**24 + 1 # We want an odd value so that 0 gets sampled. 46 | x_max = 10**10 47 | power = 6 48 | t = t = tf.linspace( 49 | tf.constant(-1, tf.float64), tf.constant(1, tf.float64), num_samples) 50 | t = tf.sign(t) * tf.abs(t)**power 51 | x = t * x_max 52 | 53 | # Compute losses for the values, then exponentiate the negative losses and 54 | # integrate with the trapezoid rule to get the partition function. 55 | losses = general.lossfun(x, alpha, np.float64(1)) 56 | y = tf.math.exp(-losses) 57 | partition = tf.reduce_sum((y[1:] + y[:-1]) * (x[1:] - x[:-1])) / 2. 58 | return partition 59 | 60 | 61 | def main(argv): 62 | if len(argv) > 1: 63 | raise app.UsageError('Too many command-line arguments.') 64 | 65 | # Parameters governing how the x coordinate of the spline will be laid out. 66 | # We will construct a spline with knots at 67 | # [0 : 1 / x_scale : x_max], 68 | # by fitting it to values sampled at 69 | # [0 : 1 / (x_scale * redundancy) : x_max] 70 | x_max = 12 71 | x_scale = 1024 72 | redundancy = 4 # Must be >= 2 for the spline to be useful. 73 | 74 | spline_spacing = 1. / (x_scale * redundancy) 75 | x_knots = np.arange( 76 | 0, x_max + spline_spacing, spline_spacing, dtype=np.float64) 77 | table = [] 78 | # We iterate over knots, and for each knot recover the alpha value 79 | # corresponding to that knot with inv_partition_spline_curve(), and then 80 | # with that alpha we accurately approximate its partition function using 81 | # numerical_base_partition_function(). 82 | for x_knot in x_knots: 83 | alpha = distribution.inv_partition_spline_curve(x_knot).numpy() 84 | partition = numerical_base_partition_function(alpha).numpy() 85 | table.append((x_knot, alpha, partition)) 86 | print(table[-1]) 87 | 88 | table = np.array(table) 89 | x = table[:, 0] 90 | alpha = table[:, 1] 91 | y_gt = np.log(table[:, 2]) 92 | 93 | # We grab the values from the true log-partition table that correpond to 94 | # knots, by looking for where x * x_scale is an integer. 95 | mask = np.abs(np.round(x * x_scale) - (x * x_scale)) <= 1e-8 96 | values = y_gt[mask] 97 | 98 | # Initialize `tangents` using a central differencing scheme. 99 | values_pad = np.concatenate([[values[0] - values[1] + values[0]], values, 100 | [values[-1] - values[-2] + values[-1]]], 0) 101 | tangents = (values_pad[2:] - values_pad[:-2]) / 2. 102 | 103 | # Construct the spline's value and tangent TF variables, constraining the last 104 | # knot to have a fixed value Z(infinity) and a tangent of zero. 105 | n = len(values) 106 | tangents = tf.Variable(tangents, tf.float64) 107 | values = tf.Variable(values, tf.float64) 108 | 109 | # Fit the spline. 110 | num_iters = 10001 111 | 112 | optimizer = tf.keras.optimizers.SGD(learning_rate=1e-9, momentum=0.99) 113 | 114 | trace = [] 115 | for ii in range(num_iters): 116 | with tf.GradientTape() as tape: 117 | tape.watch([values, tangents]) 118 | # Fix the endpoint to be a known constant with a zero tangent. 119 | i_values = tf.where( 120 | np.arange(n) == (n - 1), 121 | tf.ones_like(values) * 0.70526025442689566, values) 122 | i_tangents = tf.where( 123 | np.arange(n) == (n - 1), tf.zeros_like(tangents), tangents) 124 | i_y = cubic_spline.interpolate1d(x * x_scale, i_values, i_tangents) 125 | # We minimize the maximum residual, which makes for a very ugly 126 | # optimization problem but works well in practice. 127 | i_loss = tf.reduce_max(tf.abs(i_y - y_gt)) 128 | grads = tape.gradient(i_loss, [values, tangents]) 129 | optimizer.apply_gradients(zip(grads, [values, tangents])) 130 | trace.append(i_loss.numpy()) 131 | if (ii % 200) == 0: 132 | print('{:5d}: {:e}'.format(ii, trace[-1])) 133 | 134 | mask = alpha <= 4 135 | max_error_a4 = np.max(np.abs(i_y[mask] - y_gt[mask])) 136 | max_error = np.max(np.abs(i_y - y_gt)) 137 | print('Max Error (a <= 4): {:e}'.format(max_error_a4)) 138 | print('Max Error: {:e}'.format(max_error)) 139 | 140 | # Just a sanity-check on the error. 141 | assert max_error_a4 <= 5e-7 142 | assert max_error <= 5e-7 143 | 144 | # Save the spline to disk. 145 | np.savez( 146 | './data/partition_spline.npz', 147 | x_scale=x_scale, 148 | values=i_values.numpy(), 149 | tangents=i_tangents.numpy()) 150 | 151 | 152 | if __name__ == '__main__': 153 | app.run(main) 154 | -------------------------------------------------------------------------------- /src/robust_loss/fit_partition_spline_test.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | # Copyright 2020 The Google Research Authors. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | """Tests for fit_partition_spline.py.""" 17 | 18 | import tensorflow.compat.v2 as tf 19 | from robust_loss import distribution 20 | from robust_loss import fit_partition_spline 21 | 22 | tf.enable_v2_behavior() 23 | 24 | 25 | class FitPartitionSplineTest(tf.test.TestCase): 26 | 27 | def testNumericalPartitionIsAccurate(self): 28 | """Test _numerical_base_partition_function against some golden data.""" 29 | for (numer, denom) in [(0, 1), (1, 8), (1, 2), (1, 1), (2, 1), (8, 1)]: 30 | alpha = tf.cast(numer, tf.float64) / tf.cast(denom, tf.float64) 31 | z_true = distribution.analytical_base_partition_function(numer, denom) 32 | z = fit_partition_spline.numerical_base_partition_function(alpha) 33 | self.assertAllClose(z, z_true, atol=1e-10, rtol=1e-10) 34 | 35 | 36 | if __name__ == '__main__': 37 | tf.test.main() 38 | -------------------------------------------------------------------------------- /src/robust_loss/general.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | # Copyright 2020 The Google Research Authors. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | r"""Implements the general form of the loss. 17 | 18 | This is the simplest way of using this loss. No parameters will be tuned 19 | automatically, it's just a simple function that takes in parameters (likely 20 | hand-tuned ones) and return a loss. For an adaptive loss, look at adaptive.py 21 | or distribution.py. 22 | """ 23 | 24 | import numpy as np 25 | import tensorflow.compat.v2 as tf 26 | from robust_loss import util 27 | 28 | 29 | def lossfun(x, alpha, scale, approximate=False, epsilon=1e-6): 30 | r"""Implements the general form of the loss. 31 | 32 | This implements the rho(x, \alpha, c) function described in "A General and 33 | Adaptive Robust Loss Function", Jonathan T. Barron, 34 | https://arxiv.org/abs/1701.03077. 35 | 36 | Args: 37 | x: The residual for which the loss is being computed. x can have any shape, 38 | and alpha and scale will be broadcasted to match x's shape if necessary. 39 | Must be a tensorflow tensor or numpy array of floats. 40 | alpha: The shape parameter of the loss (\alpha in the paper), where more 41 | negative values produce a loss with more robust behavior (outliers "cost" 42 | less), and more positive values produce a loss with less robust behavior 43 | (outliers are penalized more heavily). Alpha can be any value in 44 | [-infinity, infinity], but the gradient of the loss with respect to alpha 45 | is 0 at -infinity, infinity, 0, and 2. Must be a tensorflow tensor or 46 | numpy array of floats with the same precision as `x`. Varying alpha allows 47 | for smooth interpolation between a number of discrete robust losses: 48 | alpha=-Infinity: Welsch/Leclerc Loss. 49 | alpha=-2: Geman-McClure loss. 50 | alpha=0: Cauchy/Lortentzian loss. 51 | alpha=1: Charbonnier/pseudo-Huber loss. 52 | alpha=2: L2 loss. 53 | scale: The scale parameter of the loss. When |x| < scale, the loss is an 54 | L2-like quadratic bowl, and when |x| > scale the loss function takes on a 55 | different shape according to alpha. Must be a tensorflow tensor or numpy 56 | array of single-precision floats. 57 | approximate: a bool, where if True, this function returns an approximate and 58 | faster form of the loss, as described in the appendix of the paper. This 59 | approximation holds well everywhere except as x and alpha approach zero. 60 | epsilon: A float that determines how inaccurate the "approximate" version of 61 | the loss will be. Larger values are less accurate but more numerically 62 | stable. Must be great than single-precision machine epsilon. 63 | 64 | Returns: 65 | The losses for each element of x, in the same shape as x. This is returned 66 | as a TensorFlow graph node of single precision floats. 67 | """ 68 | # `scale` and `alpha` must have the same type as `x`. 69 | float_dtype = x.dtype 70 | tf.debugging.assert_type(scale, float_dtype) 71 | tf.debugging.assert_type(alpha, float_dtype) 72 | # `scale` must be > 0. 73 | assert_ops = [tf.Assert(tf.reduce_all(tf.greater(scale, 0.)), [scale])] 74 | with tf.control_dependencies(assert_ops): 75 | # Broadcast `alpha` and `scale` to have the same shape as `x`. 76 | alpha = tf.broadcast_to(alpha, tf.shape(x)) 77 | scale = tf.broadcast_to(scale, tf.shape(x)) 78 | 79 | if approximate: 80 | # `epsilon` must be greater than single-precision machine epsilon. 81 | assert epsilon > np.finfo(np.float32).eps 82 | # Compute an approximate form of the loss which is faster, but innacurate 83 | # when x and alpha are near zero. 84 | b = tf.abs(alpha - tf.cast(2., float_dtype)) + epsilon 85 | d = tf.where( 86 | tf.greater_equal(alpha, 0.), alpha + epsilon, alpha - epsilon) 87 | loss = (b / d) * (tf.pow(tf.square(x / scale) / b + 1., 0.5 * d) - 1.) 88 | else: 89 | # Compute the exact loss. 90 | 91 | # This will be used repeatedly. 92 | squared_scaled_x = tf.square(x / scale) 93 | 94 | # The loss when alpha == 2. 95 | loss_two = 0.5 * squared_scaled_x 96 | # The loss when alpha == 0. 97 | loss_zero = util.log1p_safe(0.5 * squared_scaled_x) 98 | # The loss when alpha == -infinity. 99 | loss_neginf = -tf.math.expm1(-0.5 * squared_scaled_x) 100 | # The loss when alpha == +infinity. 101 | loss_posinf = util.expm1_safe(0.5 * squared_scaled_x) 102 | 103 | # The loss when not in one of the above special cases. 104 | machine_epsilon = tf.cast(np.finfo(np.float32).eps, float_dtype) 105 | # Clamp |2-alpha| to be >= machine epsilon so that it's safe to divide by. 106 | beta_safe = tf.maximum(machine_epsilon, tf.abs(alpha - 2.)) 107 | # Clamp |alpha| to be >= machine epsilon so that it's safe to divide by. 108 | alpha_safe = tf.where( 109 | tf.greater_equal(alpha, 0.), tf.ones_like(alpha), 110 | -tf.ones_like(alpha)) * tf.maximum(machine_epsilon, tf.abs(alpha)) 111 | loss_otherwise = (beta_safe / alpha_safe) * ( 112 | tf.pow(squared_scaled_x / beta_safe + 1., 0.5 * alpha) - 1.) 113 | 114 | # Select which of the cases of the loss to return. 115 | loss = tf.where( 116 | tf.equal(alpha, -tf.cast(float('inf'), float_dtype)), loss_neginf, 117 | tf.where( 118 | tf.equal(alpha, 0.), loss_zero, 119 | tf.where( 120 | tf.equal(alpha, 2.), loss_two, 121 | tf.where( 122 | tf.equal(alpha, tf.cast(float('inf'), float_dtype)), 123 | loss_posinf, loss_otherwise)))) 124 | 125 | return loss 126 | -------------------------------------------------------------------------------- /src/robust_loss/general_test.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | # Copyright 2020 The Google Research Authors. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | """Tests for general.py.""" 17 | 18 | from absl.testing import parameterized 19 | import numpy as np 20 | import tensorflow.compat.v2 as tf 21 | from robust_loss import general 22 | 23 | tf.enable_v2_behavior() 24 | 25 | 26 | class LossfunTest(parameterized.TestCase, tf.test.TestCase): 27 | 28 | def setUp(self): 29 | super(LossfunTest, self).setUp() 30 | np.random.seed(0) 31 | 32 | def _assert_all_close_according_to_type(self, a, b): 33 | """AssertAllClose() with tighter thresholds for float64 than float32.""" 34 | self.assertAllCloseAccordingToType( 35 | a, b, rtol=1e-15, atol=1e-15, float_rtol=1e-6, float_atol=1e-6) 36 | 37 | def _precompute_lossfun_inputs(self, float_dtype): 38 | """Precompute a loss and its derivatives for random inputs and parameters. 39 | 40 | Generates a large number of random inputs to the loss, and random 41 | shape/scale parameters for the loss function at each sample, and 42 | computes the loss and its derivative with respect to all inputs and 43 | parameters, returning everything to be used to assert various properties 44 | in our unit tests. 45 | 46 | Args: 47 | float_dtype: The float precision to be used (np.float32 or np.float64). 48 | 49 | Returns: 50 | A tuple containing: 51 | (the number (int) of samples, and the length of all following arrays, 52 | A np.array (float_dtype) of losses for each sample, 53 | A np.array (float_dtype) of residuals of each sample (the loss inputs), 54 | A np array (float_dtype) of shape parameters of each loss, 55 | A np.array (float_dtype) of scale parameters of each loss, 56 | A np.array (float_dtype) of derivatives of each loss wrt each x, 57 | A np.array (float_dtype) of derivatives of each loss wrt each alpha, 58 | A np.array (float_dtype) of derivatives of each loss wrt each scale) 59 | 60 | Typical usage example: 61 | (num_samples, loss, x, alpha, scale, d_x, d_alpha, d_scale) 62 | = self._precompute_lossfun_inputs(np.float32) 63 | """ 64 | num_samples = 100000 65 | # Normally distributed inputs. 66 | x = float_dtype(np.random.normal(size=num_samples)) 67 | 68 | # Uniformly distributed values in (-16, 3), quantized to the nearest 0.1 69 | # to ensure that we hit the special cases at 0, 2. 70 | alpha = float_dtype( 71 | np.round(np.random.uniform(-16, 3, num_samples) * 10) / 10.) 72 | # Push the sampled alphas at the extents of the range to +/- infinity, so 73 | # that we probe those cases too. 74 | alpha[alpha == 3.] = float_dtype(float('inf')) 75 | alpha[alpha == -16.] = -float_dtype(float('inf')) 76 | 77 | # Random log-normally distributed values in approx (1e-5, 100000): 78 | scale = float_dtype( 79 | np.exp(np.random.normal(size=num_samples) * 4.) + 1e-5) 80 | 81 | x, alpha, scale = [tf.convert_to_tensor(z) for z in (x, alpha, scale)] 82 | with tf.GradientTape(persistent=True) as tape: 83 | for z in (x, alpha, scale): 84 | tape.watch(z) 85 | loss = general.lossfun(x, alpha, scale) 86 | d_x, d_alpha, d_scale = [ 87 | tape.gradient(tf.reduce_sum(loss), z) for z in (x, alpha, scale) 88 | ] 89 | return (num_samples, loss, x, alpha, scale, d_x, d_alpha, d_scale) 90 | 91 | @parameterized.named_parameters(('Single', np.float32), 92 | ('Double', np.float64)) 93 | def testLossfunPreservesDtype(self, float_dtype): 94 | """Check the loss's output has the same precision as its input.""" 95 | n = 16 96 | x = float_dtype(np.random.normal(size=n)) 97 | alpha = float_dtype(np.random.normal(size=n)) 98 | scale = float_dtype(np.exp(np.random.normal(size=n))) 99 | y = general.lossfun(x, alpha, scale) 100 | self.assertDTypeEqual(y, float_dtype) 101 | 102 | @parameterized.named_parameters(('Single', np.float32), 103 | ('Double', np.float64)) 104 | def testDerivativeIsMonotonicWrtX(self, float_dtype): 105 | # Check that the loss increases monotonically with |x|. 106 | _, _, x, alpha, _, d_x, _, _ = self._precompute_lossfun_inputs(float_dtype) 107 | # This is just to suppress a warning below. 108 | d_x = tf.where(tf.math.is_finite(d_x), d_x, tf.zeros_like(d_x)) 109 | mask = np.isfinite(alpha) & ( 110 | np.abs(d_x) > (300. * np.finfo(float_dtype).eps)) 111 | self.assertAllEqual(np.sign(d_x[mask]), np.sign(x[mask])) 112 | 113 | @parameterized.named_parameters(('Single', np.float32), 114 | ('Double', np.float64)) 115 | def testLossIsNearZeroAtOrigin(self, float_dtype): 116 | # Check that the loss is near-zero when x is near-zero. 117 | _, loss, x, _, _, _, _, _ = self._precompute_lossfun_inputs(float_dtype) 118 | self.assertTrue(np.all(np.abs(loss[np.abs(x) < 1e-5]) < 1e-5)) 119 | 120 | @parameterized.named_parameters(('Single', np.float32), 121 | ('Double', np.float64)) 122 | def testLossIsQuadraticNearOrigin(self, float_dtype): 123 | # Check that the loss is well-approximated by a quadratic bowl when 124 | # |x| < scale 125 | _, loss, x, _, scale, _, _, _ = self._precompute_lossfun_inputs(float_dtype) 126 | mask = np.abs(x) < (0.5 * scale) 127 | loss_quad = 0.5 * np.square(x / scale) 128 | self.assertAllClose(loss_quad[mask], loss[mask], rtol=1e-5, atol=1e-2) 129 | 130 | @parameterized.named_parameters(('Single', np.float32), 131 | ('Double', np.float64)) 132 | def testLossIsBoundedWhenAlphaIsNegative(self, float_dtype): 133 | # Assert that loss < (alpha - 2)/alpha when alpha < 0. 134 | _, loss, _, alpha, _, _, _, _ = self._precompute_lossfun_inputs(float_dtype) 135 | mask = alpha < 0. 136 | min_val = np.finfo(float_dtype).min 137 | alpha_clipped = np.maximum(min_val, alpha[mask]) 138 | self.assertTrue( 139 | np.all(loss[mask] <= ((alpha_clipped - 2.) / alpha_clipped))) 140 | 141 | @parameterized.named_parameters(('Single', np.float32), 142 | ('Double', np.float64)) 143 | def testDerivativeIsBoundedWhenAlphaIsBelow2(self, float_dtype): 144 | # Assert that |d_x| < |x|/scale^2 when alpha <= 2. 145 | _, _, x, alpha, scale, d_x, _, _ = self._precompute_lossfun_inputs( 146 | float_dtype) 147 | mask = np.isfinite(alpha) & (alpha <= 2) 148 | self.assertTrue( 149 | np.all((np.abs(d_x[mask]) <= 150 | ((np.abs(x[mask]) + 151 | (300. * np.finfo(float_dtype).eps)) / scale[mask]**2)))) 152 | 153 | @parameterized.named_parameters(('Single', np.float32), 154 | ('Double', np.float64)) 155 | def testDerivativeIsBoundedWhenAlphaIsBelow1(self, float_dtype): 156 | # Assert that |d_x| < 1/scale when alpha <= 1. 157 | _, _, _, alpha, scale, d_x, _, _ = self._precompute_lossfun_inputs( 158 | float_dtype) 159 | mask = np.isfinite(alpha) & (alpha <= 1) 160 | self.assertTrue( 161 | np.all((np.abs(d_x[mask]) <= 162 | ((1. + (300. * np.finfo(float_dtype).eps)) / scale[mask])))) 163 | 164 | @parameterized.named_parameters(('Single', np.float32), 165 | ('Double', np.float64)) 166 | def testAlphaDerivativeIsPositive(self, float_dtype): 167 | # Assert that d_loss / d_alpha > 0. 168 | _, _, _, alpha, _, _, d_alpha, _ = self._precompute_lossfun_inputs( 169 | float_dtype) 170 | mask = np.isfinite(alpha) 171 | self.assertTrue(np.all(d_alpha[mask] > (-300. * np.finfo(float_dtype).eps))) 172 | 173 | @parameterized.named_parameters(('Single', np.float32), 174 | ('Double', np.float64)) 175 | def testScaleDerivativeIsNegative(self, float_dtype): 176 | # Assert that d_loss / d_scale < 0. 177 | _, _, _, alpha, _, _, _, d_scale = self._precompute_lossfun_inputs( 178 | float_dtype) 179 | mask = np.isfinite(alpha) 180 | self.assertTrue(np.all(d_scale[mask] < (300. * np.finfo(float_dtype).eps))) 181 | 182 | @parameterized.named_parameters(('Single', np.float32), 183 | ('Double', np.float64)) 184 | def testLossIsScaleInvariant(self, float_dtype): 185 | # Check that loss(mult * x, alpha, mult * scale) == loss(x, alpha, scale) 186 | (num_samples, loss, x, alpha, scale, _, _, _) = ( 187 | self._precompute_lossfun_inputs(float_dtype)) 188 | # Random log-normally distributed scalings in ~(0.2, 20) 189 | mult = float_dtype( 190 | np.maximum(0.2, np.exp(np.random.normal(size=num_samples)))) 191 | # Compute the scaled loss. 192 | loss_scaled = general.lossfun(mult * x, alpha, mult * scale) 193 | self.assertAllClose(loss, loss_scaled, atol=1e-4, rtol=1e-4) 194 | 195 | @parameterized.named_parameters(('Single', np.float32), 196 | ('Double', np.float64)) 197 | def testAlphaEqualsNegativeInfinity(self, float_dtype): 198 | # Check that alpha == -Infinity reproduces Welsch aka Leclerc loss. 199 | x = np.arange(-20, 20, 0.1, float_dtype) 200 | alpha = float_dtype(-float('inf')) 201 | scale = float_dtype(1.7) 202 | 203 | # Our loss. 204 | loss = general.lossfun(x, alpha, scale) 205 | 206 | # Welsch/Leclerc loss. 207 | loss_true = (1. - tf.math.exp(-0.5 * tf.square(x / scale))) 208 | 209 | self._assert_all_close_according_to_type(loss, loss_true) 210 | 211 | @parameterized.named_parameters(('Single', np.float32), 212 | ('Double', np.float64)) 213 | def testAlphaEqualsNegativeTwo(self, float_dtype): 214 | # Check that alpha == -2 reproduces Geman-McClure loss. 215 | x = np.arange(-20, 20, 0.1, float_dtype) 216 | alpha = float_dtype(-2.) 217 | scale = float_dtype(1.7) 218 | 219 | # Our loss. 220 | loss = general.lossfun(x, alpha, scale) 221 | 222 | # Geman-McClure loss. 223 | loss_true = (2. * tf.square(x / scale) / (tf.square(x / scale) + 4.)) 224 | 225 | self._assert_all_close_according_to_type(loss, loss_true) 226 | 227 | @parameterized.named_parameters(('Single', np.float32), 228 | ('Double', np.float64)) 229 | def testAlphaEqualsZero(self, float_dtype): 230 | # Check that alpha == 0 reproduces Cauchy aka Lorentzian loss. 231 | x = np.arange(-20, 20, 0.1, float_dtype) 232 | alpha = float_dtype(0.) 233 | scale = float_dtype(1.7) 234 | 235 | # Our loss. 236 | loss = general.lossfun(x, alpha, scale) 237 | 238 | # Cauchy/Lorentzian loss. 239 | loss_true = (tf.math.log(0.5 * tf.square(x / scale) + 1.)) 240 | 241 | self._assert_all_close_according_to_type(loss, loss_true) 242 | 243 | @parameterized.named_parameters(('Single', np.float32), 244 | ('Double', np.float64)) 245 | def testAlphaEqualsOne(self, float_dtype): 246 | # Check that alpha == 1 reproduces Charbonnier aka pseudo-Huber loss. 247 | x = np.arange(-20, 20, 0.1, float_dtype) 248 | alpha = float_dtype(1.) 249 | scale = float_dtype(1.7) 250 | 251 | # Our loss. 252 | loss = general.lossfun(x, alpha, scale) 253 | 254 | # Charbonnier loss. 255 | loss_true = (tf.sqrt(tf.square(x / scale) + 1.) - 1.) 256 | 257 | self._assert_all_close_according_to_type(loss, loss_true) 258 | 259 | @parameterized.named_parameters(('Single', np.float32), 260 | ('Double', np.float64)) 261 | def testAlphaEqualsTwo(self, float_dtype): 262 | # Check that alpha == 2 reproduces L2 loss. 263 | x = np.arange(-20, 20, 0.1, float_dtype) 264 | alpha = float_dtype(2.) 265 | scale = float_dtype(1.7) 266 | 267 | # Our loss. 268 | loss = general.lossfun(x, alpha, scale) 269 | 270 | # L2 Loss. 271 | loss_true = (0.5 * tf.square(x / scale)) 272 | 273 | self._assert_all_close_according_to_type(loss, loss_true) 274 | 275 | @parameterized.named_parameters(('Single', np.float32), 276 | ('Double', np.float64)) 277 | def testAlphaEqualsFour(self, float_dtype): 278 | # Check that alpha == 4 reproduces a quartic. 279 | x = np.arange(-20, 20, 0.1, float_dtype) 280 | alpha = float_dtype(4.) 281 | scale = float_dtype(1.7) 282 | 283 | # Our loss. 284 | loss = general.lossfun(x, alpha, scale) 285 | 286 | # The true loss. 287 | loss_true = ( 288 | tf.square(tf.square(x / scale)) / 8. + tf.square(x / scale) / 2.) 289 | 290 | self._assert_all_close_according_to_type(loss, loss_true) 291 | 292 | @parameterized.named_parameters(('Single', np.float32), 293 | ('Double', np.float64)) 294 | def testAlphaEqualsInfinity(self, float_dtype): 295 | # Check that alpha == Infinity takes the correct form. 296 | x = np.arange(-20, 20, 0.1, float_dtype) 297 | alpha = float_dtype(float('inf')) 298 | scale = float_dtype(1.7) 299 | 300 | # Our loss. 301 | loss = general.lossfun(x, alpha, scale) 302 | 303 | # The true loss. 304 | loss_true = (tf.math.exp(0.5 * tf.square(x / scale)) - 1.) 305 | 306 | self._assert_all_close_according_to_type(loss, loss_true) 307 | 308 | @parameterized.named_parameters(('Single', np.float32), 309 | ('Double', np.float64)) 310 | def testApproximateLossIsAccurate(self, float_dtype): 311 | # Check that the approximate loss (lossfun() with epsilon=1e-6) reasonably 312 | # approximates the true loss (lossfun() with epsilon=0.) for a range of 313 | # values of alpha (skipping alpha=0, where the approximation is poor). 314 | x = np.arange(-10, 10, 0.1, float_dtype) 315 | scale = float_dtype(1.7) 316 | for alpha in [-4, -2, -0.2, -0.01, 0.01, 0.2, 1, 1.99, 2, 2.01, 4]: 317 | alpha = float_dtype(alpha) 318 | loss = general.lossfun(x, alpha, scale) 319 | loss_approx = general.lossfun(x, alpha, scale, approximate=True) 320 | self.assertAllClose( 321 | loss, loss_approx, rtol=1e-5, atol=1e-4, msg='alpha=%g' % (alpha)) 322 | 323 | @parameterized.named_parameters(('Single', np.float32), 324 | ('Double', np.float64)) 325 | def testLossAndGradientsAreFinite(self, float_dtype): 326 | # Test that the loss and its approximation both give finite losses and 327 | # derivatives everywhere that they should for a wide range of values. 328 | for approximate in [False, True]: 329 | num_samples = 100000 330 | 331 | # Normally distributed inputs. 332 | x = float_dtype(np.random.normal(size=num_samples)) 333 | 334 | # Uniformly distributed values in (-16, 3), quantized to the nearest 335 | # 0.1 to ensure that we hit the special cases at 0, 2. 336 | alpha = float_dtype( 337 | np.round(np.random.uniform(-16, 3, num_samples) * 10) / 10.) 338 | 339 | # Random log-normally distributed values in approx (1e-5, 100000): 340 | scale = float_dtype( 341 | np.exp(np.random.normal(size=num_samples) * 4.) + 1e-5) 342 | 343 | # Compute the loss and its derivative with respect to all three inputs. 344 | x, alpha, scale = [tf.convert_to_tensor(z) for z in (x, alpha, scale)] 345 | with tf.GradientTape(persistent=True) as tape: 346 | for z in (x, alpha, scale): 347 | tape.watch(z) 348 | loss = general.lossfun(x, alpha, scale, approximate=approximate) 349 | d_x, d_alpha, d_scale = [ 350 | tape.gradient(tf.reduce_sum(loss), z) for z in (x, alpha, scale) 351 | ] 352 | 353 | for v in [loss, d_x, d_alpha, d_scale]: 354 | self.assertTrue(np.all(np.isfinite(v))) 355 | 356 | @parameterized.named_parameters(('Single', np.float32), 357 | ('Double', np.float64)) 358 | def testGradientMatchesFiniteDifferences(self, float_dtype): 359 | # Test that the loss and its approximation both return gradients that are 360 | # close to the numerical gradient from finite differences, with forward 361 | # differencing. Returning correct gradients is TensorFlow's job, so this is 362 | # just an aggressive sanity check in case some implementation detail causes 363 | # gradients to incorrectly go to zero due to quantization or stop_gradients 364 | # in some op that is used by the loss. 365 | for approximate in [False, True]: 366 | num_samples = 100000 367 | 368 | # Normally distributed inputs. 369 | x = float_dtype(np.random.normal(size=num_samples)) 370 | 371 | # Uniformly distributed values in (-16, 3), quantized to the nearest 372 | # 0.1 and then shifted by 0.05 so that we avoid the special cases at 373 | # 0 and 2 where the analytical gradient wont match finite differences. 374 | alpha = float_dtype( 375 | np.round(np.random.uniform(-16, 3, num_samples) * 10) / 10.) 376 | 377 | # Random uniformy distributed values in [0.5, 1.5] 378 | scale = float_dtype(np.random.uniform(0.5, 1.5, num_samples)) 379 | 380 | # Compute the loss and its derivative with respect to all three inputs. 381 | x, alpha, scale = [tf.convert_to_tensor(z) for z in (x, alpha, scale)] 382 | with tf.GradientTape(persistent=True) as tape: 383 | for z in (x, alpha, scale): 384 | tape.watch(z) 385 | loss = general.lossfun(x, alpha, scale, approximate=approximate) 386 | d_x, d_alpha, d_scale = [ 387 | tape.gradient(tf.reduce_sum(loss), z) for z in (x, alpha, scale) 388 | ] 389 | 390 | # Assert that the 95th percentile of errors is <= 1e-2. 391 | def assert_percentile_close(v1, v2): 392 | self.assertLessEqual(np.percentile(np.abs(v1 - v2), 95), 1e-2) 393 | 394 | step_size = float_dtype(1e-3) 395 | n_x = (general.lossfun(x + step_size, alpha, scale) - loss) / step_size 396 | n_alpha = (general.lossfun(x, alpha + step_size, scale) - 397 | loss) / step_size 398 | n_scale = (general.lossfun(x, alpha, scale + step_size) - 399 | loss) / step_size 400 | assert_percentile_close(n_x, d_x) 401 | assert_percentile_close(n_alpha, d_alpha) 402 | assert_percentile_close(n_scale, d_scale) 403 | 404 | 405 | if __name__ == '__main__': 406 | tf.test.main() 407 | -------------------------------------------------------------------------------- /src/robust_loss/partition_spline.npz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/src/robust_loss/partition_spline.npz -------------------------------------------------------------------------------- /src/robust_loss/requirements.txt: -------------------------------------------------------------------------------- 1 | tensorflow=1.15.0 2 | tensorflow_probability=0.8.0 3 | numpy=1.15.4 4 | scipy=1.1.0 5 | absl-py=0.1.9 6 | mpmath=1.1.0 7 | -------------------------------------------------------------------------------- /src/robust_loss/run.sh: -------------------------------------------------------------------------------- 1 | # Copyright 2020 The Google Research Authors. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | #!/bin/bash 16 | set -e 17 | set -x 18 | 19 | virtualenv -p python3 . 20 | source ./bin/activate 21 | 22 | pip install tensorflow 23 | pip install tensorflow-probability 24 | pip install -r robust_loss/requirements.txt 25 | pip install Pillow 26 | python -m robust_loss.adaptive_test 27 | python -m robust_loss.cubic_spline_test 28 | python -m robust_loss.distribution_test 29 | python -m robust_loss.fit_partition_spline_test 30 | python -m robust_loss.general_test 31 | python -m robust_loss.util_test 32 | python -m robust_loss.wavelet_test 33 | -------------------------------------------------------------------------------- /src/robust_loss/util.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | # Copyright 2020 The Google Research Authors. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | """Helper functions.""" 17 | 18 | import numpy as np 19 | import tensorflow.compat.v2 as tf 20 | 21 | 22 | 23 | def log_safe(x): 24 | """The same as tf.math.log(x), but clamps the input to prevent NaNs.""" 25 | return tf.math.log(tf.minimum(x, tf.cast(3e37, x.dtype))) 26 | 27 | 28 | def log1p_safe(x): 29 | """The same as tf.math.log1p(x), but clamps the input to prevent NaNs.""" 30 | return tf.math.log1p(tf.minimum(x, tf.cast(3e37, x.dtype))) 31 | 32 | 33 | def exp_safe(x): 34 | """The same as tf.math.exp(x), but clamps the input to prevent NaNs.""" 35 | return tf.math.exp(tf.minimum(x, tf.cast(87.5, x.dtype))) 36 | 37 | 38 | def expm1_safe(x): 39 | """The same as tf.math.expm1(x), but clamps the input to prevent NaNs.""" 40 | return tf.math.expm1(tf.minimum(x, tf.cast(87.5, x.dtype))) 41 | 42 | 43 | def inv_softplus(y): 44 | """The inverse of tf.nn.softplus().""" 45 | return tf.where(y > 87.5, y, tf.math.log(tf.math.expm1(y))) 46 | 47 | 48 | def logit(y): 49 | """The inverse of tf.nn.sigmoid().""" 50 | return -tf.math.log(1. / y - 1.) 51 | 52 | 53 | def affine_sigmoid(real, lo=0, hi=1): 54 | """Maps reals to (lo, hi), where 0 maps to (lo+hi)/2.""" 55 | if not lo < hi: 56 | raise ValueError('`lo` (%g) must be < `hi` (%g)' % (lo, hi)) 57 | alpha = tf.sigmoid(real) * (hi - lo) + lo 58 | return alpha 59 | 60 | 61 | def inv_affine_sigmoid(alpha, lo=0, hi=1): 62 | """The inverse of affine_sigmoid(., lo, hi).""" 63 | if not lo < hi: 64 | raise ValueError('`lo` (%g) must be < `hi` (%g)' % (lo, hi)) 65 | real = logit((alpha - lo) / (hi - lo)) 66 | return real 67 | 68 | 69 | def affine_softplus(real, lo=0, ref=1): 70 | """Maps real numbers to (lo, infinity), where 0 maps to ref.""" 71 | if not lo < ref: 72 | raise ValueError('`lo` (%g) must be < `ref` (%g)' % (lo, ref)) 73 | shift = inv_softplus(tf.cast(1., real.dtype)) 74 | scale = (ref - lo) * tf.nn.softplus(real + shift) + lo 75 | return scale 76 | 77 | 78 | def inv_affine_softplus(scale, lo=0, ref=1): 79 | """The inverse of affine_softplus(., lo, ref).""" 80 | if not lo < ref: 81 | raise ValueError('`lo` (%g) must be < `ref` (%g)' % (lo, ref)) 82 | shift = inv_softplus(tf.cast(1., scale.dtype)) 83 | real = inv_softplus((scale - lo) / (ref - lo)) - shift 84 | return real 85 | 86 | 87 | def students_t_nll(x, df, scale): 88 | """The NLL of a Generalized Student's T distribution (w/o including TFP).""" 89 | return 0.5 * ((df + 1.) * tf.math.log1p( 90 | (x / scale)**2. / df) + tf.math.log(df)) + tf.math.log( 91 | tf.abs(scale)) + tf.math.lgamma( 92 | 0.5 * df) - tf.math.lgamma(0.5 * df + 0.5) + 0.5 * np.log(np.pi) 93 | 94 | 95 | # A constant scale that makes tf.image.rgb_to_yuv() volume preserving. 96 | _VOLUME_PRESERVING_YUV_SCALE = 1.580227820074 97 | 98 | 99 | def rgb_to_syuv(rgb): 100 | """A volume preserving version of tf.image.rgb_to_yuv(). 101 | 102 | By "volume preserving" we mean that rgb_to_syuv() is in the "special linear 103 | group", or equivalently, that the Jacobian determinant of the transformation 104 | is 1. 105 | 106 | Args: 107 | rgb: A tensor whose last dimension corresponds to RGB channels and is of 108 | size 3. 109 | 110 | Returns: 111 | A scaled YUV version of the input tensor, such that this transformation is 112 | volume-preserving. 113 | """ 114 | return _VOLUME_PRESERVING_YUV_SCALE * tf.image.rgb_to_yuv(rgb) 115 | 116 | 117 | def syuv_to_rgb(yuv): 118 | """A volume preserving version of tf.image.yuv_to_rgb(). 119 | 120 | By "volume preserving" we mean that rgb_to_syuv() is in the "special linear 121 | group", or equivalently, that the Jacobian determinant of the transformation 122 | is 1. 123 | 124 | Args: 125 | yuv: A tensor whose last dimension corresponds to scaled YUV channels and is 126 | of size 3 (ie, the output of rgb_to_syuv()). 127 | 128 | Returns: 129 | An RGB version of the input tensor, such that this transformation is 130 | volume-preserving. 131 | """ 132 | return tf.image.yuv_to_rgb(yuv / _VOLUME_PRESERVING_YUV_SCALE) 133 | 134 | 135 | def image_dct(image): 136 | """Does a type-II DCT (aka "The DCT") on axes 1 and 2 of a rank-3 tensor.""" 137 | dct_y = tf.transpose( 138 | a=tf.signal.dct(image, type=2, norm='ortho'), perm=[0, 2, 1]) 139 | dct_x = tf.transpose( 140 | a=tf.signal.dct(dct_y, type=2, norm='ortho'), perm=[0, 2, 1]) 141 | return dct_x 142 | 143 | 144 | def image_idct(dct_x): 145 | """Inverts image_dct(), by performing a type-III DCT.""" 146 | dct_y = tf.signal.idct( 147 | tf.transpose(dct_x, perm=[0, 2, 1]), type=2, norm='ortho') 148 | image = tf.signal.idct( 149 | tf.transpose(dct_y, perm=[0, 2, 1]), type=2, norm='ortho') 150 | return image 151 | 152 | 153 | def compute_jacobian(f, x): 154 | """Computes the Jacobian of function `f` with respect to input `x`.""" 155 | x = tf.convert_to_tensor(x) 156 | with tf.GradientTape(persistent=True) as tape: 157 | tape.watch(x) 158 | vec = lambda x: tf.reshape(x, [-1]) 159 | jacobian = tf.stack( 160 | [vec(tape.gradient(vec(f(x))[d], x)) for d in range(tf.size(x))]) 161 | return jacobian 162 | 163 | 164 | def get_resource_as_file(path): 165 | """A uniform interface for internal/open-source files.""" 166 | 167 | class NullContextManager(object): 168 | 169 | def __init__(self, dummy_resource=None): 170 | self.dummy_resource = dummy_resource 171 | 172 | def __enter__(self): 173 | return self.dummy_resource 174 | 175 | def __exit__(self, *args): 176 | pass 177 | 178 | return NullContextManager('./' + path) 179 | 180 | 181 | def get_resource_filename(path): 182 | """A uniform interface for internal/open-source filenames.""" 183 | return './' + path 184 | -------------------------------------------------------------------------------- /src/robust_loss/util_test.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | # Copyright 2020 The Google Research Authors. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | """Tests for util.py.""" 17 | 18 | import numpy as np 19 | import tensorflow.compat.v2 as tf 20 | import tensorflow_probability as tfp 21 | from robust_loss import util 22 | 23 | tf.enable_v2_behavior() 24 | 25 | 26 | class UtilTest(tf.test.TestCase): 27 | 28 | def setUp(self): 29 | super(UtilTest, self).setUp() 30 | np.random.seed(0) 31 | 32 | def testInvSoftplusIsCorrect(self): 33 | """Test that inv_softplus() is the inverse of tf.nn.softplus().""" 34 | x = np.float32(np.exp(np.linspace(-10., 10., 1000))) 35 | x_recon = tf.nn.softplus(util.inv_softplus(x)) 36 | self.assertAllClose(x, x_recon) 37 | 38 | def testLogitIsCorrect(self): 39 | """Test that logit() is the inverse of tf.sigmoid().""" 40 | x = np.float32(np.linspace(1e-5, 1. - 1e-5, 1000)) 41 | x_recon = tf.sigmoid(util.logit(x)) 42 | self.assertAllClose(x, x_recon) 43 | 44 | def testAffineSigmoidSpansRange(self): 45 | """Check that affine_sigmoid()'s output is in [lo, hi].""" 46 | x = np.finfo(np.float32).max * np.array([-1, 1], dtype=np.float32) 47 | for _ in range(10): 48 | lo = np.random.uniform(0., 0.3) 49 | hi = np.random.uniform(0.5, 4.) 50 | y = util.affine_sigmoid(x, lo=lo, hi=hi) 51 | self.assertAllClose(y[0], lo) 52 | self.assertAllClose(y[1], hi) 53 | 54 | def testAffineSigmoidIsCentered(self): 55 | """Check that affine_sigmoid(0) == (lo+hi)/2.""" 56 | for _ in range(10): 57 | lo = np.random.uniform(0., 0.3) 58 | hi = np.random.uniform(0.5, 4.) 59 | y = util.affine_sigmoid(np.array(0.), lo=lo, hi=hi) 60 | self.assertAllClose(y, (lo + hi) * 0.5) 61 | 62 | def testAffineSoftplusSpansRange(self): 63 | """Check that affine_softplus()'s output is in [lo, infinity].""" 64 | x = np.finfo(np.float32).max * np.array([-1, 1], dtype=np.float32) 65 | for _ in range(10): 66 | lo = np.random.uniform(0., 0.1) 67 | ref = np.random.uniform(0.2, 10.) 68 | y = util.affine_softplus(x, lo=lo, ref=ref) 69 | self.assertAllClose(y[0], lo) 70 | self.assertAllGreater(y[1], 1e10) 71 | 72 | def testAffineSoftplusIsCentered(self): 73 | """Check that affine_softplus(0) == 1.""" 74 | for _ in range(10): 75 | lo = np.random.uniform(0., 0.1) 76 | ref = np.random.uniform(0.2, 10.) 77 | y = util.affine_softplus(np.array(0.), lo=lo, ref=ref) 78 | self.assertAllClose(y, ref) 79 | 80 | def testDefaultAffineSigmoidMatchesSigmoid(self): 81 | """Check that affine_sigmoid() matches tf.nn.sigmoid() by default.""" 82 | x = np.float32(np.linspace(-10., 10., 1000)) 83 | y = util.affine_sigmoid(x) 84 | y_true = tf.nn.sigmoid(x) 85 | self.assertAllClose(y, y_true, atol=1e-5, rtol=1e-3) 86 | 87 | def testDefaultAffineSigmoidRoundTrip(self): 88 | """Check that x = inv_affine_sigmoid(affine_sigmoid(x)) by default.""" 89 | x = np.float32(np.linspace(-10., 10., 1000)) 90 | y = util.affine_sigmoid(x) 91 | x_recon = util.inv_affine_sigmoid(y) 92 | self.assertAllClose(x, x_recon, atol=1e-5, rtol=1e-3) 93 | 94 | def testAffineSigmoidRoundTrip(self): 95 | """Check that x = inv_affine_sigmoid(affine_sigmoid(x)) in general.""" 96 | x = np.float32(np.linspace(-10., 10., 1000)) 97 | for _ in range(10): 98 | lo = np.random.uniform(0., 0.3) 99 | hi = np.random.uniform(0.5, 4.) 100 | y = util.affine_sigmoid(x, lo=lo, hi=hi) 101 | x_recon = util.inv_affine_sigmoid(y, lo=lo, hi=hi) 102 | self.assertAllClose(x, x_recon, atol=1e-5, rtol=1e-3) 103 | 104 | def testDefaultAffineSoftplusRoundTrip(self): 105 | """Check that x = inv_affine_softplus(affine_softplus(x)) by default.""" 106 | x = np.float32(np.linspace(-10., 10., 1000)) 107 | y = util.affine_softplus(x) 108 | x_recon = util.inv_affine_softplus(y) 109 | self.assertAllClose(x, x_recon, atol=1e-5, rtol=1e-3) 110 | 111 | def testAffineSoftplusRoundTrip(self): 112 | """Check that x = inv_affine_softplus(affine_softplus(x)) in general.""" 113 | x = np.float32(np.linspace(-10., 10., 1000)) 114 | for _ in range(10): 115 | lo = np.random.uniform(0., 0.1) 116 | ref = np.random.uniform(0.2, 10.) 117 | y = util.affine_softplus(x, lo=lo, ref=ref) 118 | x_recon = util.inv_affine_softplus(y, lo=lo, ref=ref) 119 | self.assertAllClose(x, x_recon, atol=1e-5, rtol=1e-3) 120 | 121 | def testStudentsTNllAgainstTfp(self): 122 | """Check that our Student's T NLL matches TensorFlow Probability.""" 123 | for _ in range(10): 124 | x = np.random.normal() 125 | df = np.exp(4. * np.random.normal()) 126 | scale = np.exp(4. * np.random.normal()) 127 | nll = util.students_t_nll(x, df, scale) 128 | nll_true = -tfp.distributions.StudentT( 129 | df=df, loc=tf.zeros_like(scale), scale=scale).log_prob(x) 130 | self.assertAllClose(nll, nll_true) 131 | 132 | def testRgbToSyuvPreservesVolume(self): 133 | """Tests that rgb_to_syuv() is volume preserving.""" 134 | for _ in range(4): 135 | im = np.float32(np.random.uniform(size=(1, 1, 3))) 136 | jacobian = util.compute_jacobian(util.rgb_to_syuv, im) 137 | # Assert that the determinant of the Jacobian is close to 1. 138 | det = np.linalg.det(jacobian) 139 | self.assertAllClose(det, 1., atol=1e-5, rtol=1e-5) 140 | 141 | def testRgbToSyuvRoundTrip(self): 142 | """Tests that syuv_to_rgb(rgb_to_syuv(x)) == x.""" 143 | rgb = np.float32(np.random.uniform(size=(32, 32, 3))) 144 | syuv = util.rgb_to_syuv(rgb) 145 | rgb_recon = util.syuv_to_rgb(syuv) 146 | self.assertAllClose(rgb, rgb_recon) 147 | 148 | def testSyuvIsScaledYuv(self): 149 | """Tests that rgb_to_syuv is proportional to tf.image.rgb_to_yuv().""" 150 | rgb = np.float32(np.random.uniform(size=(32, 32, 3))) 151 | syuv = util.rgb_to_syuv(rgb) 152 | yuv = tf.image.rgb_to_yuv(rgb) 153 | # Check that the ratio between `syuv` and `yuv` is nearly constant. 154 | ratio = syuv / yuv 155 | self.assertAllClose(tf.reduce_min(ratio), tf.reduce_max(ratio)) 156 | 157 | def testImageDctPreservesVolume(self): 158 | """Tests that image_dct() is volume preserving.""" 159 | for _ in range(4): 160 | im = np.float32(np.random.uniform(size=(4, 4, 2))) 161 | jacobian = util.compute_jacobian(util.image_dct, im) 162 | # Assert that the determinant of the Jacobian is close to 1. 163 | det = np.linalg.det(jacobian) 164 | self.assertAllClose(det, 1., atol=1e-5, rtol=1e-5) 165 | 166 | def testImageDctIsOrthonormal(self): 167 | """Test that = .""" 168 | for _ in range(4): 169 | im0 = np.float32(np.random.uniform(size=(4, 4, 2))) 170 | im1 = np.float32(np.random.uniform(size=(4, 4, 2))) 171 | dct_im0 = util.image_dct(im0) 172 | dct_im1 = util.image_dct(im1) 173 | prod1 = tf.reduce_sum(im0 * im1) 174 | prod2 = tf.reduce_sum(dct_im0 * dct_im1) 175 | self.assertAllClose(prod1, prod2) 176 | 177 | def testImageDctRoundTrip(self): 178 | """Tests that image_idct(image_dct(x)) == x.""" 179 | image = np.float32(np.random.uniform(size=(32, 32, 3))) 180 | image_recon = util.image_idct(util.image_dct(image)) 181 | self.assertAllClose(image, image_recon) 182 | 183 | 184 | if __name__ == '__main__': 185 | tf.test.main() 186 | -------------------------------------------------------------------------------- /src/robust_loss/wavelet_test.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | # Copyright 2020 The Google Research Authors. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | """Tests for wavelet.py.""" 17 | 18 | from absl.testing import parameterized 19 | import numpy as np 20 | import PIL.Image 21 | import scipy.io 22 | import tensorflow.compat.v2 as tf 23 | from robust_loss import util 24 | from robust_loss import wavelet 25 | 26 | tf.enable_v2_behavior() 27 | 28 | 29 | class WaveletTest(parameterized.TestCase, tf.test.TestCase): 30 | 31 | def setUp(self): 32 | super(WaveletTest, self).setUp() 33 | np.random.seed(0) 34 | 35 | def _assert_pyramids_close(self, x0, x1, epsilon): 36 | """A helper function for assering that two wavelet pyramids are close.""" 37 | if isinstance(x0, tuple) or isinstance(x0, list): 38 | assert isinstance(x1, (list, tuple)) 39 | assert len(x0) == len(x1) 40 | for y0, y1 in zip(x0, x1): 41 | self._assert_pyramids_close(y0, y1, epsilon) 42 | else: 43 | assert not isinstance(x1, (list, tuple)) 44 | self.assertAllEqual(x0.shape, x1.shape) 45 | self.assertAllClose(x0, x1, atol=epsilon, rtol=epsilon) 46 | 47 | def testPadWithOneReflectionIsCorrect(self): 48 | """Tests that pad_reflecting(p) matches tf.pad(p) when p is small.""" 49 | for _ in range(4): 50 | n = int(np.ceil(np.random.uniform() * 8)) + 1 51 | x = np.random.uniform(size=(n, n, n)) 52 | padding_below = int(np.round(np.random.uniform() * (n - 1))) 53 | padding_above = int(np.round(np.random.uniform() * (n - 1))) 54 | axis = int(np.floor(np.random.uniform() * 3.)) 55 | 56 | if axis == 0: 57 | reference = tf.pad( 58 | x, [[padding_below, padding_above], [0, 0], [0, 0]], mode='REFLECT') 59 | elif axis == 1: 60 | reference = tf.pad( 61 | x, [[0, 0], [padding_below, padding_above], [0, 0]], mode='REFLECT') 62 | elif axis == 2: 63 | reference = tf.pad( 64 | x, [[0, 0], [0, 0], [padding_below, padding_above]], mode='REFLECT') 65 | 66 | result = wavelet.pad_reflecting(x, padding_below, padding_above, axis) 67 | self.assertAllEqual(result.shape, reference.shape) 68 | self.assertAllEqual(result, reference) 69 | 70 | def testPadWithManyReflectionsIsCorrect(self): 71 | """Tests that pad_reflecting(k * p) matches tf.pad(p) applied k times.""" 72 | for _ in range(4): 73 | n = int(np.random.uniform() * 8.) + 1 74 | p = n - 1 75 | x = np.random.uniform(size=(n)) 76 | result1 = wavelet.pad_reflecting(x, p, p, 0) 77 | result2 = wavelet.pad_reflecting(x, 2 * p, 2 * p, 0) 78 | result3 = wavelet.pad_reflecting(x, 3 * p, 3 * p, 0) 79 | reference1 = tf.pad(x, [[p, p]], mode='REFLECT') 80 | reference2 = tf.pad(reference1, [[p, p]], mode='REFLECT') 81 | reference3 = tf.pad(reference2, [[p, p]], mode='REFLECT') 82 | self.assertAllEqual(result1.shape, reference1.shape) 83 | self.assertAllEqual(result1, reference1) 84 | self.assertAllEqual(result2.shape, reference2.shape) 85 | self.assertAllEqual(result2, reference2) 86 | self.assertAllEqual(result3.shape, reference3.shape) 87 | self.assertAllEqual(result3, reference3) 88 | 89 | def testPadWithManyReflectionsGolden1IsCorrect(self): 90 | """Tests pad_reflecting() against a golden example.""" 91 | n = 8 92 | p0 = 17 93 | p1 = 13 94 | x = np.arange(n) 95 | reference1 = np.concatenate( 96 | (np.arange(3, 0, -1), 97 | np.arange(n), 98 | np.arange(n - 2, 0, -1), 99 | np.arange(n), 100 | np.arange(n - 2, 0, -1), 101 | np.arange(7))) # pyformat: disable 102 | result1 = wavelet.pad_reflecting(x, p0, p1, 0) 103 | self.assertAllEqual(result1.shape, reference1.shape) 104 | self.assertAllEqual(result1, reference1) 105 | 106 | def testPadWithManyReflectionsGolden2IsCorrect(self): 107 | """Tests pad_reflecting() against a golden example.""" 108 | n = 11 109 | p0 = 15 110 | p1 = 7 111 | x = np.arange(n) 112 | reference1 = np.concatenate( 113 | (np.arange(5, n), 114 | np.arange(n - 2, 0, -1), 115 | np.arange(n), 116 | np.arange(n - 2, 2, -1))) # pyformat: disable 117 | result1 = wavelet.pad_reflecting(x, p0, p1, 0) 118 | self.assertAllEqual(result1.shape, reference1.shape) 119 | self.assertAllEqual(result1, reference1) 120 | 121 | def testAnalysisLowpassFiltersAreNormalized(self): 122 | """Tests that the analysis lowpass filter doubles the input's magnitude.""" 123 | for wavelet_type in wavelet.generate_filters(): 124 | filters = wavelet.generate_filters(wavelet_type) 125 | # The sum of the outer product of the analysis lowpass filter with itself. 126 | magnitude = np.sum(filters.analysis_lo[:, np.newaxis] * 127 | filters.analysis_lo[np.newaxis, :]) 128 | self.assertAllClose(magnitude, 2., atol=1e-10, rtol=1e-10) 129 | 130 | def testWaveletTransformationIsVolumePreserving(self): 131 | """Tests that construct() is volume preserving when size is a power of 2.""" 132 | sz = (1, 4, 4) 133 | num_levels = 2 134 | im = np.float32(np.random.uniform(0., 1., sz)) 135 | for wavelet_type in wavelet.generate_filters(): 136 | # Construct the Jacobian of construct(). 137 | def fun(z): 138 | # pylint: disable=cell-var-from-loop 139 | return wavelet.flatten(wavelet.construct(z, num_levels, wavelet_type)) 140 | 141 | jacobian = util.compute_jacobian(fun, im) 142 | # Assert that the determinant of the Jacobian is close to 1. 143 | det = np.linalg.det(jacobian) 144 | self.assertAllClose(det, 1., atol=1e-5, rtol=1e-5) 145 | 146 | def _load_golden_data(self): 147 | """Loads golden data: an RGBimage and its CDF9/7 decomposition. 148 | 149 | This golden data was produced by running the code from 150 | https://www.getreuer.info/projects/wavelet-cdf-97-implementation 151 | on a test image. 152 | 153 | Returns: 154 | A tuple containing and image, its decomposition, and its wavelet type. 155 | """ 156 | with util.get_resource_as_file( 157 | 'robust_loss/data/wavelet_golden.mat') as golden_filename: 158 | data = scipy.io.loadmat(golden_filename) 159 | im = np.float32(data['I_color']) 160 | pyr_true = data['pyr_color'][0, :].tolist() 161 | for i in range(len(pyr_true) - 1): 162 | pyr_true[i] = tuple(pyr_true[i].flatten()) 163 | pyr_true = tuple(pyr_true) 164 | wavelet_type = 'CDF9/7' 165 | return im, pyr_true, wavelet_type 166 | 167 | def testConstructMatchesGoldenData(self): 168 | """Tests construct() against golden data.""" 169 | im, pyr_true, wavelet_type = self._load_golden_data() 170 | pyr = wavelet.construct(im, len(pyr_true) - 1, wavelet_type) 171 | self._assert_pyramids_close(pyr, pyr_true, 1e-5) 172 | 173 | def testCollapseMatchesGoldenData(self): 174 | """Tests collapse() against golden data.""" 175 | im, pyr_true, wavelet_type = self._load_golden_data() 176 | recon = wavelet.collapse(pyr_true, wavelet_type) 177 | self.assertAllClose(recon, im, atol=1e-5, rtol=1e-5) 178 | 179 | def testVisualizeMatchesGoldenData(self): 180 | """Tests visualize() (and implicitly flatten()).""" 181 | _, pyr, _ = self._load_golden_data() 182 | vis = wavelet.visualize(pyr) 183 | golden_vis_filename = 'robust_loss/data/wavelet_vis_golden.png' 184 | vis_true = np.asarray( 185 | PIL.Image.open(util.get_resource_filename(golden_vis_filename))) 186 | # Allow for some slack as quantization may exaggerate some errors. 187 | self.assertAllClose(vis_true, vis, atol=1., rtol=0) 188 | 189 | def testAccurateRoundTripWithSmallRandomImages(self): 190 | """Tests that collapse(construct(x)) == x for x = [1, k, k], k in [1, 4].""" 191 | for wavelet_type in wavelet.generate_filters(): 192 | for width in range(1, 5): 193 | sz = [1, width, width] 194 | num_levels = wavelet.get_max_num_levels(sz) 195 | im = np.random.uniform(size=sz) 196 | 197 | pyr = wavelet.construct(im, num_levels, wavelet_type) 198 | recon = wavelet.collapse(pyr, wavelet_type) 199 | self.assertAllClose(recon, im, atol=1e-8, rtol=1e-8) 200 | 201 | def testAccurateRoundTripWithLargeRandomImages(self): 202 | """Tests that collapse(construct(x)) == x for large random x's.""" 203 | for wavelet_type in wavelet.generate_filters(): 204 | for _ in range(4): 205 | num_levels = np.int32(np.ceil(4 * np.random.uniform())) 206 | sz_clamp = 2**(num_levels - 1) + 1 207 | sz = np.maximum( 208 | np.int32( 209 | np.ceil(np.array([2, 32, 32]) * np.random.uniform(size=3))), 210 | np.array([0, sz_clamp, sz_clamp])) 211 | im = np.random.uniform(size=sz) 212 | pyr = wavelet.construct(im, num_levels, wavelet_type) 213 | recon = wavelet.collapse(pyr, wavelet_type) 214 | self.assertAllClose(recon, im, atol=1e-8, rtol=1e-8) 215 | 216 | def testDecompositionIsNonRedundant(self): 217 | """Test that wavelet construction is not redundant. 218 | 219 | If the wavelet decompositon is not redundant, then we should be able to 220 | 1) Construct a wavelet decomposition 221 | 2) Alter a single coefficient in the decomposition 222 | 3) Collapse that decomposition into an image and back 223 | and the two wavelet decompositions should be the same. 224 | """ 225 | for wavelet_type in wavelet.generate_filters(): 226 | for _ in range(4): 227 | # Construct an image and a wavelet decomposition of it. 228 | num_levels = np.int32(np.ceil(4 * np.random.uniform())) 229 | sz_clamp = 2**(num_levels - 1) + 1 230 | sz = np.maximum( 231 | np.int32( 232 | np.ceil(np.array([2, 32, 32]) * np.random.uniform(size=3))), 233 | np.array([0, sz_clamp, sz_clamp])) 234 | im = np.random.uniform(size=sz) 235 | pyr = wavelet.construct(im, num_levels, wavelet_type) 236 | pyr = list(pyr) 237 | 238 | # Pick a coefficient at random in the decomposition to alter. 239 | d = np.int32(np.floor(np.random.uniform() * len(pyr))) 240 | v = np.random.uniform() 241 | if d == (len(pyr) - 1): 242 | if np.prod(pyr[d].shape) > 0: 243 | c, i, j = np.int32( 244 | np.floor(np.array(np.random.uniform(size=3)) * 245 | pyr[d].shape)).tolist() 246 | pyr[d] = pyr[d].numpy() 247 | pyr[d][c, i, j] = v 248 | else: 249 | b = np.int32(np.floor(np.random.uniform() * len(pyr[d]))) 250 | if np.prod(pyr[d][b].shape) > 0: 251 | c, i, j = np.int32( 252 | np.floor(np.array(np.random.uniform(size=3)) * 253 | pyr[d][b].shape)).tolist() 254 | pyr[d] = list(pyr[d]) 255 | pyr[d][b] = pyr[d][b].numpy() 256 | pyr[d][b][c, i, j] = v 257 | 258 | # Collapse and then reconstruct the wavelet decomposition, and check 259 | # that it is unchanged. 260 | recon = wavelet.collapse(pyr, wavelet_type) 261 | pyr_again = wavelet.construct(recon, num_levels, wavelet_type) 262 | self._assert_pyramids_close(pyr, pyr_again, 1e-8) 263 | 264 | def testUpsampleAndDownsampleAreTransposes(self): 265 | """Tests that _downsample() is the transpose of _upsample().""" 266 | n = 8 267 | x = tf.convert_to_tensor(np.random.uniform(size=(1, n, 1))) 268 | 269 | for f_len in range(1, 5): 270 | f = np.random.uniform(size=f_len) 271 | for shift in [0, 1]: 272 | 273 | # We're only testing the resampling operators away from the boundaries, 274 | # as this test appears to fail in the presences of boundary conditions. 275 | # TODO(barron): Figure out what's happening and make this test more 276 | # thorough, and then set range1 = range(d), range2 = range(d//2) and 277 | # have this code depend on util.compute_jacobian(). 278 | range1 = np.arange(f_len // 2 + 1, n - (f_len // 2 + 1)) 279 | range2 = np.arange(f_len // 4, n // 2 - (f_len // 4)) 280 | 281 | y = wavelet._downsample(x, f, 0, shift) 282 | vec = lambda z: tf.reshape(z, [-1]) 283 | 284 | jacobian_down = [] 285 | with tf.GradientTape(persistent=True) as tape: 286 | tape.watch(x) 287 | for d in range2: 288 | yd = vec(wavelet._downsample(x, f, 0, shift))[d] 289 | jacobian_down.append(vec(tape.gradient(yd, x))) 290 | jacobian_down = tf.stack(jacobian_down, 1).numpy() 291 | 292 | jacobian_up = [] 293 | with tf.GradientTape(persistent=True) as tape: 294 | tape.watch(y) 295 | for d in range1: 296 | xd = vec(wavelet._upsample(y, x.shape[1:], f, 0, shift))[d] 297 | jacobian_up.append(vec(tape.gradient(xd, y))) 298 | jacobian_up = tf.stack(jacobian_up, 1).numpy() 299 | 300 | # Test that the jacobian of _downsample() is close to the transpose of 301 | # the jacobian of _upsample(). 302 | self.assertAllClose( 303 | jacobian_down[range1, :], 304 | np.transpose(jacobian_up[range2, :]), 305 | atol=1e-6, 306 | rtol=1e-6) 307 | 308 | @parameterized.named_parameters(('Single', np.float32), 309 | ('Double', np.float64)) 310 | def testConstructPreservesDtype(self, float_dtype): 311 | """Checks that construct()'s output has the same precision as its input.""" 312 | x = float_dtype(np.random.normal(size=(3, 16, 16))) 313 | for wavelet_type in wavelet.generate_filters(): 314 | y = wavelet.flatten(wavelet.construct(x, 3, wavelet_type)) 315 | self.assertDTypeEqual(y, float_dtype) 316 | 317 | @parameterized.named_parameters(('Single', np.float32), 318 | ('Double', np.float64)) 319 | def testCollapsePreservesDtype(self, float_dtype): 320 | """Checks that collapse()'s output has the same precision as its input.""" 321 | n = 16 322 | x = [] 323 | for n in [8, 4, 2]: 324 | band = [] 325 | for _ in range(3): 326 | band.append(float_dtype(np.random.normal(size=(3, n, n)))) 327 | x.append(band) 328 | x.append(float_dtype(np.random.normal(size=(3, n, n)))) 329 | for wavelet_type in wavelet.generate_filters(): 330 | y = wavelet.collapse(x, wavelet_type) 331 | self.assertDTypeEqual(y, float_dtype) 332 | 333 | def testRescaleOneIsANoOp(self): 334 | """Tests that rescale(x, 1) = x.""" 335 | im = np.random.uniform(size=(2, 32, 32)) 336 | pyr = wavelet.construct(im, 4, 'LeGall5/3') 337 | pyr_rescaled = wavelet.rescale(pyr, 1.) 338 | self._assert_pyramids_close(pyr, pyr_rescaled, 1e-8) 339 | 340 | def testRescaleDoesNotAffectTheFirstLevel(self): 341 | """Tests that rescale(x, s)[0] = x[0] for any s.""" 342 | im = np.random.uniform(size=(2, 32, 32)) 343 | pyr = wavelet.construct(im, 4, 'LeGall5/3') 344 | pyr_rescaled = wavelet.rescale(pyr, np.exp(np.random.normal())) 345 | self._assert_pyramids_close(pyr[0:1], pyr_rescaled[0:1], 1e-8) 346 | 347 | def testRescaleOneHalfIsNormalized(self): 348 | """Tests that rescale(construct(k), 0.5)[-1] = k for constant image k.""" 349 | for num_levels in range(5): 350 | k = np.random.uniform() 351 | im = k * np.ones((2, 32, 32)) 352 | pyr = wavelet.construct(im, num_levels, 'LeGall5/3') 353 | pyr_rescaled = wavelet.rescale(pyr, 0.5) 354 | self.assertAllClose( 355 | pyr_rescaled[-1], 356 | k * np.ones_like(pyr_rescaled[-1]), 357 | atol=1e-8, 358 | rtol=1e-8) 359 | 360 | def testRescaleAndUnrescaleReproducesInput(self): 361 | """Tests that rescale(rescale(x, k), 1/k) = x.""" 362 | im = np.random.uniform(size=(2, 32, 32)) 363 | scale_base = np.exp(np.random.normal()) 364 | pyr = wavelet.construct(im, 4, 'LeGall5/3') 365 | pyr_rescaled = wavelet.rescale(pyr, scale_base) 366 | pyr_recon = wavelet.rescale(pyr_rescaled, 1. / scale_base) 367 | self._assert_pyramids_close(pyr, pyr_recon, 1e-8) 368 | 369 | 370 | if __name__ == '__main__': 371 | tf.test.main() 372 | -------------------------------------------------------------------------------- /thinplate/__init__.py: -------------------------------------------------------------------------------- 1 | from thinplate.numpy import * 2 | 3 | try: 4 | import torch 5 | import thinplate.pytorch as torch 6 | except ImportError: 7 | pass 8 | 9 | __version__ = '1.0.0' -------------------------------------------------------------------------------- /thinplate/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/thinplate/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /thinplate/__pycache__/numpy.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/thinplate/__pycache__/numpy.cpython-37.pyc -------------------------------------------------------------------------------- /thinplate/numpy.py: -------------------------------------------------------------------------------- 1 | # Copyright 2018 Christoph Heindl. 2 | # 3 | # Licensed under MIT License 4 | # ============================================================ 5 | 6 | import numpy as np 7 | 8 | class TPS: 9 | @staticmethod 10 | def fit(c, lambd=0., reduced=False): 11 | n = c.shape[0] 12 | 13 | U = TPS.u(TPS.d(c, c)) 14 | K = U + np.eye(n, dtype=np.float32)*lambd 15 | 16 | P = np.ones((n, 3), dtype=np.float32) 17 | P[:, 1:] = c[:, :2] 18 | 19 | v = np.zeros(n+3, dtype=np.float32) 20 | v[:n] = c[:, -1] 21 | 22 | A = np.zeros((n+3, n+3), dtype=np.float32) 23 | A[:n, :n] = K 24 | A[:n, -3:] = P 25 | A[-3:, :n] = P.T 26 | 27 | theta = np.linalg.solve(A, v) # p has structure w,a 28 | return theta[1:] if reduced else theta 29 | 30 | @staticmethod 31 | def d(a, b): 32 | return np.sqrt(np.square(a[:, None, :2] - b[None, :, :2]).sum(-1)) 33 | 34 | @staticmethod 35 | def u(r): 36 | return r**2 * np.log(r + 1e-6) 37 | 38 | @staticmethod 39 | def z(x, c, theta): 40 | x = np.atleast_2d(x) 41 | U = TPS.u(TPS.d(x, c)) 42 | w, a = theta[:-3], theta[-3:] 43 | reduced = theta.shape[0] == c.shape[0] + 2 44 | if reduced: 45 | w = np.concatenate((-np.sum(w, keepdims=True), w)) 46 | b = np.dot(U, w) 47 | return a[0] + a[1]*x[:, 0] + a[2]*x[:, 1] + b 48 | 49 | def uniform_grid(shape): 50 | '''Uniform grid coordinates. 51 | 52 | Params 53 | ------ 54 | shape : tuple 55 | HxW defining the number of height and width dimension of the grid 56 | 57 | Returns 58 | ------- 59 | points: HxWx2 tensor 60 | Grid coordinates over [0,1] normalized image range. 61 | ''' 62 | 63 | H,W = shape[:2] 64 | c = np.empty((H, W, 2)) 65 | c[..., 0] = np.linspace(0, 1, W, dtype=np.float32) 66 | c[..., 1] = np.expand_dims(np.linspace(0, 1, H, dtype=np.float32), -1) 67 | 68 | return c 69 | 70 | def tps_theta_from_points(c_src, c_dst, reduced=False): 71 | delta = c_src - c_dst 72 | 73 | cx = np.column_stack((c_dst, delta[:, 0])) 74 | cy = np.column_stack((c_dst, delta[:, 1])) 75 | 76 | theta_dx = TPS.fit(cx, reduced=reduced) 77 | theta_dy = TPS.fit(cy, reduced=reduced) 78 | 79 | return np.stack((theta_dx, theta_dy), -1) 80 | 81 | 82 | def tps_grid(theta, c_dst, dshape): 83 | ugrid = uniform_grid(dshape) 84 | 85 | reduced = c_dst.shape[0] + 2 == theta.shape[0] 86 | 87 | dx = TPS.z(ugrid.reshape((-1, 2)), c_dst, theta[:, 0]).reshape(dshape[:2]) 88 | dy = TPS.z(ugrid.reshape((-1, 2)), c_dst, theta[:, 1]).reshape(dshape[:2]) 89 | dgrid = np.stack((dx, dy), -1) 90 | 91 | grid = dgrid + ugrid 92 | 93 | return grid # H'xW'x2 grid[i,j] in range [0..1] 94 | 95 | def tps_grid_to_remap(grid, sshape): 96 | '''Convert a dense grid to OpenCV's remap compatible maps. 97 | 98 | Params 99 | ------ 100 | grid : HxWx2 array 101 | Normalized flow field coordinates as computed by compute_densegrid. 102 | sshape : tuple 103 | Height and width of source image in pixels. 104 | 105 | 106 | Returns 107 | ------- 108 | mapx : HxW array 109 | mapy : HxW array 110 | ''' 111 | 112 | mx = (grid[:, :, 0] * sshape[1]).astype(np.float32) 113 | my = (grid[:, :, 1] * sshape[0]).astype(np.float32) 114 | 115 | return mx, my -------------------------------------------------------------------------------- /thinplate/pytorch.py: -------------------------------------------------------------------------------- 1 | # Copyright 2018 Christoph Heindl. 2 | # 3 | # Licensed under MIT License 4 | # ============================================================ 5 | 6 | import torch 7 | 8 | def tps(theta, ctrl, grid): 9 | '''Evaluate the thin-plate-spline (TPS) surface at xy locations arranged in a grid. 10 | The TPS surface is a minimum bend interpolation surface defined by a set of control points. 11 | The function value for a x,y location is given by 12 | 13 | TPS(x,y) := theta[-3] + theta[-2]*x + theta[-1]*y + \sum_t=0,T theta[t] U(x,y,ctrl[t]) 14 | 15 | This method computes the TPS value for multiple batches over multiple grid locations for 2 16 | surfaces in one go. 17 | 18 | Params 19 | ------ 20 | theta: Nx(T+3)x2 tensor, or Nx(T+2)x2 tensor 21 | Batch size N, T+3 or T+2 (reduced form) model parameters for T control points in dx and dy. 22 | ctrl: NxTx2 tensor or Tx2 tensor 23 | T control points in normalized image coordinates [0..1] 24 | grid: NxHxWx3 tensor 25 | Grid locations to evaluate with homogeneous 1 in first coordinate. 26 | 27 | Returns 28 | ------- 29 | z: NxHxWx2 tensor 30 | Function values at each grid location in dx and dy. 31 | ''' 32 | 33 | N, H, W, _ = grid.size() 34 | 35 | if ctrl.dim() == 2: 36 | ctrl = ctrl.expand(N, *ctrl.size()) 37 | 38 | T = ctrl.shape[1] 39 | 40 | diff = grid[...,1:].unsqueeze(-2) - ctrl.unsqueeze(1).unsqueeze(1) 41 | D = torch.sqrt((diff**2).sum(-1)) 42 | U = (D**2) * torch.log(D + 1e-6) 43 | 44 | w, a = theta[:, :-3, :], theta[:, -3:, :] 45 | 46 | reduced = T + 2 == theta.shape[1] 47 | if reduced: 48 | w = torch.cat((-w.sum(dim=1, keepdim=True), w), dim=1) 49 | 50 | # U is NxHxWxT 51 | b = torch.bmm(U.view(N, -1, T), w).view(N,H,W,2) 52 | # b is NxHxWx2 53 | z = torch.bmm(grid.view(N,-1,3), a).view(N,H,W,2) + b 54 | 55 | return z 56 | 57 | def tps_grid(theta, ctrl, size): 58 | '''Compute a thin-plate-spline grid from parameters for sampling. 59 | 60 | Params 61 | ------ 62 | theta: Nx(T+3)x2 tensor 63 | Batch size N, T+3 model parameters for T control points in dx and dy. 64 | ctrl: NxTx2 tensor, or Tx2 tensor 65 | T control points in normalized image coordinates [0..1] 66 | size: tuple 67 | Output grid size as NxCxHxW. C unused. This defines the output image 68 | size when sampling. 69 | 70 | Returns 71 | ------- 72 | grid : NxHxWx2 tensor 73 | Grid suitable for sampling in pytorch containing source image 74 | locations for each output pixel. 75 | ''' 76 | N, _, H, W = size 77 | 78 | grid = theta.new(N, H, W, 3) 79 | grid[:, :, :, 0] = 1. 80 | grid[:, :, :, 1] = torch.linspace(0, 1, W) 81 | grid[:, :, :, 2] = torch.linspace(0, 1, H).unsqueeze(-1) 82 | 83 | z = tps(theta, ctrl, grid) 84 | return (grid[...,1:] + z)*2-1 # [-1,1] range required by F.sample_grid 85 | 86 | def tps_sparse(theta, ctrl, xy): 87 | if xy.dim() == 2: 88 | xy = xy.expand(theta.shape[0], *xy.size()) 89 | 90 | N, M = xy.shape[:2] 91 | grid = xy.new(N, M, 3) 92 | grid[..., 0] = 1. 93 | grid[..., 1:] = xy 94 | 95 | z = tps(theta, ctrl, grid.view(N,M,1,3)) 96 | return xy + z.view(N, M, 2) 97 | 98 | def uniform_grid(shape): 99 | '''Uniformly places control points aranged in grid accross normalized image coordinates. 100 | 101 | Params 102 | ------ 103 | shape : tuple 104 | HxW defining the number of control points in height and width dimension 105 | 106 | Returns 107 | ------- 108 | points: HxWx2 tensor 109 | Control points over [0,1] normalized image range. 110 | ''' 111 | H,W = shape[:2] 112 | c = torch.zeros(H, W, 2) 113 | c[..., 0] = torch.linspace(0, 1, W) 114 | c[..., 1] = torch.linspace(0, 1, H).unsqueeze(-1) 115 | return c 116 | 117 | if __name__ == '__main__': 118 | c = torch.tensor([ 119 | [0., 0], 120 | [1., 0], 121 | [1., 1], 122 | [0, 1], 123 | ]).unsqueeze(0) 124 | theta = torch.zeros(1, 4+3, 2) 125 | size= (1,1,6,3) 126 | print(tps_grid(theta, c, size).shape) -------------------------------------------------------------------------------- /thinplate/tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Leedeng/LineCounter/80713febcb8b156384e6e7d7505e0fac8aa76bf2/thinplate/tests/__init__.py -------------------------------------------------------------------------------- /thinplate/tests/test_tps_numpy.py: -------------------------------------------------------------------------------- 1 | 2 | import numpy as np 3 | from numpy.testing import assert_allclose 4 | import thinplate as tps 5 | 6 | def test_numpy_fit(): 7 | c = np.array([ 8 | [0., 0, 0.0], 9 | [1., 0, 0.0], 10 | [1., 1, 0.0], 11 | [0, 1, 0.0], 12 | ]) 13 | 14 | theta = tps.TPS.fit(c) 15 | assert_allclose(theta, 0) 16 | assert_allclose(tps.TPS.z(c, c, theta), c[:, 2]) 17 | 18 | c = np.array([ 19 | [0., 0, 1.0], 20 | [1., 0, 1.0], 21 | [1., 1, 1.0], 22 | [0, 1, 1.0], 23 | ]) 24 | 25 | theta = tps.TPS.fit(c) 26 | assert_allclose(theta[:-3], 0) 27 | assert_allclose(theta[-3:], [1, 0, 0]) 28 | assert_allclose(tps.TPS.z(c, c, theta), c[:, 2], atol=1e-3) 29 | 30 | # reduced form 31 | theta = tps.TPS.fit(c, reduced=True) 32 | assert len(theta) == c.shape[0] + 2 33 | assert_allclose(theta[:-3], 0) 34 | assert_allclose(theta[-3:], [1, 0, 0]) 35 | assert_allclose(tps.TPS.z(c, c, theta), c[:, 2], atol=1e-3) 36 | 37 | c = np.array([ 38 | [0., 0, -.5], 39 | [1., 0, 0.5], 40 | [1., 1, 0.2], 41 | [0, 1, 0.8], 42 | ]) 43 | 44 | theta = tps.TPS.fit(c) 45 | assert_allclose(tps.TPS.z(c, c, theta), c[:, 2], atol=1e-3) 46 | 47 | def test_numpy_densegrid(): 48 | 49 | # enlarges a small rectangle to full view 50 | 51 | import cv2 52 | 53 | img = np.zeros((40, 40), dtype=np.uint8) 54 | img[10:21, 10:21] = 255 55 | 56 | c_dst = np.array([ 57 | [0., 0], 58 | [1., 0], 59 | [1, 1], 60 | [0, 1], 61 | ]) 62 | 63 | 64 | c_src = np.array([ 65 | [10., 10], 66 | [20., 10], 67 | [20, 20], 68 | [10, 20], 69 | ]) / 40. 70 | 71 | theta = tps.tps_theta_from_points(c_src, c_dst) 72 | theta_r = tps.tps_theta_from_points(c_src, c_dst, reduced=True) 73 | 74 | grid = tps.tps_grid(theta, c_dst, (20,20)) 75 | grid_r = tps.tps_grid(theta_r, c_dst, (20,20)) 76 | 77 | mapx, mapy = tps.tps_grid_to_remap(grid, img.shape) 78 | warped = cv2.remap(img, mapx, mapy, cv2.INTER_CUBIC) 79 | 80 | assert img.min() == 0. 81 | assert img.max() == 255. 82 | assert warped.shape == (20,20) 83 | assert warped.min() == 255. 84 | assert warped.max() == 255. 85 | assert np.linalg.norm(grid.reshape(-1,2) - grid_r.reshape(-1,2)) < 1e-3 86 | -------------------------------------------------------------------------------- /thinplate/tests/test_tps_pytorch.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.optim as optim 3 | import torch.nn.functional as F 4 | 5 | import numpy as np 6 | import thinplate as tps 7 | 8 | from numpy.testing import assert_allclose 9 | 10 | def test_pytorch_grid(): 11 | 12 | c_dst = np.array([ 13 | [0., 0], 14 | [1., 0], 15 | [1, 1], 16 | [0, 1], 17 | ], dtype=np.float32) 18 | 19 | 20 | c_src = np.array([ 21 | [10., 10], 22 | [20., 10], 23 | [20, 20], 24 | [10, 20], 25 | ], dtype=np.float32) / 40. 26 | 27 | theta = tps.tps_theta_from_points(c_src, c_dst) 28 | theta_r = tps.tps_theta_from_points(c_src, c_dst, reduced=True) 29 | 30 | np_grid = tps.tps_grid(theta, c_dst, (20,20)) 31 | np_grid_r = tps.tps_grid(theta_r, c_dst, (20,20)) 32 | 33 | pth_theta = torch.tensor(theta).unsqueeze(0) 34 | pth_grid = tps.torch.tps_grid(pth_theta, torch.tensor(c_dst), (1, 1, 20, 20)).squeeze().numpy() 35 | pth_grid = (pth_grid + 1) / 2 # convert [-1,1] range to [0,1] 36 | 37 | pth_theta_r = torch.tensor(theta_r).unsqueeze(0) 38 | pth_grid_r = tps.torch.tps_grid(pth_theta_r, torch.tensor(c_dst), (1, 1, 20, 20)).squeeze().numpy() 39 | pth_grid_r = (pth_grid_r + 1) / 2 # convert [-1,1] range to [0,1] 40 | 41 | assert_allclose(np_grid, pth_grid) 42 | assert_allclose(np_grid_r, pth_grid_r) 43 | assert_allclose(np_grid_r, np_grid) --------------------------------------------------------------------------------