#!/bin/sh

# prepare_images.sh
# 2026-02-21
# by Gernot Walzl

# This script prepares images for photogrammetry:
# 1. Extract the images (ether from zip or from a mp4 video)
# 2. Create background masks from images

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

CWD=$(pwd)
FPS=${FPS:-'4'}
MODEL=${MODEL:-'u2net'}
ARCHIVE="$1"

set -e

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

echo 'Creating background masks:'
rembg p -m "$MODEL" -om images masks
echo

# Naming of image masks:
# https://colmap.github.io/faq.html#mask-image-regions
for MASK in masks/*.png; do
  mv "$MASK" "masks/$(basename "$MASK" .png).jpg.png"
done

# Camera intrinsics per folder:
# https://github.com/colmap/colmap/blob/main/src/colmap/exe/feature.h#L72
for CAM_IDX in 0 1 2 3; do
  if ls images/*_${CAM_IDX}.jpg > /dev/null 2>&1; then
    mkdir "images/$CAM_IDX"
    mv images/*_${CAM_IDX}.jpg "images/$CAM_IDX"
    mkdir "masks/$CAM_IDX"
    mv masks/*_${CAM_IDX}.jpg.png "masks/$CAM_IDX"
  fi
done