import os
import sys
import rbcc
from preprocess import preprocess


def generate_img(img_dir):
    img_files = [os.path.join(img_dir, f) for f in os.listdir(img_dir) if '.jpg' in f]
    img_files.sort()

    for img_file in img_files:
        img = preprocess(img_file)
        yield img


def main(sg_file, img_dir):
    sg = rbcc.load_sg(sg_file)
    qconfig = {
        "quant_ops": {"Conv2D": {"w_axis": 0}},
        "layers": {
            "/last_layer0_6/Conv": {'bit': 16, "observer": "MinMaxObserver"},
            "/last_layer1_6/Conv": {'bit': 16, "observer": "MinMaxObserver"}
        }
    }
    qsg = rbcc.quantize(sg, dataset=generate_img(img_dir), gpu_id=0, qconfig=qconfig)
    rbcc.save_sg(qsg, sg_file.replace('.sg', '_8bit.sg'))


if __name__ == '__main__':
    main(sys.argv[1], sys.argv[2])