12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849 |
- 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()
|