#!/usr/bin/env bash
#
# fetch-kitti.sh --- download the KITTI Road/Lane benchmark used in CSCI 3240U.
#
# Usage:
#     ./fetch-kitti.sh [destination]        # default destination: ./kitti
#     ./fetch-kitti.sh --left-only [dest]   # skip the right (stereo) images
#
# You only need the stereo images for Lab 11.  Until then, --left-only is
# enough and saves you about 430 MB.
#
# Windows users: run this from WSL or Git Bash.
#
# Faisal Z. Qureshi
# faisal.qureshi@ontariotechu.ca

set -euo pipefail

BASE_URL="https://s3.eu-central-1.amazonaws.com/avg-kitti"
LEFT_ZIP="data_road.zip"
RIGHT_ZIP="data_road_right.zip"
LEFT_BYTES=470992343
RIGHT_BYTES=447295020

WANT_RIGHT=1
if [ "${1:-}" = "--left-only" ]; then
    WANT_RIGHT=0
    shift
fi
DEST="${1:-./kitti}"

say() { printf '\033[1m==>\033[0m %s\n' "$*"; }
die() { printf '\033[1;31mError:\033[0m %s\n' "$*" >&2; exit 1; }

command -v curl  >/dev/null || die "curl is required but was not found."
command -v unzip >/dev/null || die "unzip is required but was not found."

mkdir -p "$DEST"
cd "$DEST"

# ---------------------------------------------------------------- download ---
# Downloads are resumable (curl -C -), so a dropped connection is not fatal;
# just run the script again.
fetch() {
    local zip="$1" expected="$2" actual
    if [ -f "$zip" ]; then
        actual=$(wc -c < "$zip" | tr -d ' ')
        if [ "$actual" = "$expected" ]; then
            say "$zip already present and complete, skipping download."
            return
        fi
        say "$zip is incomplete ($actual of $expected bytes), resuming."
    else
        say "Downloading $zip ($(( expected / 1048576 )) MB) ..."
    fi
    curl -# -L -C - -o "$zip" "$BASE_URL/$zip"

    actual=$(wc -c < "$zip" | tr -d ' ')
    [ "$actual" = "$expected" ] || die "$zip is $actual bytes, expected $expected. Delete it and re-run."
    unzip -tq "$zip" >/dev/null || die "$zip failed its integrity check. Delete it and re-run."
}

fetch "$LEFT_ZIP" "$LEFT_BYTES"
if [ "$WANT_RIGHT" = 1 ]; then
    fetch "$RIGHT_ZIP" "$RIGHT_BYTES"
fi

# ----------------------------------------------------------------- extract ---
say "Extracting $LEFT_ZIP ..."
unzip -q -o "$LEFT_ZIP"

if [ "$WANT_RIGHT" = 1 ]; then
    # Note: data_road_right.zip has training/ and testing/ at its top level,
    # with no data_road/ prefix, so it must be extracted into a staging
    # directory and merged rather than unzipped in place.
    say "Extracting $RIGHT_ZIP ..."
    rm -rf .right-staging
    mkdir -p .right-staging
    unzip -q -o "$RIGHT_ZIP" -d .right-staging
    for split in training testing; do
        if [ -d ".right-staging/$split/image_3" ]; then
            rm -rf "data_road/$split/image_3"
            mv ".right-staging/$split/image_3" "data_road/$split/image_3"
        fi
    done
    rm -rf .right-staging
fi

# ------------------------------------------------------------------ verify ---
count() { ls "data_road/$1" 2>/dev/null | wc -l | tr -d ' '; }

say "Verifying ..."
ok=1
check() {
    local path="$1" want="$2" got
    got=$(count "$path")
    if [ "$got" = "$want" ]; then
        printf '    %-24s %4s  ok\n' "$path" "$got"
    else
        printf '    %-24s %4s  EXPECTED %s\n' "$path" "$got" "$want"
        ok=0
    fi
}

check training/image_2    289
check training/calib      289
check training/gt_image_2 384
check testing/image_2     290
check testing/calib       290
if [ "$WANT_RIGHT" = 1 ]; then
    check training/image_3 289
    check testing/image_3  290
fi

[ "$ok" = 1 ] || die "Verification failed. Remove '$DEST' and run this script again."

cat <<'EOF'

==> Done.

Layout:

    data_road/training/image_2/       left colour images
    data_road/training/image_3/       right colour images (stereo)
    data_road/training/calib/         camera calibration, one file per image
    data_road/training/gt_image_2/    ground truth (training split only)
    data_road/testing/                images and calibration, no ground truth

Ground truth is colour coded: magenta marks the positive region (road, or the
ego lane), red marks the negative region, and black marks "do not care"
pixels that are excluded from scoring.

There are two kinds of ground truth.  "*_road_*" covers the whole road
surface and exists for every training image.  "*_lane_*" covers the ego lane
only and exists for the 95 "um" (urban marked) images.

You may delete the .zip files once you are happy that everything extracted.

--------------------------------------------------------------------------
The KITTI Vision Benchmark Suite is (c) Andreas Geiger, Philip Lenz and
Raquel Urtasun, released under a Creative Commons Attribution-NonCommercial-
ShareAlike 3.0 licence.  It is used here for coursework only.  If you refer
to this data in a report, cite:

  A. Geiger, P. Lenz and R. Urtasun, "Are we ready for autonomous driving?
  The KITTI vision benchmark suite", CVPR 2012.
--------------------------------------------------------------------------
EOF