from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from scipy import misc
import sys
import os
import argparse
import tensorflow as tf
import numpy as np
import mxnet as mx
import random
import cv2
import sklearn
from sklearn.decomposition import PCA
from time import sleep
from easydict import EasyDict as edict

parser = argparse.ArgumentParser(description='face model slim')
# general
parser.add_argument('--model', default='../models/model-r34-amf/model,60', help='path to load model.')
args = parser.parse_args()

_vec = args.model.split(',')
assert len(_vec)==2
prefix = _vec[0]
epoch = int(_vec[1])
print('loading',prefix, epoch)
sym, arg_params, aux_params = mx.model.load_checkpoint(prefix, epoch)
all_layers = sym.get_internals()
sym = all_layers['fc1_output']
dellist = []
for k,v in arg_params.iteritems():
  if k.startswith('fc7'):
    dellist.append(k)
for d in dellist:
  del arg_params[d]
mx.model.save_checkpoint(prefix, 0, sym, arg_params, aux_params)