from __future__ import division
import os
import time
import cPickle
from keras.layers import Input
from keras.models import Model
from keras_csp import config, bbox_process
from keras_csp.utilsfunc import *

os.environ["CUDA_VISIBLE_DEVICES"] = '1'
C = config.Config()
C.offset = False
cache_path = 'data/cache/cityperson/val_500'
with open(cache_path, 'rb') as fid:
    val_data = cPickle.load(fid)
num_imgs = len(val_data)
print 'num of val samples: {}'.format(num_imgs)

C.size_test = (1024, 2048)
input_shape_img = (C.size_test[0], C.size_test[1], 3)
img_input = Input(shape=input_shape_img)

# define the base network (resnet here, can be MobileNet, etc)
if C.network == 'resnet50':
    from keras_csp import resnet50 as nn
    weight_path = 'data/models/net_e54_l0.00451828871423.hdf5'
    preds = nn.nn_p3p4p5(img_input, offset=C.offset, num_scale=C.num_scale, trainable=True)
    preds_tea = nn.nn_p3p4p5(img_input, offset=C.offset, num_scale=C.num_scale, trainable=True)
elif C.network == 'mobilenet':
    from keras_csp import mobilenet as nn
    weight_path = '/home/zyh/abc/CSP/data/models/mobilenet_1_0_224_tf.h5'
    preds = nn.nn_p3p4p5(img_input)
    preds_tea = nn.nn_p3p4p5(img_input)
else:
    print 'not define backbone {}'.format(C.network)
# define the network prediction
# preds = nn.nn_p3p4p5(img_input, offset=C.offset, num_scale=C.num_scale, trainable=True)
model = Model(img_input, preds)

if C.offset:
    w_path = 'output/valmodels/city/%s/off' % (C.scale)
    out_path = 'output/valresults/city/%s/off' % (C.scale)
else:
    w_path = 'output/valmodels/city/%s/nooff' % (C.scale)
    out_path = 'output/valresults/city/%s/nooff' % (C.scale)

w_path = '/home/zyh/abc/CSP/data/models/'

if not os.path.exists(out_path):
    os.makedirs(out_path)
files = sorted(os.listdir(w_path))
# get the results from epoch 51 to epoch 150
for w_ind in range(51, 151):
    for f in files:
        if f.split('_')[0] == 'net' and int(f.split('_')[1][1:]) == w_ind:
            cur_file = f
            break
        cur_file = f
    weight1 = os.path.join(w_path, cur_file)
    print 'load weights from {}'.format(weight1)
    model.load_weights(weight1, by_name=True)
    res_path = os.path.join(out_path, '%03d' % int(str(w_ind)))
    if not os.path.exists(res_path):
        os.makedirs(res_path)
    print res_path
    res_file = os.path.join(res_path, 'val_det.txt')
    res_all = []
    start_time = time.time()
    for f in range(num_imgs):
        filepath = val_data[f]['filepath']
        img = cv2.imread(filepath)
        x_rcnn = format_img(img, C)
        Y = model.predict(x_rcnn)

        if C.offset:
            boxes = bbox_process.parse_det_offset(Y, C, score=0.1, down=4)
        else:
            boxes = bbox_process.parse_det(Y, C, score=0.1, down=4, scale=C.scale)
        if len(boxes) > 0:
            f_res = np.repeat(f + 1, len(boxes), axis=0).reshape((-1, 1))
            boxes[:, [2, 3]] -= boxes[:, [0, 1]]
            res_all += np.concatenate((f_res, boxes), axis=-1).tolist()
    np.savetxt(res_file, np.array(res_all), fmt='%6f')
    print time.time() - start_time
