#!/bin/bash

__doc__="

This demonstrates an end-to-end pipeline on multispectral toydata

This walks through the entire process of fit -> predict -> evaluate and the

output if you run this should end with something like

source ~/code/geowatch/docs/source/manual/tutorial/tutorial2_msi_network.sh

"

# Define wherever you want to store results and create the directories

# if they don't already exist.

DVC_DATA_DPATH=$HOME/data/dvc-repos/toy_data_dvc

DVC_EXPT_DPATH=$HOME/data/dvc-repos/toy_expt_dvc

# Create the above directories of they don't exist

mkdir -p "$DVC_DATA_DPATH"

mkdir -p "$DVC_EXPT_DPATH"

# The user can overwrite these variables, but will default if they are not given.

NUM_TOY_TRAIN_VIDS="${NUM_TOY_TRAIN_VIDS:=100}" # If variable not set or null, use default.

NUM_TOY_VALI_VIDS="${NUM_TOY_VALI_VIDS:=5}" # If variable not set or null, use default.

NUM_TOY_TEST_VIDS="${NUM_TOY_TEST_VIDS:=2}" # If variable not set or null, use default.

# Generate toy datasets

TRAIN_FPATH=$DVC_DATA_DPATH/vidshapes_msi_train${NUM_TOY_TRAIN_VIDS}/data.kwcoco.json

VALI_FPATH=$DVC_DATA_DPATH/vidshapes_msi_vali${NUM_TOY_VALI_VIDS}/data.kwcoco.json

TEST_FPATH=$DVC_DATA_DPATH/vidshapes_msi_test${NUM_TOY_TEST_VIDS}/data.kwcoco.json

generate_data(){

kwcoco toydata --key="vidshapes${NUM_TOY_TRAIN_VIDS}-frames5-randgsize-speed0.2-msi-multisensor" \

--bundle_dpath "$DVC_DATA_DPATH/vidshapes_msi_train${NUM_TOY_TRAIN_VIDS}" --verbose=1

kwcoco toydata --key="vidshapes${NUM_TOY_VALI_VIDS}-frames5-randgsize-speed0.2-msi-multisensor" \

--bundle_dpath "$DVC_DATA_DPATH/vidshapes_msi_vali${NUM_TOY_VALI_VIDS}" --verbose=1

kwcoco toydata --key="vidshapes${NUM_TOY_TEST_VIDS}-frames6-randgsize-speed0.2-msi-multisensor" \

--bundle_dpath "$DVC_DATA_DPATH/vidshapes_msi_test${NUM_TOY_TEST_VIDS}" --verbose=1

}

print_stats(){

# Print stats

kwcoco stats "$TRAIN_FPATH" "$VALI_FPATH" "$TEST_FPATH"

geowatch stats "$TRAIN_FPATH" "$VALI_FPATH" "$TEST_FPATH"

}

# If the train file path does not exist then

if [[ ! -e "$TRAIN_FPATH" ]]; then

generate_data

print_stats

fi

__doc__="

Should look like

dset n_anns n_imgs n_videos n_cats r|g|b|disparity|gauss|B8|B11 B1|B8|B8a|B10|B11 r|g|b|flowx|flowy|distri|B10|B11

0 vidshapes_msi_train/data.kwcoco.json 80 40 8 3 12 12 16

1 vidshapes_msi_vali/data.kwcoco.json 50 25 5 3 9 10 6

2 vidshapes_msi_test/data.kwcoco.json 24 12 2 3 5 3 4

"

demo_visualize_toydata(){

python -m geowatch.cli.coco_visualize_videos \

--src "$TRAIN_FPATH" \

--channels="gauss|B11,r|g|b,B1|B8|B11" \

--viz_dpath="$DVC_DATA_DPATH/vidshapes_msi_train_viz" --animate=False

}

# Run the function if you want to visualize the data

# demo_visualize_toydata

# Define the channels we want to use

# The sensors and channels are specified by the kwcoco SensorChanSpec

# in this example the data does not contain sensor metadata, so we

# use a "*" to indicate a generic sensor.

