Commit c9adfd5f authored by Alexander Gehrke's avatar Alexander Gehrke

Add new commands for image maps

parent dbe61c1e
import argparse
import functools
import multiprocessing
import os
from ocr4all.image_map import ImageMap
from tqdm import tqdm
def main():
parser = argparse.ArgumentParser(add_help=False)
paths_args = parser.add_argument_group("Paths")
paths_args.add_argument("-i", "--input-color-map", type=str, required=True,
help="color mapping for the input images")
paths_args.add_argument("-o", "--output-color-map", type=str, required=True,
help="desired output color mapping")
paths_args.add_argument("-O", "--output-dir", type=str, required=True,
help="The output dir for the recolored masks")
paths_args.add_argument("images", metavar="IMAGE", type=str, nargs="+",
help="the images to convert")
opt_args = parser.add_argument_group("optional arguments")
opt_args.add_argument("-j", "--jobs", "--threads", metavar='THREADS', dest='threads',
type=int, default=multiprocessing.cpu_count(),
help="Number of threads to use")
opt_args.add_argument("-h", "--help", action="help", help="show this help message and exit")
args = parser.parse_args()
in_map = ImageMap.load(args.input_color_map)
out_map = ImageMap.load(args.output_color_map)
os.makedirs(args.output_dir, exist_ok=True)
convert_m = functools.partial(convert, outdir=args.output_dir,
in_map=in_map,
out_map=out_map,
)
with multiprocessing.Pool(processes=args.threads) as p:
for _ in tqdm(p.imap(convert_m, args.images), total=len(args.images), unit="img"):
pass
def convert(img_path, outdir, in_map: ImageMap, out_map: ImageMap):
labels = in_map.imread_labels(img_path)
res = out_map.to_image(labels)
res.save(os.path.join(outdir, os.path.basename(img_path)))
if __name__ == '__main__':
main()
import argparse
from ocr4all.image_map import ImageMap
def main():
parser = argparse.ArgumentParser(add_help=False)
parser.add_argument('vars', metavar='IMAGEMAP', type=str, nargs='+',
help='image map json files')
opt_args = parser.add_argument_group("optional arguments")
opt_args.add_argument("-h", "--help", action="help", help="show this help message and exit")
args = parser.parse_args()
for imap in args.vars:
map = ImageMap.load(imap)
for k, v in map.mapping.items():
print(f"\x1b[48;2;{k[0]};{k[1]};{k[2]}m \x1b[0m \x1b[38;2;{k[0]};{k[1]};{k[2]}m{v}\x1b[0m\n")
if __name__ == '__main__':
main()
......@@ -34,6 +34,16 @@ commands = {
'main': 'main',
'help': 'Generates color map'
},
'inspect-image-map': {
'script': 'ocr4all_pixel_classifier_frontend.inspect_image_map',
'main': 'main',
'help': 'Displays image map on a color-enabled terminal'
},
'convert-colors': {
'script': 'ocr4all_pixel_classifier_frontend.convert_colors',
'main': 'main',
'help': 'Convert predictions or masks from one color mapping to another'
},
'migrate-model': {
'script': 'ocr4all_pixel_classifier_frontend.migrate_model',
'main': 'main',
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment