#!/bin/bash

basedir=`pwd`

img="../../pics/147.png"

baseline_code_dir='shrec-22/segmentation-inference'
baseline_model_dir='../../baseline/models/'
baseline_model_name='pothole-segmentron-deeplabv3+-resnet101-basic-data-aug-img_size-540-540-2a-2021-11-04_11.02.00-WD-0.0001-BS-4-LR-0.0001-0.001-epoch-16-valid_loss-0.0402.pth'

team1_model_dir='../submission-1/team1/models/'
team1_manet_name='weights_smp_manet_320_effb5_cecv_01_last.pt'
team1_unet_name='weights_smp_unet_320_effb0_cecv_01_metric.pt'
team1_unetpp_name='weights_smp_unetpp_320_effb5_cecv_01_loss.pt'

declare -a out_fns

git lfs pull		# just in case...

test_imagemagick()
{
 convert pics/147.png /tmp/147.png.jpg # ; echo $?
 if [ $? -ne 0 ] ; then
	echo '*******************************************'
	echo '* Imagemagick\"s convert is not installed *'
	echo '*         Please install it with:         *'
	echo '*      sudo apt install imagemagick       *'
	echo '*******************************************'
	exit 1
 fi
}

test_imagemagick

img_dir=`mktemp --directory --suffix=temp-147`
#img_dir='/tmp/tmp.YNh47yAimbtemp-147'						# use these for testing

create_venv()
{
	venvdir=`mktemp --directory --suffix=pothole-mix-segmentation-inference-venv`
	#venvdir="/tmp/tmp.WqBb61LpIwpothole-mix-segmentation-inference-venv"	# use these for testing

	echo "Creating virtual env into: $venvdir"

	python3.8 -m venv "$venvdir"

	echo "Sourcing the environment with: source \"$venvdir/bin/activate\""
	source "$venvdir"/bin/activate
}

part_one()
{
	cd "$baseline_code_dir"

	echo "Moved to directory:" `pwd`

	mkdir models

	echo "Installing all needed packages with: pip install -r requirements.txt"

	pip install -r requirements.txt
}

baseline()
{
	## Baseline
	
	echo '****************************************'
	echo '************** BASELINE ****************'
	echo '****************************************'
	
	outdir=outdir/baseline-deeplabv3+-resnet101
	mkdir -p $outdir
	
	cd models
	ln -s "$baseline_model_dir"/"$baseline_model_name"
	cd ..
	
	classes='asphalt[0;0;0],pothole[0;0;255],crack[255;0;0]'
	
	./inference_v2_seg.py --model "$baseline_model_name" --dataset-dir ./fake-dataset --classes "$classes" --img "$img"
	
	mv *147* "$outdir"

	out_fns[0]="$baseline_code_dir/$outdir/blend*147*"
}

manet()
{
	echo '****************************************'
	echo '**************** MANET *****************'
	echo '****************************************'
	
	outdir=outdir/manet-320-effb5
	mkdir -p $outdir
	
	cd models
	ln -s "$team1_model_dir"/"$team1_manet_name"
	cd ..
	
	classes='asphalt[0;0;0],pothole[255;0;0],crack[0;0;255]'
	
	./inference_v2_seg.py --arch create_model_manet_320_effb5 --model "$team1_model_dir"/"$team1_manet_name" --dataset-dir ./fake-dataset --classes "$classes" --img "$img"
	
	mv *147* "$outdir"

	out_fns[1]="$baseline_code_dir/$outdir/blend*147*"
}


unet()
{
	echo '****************************************'
	echo '***************** UNET *****************'
	echo '****************************************'
	
	outdir=outdir/unet_320_effb0
	mkdir -p $outdir
	
	cd models
	ln -s "$team1_model_dir"/"$team1_unet_name"
	cd ..
	
	classes='asphalt[0;0;0],pothole[255;0;0],crack[0;0;255]'
	
	./inference_v2_seg.py --arch create_model_unet_320_effb0 --model "$team1_model_dir"/"$team1_unet_name" --dataset-dir ./fake-dataset --classes "$classes" --img "$img"
	
	mv *147* "$outdir"

	out_fns[2]="$baseline_code_dir/$outdir/blend*147*"
}


unetpp()
{
	echo '****************************************'
	echo '**************** UNET++ ****************'
	echo '****************************************'
	
	outdir=outdir/unetpp_320_effb5
	mkdir -p $outdir
	
	cd models
	ln -s "$team1_model_dir"/"$team1_unetpp_name"
	cd ..
	
	classes='asphalt[0;0;0],pothole[255;0;0],crack[0;0;255]'
	
	./inference_v2_seg.py --arch create_model_unetpp_320_effb5 --model "$team1_model_dir"/"$team1_unetpp_name" --dataset-dir ./fake-dataset --classes "$classes" --img "$img"
	
	mv *147* "$outdir"

	out_fns[3]="$baseline_code_dir/$outdir/blend*147*"
}

create_venv

part_one
baseline
manet
unet
unetpp

echo
echo '***********************************************************'
echo '***********************************************************'
echo "Results for \"baseline\", \"manet\", \"unet\" and \"unetpp\" are in"
echo "shrec-22/segmentation-inference/outdir/unetpp_320_effb5"
echo '***********************************************************'
echo '***********************************************************'
echo


