Repository: chenhsuanlin/inverse-compositional-STN
Branch: master
Commit: 4a2a8fc7b9a1
Files: 22
Total size: 86.3 KB
Directory structure:
gitextract_t1x_4nxr/
├── .editorconfig
├── .gitignore
├── LICENSE
├── MNIST-pytorch/
│ ├── data.py
│ ├── graph.py
│ ├── options.py
│ ├── train.py
│ ├── util.py
│ └── warp.py
├── MNIST-tensorflow/
│ ├── data.py
│ ├── graph.py
│ ├── options.py
│ ├── train.py
│ ├── util.py
│ └── warp.py
├── README.md
└── traffic-sign-tensorflow/
├── data.py
├── graph.py
├── options.py
├── train.py
├── util.py
└── warp.py
================================================
FILE CONTENTS
================================================
================================================
FILE: .editorconfig
================================================
root = true
[*]
end_of_line = lf
insert_final_newline = true
indent_style = tab
indent_size = 4
trim_trailing_whitespace = true
[*.md]
trim_trailing_whitespace = false
================================================
FILE: .gitignore
================================================
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
*$py.class
# C extensions
*.so
# Distribution / packaging
.Python
env/
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
*.egg-info/
.installed.cfg
*.egg
# PyInstaller
# Usually these files are written by a python script from a template
# before PyInstaller builds the exe, so as to inject date/other infos into it.
*.manifest
*.spec
# Installer logs
pip-log.txt
pip-delete-this-directory.txt
# Unit test / coverage reports
htmlcov/
.tox/
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*,cover
.hypothesis/
# Translations
*.mo
*.pot
# Django stuff:
*.log
local_settings.py
# Flask stuff:
instance/
.webassets-cache
# Scrapy stuff:
.scrapy
# Sphinx documentation
docs/_build/
# PyBuilder
target/
# IPython Notebook
.ipynb_checkpoints
# pyenv
.python-version
# celery beat schedule file
celerybeat-schedule
# dotenv
.env
# virtualenv
venv/
ENV/
# Spyder project settings
.spyderproject
# Rope project settings
.ropeproject
================================================
FILE: LICENSE
================================================
MIT License
Copyright (c) 2018 Chen-Hsuan Lin
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
================================================
FILE: MNIST-pytorch/data.py
================================================
import numpy as np
import scipy.linalg
import os,time
import torch
import torchvision
import warp,util
# load MNIST data
def loadMNIST(opt,path):
os.makedirs(path,exist_ok=True)
trainDataset = torchvision.datasets.MNIST(path,train=True,download=True)
testDataset = torchvision.datasets.MNIST(path,train=False,download=True)
trainData,testData = {},{}
trainData["image"] = torch.tensor([np.array(sample[0])/255.0 for sample in trainDataset],dtype=torch.float32)
testData["image"] = torch.tensor([np.array(sample[0])/255.0 for sample in testDataset],dtype=torch.float32)
trainData["label"] = torch.tensor([sample[1] for sample in trainDataset])
testData["label"] = torch.tensor([sample[1] for sample in testDataset])
return trainData,testData
# generate training batch
def genPerturbations(opt):
X = np.tile(opt.canon4pts[:,0],[opt.batchSize,1])
Y = np.tile(opt.canon4pts[:,1],[opt.batchSize,1])
O = np.zeros([opt.batchSize,4],dtype=np.float32)
I = np.ones([opt.batchSize,4],dtype=np.float32)
dX = np.random.randn(opt.batchSize,4)*opt.pertScale \
+np.random.randn(opt.batchSize,1)*opt.transScale
dY = np.random.randn(opt.batchSize,4)*opt.pertScale \
+np.random.randn(opt.batchSize,1)*opt.transScale
dX,dY = dX.astype(np.float32),dY.astype(np.float32)
# fit warp parameters to generated displacements
if opt.warpType=="homography":
A = np.concatenate([np.stack([X,Y,I,O,O,O,-X*(X+dX),-Y*(X+dX)],axis=-1),
np.stack([O,O,O,X,Y,I,-X*(Y+dY),-Y*(Y+dY)],axis=-1)],axis=1)
b = np.expand_dims(np.concatenate([X+dX,Y+dY],axis=1),axis=-1)
pPert = np.matmul(np.linalg.inv(A),b).squeeze()
pPert -= np.array([1,0,0,0,1,0,0,0])
else:
if opt.warpType=="translation":
J = np.concatenate([np.stack([I,O],axis=-1),
np.stack([O,I],axis=-1)],axis=1)
if opt.warpType=="similarity":
J = np.concatenate([np.stack([X,Y,I,O],axis=-1),
np.stack([-Y,X,O,I],axis=-1)],axis=1)
if opt.warpType=="affine":
J = np.concatenate([np.stack([X,Y,I,O,O,O],axis=-1),
np.stack([O,O,O,X,Y,I],axis=-1)],axis=1)
dXY = np.expand_dims(np.concatenate([dX,dY],axis=1),axis=-1)
Jtransp = np.transpose(J,axes=[0,2,1])
pPert = np.matmul(np.linalg.inv(np.matmul(Jtransp,J)),np.matmul(Jtransp,dXY)).squeeze()
pInit = torch.from_numpy(pPert).cuda()
return pInit
# make training batch
def makeBatch(opt,data):
N = len(data["image"])
randIdx = np.random.randint(N,size=[opt.batchSize])
batch = {
"image": data["image"][randIdx].cuda(),
"label": data["label"][randIdx].cuda(),
}
return batch
# evaluation on test set
def evalTest(opt,data,geometric,classifier):
geometric.eval()
classifier.eval()
N = len(data["image"])
batchN = int(np.ceil(N/opt.batchSize))
warped = [{},{}]
count = 0
for b in range(batchN):
# use some dummy data (0) as batch filler if necessary
if b!=batchN-1:
realIdx = np.arange(opt.batchSize*b,opt.batchSize*(b+1))
else:
realIdx = np.arange(opt.batchSize*b,N)
idx = np.zeros([opt.batchSize],dtype=int)
idx[:len(realIdx)] = realIdx
# make training batch
image = data["image"][idx].cuda()
label = data["label"][idx].cuda()
image.data.unsqueeze_(dim=1)
# generate perturbation
pInit = genPerturbations(opt)
pInitMtrx = warp.vec2mtrx(opt,pInit)
imagePert = warp.transformImage(opt,image,pInitMtrx)
imageWarpAll = geometric(opt,image,pInit) if opt.netType=="IC-STN" else geometric(opt,imagePert)
imageWarp = imageWarpAll[-1]
output = classifier(opt,imageWarp)
_,pred = output.max(dim=1)
count += int((pred==label).sum().cpu().numpy())
if opt.netType=="STN" or opt.netType=="IC-STN":
imgPert = imagePert.detach().cpu().numpy()
imgWarp = imageWarp.detach().cpu().numpy()
for i in range(len(realIdx)):
l = data["label"][idx[i]].item()
if l not in warped[0]: warped[0][l] = []
if l not in warped[1]: warped[1][l] = []
warped[0][l].append(imgPert[i])
warped[1][l].append(imgWarp[i])
accuracy = float(count)/N
if opt.netType=="STN" or opt.netType=="IC-STN":
mean = [np.array([np.mean(warped[0][l],axis=0) for l in warped[0]]),
np.array([np.mean(warped[1][l],axis=0) for l in warped[1]])]
var = [np.array([np.var(warped[0][l],axis=0) for l in warped[0]]),
np.array([np.var(warped[1][l],axis=0) for l in warped[1]])]
else: mean,var = None,None
geometric.train()
classifier.train()
return accuracy,mean,var
================================================
FILE: MNIST-pytorch/graph.py
================================================
import numpy as np
import torch
import time
import data,warp,util
# build classification network
class FullCNN(torch.nn.Module):
def __init__(self,opt):
super(FullCNN,self).__init__()
self.inDim = 1
def conv2Layer(outDim):
conv = torch.nn.Conv2d(self.inDim,outDim,kernel_size=[3,3],stride=1,padding=0)
self.inDim = outDim
return conv
def linearLayer(outDim):
fc = torch.nn.Linear(self.inDim,outDim)
self.inDim = outDim
return fc
def maxpoolLayer(): return torch.nn.MaxPool2d([2,2],stride=2)
self.conv2Layers = torch.nn.Sequential(
conv2Layer(3),torch.nn.ReLU(True),
conv2Layer(6),torch.nn.ReLU(True),maxpoolLayer(),
conv2Layer(9),torch.nn.ReLU(True),
conv2Layer(12),torch.nn.ReLU(True)
)
self.inDim *= 8**2
self.linearLayers = torch.nn.Sequential(
linearLayer(48),torch.nn.ReLU(True),
linearLayer(opt.labelN)
)
initialize(opt,self,opt.stdC)
def forward(self,opt,image):
feat = image
feat = self.conv2Layers(feat).reshape(opt.batchSize,-1)
feat = self.linearLayers(feat)
output = feat
return output
# build classification network
class CNN(torch.nn.Module):
def __init__(self,opt):
super(CNN,self).__init__()
self.inDim = 1
def conv2Layer(outDim):
conv = torch.nn.Conv2d(self.inDim,outDim,kernel_size=[9,9],stride=1,padding=0)
self.inDim = outDim
return conv
def linearLayer(outDim):
fc = torch.nn.Linear(self.inDim,outDim)
self.inDim = outDim
return fc
def maxpoolLayer(): return torch.nn.MaxPool2d([2,2],stride=2)
self.conv2Layers = torch.nn.Sequential(
conv2Layer(3),torch.nn.ReLU(True)
)
self.inDim *= 20**2
self.linearLayers = torch.nn.Sequential(
linearLayer(opt.labelN)
)
initialize(opt,self,opt.stdC)
def forward(self,opt,image):
feat = image
feat = self.conv2Layers(feat).reshape(opt.batchSize,-1)
feat = self.linearLayers(feat)
output = feat
return output
# an identity class to skip geometric predictors
class Identity(torch.nn.Module):
def __init__(self): super(Identity,self).__init__()
def forward(self,opt,feat): return [feat]
# build Spatial Transformer Network
class STN(torch.nn.Module):
def __init__(self,opt):
super(STN,self).__init__()
self.inDim = 1
def conv2Layer(outDim):
conv = torch.nn.Conv2d(self.inDim,outDim,kernel_size=[7,7],stride=1,padding=0)
self.inDim = outDim
return conv
def linearLayer(outDim):
fc = torch.nn.Linear(self.inDim,outDim)
self.inDim = outDim
return fc
def maxpoolLayer(): return torch.nn.MaxPool2d([2,2],stride=2)
self.conv2Layers = torch.nn.Sequential(
conv2Layer(4),torch.nn.ReLU(True),
conv2Layer(8),torch.nn.ReLU(True),maxpoolLayer()
)
self.inDim *= 8**2
self.linearLayers = torch.nn.Sequential(
linearLayer(48),torch.nn.ReLU(True),
linearLayer(opt.warpDim)
)
initialize(opt,self,opt.stdGP,last0=True)
def forward(self,opt,image):
imageWarpAll = [image]
feat = image
feat = self.conv2Layers(feat).reshape(opt.batchSize,-1)
feat = self.linearLayers(feat)
p = feat
pMtrx = warp.vec2mtrx(opt,p)
imageWarp = warp.transformImage(opt,image,pMtrx)
imageWarpAll.append(imageWarp)
return imageWarpAll
# build Inverse Compositional STN
class ICSTN(torch.nn.Module):
def __init__(self,opt):
super(ICSTN,self).__init__()
self.inDim = 1
def conv2Layer(outDim):
conv = torch.nn.Conv2d(self.inDim,outDim,kernel_size=[7,7],stride=1,padding=0)
self.inDim = outDim
return conv
def linearLayer(outDim):
fc = torch.nn.Linear(self.inDim,outDim)
self.inDim = outDim
return fc
def maxpoolLayer(): return torch.nn.MaxPool2d([2,2],stride=2)
self.conv2Layers = torch.nn.Sequential(
conv2Layer(4),torch.nn.ReLU(True),
conv2Layer(8),torch.nn.ReLU(True),maxpoolLayer()
)
self.inDim *= 8**2
self.linearLayers = torch.nn.Sequential(
linearLayer(48),torch.nn.ReLU(True),
linearLayer(opt.warpDim)
)
initialize(opt,self,opt.stdGP,last0=True)
def forward(self,opt,image,p):
imageWarpAll = []
for l in range(opt.warpN):
pMtrx = warp.vec2mtrx(opt,p)
imageWarp = warp.transformImage(opt,image,pMtrx)
imageWarpAll.append(imageWarp)
feat = imageWarp
feat = self.conv2Layers(feat).reshape(opt.batchSize,-1)
feat = self.linearLayers(feat)
dp = feat
p = warp.compose(opt,p,dp)
pMtrx = warp.vec2mtrx(opt,p)
imageWarp = warp.transformImage(opt,image,pMtrx)
imageWarpAll.append(imageWarp)
return imageWarpAll
# initialize weights/biases
def initialize(opt,model,stddev,last0=False):
for m in model.conv2Layers:
if isinstance(m,torch.nn.Conv2d):
m.weight.data.normal_(0,stddev)
m.bias.data.normal_(0,stddev)
for m in model.linearLayers:
if isinstance(m,torch.nn.Linear):
if last0 and m is model.linearLayers[-1]:
m.weight.data.zero_()
m.bias.data.zero_()
else:
m.weight.data.normal_(0,stddev)
m.bias.data.normal_(0,stddev)
================================================
FILE: MNIST-pytorch/options.py
================================================
import numpy as np
import argparse
import warp
import util
import torch
def set(training):
# parse input arguments
parser = argparse.ArgumentParser()
parser.add_argument("netType", choices=["CNN","STN","IC-STN"], help="type of network")
parser.add_argument("--group", default="0", help="name for group")
parser.add_argument("--model", default="test", help="name for model instance")
parser.add_argument("--size", default="28x28", help="image resolution")
parser.add_argument("--warpType", default="homography", help="type of warp function on images",
choices=["translation","similarity","affine","homography"])
parser.add_argument("--warpN", type=int, default=4, help="number of recurrent transformations (for IC-STN)")
parser.add_argument("--stdC", type=float, default=0.1, help="initialization stddev (classification network)")
parser.add_argument("--stdGP", type=float, default=0.1, help="initialization stddev (geometric predictor)")
parser.add_argument("--pertScale", type=float, default=0.25, help="initial perturbation scale")
parser.add_argument("--transScale", type=float, default=0.25, help="initial translation scale")
if training: # training
parser.add_argument("--port", type=int, default=8097, help="port number for visdom visualization")
parser.add_argument("--batchSize", type=int, default=100, help="batch size for SGD")
parser.add_argument("--lrC", type=float, default=1e-2, help="learning rate (classification network)")
parser.add_argument("--lrGP", type=float, default=None, help="learning rate (geometric predictor)")
parser.add_argument("--lrDecay", type=float, default=1.0, help="learning rate decay")
parser.add_argument("--lrStep", type=int, default=100000, help="learning rate decay step size")
parser.add_argument("--fromIt", type=int, default=0, help="resume training from iteration number")
parser.add_argument("--toIt", type=int, default=500000, help="run training to iteration number")
else: # evaluation
parser.add_argument("--batchSize", type=int, default=1, help="batch size for evaluation")
opt = parser.parse_args()
if opt.lrGP is None: opt.lrGP = 0 if opt.netType=="CNN" else \
1e-2 if opt.netType=="STN" else \
1e-4 if opt.netType=="IC-STN" else None
# --- below are automatically set ---
assert(torch.cuda.is_available()) # support only training on GPU for now
torch.set_default_tensor_type("torch.cuda.FloatTensor")
opt.training = training
opt.H,opt.W = [int(x) for x in opt.size.split("x")]
opt.visBlockSize = int(np.floor(np.sqrt(opt.batchSize)))
opt.warpDim = 2 if opt.warpType == "translation" else \
4 if opt.warpType == "similarity" else \
6 if opt.warpType == "affine" else \
8 if opt.warpType == "homography" else None
opt.labelN = 10
opt.canon4pts = np.array([[-1,-1],[-1,1],[1,1],[1,-1]],dtype=np.float32)
opt.image4pts = np.array([[0,0],[0,opt.H-1],[opt.W-1,opt.H-1],[opt.W-1,0]],dtype=np.float32)
opt.refMtrx = np.eye(3).astype(np.float32)
if opt.netType=="STN": opt.warpN = 1
print("({0}) {1}".format(
util.toGreen("{0}".format(opt.group)),
util.toGreen("{0}".format(opt.model))))
print("------------------------------------------")
print("network type: {0}, recurrent warps: {1}".format(
util.toYellow("{0}".format(opt.netType)),
util.toYellow("{0}".format(opt.warpN if opt.netType=="IC-STN" else "X"))))
print("batch size: {0}, image size: {1}x{2}".format(
util.toYellow("{0}".format(opt.batchSize)),
util.toYellow("{0}".format(opt.H)),
util.toYellow("{0}".format(opt.W))))
print("warpScale: (pert) {0} (trans) {1}".format(
util.toYellow("{0}".format(opt.pertScale)),
util.toYellow("{0}".format(opt.transScale))))
if training:
print("[geometric predictor] stddev={0}, lr={1}".format(
util.toYellow("{0:.0e}".format(opt.stdGP)),
util.toYellow("{0:.0e}".format(opt.lrGP))))
print("[classification network] stddev={0}, lr={1}".format(
util.toYellow("{0:.0e}".format(opt.stdC)),
util.toYellow("{0:.0e}".format(opt.lrC))))
print("------------------------------------------")
if training:
print(util.toMagenta("training model ({0}) {1}...".format(opt.group,opt.model)))
return opt
================================================
FILE: MNIST-pytorch/train.py
================================================
import numpy as np
import time,os,sys
import argparse
import util
print(util.toYellow("======================================================="))
print(util.toYellow("train.py (training on MNIST)"))
print(util.toYellow("======================================================="))
import torch
import data,graph,warp,util
import options
print(util.toMagenta("setting configurations..."))
opt = options.set(training=True)
# create directories for model output
util.mkdir("models_{0}".format(opt.group))
print(util.toMagenta("building network..."))
with torch.cuda.device(0):
# ------ build network ------
if opt.netType=="CNN":
geometric = graph.Identity()
classifier = graph.FullCNN(opt)
elif opt.netType=="STN":
geometric = graph.STN(opt)
classifier = graph.CNN(opt)
elif opt.netType=="IC-STN":
geometric = graph.ICSTN(opt)
classifier = graph.CNN(opt)
# ------ define loss ------
loss = torch.nn.CrossEntropyLoss()
# ------ optimizer ------
optimList = [{ "params": geometric.parameters(), "lr": opt.lrGP },
{ "params": classifier.parameters(), "lr": opt.lrC }]
optim = torch.optim.SGD(optimList)
# load data
print(util.toMagenta("loading MNIST dataset..."))
trainData,testData = data.loadMNIST(opt,"data")
# visdom visualizer
vis = util.Visdom(opt)
print(util.toYellow("======= TRAINING START ======="))
timeStart = time.time()
# start session
with torch.cuda.device(0):
geometric.train()
classifier.train()
if opt.fromIt!=0:
util.restoreModel(opt,geometric,classifier,opt.fromIt)
print(util.toMagenta("resuming from iteration {0}...".format(opt.fromIt)))
print(util.toMagenta("start training..."))
# training loop
for i in range(opt.fromIt,opt.toIt):
lrGP = opt.lrGP*opt.lrDecay**(i//opt.lrStep)
lrC = opt.lrC*opt.lrDecay**(i//opt.lrStep)
# make training batch
batch = data.makeBatch(opt,trainData)
image = batch["image"].unsqueeze(dim=1)
label = batch["label"]
# generate perturbation
pInit = data.genPerturbations(opt)
pInitMtrx = warp.vec2mtrx(opt,pInit)
# forward/backprop through network
optim.zero_grad()
imagePert = warp.transformImage(opt,image,pInitMtrx)
imageWarpAll = geometric(opt,image,pInit) if opt.netType=="IC-STN" else geometric(opt,imagePert)
imageWarp = imageWarpAll[-1]
output = classifier(opt,imageWarp)
train_loss = loss(output,label)
train_loss.backward()
# run one step
optim.step()
if (i+1)%100==0:
print("it. {0}/{1} lr={3}(GP),{4}(C), loss={5}, time={2}"
.format(util.toCyan("{0}".format(i+1)),
opt.toIt,
util.toGreen("{0:.2f}".format(time.time()-timeStart)),
util.toYellow("{0:.0e}".format(lrGP)),
util.toYellow("{0:.0e}".format(lrC)),
util.toRed("{0:.4f}".format(train_loss))))
if (i+1)%200==0: vis.trainLoss(opt,i+1,train_loss)
if (i+1)%1000==0:
# evaluate on test set
testAcc,testMean,testVar = data.evalTest(opt,testData,geometric,classifier)
testError = (1-testAcc)*100
vis.testLoss(opt,i+1,testError)
if opt.netType=="STN" or opt.netType=="IC-STN":
vis.meanVar(opt,testMean,testVar)
if (i+1)%10000==0:
util.saveModel(opt,geometric,classifier,i+1)
print(util.toGreen("model saved: {0}/{1}, it.{2}".format(opt.group,opt.model,i+1)))
print(util.toYellow("======= TRAINING DONE ======="))
================================================
FILE: MNIST-pytorch/util.py
================================================
import numpy as np
import scipy.misc
import torch
import os
import termcolor
import visdom
def mkdir(path):
if not os.path.exists(path): os.mkdir(path)
def imread(fname):
return scipy.misc.imread(fname)/255.0
def imsave(fname,array):
scipy.misc.toimage(array,cmin=0.0,cmax=1.0).save(fname)
# convert to colored strings
def toRed(content): return termcolor.colored(content,"red",attrs=["bold"])
def toGreen(content): return termcolor.colored(content,"green",attrs=["bold"])
def toBlue(content): return termcolor.colored(content,"blue",attrs=["bold"])
def toCyan(content): return termcolor.colored(content,"cyan",attrs=["bold"])
def toYellow(content): return termcolor.colored(content,"yellow",attrs=["bold"])
def toMagenta(content): return termcolor.colored(content,"magenta",attrs=["bold"])
# restore model
def restoreModel(opt,geometric,classifier,it):
geometric.load_state_dict(torch.load("models_{0}/{1}_it{2}_GP.npy".format(opt.group,opt.model,it)))
classifier.load_state_dict(torch.load("models_{0}/{1}_it{2}_C.npy".format(opt.group,opt.model,it)))
# save model
def saveModel(opt,geometric,classifier,it):
torch.save(geometric.state_dict(),"models_{0}/{1}_it{2}_GP.npy".format(opt.group,opt.model,it))
torch.save(classifier.state_dict(),"models_{0}/{1}_it{2}_C.npy".format(opt.group,opt.model,it))
class Visdom():
def __init__(self,opt):
self.vis = visdom.Visdom(port=opt.port,use_incoming_socket=False)
self.trainLossInit = True
self.testLossInit = True
self.meanVarInit = True
def tileImages(self,opt,images,H,W,HN,WN):
assert(len(images)==HN*WN)
images = images.reshape([HN,WN,-1,H,W])
images = [list(i) for i in images]
imageBlocks = np.concatenate([np.concatenate(row,axis=2) for row in images],axis=1)
return imageBlocks
def trainLoss(self,opt,it,loss):
loss = float(loss.detach().cpu().numpy())
if self.trainLossInit:
self.vis.line(Y=np.array([loss]),X=np.array([it]),win="{0}_trainloss".format(opt.model),
opts={ "title": "{0} (TRAIN_loss)".format(opt.model) })
self.trainLossInit = False
else: self.vis.line(Y=np.array([loss]),X=np.array([it]),win=opt.model+"_trainloss",update="append")
def testLoss(self,opt,it,loss):
if self.testLossInit:
self.vis.line(Y=np.array([loss]),X=np.array([it]),win="{0}_testloss".format(opt.model),
opts={ "title": "{0} (TEST_error)".format(opt.model) })
self.testLossInit = False
else: self.vis.line(Y=np.array([loss]),X=np.array([it]),win=opt.model+"_testloss",update="append")
def meanVar(self,opt,mean,var):
mean = [self.tileImages(opt,m,opt.H,opt.W,1,10) for m in mean]
var = [self.tileImages(opt,v,opt.H,opt.W,1,10)*3 for v in var]
self.vis.image(mean[0].clip(0,1),win="{0}_meaninit".format(opt.model), opts={ "title": "{0} (TEST_mean_init)".format(opt.model) })
self.vis.image(mean[1].clip(0,1),win="{0}_meanwarped".format(opt.model), opts={ "title": "{0} (TEST_mean_warped)".format(opt.model) })
self.vis.image(var[0].clip(0,1),win="{0}_varinit".format(opt.model), opts={ "title": "{0} (TEST_var_init)".format(opt.model) })
self.vis.image(var[1].clip(0,1),win="{0}_varwarped".format(opt.model), opts={ "title": "{0} (TEST_var_warped)".format(opt.model) })
================================================
FILE: MNIST-pytorch/warp.py
================================================
import numpy as np
import scipy.linalg
import torch
import util
# fit (affine) warp between two sets of points
def fit(Xsrc,Xdst):
ptsN = len(Xsrc)
X,Y,U,V,O,I = Xsrc[:,0],Xsrc[:,1],Xdst[:,0],Xdst[:,1],np.zeros([ptsN]),np.ones([ptsN])
A = np.concatenate((np.stack([X,Y,I,O,O,O],axis=1),
np.stack([O,O,O,X,Y,I],axis=1)),axis=0)
b = np.concatenate((U,V),axis=0)
p1,p2,p3,p4,p5,p6 = scipy.linalg.lstsq(A,b)[0].squeeze()
pMtrx = np.array([[p1,p2,p3],[p4,p5,p6],[0,0,1]],dtype=torch.float32)
return pMtrx
# compute composition of warp parameters
def compose(opt,p,dp):
pMtrx = vec2mtrx(opt,p)
dpMtrx = vec2mtrx(opt,dp)
pMtrxNew = dpMtrx.matmul(pMtrx)
pMtrxNew = pMtrxNew/pMtrxNew[:,2:3,2:3]
pNew = mtrx2vec(opt,pMtrxNew)
return pNew
# compute inverse of warp parameters
def inverse(opt,p):
pMtrx = vec2mtrx(opt,p)
pInvMtrx = pMtrx.inverse()
pInv = mtrx2vec(opt,pInvMtrx)
return pInv
# convert warp parameters to matrix
def vec2mtrx(opt,p):
O = torch.zeros(opt.batchSize,dtype=torch.float32).cuda()
I = torch.ones(opt.batchSize,dtype=torch.float32).cuda()
if opt.warpType=="translation":
tx,ty = torch.unbind(p,dim=1)
pMtrx = torch.stack([torch.stack([I,O,tx],dim=-1),
torch.stack([O,I,ty],dim=-1),
torch.stack([O,O,I],dim=-1)],dim=1)
if opt.warpType=="similarity":
pc,ps,tx,ty = torch.unbind(p,dim=1)
pMtrx = torch.stack([torch.stack([I+pc,-ps,tx],dim=-1),
torch.stack([ps,I+pc,ty],dim=-1),
torch.stack([O,O,I],dim=-1)],dim=1)
if opt.warpType=="affine":
p1,p2,p3,p4,p5,p6 = torch.unbind(p,dim=1)
pMtrx = torch.stack([torch.stack([I+p1,p2,p3],dim=-1),
torch.stack([p4,I+p5,p6],dim=-1),
torch.stack([O,O,I],dim=-1)],dim=1)
if opt.warpType=="homography":
p1,p2,p3,p4,p5,p6,p7,p8 = torch.unbind(p,dim=1)
pMtrx = torch.stack([torch.stack([I+p1,p2,p3],dim=-1),
torch.stack([p4,I+p5,p6],dim=-1),
torch.stack([p7,p8,I],dim=-1)],dim=1)
return pMtrx
# convert warp matrix to parameters
def mtrx2vec(opt,pMtrx):
[row0,row1,row2] = torch.unbind(pMtrx,dim=1)
[e00,e01,e02] = torch.unbind(row0,dim=1)
[e10,e11,e12] = torch.unbind(row1,dim=1)
[e20,e21,e22] = torch.unbind(row2,dim=1)
if opt.warpType=="translation": p = torch.stack([e02,e12],dim=1)
if opt.warpType=="similarity": p = torch.stack([e00-1,e10,e02,e12],dim=1)
if opt.warpType=="affine": p = torch.stack([e00-1,e01,e02,e10,e11-1,e12],dim=1)
if opt.warpType=="homography": p = torch.stack([e00-1,e01,e02,e10,e11-1,e12,e20,e21],dim=1)
return p
# warp the image
def transformImage(opt,image,pMtrx):
refMtrx = torch.from_numpy(opt.refMtrx).cuda()
refMtrx = refMtrx.repeat(opt.batchSize,1,1)
transMtrx = refMtrx.matmul(pMtrx)
# warp the canonical coordinates
X,Y = np.meshgrid(np.linspace(-1,1,opt.W),np.linspace(-1,1,opt.H))
X,Y = X.flatten(),Y.flatten()
XYhom = np.stack([X,Y,np.ones_like(X)],axis=1).T
XYhom = np.tile(XYhom,[opt.batchSize,1,1]).astype(np.float32)
XYhom = torch.from_numpy(XYhom).cuda()
XYwarpHom = transMtrx.matmul(XYhom)
XwarpHom,YwarpHom,ZwarpHom = torch.unbind(XYwarpHom,dim=1)
Xwarp = (XwarpHom/(ZwarpHom+1e-8)).reshape(opt.batchSize,opt.H,opt.W)
Ywarp = (YwarpHom/(ZwarpHom+1e-8)).reshape(opt.batchSize,opt.H,opt.W)
grid = torch.stack([Xwarp,Ywarp],dim=-1)
# sampling with bilinear interpolation
imageWarp = torch.nn.functional.grid_sample(image,grid,mode="bilinear")
return imageWarp
================================================
FILE: MNIST-tensorflow/data.py
================================================
import numpy as np
import scipy.linalg
import os,time
import tensorflow as tf
import warp
# load MNIST data
def loadMNIST(fname):
if not os.path.exists(fname):
# download and preprocess MNIST dataset
from tensorflow.examples.tutorials.mnist import input_data
mnist = input_data.read_data_sets("MNIST_data/",one_hot=True)
trainData,validData,testData = {},{},{}
trainData["image"] = mnist.train.images.reshape([-1,28,28]).astype(np.float32)
validData["image"] = mnist.validation.images.reshape([-1,28,28]).astype(np.float32)
testData["image"] = mnist.test.images.reshape([-1,28,28]).astype(np.float32)
trainData["label"] = np.argmax(mnist.train.labels.astype(np.float32),axis=1)
validData["label"] = np.argmax(mnist.validation.labels.astype(np.float32),axis=1)
testData["label"] = np.argmax(mnist.test.labels.astype(np.float32),axis=1)
os.makedirs(os.path.dirname(fname))
np.savez(fname,train=trainData,valid=validData,test=testData)
os.system("rm -rf MNIST_data")
MNIST = np.load(fname)
trainData = MNIST["train"].item()
validData = MNIST["valid"].item()
testData = MNIST["test"].item()
return trainData,validData,testData
# generate training batch
def genPerturbations(opt):
with tf.name_scope("genPerturbations"):
X = np.tile(opt.canon4pts[:,0],[opt.batchSize,1])
Y = np.tile(opt.canon4pts[:,1],[opt.batchSize,1])
dX = tf.random_normal([opt.batchSize,4])*opt.pertScale \
+tf.random_normal([opt.batchSize,1])*opt.transScale
dY = tf.random_normal([opt.batchSize,4])*opt.pertScale \
+tf.random_normal([opt.batchSize,1])*opt.transScale
O = np.zeros([opt.batchSize,4],dtype=np.float32)
I = np.ones([opt.batchSize,4],dtype=np.float32)
# fit warp parameters to generated displacements
if opt.warpType=="homography":
A = tf.concat([tf.stack([X,Y,I,O,O,O,-X*(X+dX),-Y*(X+dX)],axis=-1),
tf.stack([O,O,O,X,Y,I,-X*(Y+dY),-Y*(Y+dY)],axis=-1)],1)
b = tf.expand_dims(tf.concat([X+dX,Y+dY],1),-1)
pPert = tf.matrix_solve(A,b)[:,:,0]
pPert -= tf.to_float([[1,0,0,0,1,0,0,0]])
else:
if opt.warpType=="translation":
J = np.concatenate([np.stack([I,O],axis=-1),
np.stack([O,I],axis=-1)],axis=1)
if opt.warpType=="similarity":
J = np.concatenate([np.stack([X,Y,I,O],axis=-1),
np.stack([-Y,X,O,I],axis=-1)],axis=1)
if opt.warpType=="affine":
J = np.concatenate([np.stack([X,Y,I,O,O,O],axis=-1),
np.stack([O,O,O,X,Y,I],axis=-1)],axis=1)
dXY = tf.expand_dims(tf.concat([dX,dY],1),-1)
pPert = tf.matrix_solve_ls(J,dXY)[:,:,0]
return pPert
# make training batch
def makeBatch(opt,data,PH):
N = len(data["image"])
randIdx = np.random.randint(N,size=[opt.batchSize])
# put data in placeholders
[image,label] = PH
batch = {
image: data["image"][randIdx],
label: data["label"][randIdx],
}
return batch
# evaluation on test set
def evalTest(opt,sess,data,PH,prediction,imagesEval=[]):
N = len(data["image"])
# put data in placeholders
[image,label] = PH
batchN = int(np.ceil(N/opt.batchSize))
warped = [{},{}]
count = 0
for b in range(batchN):
# use some dummy data (0) as batch filler if necessary
if b!=batchN-1:
realIdx = np.arange(opt.batchSize*b,opt.batchSize*(b+1))
else:
realIdx = np.arange(opt.batchSize*b,N)
idx = np.zeros([opt.batchSize],dtype=int)
idx[:len(realIdx)] = realIdx
batch = {
image: data["image"][idx],
label: data["label"][idx],
}
evalList = sess.run([prediction]+imagesEval,feed_dict=batch)
pred = evalList[0]
count += pred[:len(realIdx)].sum()
if opt.netType=="STN" or opt.netType=="IC-STN":
imgs = evalList[1:]
for i in range(len(realIdx)):
l = data["label"][idx[i]]
if l not in warped[0]: warped[0][l] = []
if l not in warped[1]: warped[1][l] = []
warped[0][l].append(imgs[0][i])
warped[1][l].append(imgs[1][i])
accuracy = float(count)/N
if opt.netType=="STN" or opt.netType=="IC-STN":
mean = [np.array([np.mean(warped[0][l],axis=0) for l in warped[0]]),
np.array([np.mean(warped[1][l],axis=0) for l in warped[1]])]
var = [np.array([np.var(warped[0][l],axis=0) for l in warped[0]]),
np.array([np.var(warped[1][l],axis=0) for l in warped[1]])]
else: mean,var = None,None
return accuracy,mean,var
================================================
FILE: MNIST-tensorflow/graph.py
================================================
import numpy as np
import tensorflow as tf
import time
import data,warp,util
# build classification network
def fullCNN(opt,image):
def conv2Layer(opt,feat,outDim):
weight,bias = createVariable(opt,[3,3,int(feat.shape[-1]),outDim],stddev=opt.stdC)
conv = tf.nn.conv2d(feat,weight,strides=[1,1,1,1],padding="VALID")+bias
return conv
def linearLayer(opt,feat,outDim):
weight,bias = createVariable(opt,[int(feat.shape[-1]),outDim],stddev=opt.stdC)
fc = tf.matmul(feat,weight)+bias
return fc
with tf.variable_scope("classifier"):
feat = image
with tf.variable_scope("conv1"):
feat = conv2Layer(opt,feat,3)
feat = tf.nn.relu(feat)
with tf.variable_scope("conv2"):
feat = conv2Layer(opt,feat,6)
feat = tf.nn.relu(feat)
feat = tf.nn.max_pool(feat,ksize=[1,2,2,1],strides=[1,2,2,1],padding="VALID")
with tf.variable_scope("conv3"):
feat = conv2Layer(opt,feat,9)
feat = tf.nn.relu(feat)
with tf.variable_scope("conv4"):
feat = conv2Layer(opt,feat,12)
feat = tf.nn.relu(feat)
feat = tf.reshape(feat,[opt.batchSize,-1])
with tf.variable_scope("fc5"):
feat = linearLayer(opt,feat,48)
feat = tf.nn.relu(feat)
with tf.variable_scope("fc6"):
feat = linearLayer(opt,feat,opt.labelN)
output = feat
return output
# build classification network
def CNN(opt,image):
def conv2Layer(opt,feat,outDim):
weight,bias = createVariable(opt,[9,9,int(feat.shape[-1]),outDim],stddev=opt.stdC)
conv = tf.nn.conv2d(feat,weight,strides=[1,1,1,1],padding="VALID")+bias
return conv
def linearLayer(opt,feat,outDim):
weight,bias = createVariable(opt,[int(feat.shape[-1]),outDim],stddev=opt.stdC)
fc = tf.matmul(feat,weight)+bias
return fc
with tf.variable_scope("classifier"):
feat = image
with tf.variable_scope("conv1"):
feat = conv2Layer(opt,feat,3)
feat = tf.nn.relu(feat)
feat = tf.reshape(feat,[opt.batchSize,-1])
with tf.variable_scope("fc2"):
feat = linearLayer(opt,feat,opt.labelN)
output = feat
return output
# build Spatial Transformer Network
def STN(opt,image):
def conv2Layer(opt,feat,outDim):
weight,bias = createVariable(opt,[7,7,int(feat.shape[-1]),outDim],stddev=opt.stdGP)
conv = tf.nn.conv2d(feat,weight,strides=[1,1,1,1],padding="VALID")+bias
return conv
def linearLayer(opt,feat,outDim,final=False):
weight,bias = createVariable(opt,[int(feat.shape[-1]),outDim],stddev=0.0 if final else opt.stdGP)
fc = tf.matmul(feat,weight)+bias
return fc
imageWarpAll = [image]
with tf.variable_scope("geometric"):
feat = image
with tf.variable_scope("conv1"):
feat = conv2Layer(opt,feat,4)
feat = tf.nn.relu(feat)
with tf.variable_scope("conv2"):
feat = conv2Layer(opt,feat,8)
feat = tf.nn.relu(feat)
feat = tf.nn.max_pool(feat,ksize=[1,2,2,1],strides=[1,2,2,1],padding="VALID")
feat = tf.reshape(feat,[opt.batchSize,-1])
with tf.variable_scope("fc3"):
feat = linearLayer(opt,feat,48)
feat = tf.nn.relu(feat)
with tf.variable_scope("fc4"):
feat = linearLayer(opt,feat,opt.warpDim,final=True)
p = feat
pMtrx = warp.vec2mtrx(opt,p)
imageWarp = warp.transformImage(opt,image,pMtrx)
imageWarpAll.append(imageWarp)
return imageWarpAll
# build Inverse Compositional STN
def ICSTN(opt,image,p):
def conv2Layer(opt,feat,outDim):
weight,bias = createVariable(opt,[7,7,int(feat.shape[-1]),outDim],stddev=opt.stdGP)
conv = tf.nn.conv2d(feat,weight,strides=[1,1,1,1],padding="VALID")+bias
return conv
def linearLayer(opt,feat,outDim,final=False):
weight,bias = createVariable(opt,[int(feat.shape[-1]),outDim],stddev=0.0 if final else opt.stdGP)
fc = tf.matmul(feat,weight)+bias
return fc
imageWarpAll = []
for l in range(opt.warpN):
with tf.variable_scope("geometric",reuse=l>0):
pMtrx = warp.vec2mtrx(opt,p)
imageWarp = warp.transformImage(opt,image,pMtrx)
imageWarpAll.append(imageWarp)
feat = imageWarp
with tf.variable_scope("conv1"):
feat = conv2Layer(opt,feat,4)
feat = tf.nn.relu(feat)
with tf.variable_scope("conv2"):
feat = conv2Layer(opt,feat,8)
feat = tf.nn.relu(feat)
feat = tf.nn.max_pool(feat,ksize=[1,2,2,1],strides=[1,2,2,1],padding="VALID")
feat = tf.reshape(feat,[opt.batchSize,-1])
with tf.variable_scope("fc3"):
feat = linearLayer(opt,feat,48)
feat = tf.nn.relu(feat)
with tf.variable_scope("fc4"):
feat = linearLayer(opt,feat,opt.warpDim,final=True)
dp = feat
p = warp.compose(opt,p,dp)
pMtrx = warp.vec2mtrx(opt,p)
imageWarp = warp.transformImage(opt,image,pMtrx)
imageWarpAll.append(imageWarp)
return imageWarpAll
# auxiliary function for creating weight and bias
def createVariable(opt,weightShape,biasShape=None,stddev=None):
if biasShape is None: biasShape = [weightShape[-1]]
weight = tf.get_variable("weight",shape=weightShape,dtype=tf.float32,
initializer=tf.random_normal_initializer(stddev=stddev))
bias = tf.get_variable("bias",shape=biasShape,dtype=tf.float32,
initializer=tf.random_normal_initializer(stddev=stddev))
return weight,bias
================================================
FILE: MNIST-tensorflow/options.py
================================================
import numpy as np
import argparse
import warp
import util
def set(training):
# parse input arguments
parser = argparse.ArgumentParser()
parser.add_argument("netType", choices=["CNN","STN","IC-STN"], help="type of network")
parser.add_argument("--group", default="0", help="name for group")
parser.add_argument("--model", default="test", help="name for model instance")
parser.add_argument("--size", default="28x28", help="image resolution")
parser.add_argument("--warpType", default="homography", help="type of warp function on images",
choices=["translation","similarity","affine","homography"])
parser.add_argument("--warpN", type=int, default=4, help="number of recurrent transformations (for IC-STN)")
parser.add_argument("--stdC", type=float, default=0.1, help="initialization stddev (classification network)")
parser.add_argument("--stdGP", type=float, default=0.1, help="initialization stddev (geometric predictor)")
parser.add_argument("--pertScale", type=float, default=0.25, help="initial perturbation scale")
parser.add_argument("--transScale", type=float, default=0.25, help="initial translation scale")
if training: # training
parser.add_argument("--batchSize", type=int, default=100, help="batch size for SGD")
parser.add_argument("--lrC", type=float, default=1e-2, help="learning rate (classification network)")
parser.add_argument("--lrCdecay", type=float, default=1.0, help="learning rate decay (classification network)")
parser.add_argument("--lrCstep", type=int, default=100000, help="learning rate decay step size (classification network)")
parser.add_argument("--lrGP", type=float, default=None, help="learning rate (geometric predictor)")
parser.add_argument("--lrGPdecay", type=float, default=1.0, help="learning rate decay (geometric predictor)")
parser.add_argument("--lrGPstep", type=int, default=100000, help="learning rate decay step size (geometric predictor)")
parser.add_argument("--fromIt", type=int, default=0, help="resume training from iteration number")
parser.add_argument("--toIt", type=int, default=500000, help="run training to iteration number")
else: # evaluation
parser.add_argument("--batchSize", type=int, default=1, help="batch size for evaluation")
opt = parser.parse_args()
if opt.lrGP is None: opt.lrGP = 0 if opt.netType=="CNN" else \
1e-2 if opt.netType=="STN" else \
1e-4 if opt.netType=="IC-STN" else None
# --- below are automatically set ---
opt.training = training
opt.H,opt.W = [int(x) for x in opt.size.split("x")]
opt.visBlockSize = int(np.floor(np.sqrt(opt.batchSize)))
opt.warpDim = 2 if opt.warpType == "translation" else \
4 if opt.warpType == "similarity" else \
6 if opt.warpType == "affine" else \
8 if opt.warpType == "homography" else None
opt.labelN = 10
opt.canon4pts = np.array([[-1,-1],[-1,1],[1,1],[1,-1]],dtype=np.float32)
opt.image4pts = np.array([[0,0],[0,opt.H-1],[opt.W-1,opt.H-1],[opt.W-1,0]],dtype=np.float32)
opt.refMtrx = warp.fit(Xsrc=opt.canon4pts,Xdst=opt.image4pts)
if opt.netType=="STN": opt.warpN = 1
print("({0}) {1}".format(
util.toGreen("{0}".format(opt.group)),
util.toGreen("{0}".format(opt.model))))
print("------------------------------------------")
print("network type: {0}, recurrent warps: {1}".format(
util.toYellow("{0}".format(opt.netType)),
util.toYellow("{0}".format(opt.warpN if opt.netType=="IC-STN" else "X"))))
print("batch size: {0}, image size: {1}x{2}".format(
util.toYellow("{0}".format(opt.batchSize)),
util.toYellow("{0}".format(opt.H)),
util.toYellow("{0}".format(opt.W))))
print("warpScale: (pert) {0} (trans) {1}".format(
util.toYellow("{0}".format(opt.pertScale)),
util.toYellow("{0}".format(opt.transScale))))
if training:
print("[geometric predictor] stddev={0}, lr={1}".format(
util.toYellow("{0:.0e}".format(opt.stdGP)),
util.toYellow("{0:.0e}".format(opt.lrGP))))
print("[classification network] stddev={0}, lr={1}".format(
util.toYellow("{0:.0e}".format(opt.stdC)),
util.toYellow("{0:.0e}".format(opt.lrC))))
print("------------------------------------------")
if training:
print(util.toMagenta("training model ({0}) {1}...".format(opt.group,opt.model)))
return opt
================================================
FILE: MNIST-tensorflow/train.py
================================================
import numpy as np
import time,os,sys
import argparse
import util
print(util.toYellow("======================================================="))
print(util.toYellow("train.py (training on MNIST)"))
print(util.toYellow("======================================================="))
import tensorflow as tf
import data,graph,warp,util
import options
print(util.toMagenta("setting configurations..."))
opt = options.set(training=True)
# create directories for model output
util.mkdir("models_{0}".format(opt.group))
print(util.toMagenta("building graph..."))
tf.reset_default_graph()
# build graph
with tf.device("/gpu:0"):
# ------ define input data ------
image = tf.placeholder(tf.float32,shape=[opt.batchSize,opt.H,opt.W])
label = tf.placeholder(tf.int64,shape=[opt.batchSize])
PH = [image,label]
# ------ generate perturbation ------
pInit = data.genPerturbations(opt)
pInitMtrx = warp.vec2mtrx(opt,pInit)
# ------ build network ------
image = tf.expand_dims(image,axis=-1)
imagePert = warp.transformImage(opt,image,pInitMtrx)
if opt.netType=="CNN":
output = graph.fullCNN(opt,imagePert)
elif opt.netType=="STN":
imageWarpAll = graph.STN(opt,imagePert)
imageWarp = imageWarpAll[-1]
output = graph.CNN(opt,imageWarp)
elif opt.netType=="IC-STN":
imageWarpAll = graph.ICSTN(opt,image,pInit)
imageWarp = imageWarpAll[-1]
output = graph.CNN(opt,imageWarp)
softmax = tf.nn.softmax(output)
labelOnehot = tf.one_hot(label,opt.labelN)
prediction = tf.equal(tf.argmax(softmax,1),label)
# ------ define loss ------
softmaxLoss = tf.nn.softmax_cross_entropy_with_logits(logits=output,labels=labelOnehot)
loss = tf.reduce_mean(softmaxLoss)
# ------ optimizer ------
lrGP_PH,lrC_PH = tf.placeholder(tf.float32,shape=[]),tf.placeholder(tf.float32,shape=[])
optim = util.setOptimizer(opt,loss,lrGP_PH,lrC_PH)
# ------ generate summaries ------
summaryImageTrain = []
summaryImageTest = []
if opt.netType=="STN" or opt.netType=="IC-STN":
for l in range(opt.warpN+1):
summaryImageTrain.append(util.imageSummary(opt,imageWarpAll[l],"TRAIN_warp{0}".format(l),opt.H,opt.W))
summaryImageTest.append(util.imageSummary(opt,imageWarpAll[l],"TEST_warp{0}".format(l),opt.H,opt.W))
summaryImageTrain = tf.summary.merge(summaryImageTrain)
summaryImageTest = tf.summary.merge(summaryImageTest)
summaryLossTrain = tf.summary.scalar("TRAIN_loss",loss)
testErrorPH = tf.placeholder(tf.float32,shape=[])
testImagePH = tf.placeholder(tf.float32,shape=[opt.labelN,opt.H,opt.W,1])
summaryErrorTest = tf.summary.scalar("TEST_error",testErrorPH)
if opt.netType=="STN" or opt.netType=="IC-STN":
summaryMeanTest0 = util.imageSummaryMeanVar(opt,testImagePH,"TEST_mean_init",opt.H,opt.W)
summaryMeanTest1 = util.imageSummaryMeanVar(opt,testImagePH,"TEST_mean_warped",opt.H,opt.W)
summaryVarTest0 = util.imageSummaryMeanVar(opt,testImagePH*3,"TEST_var_init",opt.H,opt.W)
summaryVarTest1 = util.imageSummaryMeanVar(opt,testImagePH*3,"TEST_var_warped",opt.H,opt.W)
# load data
print(util.toMagenta("loading MNIST dataset..."))
trainData,validData,testData = data.loadMNIST("data/MNIST.npz")
# prepare model saver/summary writer
saver = tf.train.Saver(max_to_keep=20)
summaryWriter = tf.summary.FileWriter("summary_{0}/{1}".format(opt.group,opt.model))
print(util.toYellow("======= TRAINING START ======="))
timeStart = time.time()
# start session
tfConfig = tf.ConfigProto(allow_soft_placement=True)
tfConfig.gpu_options.allow_growth = True
with tf.Session(config=tfConfig) as sess:
sess.run(tf.global_variables_initializer())
summaryWriter.add_graph(sess.graph)
if opt.fromIt!=0:
util.restoreModel(opt,sess,saver,opt.fromIt)
print(util.toMagenta("resuming from iteration {0}...".format(opt.fromIt)))
print(util.toMagenta("start training..."))
# training loop
for i in range(opt.fromIt,opt.toIt):
lrGP = opt.lrGP*opt.lrGPdecay**(i//opt.lrGPstep)
lrC = opt.lrC*opt.lrCdecay**(i//opt.lrCstep)
# make training batch
batch = data.makeBatch(opt,trainData,PH)
batch[lrGP_PH] = lrGP
batch[lrC_PH] = lrC
# run one step
_,l = sess.run([optim,loss],feed_dict=batch)
if (i+1)%100==0:
print("it. {0}/{1} lr={3}(GP),{4}(C), loss={5}, time={2}"
.format(util.toCyan("{0}".format(i+1)),
opt.toIt,
util.toGreen("{0:.2f}".format(time.time()-timeStart)),
util.toYellow("{0:.0e}".format(lrGP)),
util.toYellow("{0:.0e}".format(lrC)),
util.toRed("{0:.4f}".format(l))))
if (i+1)%100==0:
summaryWriter.add_summary(sess.run(summaryLossTrain,feed_dict=batch),i+1)
if (i+1)%500==0 and (opt.netType=="STN" or opt.netType=="IC-STN"):
summaryWriter.add_summary(sess.run(summaryImageTrain,feed_dict=batch),i+1)
summaryWriter.add_summary(sess.run(summaryImageTest,feed_dict=batch),i+1)
if (i+1)%1000==0:
# evaluate on test set
if opt.netType=="STN" or opt.netType=="IC-STN":
testAcc,testMean,testVar = data.evalTest(opt,sess,testData,PH,prediction,imagesEval=[imagePert,imageWarp])
else:
testAcc,_,_ = data.evalTest(opt,sess,testData,PH,prediction)
testError = (1-testAcc)*100
summaryWriter.add_summary(sess.run(summaryErrorTest,feed_dict={testErrorPH:testError}),i+1)
if opt.netType=="STN" or opt.netType=="IC-STN":
summaryWriter.add_summary(sess.run(summaryMeanTest0,feed_dict={testImagePH:testMean[0]}),i+1)
summaryWriter.add_summary(sess.run(summaryMeanTest1,feed_dict={testImagePH:testMean[1]}),i+1)
summaryWriter.add_summary(sess.run(summaryVarTest0,feed_dict={testImagePH:testVar[0]}),i+1)
summaryWriter.add_summary(sess.run(summaryVarTest1,feed_dict={testImagePH:testVar[1]}),i+1)
if (i+1)%10000==0:
util.saveModel(opt,sess,saver,i+1)
print(util.toGreen("model saved: {0}/{1}, it.{2}".format(opt.group,opt.model,i+1)))
print(util.toYellow("======= TRAINING DONE ======="))
================================================
FILE: MNIST-tensorflow/util.py
================================================
import numpy as np
import scipy.misc
import tensorflow as tf
import os
import termcolor
def mkdir(path):
if not os.path.exists(path): os.mkdir(path)
def imread(fname):
return scipy.misc.imread(fname)/255.0
def imsave(fname,array):
scipy.misc.toimage(array,cmin=0.0,cmax=1.0).save(fname)
# convert to colored strings
def toRed(content): return termcolor.colored(content,"red",attrs=["bold"])
def toGreen(content): return termcolor.colored(content,"green",attrs=["bold"])
def toBlue(content): return termcolor.colored(content,"blue",attrs=["bold"])
def toCyan(content): return termcolor.colored(content,"cyan",attrs=["bold"])
def toYellow(content): return termcolor.colored(content,"yellow",attrs=["bold"])
def toMagenta(content): return termcolor.colored(content,"magenta",attrs=["bold"])
# make image summary from image batch
def imageSummary(opt,image,tag,H,W):
blockSize = opt.visBlockSize
imageOne = tf.batch_to_space(image[:blockSize**2],crops=[[0,0],[0,0]],block_size=blockSize)
imagePermute = tf.reshape(imageOne,[H,blockSize,W,blockSize,-1])
imageTransp = tf.transpose(imagePermute,[1,0,3,2,4])
imageBlocks = tf.reshape(imageTransp,[1,H*blockSize,W*blockSize,-1])
imageBlocks = tf.cast(imageBlocks*255,tf.uint8)
summary = tf.summary.image(tag,imageBlocks)
return summary
# make image summary from image batch (mean/variance)
def imageSummaryMeanVar(opt,image,tag,H,W):
imageOne = tf.batch_to_space_nd(image,crops=[[0,0],[0,0]],block_shape=[1,10])
imagePermute = tf.reshape(imageOne,[H,1,W,10,-1])
imageTransp = tf.transpose(imagePermute,[1,0,3,2,4])
imageBlocks = tf.reshape(imageTransp,[1,H*1,W*10,-1])
imageBlocks = tf.cast(imageBlocks*255,tf.uint8)
summary = tf.summary.image(tag,imageBlocks)
return summary
# set optimizer for different learning rates
def setOptimizer(opt,loss,lrGP,lrC):
varsGP = [v for v in tf.global_variables() if "geometric" in v.name]
varsC = [v for v in tf.global_variables() if "classifier" in v.name]
gradC = tf.gradients(loss,varsC)
optimC = tf.train.GradientDescentOptimizer(lrC).apply_gradients(zip(gradC,varsC))
if len(varsGP)>0:
gradGP = tf.gradients(loss,varsGP)
optimGP = tf.train.GradientDescentOptimizer(lrGP).apply_gradients(zip(gradGP,varsGP))
optim = tf.group(optimC,optimGP)
else:
optim = optimC
return optim
# restore model
def restoreModel(opt,sess,saver,it):
saver.restore(sess,"models_{0}/{1}_it{2}.ckpt".format(opt.group,opt.model,it,opt.warpN))
# save model
def saveModel(opt,sess,saver,it):
saver.save(sess,"models_{0}/{1}_it{2}.ckpt".format(opt.group,opt.model,it,opt.warpN))
================================================
FILE: MNIST-tensorflow/warp.py
================================================
import numpy as np
import scipy.linalg
import tensorflow as tf
# fit (affine) warp between two sets of points
def fit(Xsrc,Xdst):
ptsN = len(Xsrc)
X,Y,U,V,O,I = Xsrc[:,0],Xsrc[:,1],Xdst[:,0],Xdst[:,1],np.zeros([ptsN]),np.ones([ptsN])
A = np.concatenate((np.stack([X,Y,I,O,O,O],axis=1),
np.stack([O,O,O,X,Y,I],axis=1)),axis=0)
b = np.concatenate((U,V),axis=0)
p1,p2,p3,p4,p5,p6 = scipy.linalg.lstsq(A,b)[0].squeeze()
pMtrx = np.array([[p1,p2,p3],[p4,p5,p6],[0,0,1]],dtype=np.float32)
return pMtrx
# compute composition of warp parameters
def compose(opt,p,dp):
with tf.name_scope("compose"):
pMtrx = vec2mtrx(opt,p)
dpMtrx = vec2mtrx(opt,dp)
pMtrxNew = tf.matmul(dpMtrx,pMtrx)
pMtrxNew /= pMtrxNew[:,2:3,2:3]
pNew = mtrx2vec(opt,pMtrxNew)
return pNew
# compute inverse of warp parameters
def inverse(opt,p):
with tf.name_scope("inverse"):
pMtrx = vec2mtrx(opt,p)
pInvMtrx = tf.matrix_inverse(pMtrx)
pInv = mtrx2vec(opt,pInvMtrx)
return pInv
# convert warp parameters to matrix
def vec2mtrx(opt,p):
with tf.name_scope("vec2mtrx"):
O = tf.zeros([opt.batchSize])
I = tf.ones([opt.batchSize])
if opt.warpType=="translation":
tx,ty = tf.unstack(p,axis=1)
pMtrx = tf.transpose(tf.stack([[I,O,tx],[O,I,ty],[O,O,I]]),perm=[2,0,1])
if opt.warpType=="similarity":
pc,ps,tx,ty = tf.unstack(p,axis=1)
pMtrx = tf.transpose(tf.stack([[I+pc,-ps,tx],[ps,I+pc,ty],[O,O,I]]),perm=[2,0,1])
if opt.warpType=="affine":
p1,p2,p3,p4,p5,p6,p7,p8 = tf.unstack(p,axis=1)
pMtrx = tf.transpose(tf.stack([[I+p1,p2,p3],[p4,I+p5,p6],[O,O,I]]),perm=[2,0,1])
if opt.warpType=="homography":
p1,p2,p3,p4,p5,p6,p7,p8 = tf.unstack(p,axis=1)
pMtrx = tf.transpose(tf.stack([[I+p1,p2,p3],[p4,I+p5,p6],[p7,p8,I]]),perm=[2,0,1])
return pMtrx
# convert warp matrix to parameters
def mtrx2vec(opt,pMtrx):
with tf.name_scope("mtrx2vec"):
[row0,row1,row2] = tf.unstack(pMtrx,axis=1)
[e00,e01,e02] = tf.unstack(row0,axis=1)
[e10,e11,e12] = tf.unstack(row1,axis=1)
[e20,e21,e22] = tf.unstack(row2,axis=1)
if opt.warpType=="translation": p = tf.stack([e02,e12],axis=1)
if opt.warpType=="similarity": p = tf.stack([e00-1,e10,e02,e12],axis=1)
if opt.warpType=="affine": p = tf.stack([e00-1,e01,e02,e10,e11-1,e12],axis=1)
if opt.warpType=="homography": p = tf.stack([e00-1,e01,e02,e10,e11-1,e12,e20,e21],axis=1)
return p
# warp the image
def transformImage(opt,image,pMtrx):
with tf.name_scope("transformImage"):
refMtrx = tf.tile(tf.expand_dims(opt.refMtrx,axis=0),[opt.batchSize,1,1])
transMtrx = tf.matmul(refMtrx,pMtrx)
# warp the canonical coordinates
X,Y = np.meshgrid(np.linspace(-1,1,opt.W),np.linspace(-1,1,opt.H))
X,Y = X.flatten(),Y.flatten()
XYhom = np.stack([X,Y,np.ones_like(X)],axis=1).T
XYhom = np.tile(XYhom,[opt.batchSize,1,1]).astype(np.float32)
XYwarpHom = tf.matmul(transMtrx,XYhom)
XwarpHom,YwarpHom,ZwarpHom = tf.unstack(XYwarpHom,axis=1)
Xwarp = tf.reshape(XwarpHom/(ZwarpHom+1e-8),[opt.batchSize,opt.H,opt.W])
Ywarp = tf.reshape(YwarpHom/(ZwarpHom+1e-8),[opt.batchSize,opt.H,opt.W])
# get the integer sampling coordinates
Xfloor,Xceil = tf.floor(Xwarp),tf.ceil(Xwarp)
Yfloor,Yceil = tf.floor(Ywarp),tf.ceil(Ywarp)
XfloorInt,XceilInt = tf.to_int32(Xfloor),tf.to_int32(Xceil)
YfloorInt,YceilInt = tf.to_int32(Yfloor),tf.to_int32(Yceil)
imageIdx = np.tile(np.arange(opt.batchSize).reshape([opt.batchSize,1,1]),[1,opt.H,opt.W])
imageVec = tf.reshape(image,[-1,int(image.shape[-1])])
imageVecOut = tf.concat([imageVec,tf.zeros([1,int(image.shape[-1])])],axis=0)
idxUL = (imageIdx*opt.H+YfloorInt)*opt.W+XfloorInt
idxUR = (imageIdx*opt.H+YfloorInt)*opt.W+XceilInt
idxBL = (imageIdx*opt.H+YceilInt)*opt.W+XfloorInt
idxBR = (imageIdx*opt.H+YceilInt)*opt.W+XceilInt
idxOutside = tf.fill([opt.batchSize,opt.H,opt.W],opt.batchSize*opt.H*opt.W)
def insideImage(Xint,Yint):
return (Xint>=0)&(Xint<opt.W)&(Yint>=0)&(Yint<opt.H)
idxUL = tf.where(insideImage(XfloorInt,YfloorInt),idxUL,idxOutside)
idxUR = tf.where(insideImage(XceilInt,YfloorInt),idxUR,idxOutside)
idxBL = tf.where(insideImage(XfloorInt,YceilInt),idxBL,idxOutside)
idxBR = tf.where(insideImage(XceilInt,YceilInt),idxBR,idxOutside)
# bilinear interpolation
Xratio = tf.reshape(Xwarp-Xfloor,[opt.batchSize,opt.H,opt.W,1])
Yratio = tf.reshape(Ywarp-Yfloor,[opt.batchSize,opt.H,opt.W,1])
imageUL = tf.to_float(tf.gather(imageVecOut,idxUL))*(1-Xratio)*(1-Yratio)
imageUR = tf.to_float(tf.gather(imageVecOut,idxUR))*(Xratio)*(1-Yratio)
imageBL = tf.to_float(tf.gather(imageVecOut,idxBL))*(1-Xratio)*(Yratio)
imageBR = tf.to_float(tf.gather(imageVecOut,idxBR))*(Xratio)*(Yratio)
imageWarp = imageUL+imageUR+imageBL+imageBR
return imageWarp
================================================
FILE: README.md
================================================
## Inverse Compositional Spatial Transformer Networks
[Chen-Hsuan Lin](https://chenhsuanlin.bitbucket.io/)
and [Simon Lucey](http://www.simonlucey.com/)
IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2017 (**oral presentation**)
Project page: https://chenhsuanlin.bitbucket.io/inverse-compositional-STN
Paper: https://chenhsuanlin.bitbucket.io/inverse-compositional-STN/paper.pdf
Poster: https://chenhsuanlin.bitbucket.io/inverse-compositional-STN/poster.pdf
arXiv preprint: https://arxiv.org/abs/1612.03897
<p align="center"><img src="https://www.andrew.cmu.edu/user/chenhsul/images/ICSTN2.png" width=600 height=250></p>
We provide TensorFlow code for the following experiments:
- MNIST classification
- traffic sign classification
**[NEW!]** The PyTorch implementation of the MNIST experiment is now up!
--------------------------------------
## TensorFlow
### Prerequisites
This code is developed with Python3 (`python3`) but it is also compatible with Python2.7 (`python`). TensorFlow r1.0+ is required. The dependencies can install by running
```
pip3 install --upgrade numpy scipy termcolor matplotlib tensorflow-gpu
```
If you're using Python2.7, use `pip2` instead; if you don't have sudo access, add the `--user` flag.
### Running the code
The training code can be executed via the command
```
python3 train.py <netType> [(options)]
```
`<netType>` should be one of the following:
1. `CNN` - standard convolutional neural network
2. `STN` - Spatial Transformer Network (STN)
3. `IC-STN` - Inverse Compositional Spatial Transformer Network (IC-STN)
The list of optional arguments can be found by executing `python3 train.py --help`.
The default training settings in this released code is slightly different from that in the paper; it is stabler and optimizes the networks better.
When the code is run for the first time, the datasets will be automatically downloaded and preprocessed.
The checkpoints are saved in the automatically created directory `model_GROUP`; summaries are saved in `summary_GROUP`.
### Visualizing the results
We've included code to visualize the training over TensorBoard. To execute, run
```
tensorboard --logdir=summary_GROUP --port=6006
```
We provide three types of data visualization:
1. **SCALARS**: training/test error over iterations
2. **IMAGES**: alignment results and mean/variance appearances
3. **GRAPH**: network architecture
--------------------------------------
## PyTorch
The PyTorch version of the code is stil under active development. The training speed is currently slower than the TensorFlow version. Suggestions on improvements are welcome! :)
### Prerequisites
This code is developed with Python3 (`python3`). It has not been tested with Python2.7 yet. PyTorch 0.2.0+ is required. Please see http://pytorch.org/ for installation instructions.
Visdom is also required; it can be installed by running
```
pip3 install --upgrade visdom
```
If you don't have sudo access, add the `--user` flag.
### Running the code
First, start a Visdom server by running
```
python3 -m visdom.server -port=7000
```
The training code can be executed via the command (using the same port number)
```
python3 train.py <netType> --port=7000 [(options)]
```
`<netType>` should be one of the following:
1. `CNN` - standard convolutional neural network
2. `STN` - Spatial Transformer Network (STN)
3. `IC-STN` - Inverse Compositional Spatial Transformer Network (IC-STN)
The list of optional arguments can be found by executing `python3 train.py --help`.
The default training settings in this released code is slightly different from that in the paper; it is stabler and optimizes the networks better.
When the code is run for the first time, the datasets will be automatically downloaded and preprocessed.
The checkpoints are saved in the automatically created directory `model_GROUP`; summaries are saved in `summary_GROUP`.
### Visualizing the results
We provide three types of data visualization on Visdom:
1. Training/test error over iterations
2. Alignment results and mean/variance appearances
--------------------------------------
If you find our code useful for your research, please cite
```
@inproceedings{lin2017inverse,
title={Inverse Compositional Spatial Transformer Networks},
author={Lin, Chen-Hsuan and Lucey, Simon},
booktitle={IEEE Conference on Computer Vision and Pattern Recognition ({CVPR})},
year={2017}
}
```
Please contact me (chlin@cmu.edu) if you have any questions!
================================================
FILE: traffic-sign-tensorflow/data.py
================================================
import numpy as np
import scipy.linalg,scipy.misc
import os,time
import tensorflow as tf
import matplotlib.pyplot as plt
import csv
import warp
# load GTSRB data
def loadGTSRB(opt,fname):
if not os.path.exists(fname):
# download and preprocess GTSRB dataset
os.makedirs(os.path.dirname(fname))
os.system("wget -O data/GTSRB_Final_Training_Images.zip http://benchmark.ini.rub.de/Dataset/GTSRB_Final_Training_Images.zip")
os.system("wget -O data/GTSRB_Final_Test_Images.zip http://benchmark.ini.rub.de/Dataset/GTSRB_Final_Test_Images.zip")
os.system("wget -O data/GTSRB_Final_Test_GT.zip http://benchmark.ini.rub.de/Dataset/GTSRB_Final_Test_GT.zip")
os.system("cd data && unzip GTSRB_Final_Training_Images.zip")
os.system("cd data && unzip GTSRB_Final_Test_Images.zip")
os.system("cd data && unzip GTSRB_Final_Test_GT.zip")
# training data
print("preparing training data...")
images,bboxes,labels = [],[],[]
for c in range(43):
prefix = "data/GTSRB/Final_Training/Images/{0:05d}".format(c)
with open("{0}/GT-{1:05d}.csv".format(prefix,c)) as file:
reader = csv.reader(file,delimiter=";")
next(reader)
for line in reader:
img = plt.imread(prefix+"/"+line[0])
rawH,rawW = img.shape[0],img.shape[1]
scaleH,scaleW = float(opt.fullH)/rawH,float(opt.fullW)/rawW
imgResize = scipy.misc.imresize(img,(opt.fullH,opt.fullW,3))
images.append(imgResize)
bboxes.append([float(line[3])*scaleW,float(line[4])*scaleH,
float(line[5])*scaleW,float(line[6])*scaleH])
labels.append(int(line[7]))
trainData = {
"image": np.array(images),
"bbox": np.array(bboxes),
"label": np.array(labels)
}
# test data
print("preparing test data...")
images,bboxes,labels = [],[],[]
prefix = "data/GTSRB/Final_Test/Images/"
with open("data/GT-final_test.csv") as file:
reader = csv.reader(file,delimiter=";")
next(reader)
for line in reader:
img = plt.imread(prefix+"/"+line[0])
rawH,rawW = img.shape[0],img.shape[1]
scaleH,scaleW = float(opt.fullH)/rawH,float(opt.fullW)/rawW
imgResize = scipy.misc.imresize(img,(opt.fullH,opt.fullW,3))
images.append(imgResize)
bboxes.append([float(line[3])*scaleW,float(line[4])*scaleH,
float(line[5])*scaleW,float(line[6])*scaleH])
labels.append(int(line[7]))
testData = {
"image": np.array(images),
"bbox": np.array(bboxes),
"label": np.array(labels)
}
np.savez(fname,train=trainData,test=testData)
os.system("rm -rf data/*.zip")
GTSRB = np.load(fname)
trainData = GTSRB["train"].item()
testData = GTSRB["test"].item()
return trainData,testData
# generate training batch
def genPerturbations(opt):
with tf.name_scope("genPerturbations"):
X = np.tile(opt.canon4pts[:,0],[opt.batchSize,1])
Y = np.tile(opt.canon4pts[:,1],[opt.batchSize,1])
dX = tf.random_normal([opt.batchSize,4])*opt.pertScale \
+tf.random_normal([opt.batchSize,1])*opt.transScale
dY = tf.random_normal([opt.batchSize,4])*opt.pertScale \
+tf.random_normal([opt.batchSize,1])*opt.transScale
O = np.zeros([opt.batchSize,4],dtype=np.float32)
I = np.ones([opt.batchSize,4],dtype=np.float32)
# fit warp parameters to generated displacements
if opt.warpType=="homography":
A = tf.concat([tf.stack([X,Y,I,O,O,O,-X*(X+dX),-Y*(X+dX)],axis=-1),
tf.stack([O,O,O,X,Y,I,-X*(Y+dY),-Y*(Y+dY)],axis=-1)],1)
b = tf.expand_dims(tf.concat([X+dX,Y+dY],1),-1)
pPert = tf.matrix_solve(A,b)[:,:,0]
pPert -= tf.to_float([[1,0,0,0,1,0,0,0]])
else:
if opt.warpType=="translation":
J = np.concatenate([np.stack([I,O],axis=-1),
np.stack([O,I],axis=-1)],axis=1)
if opt.warpType=="similarity":
J = np.concatenate([np.stack([X,Y,I,O],axis=-1),
np.stack([-Y,X,O,I],axis=-1)],axis=1)
if opt.warpType=="affine":
J = np.concatenate([np.stack([X,Y,I,O,O,O],axis=-1),
np.stack([O,O,O,X,Y,I],axis=-1)],axis=1)
dXY = tf.expand_dims(tf.concat([dX,dY],1),-1)
pPert = tf.matrix_solve_ls(J,dXY)[:,:,0]
return pPert
# make training batch
def makeBatch(opt,data,PH):
N = len(data["image"])
randIdx = np.random.randint(N,size=[opt.batchSize])
# put data in placeholders
[image,label] = PH
batch = {
image: data["image"][randIdx]/255.0,
label: data["label"][randIdx],
}
return batch
# evaluation on test set
def evalTest(opt,sess,data,PH,prediction,imagesEval=[]):
N = len(data["image"])
# put data in placeholders
[image,label] = PH
batchN = int(np.ceil(N/opt.batchSize))
warped = [{},{}]
count = 0
for b in range(batchN):
# use some dummy data (0) as batch filler if necessary
if b!=batchN-1:
realIdx = np.arange(opt.batchSize*b,opt.batchSize*(b+1))
else:
realIdx = np.arange(opt.batchSize*b,N)
idx = np.zeros([opt.batchSize],dtype=int)
idx[:len(realIdx)] = realIdx
batch = {
image: data["image"][idx]/255.0,
label: data["label"][idx],
}
evalList = sess.run([prediction]+imagesEval,feed_dict=batch)
pred = evalList[0]
count += pred[:len(realIdx)].sum()
if len(imagesEval)>0:
imgs = evalList[1:]
for i in range(len(realIdx)):
if data["label"][idx[i]] not in warped[0]: warped[0][data["label"][idx[i]]] = []
if data["label"][idx[i]] not in warped[1]: warped[1][data["label"][idx[i]]] = []
warped[0][data["label"][idx[i]]].append(imgs[0][i])
warped[1][data["label"][idx[i]]].append(imgs[1][i])
accuracy = float(count)/N
if len(imagesEval)>0:
mean = [np.array([np.mean(warped[0][l],axis=0) for l in warped[0]]),
np.array([np.mean(warped[1][l],axis=0) for l in warped[1]])]
var = [np.array([np.var(warped[0][l],axis=0) for l in warped[0]]),
np.array([np.var(warped[1][l],axis=0) for l in warped[1]])]
else: mean,var = None,None
return accuracy,mean,var
================================================
FILE: traffic-sign-tensorflow/graph.py
================================================
import numpy as np
import tensorflow as tf
import time
import data,warp,util
# build classification network
def fullCNN(opt,image):
def conv2Layer(opt,feat,outDim):
weight,bias = createVariable(opt,[7,7,int(feat.shape[-1]),outDim],stddev=opt.stdC)
conv = tf.nn.conv2d(feat,weight,strides=[1,1,1,1],padding="VALID")+bias
return conv
def linearLayer(opt,feat,outDim):
weight,bias = createVariable(opt,[int(feat.shape[-1]),outDim],stddev=opt.stdC)
fc = tf.matmul(feat,weight)+bias
return fc
with tf.variable_scope("classifier"):
feat = image
with tf.variable_scope("conv1"):
feat = conv2Layer(opt,feat,6)
feat = tf.nn.relu(feat)
with tf.variable_scope("conv2"):
feat = conv2Layer(opt,feat,12)
feat = tf.nn.relu(feat)
feat = tf.nn.max_pool(feat,ksize=[1,2,2,1],strides=[1,2,2,1],padding="VALID")
with tf.variable_scope("conv3"):
feat = conv2Layer(opt,feat,24)
feat = tf.nn.relu(feat)
feat = tf.reshape(feat,[opt.batchSize,-1])
with tf.variable_scope("fc4"):
feat = linearLayer(opt,feat,200)
feat = tf.nn.relu(feat)
with tf.variable_scope("fc5"):
feat = linearLayer(opt,feat,opt.labelN)
output = feat
return output
# build classification network
def CNN(opt,image):
def conv2Layer(opt,feat,outDim):
weight,bias = createVariable(opt,[7,7,int(feat.shape[-1]),outDim],stddev=opt.stdC)
conv = tf.nn.conv2d(feat,weight,strides=[1,1,1,1],padding="VALID")+bias
return conv
def linearLayer(opt,feat,outDim):
weight,bias = createVariable(opt,[int(feat.shape[-1]),outDim],stddev=opt.stdC)
fc = tf.matmul(feat,weight)+bias
return fc
with tf.variable_scope("classifier"):
feat = image
with tf.variable_scope("conv1"):
feat = conv2Layer(opt,feat,6)
feat = tf.nn.relu(feat)
with tf.variable_scope("conv2"):
feat = conv2Layer(opt,feat,12)
feat = tf.nn.relu(feat)
feat = tf.nn.max_pool(feat,ksize=[1,2,2,1],strides=[1,2,2,1],padding="VALID")
feat = tf.reshape(feat,[opt.batchSize,-1])
with tf.variable_scope("fc3"):
feat = linearLayer(opt,feat,opt.labelN)
output = feat
return output
# build Spatial Transformer Network
def STN(opt,image):
def conv2Layer(opt,feat,outDim):
weight,bias = createVariable(opt,[7,7,int(feat.shape[-1]),outDim],stddev=opt.stdGP)
conv = tf.nn.conv2d(feat,weight,strides=[1,1,1,1],padding="VALID")+bias
return conv
def linearLayer(opt,feat,outDim):
weight,bias = createVariable(opt,[int(feat.shape[-1]),outDim],stddev=opt.stdGP)
fc = tf.matmul(feat,weight)+bias
return fc
imageWarpAll = [image]
with tf.variable_scope("geometric"):
feat = image
with tf.variable_scope("conv1"):
feat = conv2Layer(opt,feat,6)
feat = tf.nn.relu(feat)
with tf.variable_scope("conv2"):
feat = conv2Layer(opt,feat,24)
feat = tf.nn.relu(feat)
feat = tf.reshape(feat,[opt.batchSize,-1])
with tf.variable_scope("fc3"):
feat = linearLayer(opt,feat,opt.warpDim)
p = feat
pMtrx = warp.vec2mtrx(opt,p)
imageWarp = warp.transformImage(opt,image,pMtrx)
imageWarpAll.append(imageWarp)
return imageWarpAll
# build Inverse Compositional STN
def ICSTN(opt,imageFull,p):
def conv2Layer(opt,feat,outDim):
weight,bias = createVariable(opt,[7,7,int(feat.shape[-1]),outDim],stddev=opt.stdGP)
conv = tf.nn.conv2d(feat,weight,strides=[1,1,1,1],padding="VALID")+bias
return conv
def linearLayer(opt,feat,outDim):
weight,bias = createVariable(opt,[int(feat.shape[-1]),outDim],stddev=opt.stdGP)
fc = tf.matmul(feat,weight)+bias
return fc
imageWarpAll = []
for l in range(opt.warpN):
with tf.variable_scope("geometric",reuse=l>0):
pMtrx = warp.vec2mtrx(opt,p)
imageWarp = warp.transformCropImage(opt,imageFull,pMtrx)
imageWarpAll.append(imageWarp)
feat = imageWarp
with tf.variable_scope("conv1"):
feat = conv2Layer(opt,feat,6)
feat = tf.nn.relu(feat)
with tf.variable_scope("conv2"):
feat = conv2Layer(opt,feat,24)
feat = tf.nn.relu(feat)
feat = tf.reshape(feat,[opt.batchSize,-1])
with tf.variable_scope("fc3"):
feat = linearLayer(opt,feat,opt.warpDim)
dp = feat
p = warp.compose(opt,p,dp)
pMtrx = warp.vec2mtrx(opt,p)
imageWarp = warp.transformCropImage(opt,imageFull,pMtrx)
imageWarpAll.append(imageWarp)
return imageWarpAll
# auxiliary function for creating weight and bias
def createVariable(opt,weightShape,biasShape=None,stddev=None):
if biasShape is None: biasShape = [weightShape[-1]]
weight = tf.get_variable("weight",shape=weightShape,dtype=tf.float32,
initializer=tf.random_normal_initializer(stddev=stddev))
bias = tf.get_variable("bias",shape=biasShape,dtype=tf.float32,
initializer=tf.random_normal_initializer(stddev=stddev))
return weight,bias
================================================
FILE: traffic-sign-tensorflow/options.py
================================================
import numpy as np
import argparse
import warp
import util
def set(training):
# parse input arguments
parser = argparse.ArgumentParser()
parser.add_argument("netType", choices=["CNN","STN","IC-STN"], help="type of network")
parser.add_argument("--group", default="0", help="name for group")
parser.add_argument("--model", default="test", help="name for model instance")
parser.add_argument("--size", default="36x36", help="image resolution")
parser.add_argument("--sizeFull", default="50x50", help="full image resolution")
parser.add_argument("--warpType", default="homography", help="type of warp function on images",
choices=["translation","similarity","affine","homography"])
parser.add_argument("--warpN", type=int, default=4, help="number of recurrent transformations (for IC-STN)")
parser.add_argument("--stdC", type=float, default=0.01, help="initialization stddev (classification network)")
parser.add_argument("--stdGP", type=float, default=0.001, help="initialization stddev (geometric predictor)")
parser.add_argument("--pertScale", type=float, default=0.25, help="initial perturbation scale")
parser.add_argument("--transScale", type=float, default=0.25, help="initial translation scale")
if training: # training
parser.add_argument("--batchSize", type=int, default=100, help="batch size for SGD")
parser.add_argument("--lrC", type=float, default=1e-2, help="learning rate (classification network)")
parser.add_argument("--lrCdecay", type=float, default=0.1, help="learning rate decay (classification network)")
parser.add_argument("--lrCstep", type=int, default=500000, help="learning rate decay step size (classification network)")
parser.add_argument("--lrGP", type=float, default=None, help="learning rate (geometric predictor)")
parser.add_argument("--lrGPdecay", type=float, default=0.1, help="learning rate decay (geometric predictor)")
parser.add_argument("--lrGPstep", type=int, default=500000, help="learning rate decay step size (geometric predictor)")
parser.add_argument("--fromIt", type=int, default=0, help="resume training from iteration number")
parser.add_argument("--toIt", type=int, default=1000000,help="run training to iteration number")
else: # evaluation
parser.add_argument("--batchSize", type=int, default=1, help="batch size for evaluation")
opt = parser.parse_args()
if opt.lrGP is None: opt.lrGP = 0 if opt.netType=="CNN" else \
1e-3 if opt.netType=="STN" else \
3e-5 if opt.netType=="IC-STN" else None
# --- below are automatically set ---
opt.training = training
opt.H,opt.W = [int(x) for x in opt.size.split("x")]
opt.fullH,opt.fullW = [int(x) for x in opt.sizeFull.split("x")]
opt.visBlockSize = int(np.floor(np.sqrt(opt.batchSize)))
opt.warpDim = 2 if opt.warpType == "translation" else \
4 if opt.warpType == "similarity" else \
6 if opt.warpType == "affine" else \
8 if opt.warpType == "homography" else None
opt.labelN = 43
opt.canon4pts = np.array([[-1,-1],[-1,1],[1,1],[1,-1]],dtype=np.float32)
opt.image4pts = np.array([[0,0],[0,opt.H-1],[opt.W-1,opt.H-1],[opt.W-1,0]],dtype=np.float32)
opt.bbox = [int(opt.fullW/2-opt.W/2),int(opt.fullH/2-opt.H/2),int(opt.fullW/2+opt.W/2),int(opt.fullH/2+opt.H/2)]
opt.bbox4pts = np.array([[opt.bbox[0],opt.bbox[1]],[opt.bbox[0],opt.bbox[3]],
[opt.bbox[2],opt.bbox[3]],[opt.bbox[2],opt.bbox[1]]],dtype=np.float32)
opt.refMtrx = warp.fit(Xsrc=opt.canon4pts,Xdst=opt.image4pts)
opt.bboxRefMtrx = warp.fit(Xsrc=opt.canon4pts,Xdst=opt.bbox4pts)
if opt.netType=="STN": opt.warpN = 1
print("({0}) {1}".format(
util.toGreen("{0}".format(opt.group)),
util.toGreen("{0}".format(opt.model))))
print("------------------------------------------")
print("network type: {0}, recurrent warps: {1}".format(
util.toYellow("{0}".format(opt.netType)),
util.toYellow("{0}".format(opt.warpN if opt.netType=="IC-STN" else "X"))))
print("batch size: {0}, image size: {1}x{2}".format(
util.toYellow("{0}".format(opt.batchSize)),
util.toYellow("{0}".format(opt.H)),
util.toYellow("{0}".format(opt.W))))
print("warpScale: (pert) {0} (trans) {1}".format(
util.toYellow("{0}".format(opt.pertScale)),
util.toYellow("{0}".format(opt.transScale))))
if training:
print("[geometric predictor] stddev={0}, lr={1}".format(
util.toYellow("{0:.0e}".format(opt.stdGP)),
util.toYellow("{0:.0e}".format(opt.lrGP))))
print("[classification network] stddev={0}, lr={1}".format(
util.toYellow("{0:.0e}".format(opt.stdC)),
util.toYellow("{0:.0e}".format(opt.lrC))))
print("------------------------------------------")
if training:
print(util.toMagenta("training model ({0}) {1}...".format(opt.group,opt.model)))
return opt
================================================
FILE: traffic-sign-tensorflow/train.py
================================================
import numpy as np
import time,os,sys
import argparse
import util
print(util.toYellow("======================================================="))
print(util.toYellow("train.py (training on MNIST)"))
print(util.toYellow("======================================================="))
import tensorflow as tf
import data,graph,warp,util
import options
print(util.toMagenta("setting configurations..."))
opt = options.set(training=True)
# create directories for model output
util.mkdir("models_{0}".format(opt.group))
print(util.toMagenta("building graph..."))
tf.reset_default_graph()
# build graph
with tf.device("/gpu:0"):
# ------ define input data ------
imageFull = tf.placeholder(tf.float32,shape=[opt.batchSize,opt.fullH,opt.fullW,3])
imageMean,imageVar = tf.nn.moments(imageFull,axes=[1,2],keep_dims=True)
imageFullNormalize = (imageFull-imageMean)/tf.sqrt(imageVar)
label = tf.placeholder(tf.int64,shape=[opt.batchSize])
PH = [imageFull,label]
# ------ generate perturbation ------
pInit = data.genPerturbations(opt)
pInitMtrx = warp.vec2mtrx(opt,pInit)
# ------ build network ------
imagePert = warp.transformCropImage(opt,imageFullNormalize,pInitMtrx)
imagePertRescale = imagePert*tf.sqrt(imageVar)+imageMean
if opt.netType=="CNN":
output = graph.fullCNN(opt,imagePert)
elif opt.netType=="STN":
imageWarpAll = graph.STN(opt,imagePert)
imageWarp = imageWarpAll[-1]
output = graph.CNN(opt,imageWarp)
imageWarpRescale = imageWarp*tf.sqrt(imageVar)+imageMean
elif opt.netType=="IC-STN":
imageWarpAll = graph.ICSTN(opt,imageFullNormalize,pInit)
imageWarp = imageWarpAll[-1]
output = graph.CNN(opt,imageWarp)
imageWarpRescale = imageWarp*tf.sqrt(imageVar)+imageMean
softmax = tf.nn.softmax(output)
labelOnehot = tf.one_hot(label,opt.labelN)
prediction = tf.equal(tf.argmax(softmax,1),label)
# ------ define loss ------
softmaxLoss = tf.nn.softmax_cross_entropy_with_logits(logits=output,labels=labelOnehot)
loss = tf.reduce_mean(softmaxLoss)
# ------ optimizer ------
lrGP_PH,lrC_PH = tf.placeholder(tf.float32,shape=[]),tf.placeholder(tf.float32,shape=[])
optim = util.setOptimizer(opt,loss,lrGP_PH,lrC_PH)
# ------ generate summaries ------
summaryImageTrain = []
summaryImageTest = []
if opt.netType=="STN" or opt.netType=="IC-STN":
for l in range(opt.warpN+1):
summaryImageTrain.append(util.imageSummary(opt,imageWarpAll[l]*tf.sqrt(imageVar)+imageMean,"TRAIN_warp{0}".format(l),opt.H,opt.W))
summaryImageTest.append(util.imageSummary(opt,imageWarpAll[l]*tf.sqrt(imageVar)+imageMean,"TEST_warp{0}".format(l),opt.H,opt.W))
summaryImageTrain = tf.summary.merge(summaryImageTrain)
summaryImageTest = tf.summary.merge(summaryImageTest)
summaryLossTrain = tf.summary.scalar("TRAIN_loss",loss)
testErrorPH = tf.placeholder(tf.float32,shape=[])
testImagePH = tf.placeholder(tf.float32,shape=[opt.labelN,opt.H,opt.W,3])
summaryErrorTest = tf.summary.scalar("TEST_error",testErrorPH)
if opt.netType=="STN" or opt.netType=="IC-STN":
summaryMeanTest0 = util.imageSummaryMeanVar(opt,testImagePH,"TEST_mean_init",opt.H,opt.W)
summaryMeanTest1 = util.imageSummaryMeanVar(opt,testImagePH,"TEST_mean_warped",opt.H,opt.W)
summaryVarTest0 = util.imageSummaryMeanVar(opt,testImagePH,"TEST_var_init",opt.H,opt.W)
summaryVarTest1 = util.imageSummaryMeanVar(opt,testImagePH,"TEST_var_warped",opt.H,opt.W)
# load data
print(util.toMagenta("loading GTSRB dataset..."))
trainData,testData = data.loadGTSRB(opt,"data/GTSRB.npz")
# prepare model saver/summary writer
saver = tf.train.Saver(max_to_keep=20)
summaryWriter = tf.summary.FileWriter("summary_{0}/{1}".format(opt.group,opt.model))
print(util.toYellow("======= TRAINING START ======="))
timeStart = time.time()
# start session
tfConfig = tf.ConfigProto(allow_soft_placement=True)
tfConfig.gpu_options.allow_growth = True
with tf.Session(config=tfConfig) as sess:
sess.run(tf.global_variables_initializer())
summaryWriter.add_graph(sess.graph)
if opt.fromIt!=0:
util.restoreModel(opt,sess,saver,opt.fromIt)
print(util.toMagenta("resuming from iteration {0}...".format(opt.fromIt)))
print(util.toMagenta("start training..."))
# training loop
for i in range(opt.fromIt,opt.toIt):
lrGP = opt.lrGP*opt.lrGPdecay**(i//opt.lrGPstep)
lrC = opt.lrC*opt.lrCdecay**(i//opt.lrCstep)
# make training batch
batch = data.makeBatch(opt,trainData,PH)
batch[lrGP_PH] = lrGP
batch[lrC_PH] = lrC
# run one step
_,l = sess.run([optim,loss],feed_dict=batch)
if (i+1)%100==0:
print("it. {0}/{1} lr={3}(GP),{4}(C), loss={5}, time={2}"
.format(util.toCyan("{0}".format(i+1)),
opt.toIt,
util.toGreen("{0:.2f}".format(time.time()-timeStart)),
util.toYellow("{0:.0e}".format(lrGP)),
util.toYellow("{0:.0e}".format(lrC)),
util.toRed("{0:.4f}".format(l))))
if (i+1)%100==0:
summaryWriter.add_summary(sess.run(summaryLossTrain,feed_dict=batch),i+1)
if (i+1)%500==0 and (opt.netType=="STN" or opt.netType=="IC-STN"):
summaryWriter.add_summary(sess.run(summaryImageTrain,feed_dict=batch),i+1)
summaryWriter.add_summary(sess.run(summaryImageTest,feed_dict=batch),i+1)
if (i+1)%1000==0:
# evaluate on test set
if opt.netType=="STN" or opt.netType=="IC-STN":
testAcc,testMean,testVar = data.evalTest(opt,sess,testData,PH,prediction,imagesEval=[imagePert,imageWarp])
else:
testAcc,_,_ = data.evalTest(opt,sess,testData,PH,prediction)
testError = (1-testAcc)*100
summaryWriter.add_summary(sess.run(summaryErrorTest,feed_dict={testErrorPH:testError}),i+1)
if opt.netType=="STN" or opt.netType=="IC-STN":
summaryWriter.add_summary(sess.run(summaryMeanTest0,feed_dict={testImagePH:testMean[0]}),i+1)
summaryWriter.add_summary(sess.run(summaryMeanTest1,feed_dict={testImagePH:testMean[1]}),i+1)
summaryWriter.add_summary(sess.run(summaryVarTest0,feed_dict={testImagePH:testVar[0]}),i+1)
summaryWriter.add_summary(sess.run(summaryVarTest1,feed_dict={testImagePH:testVar[1]}),i+1)
if (i+1)%10000==0:
util.saveModel(opt,sess,saver,i+1)
print(util.toGreen("model saved: {0}/{1}, it.{2}".format(opt.group,opt.model,i+1)))
print(util.toYellow("======= TRAINING DONE ======="))
================================================
FILE: traffic-sign-tensorflow/util.py
================================================
import numpy as np
import scipy.misc
import tensorflow as tf
import os
import termcolor
def mkdir(path):
if not os.path.exists(path): os.mkdir(path)
def imread(fname):
return scipy.misc.imread(fname)/255.0
def imsave(fname,array):
scipy.misc.toimage(array,cmin=0.0,cmax=1.0).save(fname)
# convert to colored strings
def toRed(content): return termcolor.colored(content,"red",attrs=["bold"])
def toGreen(content): return termcolor.colored(content,"green",attrs=["bold"])
def toBlue(content): return termcolor.colored(content,"blue",attrs=["bold"])
def toCyan(content): return termcolor.colored(content,"cyan",attrs=["bold"])
def toYellow(content): return termcolor.colored(content,"yellow",attrs=["bold"])
def toMagenta(content): return termcolor.colored(content,"magenta",attrs=["bold"])
# make image summary from image batch
def imageSummary(opt,image,tag,H,W):
blockSize = opt.visBlockSize
imageOne = tf.batch_to_space(image[:blockSize**2],crops=[[0,0],[0,0]],block_size=blockSize)
imagePermute = tf.reshape(imageOne,[H,blockSize,W,blockSize,-1])
imageTransp = tf.transpose(imagePermute,[1,0,3,2,4])
imageBlocks = tf.reshape(imageTransp,[1,H*blockSize,W*blockSize,-1])
imageBlocks = tf.cast(imageBlocks*255,tf.uint8)
summary = tf.summary.image(tag,imageBlocks)
return summary
# make image summary from image batch (mean/variance)
def imageSummaryMeanVar(opt,image,tag,H,W):
image = tf.concat([image,np.zeros([2,H,W,3])],axis=0)
imageOne = tf.batch_to_space_nd(image,crops=[[0,0],[0,0]],block_shape=[5,9])
imagePermute = tf.reshape(imageOne,[H,5,W,9,-1])
imageTransp = tf.transpose(imagePermute,[1,0,3,2,4])
imageBlocks = tf.reshape(imageTransp,[1,H*5,W*9,-1])
# imageBlocks = tf.cast(imageBlocks*255,tf.uint8)
summary = tf.summary.image(tag,imageBlocks)
return summary
# set optimizer for different learning rates
def setOptimizer(opt,loss,lrGP,lrC):
varsGP = [v for v in tf.global_variables() if "geometric" in v.name]
varsC = [v for v in tf.global_variables() if "classifier" in v.name]
gradC = tf.gradients(loss,varsC)
optimC = tf.train.GradientDescentOptimizer(lrC).apply_gradients(zip(gradC,varsC))
if len(varsGP)>0:
gradGP = tf.gradients(loss,varsGP)
optimGP = tf.train.GradientDescentOptimizer(lrGP).apply_gradients(zip(gradGP,varsGP))
optim = tf.group(optimC,optimGP)
else:
optim = optimC
return optim
# restore model
def restoreModel(opt,sess,saver,it):
saver.restore(sess,"models_{0}/{1}_it{2}.ckpt".format(opt.group,opt.model,it,opt.warpN))
# save model
def saveModel(opt,sess,saver,it):
saver.save(sess,"models_{0}/{1}_it{2}.ckpt".format(opt.group,opt.model,it,opt.warpN))
================================================
FILE: traffic-sign-tensorflow/warp.py
================================================
import numpy as np
import scipy.linalg
import tensorflow as tf
# fit (affine) warp between two sets of points
def fit(Xsrc,Xdst):
ptsN = len(Xsrc)
X,Y,U,V,O,I = Xsrc[:,0],Xsrc[:,1],Xdst[:,0],Xdst[:,1],np.zeros([ptsN]),np.ones([ptsN])
A = np.concatenate((np.stack([X,Y,I,O,O,O],axis=1),
np.stack([O,O,O,X,Y,I],axis=1)),axis=0)
b = np.concatenate((U,V),axis=0)
p1,p2,p3,p4,p5,p6 = scipy.linalg.lstsq(A,b)[0].squeeze()
pMtrx = np.array([[p1,p2,p3],[p4,p5,p6],[0,0,1]],dtype=np.float32)
return pMtrx
# compute composition of warp parameters
def compose(opt,p,dp):
with tf.name_scope("compose"):
pMtrx = vec2mtrx(opt,p)
dpMtrx = vec2mtrx(opt,dp)
pMtrxNew = tf.matmul(dpMtrx,pMtrx)
pMtrxNew /= pMtrxNew[:,2:3,2:3]
pNew = mtrx2vec(opt,pMtrxNew)
return pNew
# compute inverse of warp parameters
def inverse(opt,p):
with tf.name_scope("inverse"):
pMtrx = vec2mtrx(opt,p)
pInvMtrx = tf.matrix_inverse(pMtrx)
pInv = mtrx2vec(opt,pInvMtrx)
return pInv
# convert warp parameters to matrix
def vec2mtrx(opt,p):
with tf.name_scope("vec2mtrx"):
O = tf.zeros([opt.batchSize])
I = tf.ones([opt.batchSize])
if opt.warpType=="translation":
tx,ty = tf.unstack(p,axis=1)
pMtrx = tf.transpose(tf.stack([[I,O,tx],[O,I,ty],[O,O,I]]),perm=[2,0,1])
if opt.warpType=="similarity":
pc,ps,tx,ty = tf.unstack(p,axis=1)
pMtrx = tf.transpose(tf.stack([[I+pc,-ps,tx],[ps,I+pc,ty],[O,O,I]]),perm=[2,0,1])
if opt.warpType=="affine":
p1,p2,p3,p4,p5,p6,p7,p8 = tf.unstack(p,axis=1)
pMtrx = tf.transpose(tf.stack([[I+p1,p2,p3],[p4,I+p5,p6],[O,O,I]]),perm=[2,0,1])
if opt.warpType=="homography":
p1,p2,p3,p4,p5,p6,p7,p8 = tf.unstack(p,axis=1)
pMtrx = tf.transpose(tf.stack([[I+p1,p2,p3],[p4,I+p5,p6],[p7,p8,I]]),perm=[2,0,1])
return pMtrx
# convert warp matrix to parameters
def mtrx2vec(opt,pMtrx):
with tf.name_scope("mtrx2vec"):
[row0,row1,row2] = tf.unstack(pMtrx,axis=1)
[e00,e01,e02] = tf.unstack(row0,axis=1)
[e10,e11,e12] = tf.unstack(row1,axis=1)
[e20,e21,e22] = tf.unstack(row2,axis=1)
if opt.warpType=="translation": p = tf.stack([e02,e12],axis=1)
if opt.warpType=="similarity": p = tf.stack([e00-1,e10,e02,e12],axis=1)
if opt.warpType=="affine": p = tf.stack([e00-1,e01,e02,e10,e11-1,e12],axis=1)
if opt.warpType=="homography": p = tf.stack([e00-1,e01,e02,e10,e11-1,e12,e20,e21],axis=1)
return p
# warp the image
def transformImage(opt,image,pMtrx):
with tf.name_scope("transformImage"):
refMtrx = tf.tile(tf.expand_dims(opt.refMtrx,axis=0),[opt.batchSize,1,1])
transMtrx = tf.matmul(refMtrx,pMtrx)
# warp the canonical coordinates
X,Y = np.meshgrid(np.linspace(-1,1,opt.W),np.linspace(-1,1,opt.H))
X,Y = X.flatten(),Y.flatten()
XYhom = np.stack([X,Y,np.ones_like(X)],axis=1).T
XYhom = np.tile(XYhom,[opt.batchSize,1,1]).astype(np.float32)
XYwarpHom = tf.matmul(transMtrx,XYhom)
XwarpHom,YwarpHom,ZwarpHom = tf.unstack(XYwarpHom,axis=1)
Xwarp = tf.reshape(XwarpHom/(ZwarpHom+1e-8),[opt.batchSize,opt.H,opt.W])
Ywarp = tf.reshape(YwarpHom/(ZwarpHom+1e-8),[opt.batchSize,opt.H,opt.W])
# get the integer sampling coordinates
Xfloor,Xceil = tf.floor(Xwarp),tf.ceil(Xwarp)
Yfloor,Yceil = tf.floor(Ywarp),tf.ceil(Ywarp)
XfloorInt,XceilInt = tf.to_int32(Xfloor),tf.to_int32(Xceil)
YfloorInt,YceilInt = tf.to_int32(Yfloor),tf.to_int32(Yceil)
imageIdx = np.tile(np.arange(opt.batchSize).reshape([opt.batchSize,1,1]),[1,opt.H,opt.W])
imageVec = tf.reshape(image,[-1,int(image.shape[-1])])
imageVecOut = tf.concat([imageVec,tf.zeros([1,int(image.shape[-1])])],axis=0)
idxUL = (imageIdx*opt.H+YfloorInt)*opt.W+XfloorInt
idxUR = (imageIdx*opt.H+YfloorInt)*opt.W+XceilInt
idxBL = (imageIdx*opt.H+YceilInt)*opt.W+XfloorInt
idxBR = (imageIdx*opt.H+YceilInt)*opt.W+XceilInt
idxOutside = tf.fill([opt.batchSize,opt.H,opt.W],opt.batchSize*opt.H*opt.W)
def insideImage(Xint,Yint):
return (Xint>=0)&(Xint<opt.W)&(Yint>=0)&(Yint<opt.H)
idxUL = tf.where(insideImage(XfloorInt,YfloorInt),idxUL,idxOutside)
idxUR = tf.where(insideImage(XceilInt,YfloorInt),idxUR,idxOutside)
idxBL = tf.where(insideImage(XfloorInt,YceilInt),idxBL,idxOutside)
idxBR = tf.where(insideImage(XceilInt,YceilInt),idxBR,idxOutside)
# bilinear interpolation
Xratio = tf.reshape(Xwarp-Xfloor,[opt.batchSize,opt.H,opt.W,1])
Yratio = tf.reshape(Ywarp-Yfloor,[opt.batchSize,opt.H,opt.W,1])
imageUL = tf.to_float(tf.gather(imageVecOut,idxUL))*(1-Xratio)*(1-Yratio)
imageUR = tf.to_float(tf.gather(imageVecOut,idxUR))*(Xratio)*(1-Yratio)
imageBL = tf.to_float(tf.gather(imageVecOut,idxBL))*(1-Xratio)*(Yratio)
imageBR = tf.to_float(tf.gather(imageVecOut,idxBR))*(Xratio)*(Yratio)
imageWarp = imageUL+imageUR+imageBL+imageBR
return imageWarp
# warp the image
def transformCropImage(opt,imageFull,pMtrx):
with tf.name_scope("transformImage"):
refMtrx = tf.tile(tf.expand_dims(opt.bboxRefMtrx,axis=0),[opt.batchSize,1,1])
transMtrx = tf.matmul(refMtrx,pMtrx)
# warp the canonical coordinates
X,Y = np.meshgrid(np.linspace(-1,1,opt.W),np.linspace(-1,1,opt.H))
X,Y = X.flatten(),Y.flatten()
XYhom = np.stack([X,Y,np.ones_like(X)],axis=1).T
XYhom = np.tile(XYhom,[opt.batchSize,1,1]).astype(np.float32)
XYwarpHom = tf.matmul(transMtrx,XYhom)
XwarpHom,YwarpHom,ZwarpHom = tf.unstack(XYwarpHom,axis=1)
Xwarp = tf.reshape(XwarpHom/(ZwarpHom+1e-8),[opt.batchSize,opt.H,opt.W])
Ywarp = tf.reshape(YwarpHom/(ZwarpHom+1e-8),[opt.batchSize,opt.H,opt.W])
# get the integer sampling coordinates
Xfloor,Xceil = tf.floor(Xwarp),tf.ceil(Xwarp)
Yfloor,Yceil = tf.floor(Ywarp),tf.ceil(Ywarp)
XfloorInt,XceilInt = tf.to_int32(Xfloor),tf.to_int32(Xceil)
YfloorInt,YceilInt = tf.to_int32(Yfloor),tf.to_int32(Yceil)
imageIdx = np.tile(np.arange(opt.batchSize).reshape([opt.batchSize,1,1]),[1,opt.H,opt.W])
imageVec = tf.reshape(imageFull,[-1,int(imageFull.shape[-1])])
imageVecOut = tf.concat([imageVec,tf.zeros([1,int(imageFull.shape[-1])])],axis=0)
idxUL = (imageIdx*opt.fullH+YfloorInt)*opt.fullW+XfloorInt
idxUR = (imageIdx*opt.fullH+YfloorInt)*opt.fullW+XceilInt
idxBL = (imageIdx*opt.fullH+YceilInt)*opt.fullW+XfloorInt
idxBR = (imageIdx*opt.fullH+YceilInt)*opt.fullW+XceilInt
idxOutside = tf.fill([opt.batchSize,opt.H,opt.W],opt.batchSize*opt.fullH*opt.fullW)
def insideImage(Xint,Yint):
return (Xint>=0)&(Xint<opt.fullW)&(Yint>=0)&(Yint<opt.fullH)
idxUL = tf.where(insideImage(XfloorInt,YfloorInt),idxUL,idxOutside)
idxUR = tf.where(insideImage(XceilInt,YfloorInt),idxUR,idxOutside)
idxBL = tf.where(insideImage(XfloorInt,YceilInt),idxBL,idxOutside)
idxBR = tf.where(insideImage(XceilInt,YceilInt),idxBR,idxOutside)
# bilinear interpolation
Xratio = tf.reshape(Xwarp-Xfloor,[opt.batchSize,opt.H,opt.W,1])
Yratio = tf.reshape(Ywarp-Yfloor,[opt.batchSize,opt.H,opt.W,1])
imageUL = tf.to_float(tf.gather(imageVecOut,idxUL))*(1-Xratio)*(1-Yratio)
imageUR = tf.to_float(tf.gather(imageVecOut,idxUR))*(Xratio)*(1-Yratio)
imageBL = tf.to_float(tf.gather(imageVecOut,idxBL))*(1-Xratio)*(Yratio)
imageBR = tf.to_float(tf.gather(imageVecOut,idxBR))*(Xratio)*(Yratio)
imageWarp = imageUL+imageUR+imageBL+imageBR
return imageWarp
gitextract_t1x_4nxr/
├── .editorconfig
├── .gitignore
├── LICENSE
├── MNIST-pytorch/
│ ├── data.py
│ ├── graph.py
│ ├── options.py
│ ├── train.py
│ ├── util.py
│ └── warp.py
├── MNIST-tensorflow/
│ ├── data.py
│ ├── graph.py
│ ├── options.py
│ ├── train.py
│ ├── util.py
│ └── warp.py
├── README.md
└── traffic-sign-tensorflow/
├── data.py
├── graph.py
├── options.py
├── train.py
├── util.py
└── warp.py
SYMBOL INDEX (105 symbols across 15 files)
FILE: MNIST-pytorch/data.py
function loadMNIST (line 10) | def loadMNIST(opt,path):
function genPerturbations (line 22) | def genPerturbations(opt):
function makeBatch (line 56) | def makeBatch(opt,data):
function evalTest (line 66) | def evalTest(opt,data,geometric,classifier):
FILE: MNIST-pytorch/graph.py
class FullCNN (line 7) | class FullCNN(torch.nn.Module):
method __init__ (line 8) | def __init__(self,opt):
method forward (line 32) | def forward(self,opt,image):
class CNN (line 40) | class CNN(torch.nn.Module):
method __init__ (line 41) | def __init__(self,opt):
method forward (line 61) | def forward(self,opt,image):
class Identity (line 69) | class Identity(torch.nn.Module):
method __init__ (line 70) | def __init__(self): super(Identity,self).__init__()
method forward (line 71) | def forward(self,opt,feat): return [feat]
class STN (line 74) | class STN(torch.nn.Module):
method __init__ (line 75) | def __init__(self,opt):
method forward (line 97) | def forward(self,opt,image):
class ICSTN (line 109) | class ICSTN(torch.nn.Module):
method __init__ (line 110) | def __init__(self,opt):
method forward (line 132) | def forward(self,opt,image,p):
function initialize (line 149) | def initialize(opt,model,stddev,last0=False):
FILE: MNIST-pytorch/options.py
function set (line 7) | def set(training):
FILE: MNIST-pytorch/util.py
function mkdir (line 8) | def mkdir(path):
function imread (line 10) | def imread(fname):
function imsave (line 12) | def imsave(fname,array):
function toRed (line 16) | def toRed(content): return termcolor.colored(content,"red",attrs=["bold"])
function toGreen (line 17) | def toGreen(content): return termcolor.colored(content,"green",attrs=["b...
function toBlue (line 18) | def toBlue(content): return termcolor.colored(content,"blue",attrs=["bol...
function toCyan (line 19) | def toCyan(content): return termcolor.colored(content,"cyan",attrs=["bol...
function toYellow (line 20) | def toYellow(content): return termcolor.colored(content,"yellow",attrs=[...
function toMagenta (line 21) | def toMagenta(content): return termcolor.colored(content,"magenta",attrs...
function restoreModel (line 24) | def restoreModel(opt,geometric,classifier,it):
function saveModel (line 28) | def saveModel(opt,geometric,classifier,it):
class Visdom (line 32) | class Visdom():
method __init__ (line 33) | def __init__(self,opt):
method tileImages (line 38) | def tileImages(self,opt,images,H,W,HN,WN):
method trainLoss (line 44) | def trainLoss(self,opt,it,loss):
method testLoss (line 51) | def testLoss(self,opt,it,loss):
method meanVar (line 57) | def meanVar(self,opt,mean,var):
FILE: MNIST-pytorch/warp.py
function fit (line 8) | def fit(Xsrc,Xdst):
function compose (line 19) | def compose(opt,p,dp):
function inverse (line 28) | def inverse(opt,p):
function vec2mtrx (line 35) | def vec2mtrx(opt,p):
function mtrx2vec (line 61) | def mtrx2vec(opt,pMtrx):
function transformImage (line 73) | def transformImage(opt,image,pMtrx):
FILE: MNIST-tensorflow/data.py
function loadMNIST (line 9) | def loadMNIST(fname):
function genPerturbations (line 31) | def genPerturbations(opt):
function makeBatch (line 63) | def makeBatch(opt,data,PH):
function evalTest (line 75) | def evalTest(opt,sess,data,PH,prediction,imagesEval=[]):
FILE: MNIST-tensorflow/graph.py
function fullCNN (line 7) | def fullCNN(opt,image):
function CNN (line 41) | def CNN(opt,image):
function STN (line 62) | def STN(opt,image):
function ICSTN (line 94) | def ICSTN(opt,image,p):
function createVariable (line 131) | def createVariable(opt,weightShape,biasShape=None,stddev=None):
FILE: MNIST-tensorflow/options.py
function set (line 6) | def set(training):
FILE: MNIST-tensorflow/util.py
function mkdir (line 7) | def mkdir(path):
function imread (line 9) | def imread(fname):
function imsave (line 11) | def imsave(fname,array):
function toRed (line 15) | def toRed(content): return termcolor.colored(content,"red",attrs=["bold"])
function toGreen (line 16) | def toGreen(content): return termcolor.colored(content,"green",attrs=["b...
function toBlue (line 17) | def toBlue(content): return termcolor.colored(content,"blue",attrs=["bol...
function toCyan (line 18) | def toCyan(content): return termcolor.colored(content,"cyan",attrs=["bol...
function toYellow (line 19) | def toYellow(content): return termcolor.colored(content,"yellow",attrs=[...
function toMagenta (line 20) | def toMagenta(content): return termcolor.colored(content,"magenta",attrs...
function imageSummary (line 23) | def imageSummary(opt,image,tag,H,W):
function imageSummaryMeanVar (line 34) | def imageSummaryMeanVar(opt,image,tag,H,W):
function setOptimizer (line 44) | def setOptimizer(opt,loss,lrGP,lrC):
function restoreModel (line 58) | def restoreModel(opt,sess,saver,it):
function saveModel (line 61) | def saveModel(opt,sess,saver,it):
FILE: MNIST-tensorflow/warp.py
function fit (line 6) | def fit(Xsrc,Xdst):
function compose (line 17) | def compose(opt,p,dp):
function inverse (line 27) | def inverse(opt,p):
function vec2mtrx (line 35) | def vec2mtrx(opt,p):
function mtrx2vec (line 54) | def mtrx2vec(opt,pMtrx):
function transformImage (line 67) | def transformImage(opt,image,pMtrx):
FILE: traffic-sign-tensorflow/data.py
function loadGTSRB (line 11) | def loadGTSRB(opt,fname):
function genPerturbations (line 72) | def genPerturbations(opt):
function makeBatch (line 104) | def makeBatch(opt,data,PH):
function evalTest (line 116) | def evalTest(opt,sess,data,PH,prediction,imagesEval=[]):
FILE: traffic-sign-tensorflow/graph.py
function fullCNN (line 7) | def fullCNN(opt,image):
function CNN (line 38) | def CNN(opt,image):
function STN (line 63) | def STN(opt,image):
function ICSTN (line 91) | def ICSTN(opt,imageFull,p):
function createVariable (line 124) | def createVariable(opt,weightShape,biasShape=None,stddev=None):
FILE: traffic-sign-tensorflow/options.py
function set (line 6) | def set(training):
FILE: traffic-sign-tensorflow/util.py
function mkdir (line 7) | def mkdir(path):
function imread (line 9) | def imread(fname):
function imsave (line 11) | def imsave(fname,array):
function toRed (line 15) | def toRed(content): return termcolor.colored(content,"red",attrs=["bold"])
function toGreen (line 16) | def toGreen(content): return termcolor.colored(content,"green",attrs=["b...
function toBlue (line 17) | def toBlue(content): return termcolor.colored(content,"blue",attrs=["bol...
function toCyan (line 18) | def toCyan(content): return termcolor.colored(content,"cyan",attrs=["bol...
function toYellow (line 19) | def toYellow(content): return termcolor.colored(content,"yellow",attrs=[...
function toMagenta (line 20) | def toMagenta(content): return termcolor.colored(content,"magenta",attrs...
function imageSummary (line 23) | def imageSummary(opt,image,tag,H,W):
function imageSummaryMeanVar (line 34) | def imageSummaryMeanVar(opt,image,tag,H,W):
function setOptimizer (line 45) | def setOptimizer(opt,loss,lrGP,lrC):
function restoreModel (line 59) | def restoreModel(opt,sess,saver,it):
function saveModel (line 62) | def saveModel(opt,sess,saver,it):
FILE: traffic-sign-tensorflow/warp.py
function fit (line 6) | def fit(Xsrc,Xdst):
function compose (line 17) | def compose(opt,p,dp):
function inverse (line 27) | def inverse(opt,p):
function vec2mtrx (line 35) | def vec2mtrx(opt,p):
function mtrx2vec (line 54) | def mtrx2vec(opt,pMtrx):
function transformImage (line 67) | def transformImage(opt,image,pMtrx):
function transformCropImage (line 110) | def transformCropImage(opt,imageFull,pMtrx):
Condensed preview — 22 files, each showing path, character count, and a content snippet. Download the .json file or copy for the full structured content (97K chars).
[
{
"path": ".editorconfig",
"chars": 170,
"preview": "root = true\n\n[*]\nend_of_line = lf\ninsert_final_newline = true\nindent_style = tab\nindent_size = 4\ntrim_trailing_whitespac"
},
{
"path": ".gitignore",
"chars": 1045,
"preview": "# Byte-compiled / optimized / DLL files\n__pycache__/\n*.py[cod]\n*$py.class\n\n# C extensions\n*.so\n\n# Distribution / packagi"
},
{
"path": "LICENSE",
"chars": 1071,
"preview": "MIT License\n\nCopyright (c) 2018 Chen-Hsuan Lin\n\nPermission is hereby granted, free of charge, to any person obtaining a "
},
{
"path": "MNIST-pytorch/data.py",
"chars": 4356,
"preview": "import numpy as np\nimport scipy.linalg\nimport os,time\nimport torch\nimport torchvision\n\nimport warp,util\n\n# load MNIST da"
},
{
"path": "MNIST-pytorch/graph.py",
"chars": 4861,
"preview": "import numpy as np\nimport torch\nimport time\nimport data,warp,util\n\n# build classification network\nclass FullCNN(torch.nn"
},
{
"path": "MNIST-pytorch/options.py",
"chars": 4230,
"preview": "import numpy as np\nimport argparse\nimport warp\nimport util\nimport torch\n\ndef set(training):\n\n\t# parse input arguments\n\tp"
},
{
"path": "MNIST-pytorch/train.py",
"chars": 3373,
"preview": "import numpy as np\r\nimport time,os,sys\r\nimport argparse\r\nimport util\r\n\r\nprint(util.toYellow(\"==========================="
},
{
"path": "MNIST-pytorch/util.py",
"chars": 3194,
"preview": "import numpy as np\nimport scipy.misc\nimport torch\nimport os\nimport termcolor\nimport visdom\n\ndef mkdir(path):\n\tif not os."
},
{
"path": "MNIST-pytorch/warp.py",
"chars": 3391,
"preview": "import numpy as np\nimport scipy.linalg\nimport torch\n\nimport util\n\n# fit (affine) warp between two sets of points \ndef fi"
},
{
"path": "MNIST-tensorflow/data.py",
"chars": 4227,
"preview": "import numpy as np\nimport scipy.linalg\nimport os,time\nimport tensorflow as tf\n\nimport warp\n\n# load MNIST data\ndef loadMN"
},
{
"path": "MNIST-tensorflow/graph.py",
"chars": 5009,
"preview": "import numpy as np\nimport tensorflow as tf\nimport time\nimport data,warp,util\n\n# build classification network\ndef fullCNN"
},
{
"path": "MNIST-tensorflow/options.py",
"chars": 4288,
"preview": "import numpy as np\nimport argparse\nimport warp\nimport util\n\ndef set(training):\n\n\t# parse input arguments\n\tparser = argpa"
},
{
"path": "MNIST-tensorflow/train.py",
"chars": 5942,
"preview": "import numpy as np\r\nimport time,os,sys\r\nimport argparse\r\nimport util\r\n\r\nprint(util.toYellow(\"==========================="
},
{
"path": "MNIST-tensorflow/util.py",
"chars": 2576,
"preview": "import numpy as np\nimport scipy.misc\nimport tensorflow as tf\nimport os\nimport termcolor\n\ndef mkdir(path):\n\tif not os.pat"
},
{
"path": "MNIST-tensorflow/warp.py",
"chars": 4731,
"preview": "import numpy as np\nimport scipy.linalg\nimport tensorflow as tf\n\n# fit (affine) warp between two sets of points \ndef fit("
},
{
"path": "README.md",
"chars": 4552,
"preview": "## Inverse Compositional Spatial Transformer Networks\n[Chen-Hsuan Lin](https://chenhsuanlin.bitbucket.io/)\nand [Simon Lu"
},
{
"path": "traffic-sign-tensorflow/data.py",
"chars": 5747,
"preview": "import numpy as np\nimport scipy.linalg,scipy.misc\nimport os,time\nimport tensorflow as tf\nimport matplotlib.pyplot as plt"
},
{
"path": "traffic-sign-tensorflow/graph.py",
"chars": 4677,
"preview": "import numpy as np\nimport tensorflow as tf\nimport time\nimport data,warp,util\n\n# build classification network\ndef fullCNN"
},
{
"path": "traffic-sign-tensorflow/options.py",
"chars": 4780,
"preview": "import numpy as np\nimport argparse\nimport warp\nimport util\n\ndef set(training):\n\n\t# parse input arguments\n\tparser = argpa"
},
{
"path": "traffic-sign-tensorflow/train.py",
"chars": 6312,
"preview": "import numpy as np\r\nimport time,os,sys\r\nimport argparse\r\nimport util\r\n\r\nprint(util.toYellow(\"==========================="
},
{
"path": "traffic-sign-tensorflow/util.py",
"chars": 2630,
"preview": "import numpy as np\nimport scipy.misc\nimport tensorflow as tf\nimport os\nimport termcolor\n\ndef mkdir(path):\n\tif not os.pat"
},
{
"path": "traffic-sign-tensorflow/warp.py",
"chars": 7171,
"preview": "import numpy as np\nimport scipy.linalg\nimport tensorflow as tf\n\n# fit (affine) warp between two sets of points \ndef fit("
}
]
About this extraction
This page contains the full source code of the chenhsuanlin/inverse-compositional-STN GitHub repository, extracted and formatted as plain text for AI agents and large language models (LLMs). The extraction includes 22 files (86.3 KB), approximately 27.3k tokens, and a symbol index with 105 extracted functions, classes, methods, constants, and types. Use this with OpenClaw, Claude, ChatGPT, Cursor, Windsurf, or any other AI tool that accepts text input. You can copy the full output to your clipboard or download it as a .txt file.
Extracted by GitExtract — free GitHub repo to text converter for AI. Built by Nikandr Surkov.