#!/bin/sh

# prepare_images.sh
# 2026-01-31
# by Gernot Walzl

# This script prepares images for photogrammetry:
# 1. Extract the images (ether from zip or from a mp4 video)
# 2. Crop the images
# 3. Replace the background in images with black

# Available models for rembg:
# https://github.com/danielgatis/rembg?tab=readme-ov-file#available-models

CWD=$(pwd)
CROP_GEOM=${CROP_GEOM:-'2160x2160+840+0'}
FPS=${FPS:-'4'}
MODEL=${MODEL:-'u2net'}
ARCHIVE="$1"

set -e

if [ ! -r "$ARCHIVE" ]; then
  echo 'Usage example:'
  echo "CROP_GEOM='2592x2592+1008+0' MODEL='u2net_human_seg' $0 MyProject_2026-01-31.zip"
  exit 1
fi

PROJECT=$(basename "$ARCHIVE" | rev | cut -c5- | rev)
mkdir "$PROJECT"
cd "$PROJECT" || exit 1

echo 'Extracting images:'
mkdir orig
ARCHIVE_TYPE=$(echo "$ARCHIVE" | rev | cut -c-3 | rev)
if [ "$ARCHIVE_TYPE" = "zip" ]; then
  mkdir tmp
  unzip "$CWD/$ARCHIVE" -d tmp
  find tmp -type f \( -name '*.jpg' -o -name '*.png' \) \
    -exec mv '{}' orig \;
  rm -rf tmp
elif [ "$ARCHIVE_TYPE" = "mp4" ]; then
  ffmpeg -i "$CWD/$ARCHIVE" -r "$FPS" -qmin 1 -qscale:v 1 orig/img_%04d.jpg
else
  exit 1
fi
echo

echo 'Cropping images:'
mkdir cropped
find orig -type f \( -name '*.jpg' -o -name '*.png' \) \
  | while read IMG; do
      IMG_CROPPED="cropped/$(basename "$IMG")"
      convert -verbose "$IMG" -crop "$CROP_GEOM" "$IMG_CROPPED"
    done
echo

echo 'Removing background:'
rembg p -m "$MODEL" -bgc 0 0 0 255 cropped images
echo