import ubelt as ub
import scriptconfig as scfg
import os
[docs]
def main(cmdline=False, **kwargs):
"""
Example:
>>> from geowatch.cli.coco_reformat_channels import * # NOQA
>>> import kwcoco
>>> import kwarray
>>> import shutil
>>> # Make a dataset we can modify inplace
>>> orig_dset = kwcoco.CocoDataset.demo('vidshapes1-msi')
>>> orig_dset.reroot(absolute=False)
>>> orig_bundle = ub.Path(orig_dset.bundle_dpath)
>>> new_bundle = ub.Path.appdir('kwcoco/tests/test_reformat_channels')
>>> new_bundle.delete().ensuredir()
>>> shutil.copytree(orig_bundle, new_bundle, dirs_exist_ok=True)
>>> new_fpath = new_bundle / 'data.kwcoco.json'
>>> orig_dset.dump(new_fpath)
>>> new_dset = kwcoco.CocoDataset(new_fpath)
>>> gid = 1
>>> coco_img = new_dset.coco_image(gid)
>>> rng = kwarray.ensure_rng(432)
>>> imdata = rng.rand(128, 128)
>>> coco_img.add_auxiliary_item('test_prediction.tif', channels='salient', imdata=imdata, imwrite=True)
>>> new_dset.dump(new_dset.fpath)
>>> # now reformat this new dataset
>>> orig_pred = coco_img.imdelay('salient').finalize()
>>> kwargs = {
>>> 'src': new_dset.fpath,
>>> 'quantize': 'salient',
>>> #'remove': 'B11',
>>> }
>>> cmdline = False
>>> main(cmdline=False, **kwargs)
>>> reformatted_dset = kwcoco.CocoDataset(new_fpath)
>>> assert 'quantization' in reformatted_dset.imgs[1]['auxiliary'][-1]
>>> new_coco_img = reformatted_dset.coco_image(gid)
>>> import numpy as np
>>> new_pred1 = np.nan_to_num(new_coco_img.imdelay('salient').finalize())
>>> #assert np.allclose(new_pred1, new_pred2)
>>> #new_pred2 = new_coco_img.imdelay('salient').finalize(dequantize=False)
>>> #assert new_pred2.dtype.kind == 'i'
"""
config = CocoReformatChannels.cli(data=kwargs, cmdline=cmdline, strict=True)
import kwcoco
print('config = {}'.format(ub.urepr(config, nl=1)))
dset = kwcoco.CocoDataset.coerce(config['src'])
to_quantize = kwcoco.ChannelSpec.coerce('' if config['quantize'] is None else config['quantize']).fuse()
to_remove = kwcoco.ChannelSpec.coerce('' if config['remove'] is None else config['remove']).fuse()
tasks = []
for coco_img in dset.images().coco_images:
for stream in coco_img.channels.streams():
task = {}
has_quantize = stream & to_quantize
has_remove = stream & to_remove
obj = coco_img.find_asset_obj(stream)
old_quantization = obj.get('quantization', None)
if has_remove.numel():
has_quantize = has_quantize - has_remove
if has_remove is not None and has_remove.numel():
task['has_remove'] = has_remove
if old_quantization is None:
# Dont redo quantization
if has_quantize.numel():
total_modified = has_quantize.as_oset() | has_remove.as_oset()
task['has_quantize'] = has_quantize
if len(total_modified) != len(stream):
missing = stream - total_modified
msg = f'Must quantize or remove entire stream: {stream=!s}, {has_quantize=!s}, {has_remove=!s}, {missing=!s}'
task['error'] = msg
if task:
task['gid'] = coco_img.img['id']
task['stream'] = stream
tasks.append(task)
if len(tasks) == 0:
print('No modifications need to be made')
else:
print('Found {} tasks'.format(len(tasks)))
errors = [t['error'] for t in tasks if t.get('error', None)]
if errors:
for msg in errors:
print(msg)
raise ValueError('Cannot modify')
gid_to_tasks = ub.group_items(tasks, lambda x: x['gid'])
bundle_dpath = ub.Path(dset.bundle_dpath)
jobs = ub.JobPool('process', max_workers=config['workers'])
for gid, sub_tasks in ub.ProgIter(list(gid_to_tasks.items()), desc='submit reformat jobs'):
coco_img = dset.coco_image(gid)
for task in sub_tasks:
stream = task['stream']
has_remove = task.get('has_remove', None)
has_quantize = task.get('has_quantize', None)
obj = coco_img.find_asset_obj(stream)
job = jobs.submit(reformat_obj, obj, bundle_dpath, has_remove,
has_quantize)
job.obj = obj
finalize_tasks = []
for job in jobs.as_completed(desc='reformat images'):
fpath, new_fpath, new_obj = job.result()
# Update the coco object in memory
job.obj.update(new_obj)
finalize_tasks.append({
'dst': fpath,
'src': new_fpath,
})
tmp_fpath = ub.Path(dset.fpath).augment(tail='.movtmp')
dset.dump(tmp_fpath, newlines=True)
# Point of no return
# Finalize everything by executing move
os.rename(tmp_fpath, dset.fpath)
for task in ub.ProgIter(finalize_tasks, desc='finalize image move'):
os.rename(task['src'], task['dst'])
[docs]
def schedule_quantization():
# TODO: can remove this code, or move it to a "dev" or "scripts" dir.
# Temporary job
import geowatch
dvc_dpath = ub.Path(str(geowatch.find_dvc_dpath()) + '-hdd')
dvc_dpath / 'models/fusion/'
pred_globpat = dvc_dpath / 'models/fusion/Drop2-Aligned-TA1-2022-02-15/*/pred_*/*/pred.kwcoco.json'
# pred_globpat = dvc_dpath / 'models/fusion/Drop2-Aligned-TA1-2022-02-15/*/*/*/pred.kwcoco.json'
import glob
pred_fpaths = glob.glob(str(pred_globpat))
import cmd_queue
queue = cmd_queue.Queue.create(backend='tmux', name='quantize-jobs', size=8)
for pred_fpath in pred_fpaths:
queue.submit(ub.codeblock(
rf'''
python -m geowatch.cli.coco_reformat_channels \
--src={pred_fpath} \
--quantize="not_salient|salient,Site Preparation|Active Construction|Post Construction|No Activity" \
--remove="not_salient,negative,ignore,negative,positive,background|Unknown" \
--workers=0
'''))
queue.rprint()
queue.run()
queue.monitor()
if __name__ == '__main__':
"""
CommandLine:
python ~/code/watch/geowatch/cli/coco_combine_features.py
"""
main(cmdline=True)