import argparse import json from detection.detection import detect, discretize, print_results from detection.refine import simulate, save from os.path import splitext, basename, join def main(): ap = argparse.ArgumentParser() ap.add_argument("-i", "--input", required=True, help="Uses this input as image for detection") ap.add_argument("-s", "--scale", type=float, default=0.10, help="Scales images by this factor (default: x0.1)") ap.add_argument("-c", "--config", type=str, default="detection/config.json", help="Loads config from this file (default: detection/config.json)") ap.add_argument("--save", type=str, default="save/", help="Pass the directory where saves are stored. (default: save/)") ap.add_argument("--hide", action='store_true', default=False, help="Hide images if parameter is set") ap.add_argument("--refine", type=str, help="Pass a json file to refine model") args = ap.parse_args() with open(args.config, 'r') as file: config = json.load(file) if args.refine is not None: with open(args.refine, 'r') as file: refine = json.load(file) else: refine = None orig, blob_mask, blob, food_mask, food_img = detect(args.input, config) dsc_img, dsc_blob, dsc_food_list = discretize(blob, food_mask, config['Discrete Width'], config['Discrete Height']) if args.save is not None: filename = splitext(basename(args.input))[0] + "-detect" file_path = join(args.save, filename) board, player, img = simulate(dsc_img, dsc_blob, dsc_food_list, config, refine) save(file_path, board, player, img) # Prepare file_path for details if any to save file_path += "-details" else: file_path = None print_results(orig, blob_mask, blob, food_mask, food_img, dsc_img, args.scale, file_path, args.hide) if __name__ == "__main__": main()