# A colon ":" separates channels from the sensors.

# A pipe "|" indicates channels are early fused

# A "," indicates groups of early fused channels are late fused Multiple

# sensors can be specified to the left of the channels and will distribute over

# commas.

# Note: the channel names dont really mean anything. The demo data generated

# for them is random / arbitrary. They simply simulate a case where you do have

# a meaningful multiple-channel dataset you want to learn with.

CHANNELS="(*):(disparity|gauss,X.2|Y:2:6,B1|B8a,flowx|flowy|distri)"

echo "CHANNELS = $CHANNELS"

echo "

In Tutorial #2 we expand the complexity of the 'fit' command compared to the

previous tutorial. We define some of the important parameters as environment

variables to illustrate relationships between them.

"

# Fit

DATASET_CODE=ToyDataMSI

WORKDIR=$DVC_EXPT_DPATH/training/$HOSTNAME/$USER

EXPERIMENT_NAME=ToyDataMSI_Demo_V001

DEFAULT_ROOT_DIR=$WORKDIR/$DATASET_CODE/runs/$EXPERIMENT_NAME

TARGET_LR=3e-4

WEIGHT_DECAY=$(python -c "print($TARGET_LR * 0.01)")

# Try to set the accelerator to "GPU", but fallback to cpu if nvidia-smi is not found

if which nvidia-smi ; then

echo "Detected an NVIDIA GPU"

export ACCELERATOR=gpu

# Set CUDA_VISIBLE_DEVICES to a comma separated list of the GPU #'s you want to

# use (index may start at 0).

export CUDA_VISIBLE_DEVICES=0

else

echo "No GPU detected, falling back to CPU"

export ACCELERATOR=cpu

export CUDA_VISIBLE_DEVICES=

fi

# However, if you have only have a cpu, the tutorial will still run.

# Comment out the above two commands and enable the command below.

# export ACCELERATOR=cpu

