import os
import sys
import rbcc
from preprocess import preprocess, Cityscapes


def generate_img(img_dir):
    data = Cityscapes(img_dir, split='val', target_type='semantic')

    for img_path in data.images:
        img = preprocess(img_path)
        yield img


def main(sg_file, img_dir):
    sg = rbcc.load_sg(sg_file)
    qsg = rbcc.quantize(sg, dataset=generate_img(img_dir), gpu_id=0, max_num=50)
    rbcc.save_sg(qsg, sg_file.replace('.sg', '_8bit.sg'))


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