part_two()
{
	cd "$basedir"
	cd shrec-22/submission-2/team2
	
	echo "Moved to directory:" `pwd`
	
	echo "Installing antlr4-python3-runtime because it usually fails at the first attempt..."
	pip install antlr4-python3-runtime
	
	echo "Installing all needed Team 2 packages with: pip install -e ."
	pip install -e .
	
	echo "Copying image 147.png to $img_dir"
	cp "$basedir/pics/147.png" "$img_dir"
	cp "$basedir/pics/147.png" "$img_dir/147-bis.png"		# don't ask me why, with just one img, it doesn't work -_-
}

segformer()
{
	# ## **Run 0: SegFormer**
	# 
	# - On images
	# ``` bash
	# python configs/segmentation/infer.py \
	#         -c ckpt/0.segformer/test.yaml \
	#         -o data.dataset.args.image_dir=$IMAGE_DIR
	# ```
	python configs/segmentation/infer.py -c ckpt/0.segformer/test.yaml -o data.dataset.args.image_dir=$img_dir
}

deeplab()
{
	# ## **Run 1: Efficient DeeplabV3+**
	# 
	# - On images
	# ``` bash
	# python configs/segmentation/infer.py \
	#         -c ckpt/1.deeplabv3plus/test.yaml \
	#         -o data.dataset.args.image_dir=$IMAGE_DIR
	# ```
	python configs/segmentation/infer.py -c ckpt/1.deeplabv3plus/test.yaml -o data.dataset.args.image_dir=$img_dir
}

maskedsoftcps()
{
	# ## **Run 2: Masked Soft Cross Pseudo Supervision**
	# 
	# - On images
	# ``` bash
	# python configs/cps/infer.py \
	#         -c ckpt/2.maskedsoftcps-dlunet/test.yaml \
	#         -o data.dataset.args.image_dir=$IMAGE_DIR
	# ```
	python configs/cps/infer.py -c ckpt/2.maskedsoftcps-dlunet/test.yaml -o data.dataset.args.image_dir=$img_dir
}

part_two
segformer
deeplab
maskedsoftcps

cd $basedir

show_outfiles()
{
	outfiles=`ls shrec-22/submission-2/team2/outputs/*/images/overlays/147.png`
	
	echo '****************************'
	echo '* You can find the output in'
	echo '****************************'
	
	idx=4
	for i in $outfiles
	do
		echo $i
		out_fns[$idx]="$i"
		idx=$((idx+1))
	
	done
	echo '****************************'
	
	for fn in "${out_fns[@]}"
	do
		echo "fn: $fn"
	done
}

overlay_gt()
{
	echo "Copying GT mask for final overlaying and stitching of generated images..."
	
	cp "$basedir/pics/147-gt-mask.png" "$img_dir"
	
	#convert \( "$img_dir/147-gt-mask.png" -fill "#FF0000" -opaque "#00FF00" -modulate 75,100,100 -transparent black -background transparent \) "$img_dir/147-gt-mask-transparent.png"
	convert \( "$img_dir/147-gt-mask.png" -fill "#FF3030" -opaque "#00FF00" -transparent black -background transparent \) "$img_dir/147-gt-mask-transparent.png"
	#convert "$img_dir/147-gt-mask-transparent.png" "$img_dir/147.png" -compose dstover -composite "$img_dir/147-overlay-with-gt-mask.jpg"
	convert \( "$img_dir/147-gt-mask-transparent.png" -alpha set -channel A -evaluate set 60% -fuzz 40% -transparent black -background transparent \) "$img_dir/147.png" -compose dstover -composite "$img_dir/147-overlay-with-gt-mask.jpg"
	
	out_fns[7]="$img_dir/147-overlay-with-gt-mask.jpg"
}

composite_fig_6()
{
	echo "Running: convert \"${out_fns[0]}\" \"${out_fns[1]}\" \"${out_fns[2]}\" \"${out_fns[3]}\" +append \"$img_dir/upper-half.jpg\""
	convert "${out_fns[0]}" "${out_fns[1]}" "${out_fns[2]}" "${out_fns[3]}" +append "$img_dir/upper-half.jpg"
	echo "Running: convert \"${out_fns[4]}\" \"${out_fns[5]}\" \"${out_fns[6]}\" \"${out_fns[7]}\" +append \"$img_dir/lower-half.jpg\""
	convert "${out_fns[4]}" "${out_fns[5]}" "${out_fns[6]}" "${out_fns[7]}" +append "$img_dir/lower-half.jpg"
	echo "Running: convert \"$img_dir/upper-half.jpg\" \"$img_dir/lower-half.jpg\" +append \"$img_dir/fig-6.jpg\""
	convert "$img_dir/upper-half.jpg" "$img_dir/lower-half.jpg" -append "$img_dir/fig-6.jpg"

	echo '*******************************************************'
	echo '*******************************************************'
	echo '*******************************************************'
	echo "Copying $img_dir/fig-6.jpg to $basedir..."
	echo '*******************************************************'
	echo '*******************************************************'
	echo '*******************************************************'
	cp "$img_dir/fig-6.jpg" "$basedir"
}

show_outfiles
overlay_gt
composite_fig_6

echo "All done."

exit