DEVICES=$(python -c "if 1:

import os

if os.environ.get('ACCELERATOR', '') == 'cpu':

print('1')

else:

n = len(os.environ.get('CUDA_VISIBLE_DEVICES', '').split(','))

print(','.join(list(map(str, range(n)))) + ',')

")

STRATEGY=$(python -c "if 1:

import os

if os.environ.get('ACCELERATOR', '') == 'cpu':

print('auto')

else:

n = len(os.environ.get('CUDA_VISIBLE_DEVICES', '').split(','))

print('ddp' if n > 1 else 'auto')

")

DDP_WORKAROUND=$(python -c "if 1:

import os

n = len(os.environ.get('CUDA_VISIBLE_DEVICES', '').split(','))

print(int(n > 1))

")

echo "

CUDA_VISIBLE_DEVICES = $CUDA_VISIBLE_DEVICES

ACCELERATOR = $ACCELERATOR

DEVICES = $DEVICES

STRATEGY = $STRATEGY

DDP_WORKAROUND = $DDP_WORKAROUND

"

# The following parameters are related and impose constraints on each other.

MAX_STEPS=400

MAX_EPOCHS=20

TRAIN_BATCHES_PER_EPOCH=20

ACCUMULATE_GRAD_BATCHES=1

BATCH_SIZE=2

TRAIN_ITEMS_PER_EPOCH=$(python -c "print($TRAIN_BATCHES_PER_EPOCH * $BATCH_SIZE)")

# Thus we expose this recommendation script to suggest changes to that satisfy

# the constraints. (Its ok if it isnt perfect)

python -m geowatch.cli.experimental.recommend_size_adjustments \

--MAX_STEPS=$MAX_STEPS \

--MAX_EPOCHS=$MAX_EPOCHS \

--BATCH_SIZE=$BATCH_SIZE \

--ACCUMULATE_GRAD_BATCHES=$ACCUMULATE_GRAD_BATCHES \

--TRAIN_BATCHES_PER_EPOCH="$TRAIN_BATCHES_PER_EPOCH" \

--TRAIN_ITEMS_PER_EPOCH="$TRAIN_ITEMS_PER_EPOCH"

# Now that we have important variables defined in the environment we can run

# the actual fit command.

DDP_WORKAROUND=$DDP_WORKAROUND python -m geowatch.tasks.fusion fit --config "

data:

num_workers : 4

train_dataset : $TRAIN_FPATH

vali_dataset : $VALI_FPATH

channels : '$CHANNELS'

time_steps : 5

chip_dims : 128

batch_size : $BATCH_SIZE

model:

class_path: MultimodalTransformer

init_args:

name : $EXPERIMENT_NAME

arch_name : smt_it_stm_p8

window_size : 8

dropout : 0.1

global_saliency_weight: 1.0

global_class_weight: 1.0

global_change_weight: 0.0

optimizer:

class_path: torch.optim.Adam

init_args:

lr: $TARGET_LR

weight_decay: $WEIGHT_DECAY

betas:

- 0.9

- 0.99

trainer:

accumulate_grad_batches: $ACCUMULATE_GRAD_BATCHES

default_root_dir : $DEFAULT_ROOT_DIR

accelerator : $ACCELERATOR

devices : $DEVICES

strategy : $STRATEGY

limit_train_batches : $TRAIN_BATCHES_PER_EPOCH

check_val_every_n_epoch: 1

enable_checkpointing: true

enable_model_summary: true

log_every_n_steps: 5

logger: true

max_steps: $MAX_STEPS

num_sanity_val_steps: 0

initializer:

init: noop

"

# This is the default final package that is done at the end of training.

# See comments inside of demo_force_repackage

PACKAGE_FPATH="$DEFAULT_ROOT_DIR"/final_package.pt

demo_force_repackage(){

# TODO: the new watch.mlops tool will need to handle this soon.

# NOTE: The above fit script might not produce the "best" checkpoint as the

# final output package. To evaluate a different checkpoint it must first be

# packaged. (note this can be done while training is running so

# intermediate checkpoints can be evaluated while the model is still

# learning). The following is logic for how to "package" a single checkpoint.

#

# Look at all checkpoints

ls "$DEFAULT_ROOT_DIR"/*/*/checkpoints/*.ckpt

# Grab the latest checkpoint (this is an arbitrary choice)

CHECKPOINT_FPATH=$(find "$DEFAULT_ROOT_DIR" -iname "*.ckpt" | tail -n 1)

echo "CHECKPOINT_FPATH = $CHECKPOINT_FPATH"

# Repackage a particular checkpoint as a torch.package .pt file.

python -m geowatch.tasks.fusion.repackage repackage "$CHECKPOINT_FPATH"

# Redefine package fpath variable to be that checkpoint

PACKAGE_FPATH=$(python -m geowatch.tasks.fusion.repackage repackage "$CHECKPOINT_FPATH" | tail -n 1)

echo "PACKAGE_FPATH = $PACKAGE_FPATH"

}

# Define where we will output predictions / evaluations

# Later we will see how mlops.schedule_evaluation will do this for you, but in

# this tutorial we are going to run prediction and evaluation manually.

PRED_FPATH=$DEFAULT_ROOT_DIR/predictions/pred.kwcoco.json

EVAL_DPATH=$DEFAULT_ROOT_DIR/evaluation

# Predict using one of the packaged models

python -m geowatch.tasks.fusion.predict \

--package_fpath="$PACKAGE_FPATH" \

--test_dataset="$TEST_FPATH" \

--pred_dataset="$PRED_FPATH" \

# Dump stats of truth vs prediction.

# We should see soft segmentation masks in pred, but not in truth

kwcoco stats "$TEST_FPATH" "$PRED_FPATH"

geowatch stats "$TEST_FPATH" "$PRED_FPATH"

# Visualize pixel predictions with a raw band, predicted saliency, and predicted class.

geowatch visualize "$PRED_FPATH" --channels='B11,salient,star|superstar|eff' --smart=True

# Evaluate the predictions

python -m geowatch.tasks.fusion.evaluate \

--true_dataset="$TEST_FPATH" \

--pred_dataset="$PRED_FPATH" \

--eval_dpath="$EVAL_DPATH"