import caffe
import surgery, score
import sys
import numpy as np
import os

import setproctitle
setproctitle.setproctitle(os.path.basename(os.getcwd()))

weights = '../voc-fcn32s/fcn32s-heavy-pascal.caffemodel'

# init
caffe.set_device(int(sys.argv[1]))
caffe.set_mode_gpu()

solver = caffe.SGDSolver('solver.prototxt')
solver.net.copy_from(weights)

# surgeries

interp_layers = [k for k in solver.net.params.keys() if 'up' in k]
# interp_layers: ['upscore2', 'upscore16']
surgery.interp(solver.net, interp_layers)

# scoring
val = np.loadtxt('/data02/bioinf/treml/pascal/VOC2011/ImageSets/Segmentation/seg11valid.txt', dtype=str)

#for _ in range(25):
#    solver.step(4000)
#    score.seg_tests(solver, False, val, layer='score')
