Run Inference on cluster

bash_scripts/run_tracking.sh

@@ -0,0 +1,101 @@
+#!/bin/bash
+
+#SBATCH -p gpu # a100 # partition
+#SBATCH --gres=gpu:1
+#SBATCH -N 1   # number of nodes
+#SBATCH --ntasks-per-node 8 # 2 # max number of tasks per node
+#SBATCH --mem 64G # memory pool for all cores
+#SBATCH -t 3-00:00 # time (D-HH:MM)
+#SBATCH -o slurm.%A.%N.out
+#SBATCH -e slurm.%A.%N.err
+#SBATCH --mail-type=ALL
+#SBATCH [email protected]
+
+# ---------------------
+# Source bashrc
+# ----------------------
+# Otherwise `which python` points to the miniconda module's Python
+source ~/.bashrc
+
+# memory
+# see https://pytorch.org/docs/stable/notes/cuda.html#environment-variables
+PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+
+# -----------------------------
+# Error settings for bash
+# -----------------------------
+# see https://wizardzines.com/comics/bash-errors/
+set -e  # do not continue after errors
+set -u  # throw error if variable is unset
+set -o pipefail  # make the pipe fail if any part of it fails
+
+# ---------------------
+# Define variables
+# ----------------------
+
+# video and inference config
+VIDEO_PATH=/ceph/zoo/users/sminano/crabs_tracks_label/04.09.2023-04-Right_RE_test/04.09.2023-04-Right_RE_test.mp4
+CONFIG_FILE=/ceph/zoo/users/sminano/cluster_tracking_config.yaml
+
+# checkpoint
+TRAINED_MODEL_PATH=/ceph/zoo/users/sminano/ml-runs-all/ml_runs-nikkna-copy/243676951438603508/8dbe61069f17453a87c27b4f61f6e681/checkpoints/last.ckpt
+
+
+# output directory
+OUTPUT_DIR=/ceph/zoo/users/sminano/crabs_track_output
+
+# ground truth is available
+GT_PATH=/ceph/zoo/users/sminano/crabs_tracks_label/04.09.2023-04-Right_RE_test/04.09.2023-04-Right_RE_test_corrected_ST_csv.csv
+
+# version of the codebase
+GIT_BRANCH=main
+
+# -----------------------------
+# Create virtual environment
+# -----------------------------
+module load miniconda
+
+# Define a environment for each job in the
+# temporary directory of the compute node
+ENV_NAME=crabs-dev-$SLURM_JOB_ID
+ENV_PREFIX=$TMPDIR/$ENV_NAME
+
+# create environment
+conda create \
+   --prefix $ENV_PREFIX \
+   -y \
+   python=3.10
+
+# activate environment
+conda activate $ENV_PREFIX
+
+# install crabs package in virtual env
+python -m pip install git+https://github.com/SainsburyWellcomeCentre/crabs-exploration.git@$GIT_BRANCH
+
+
+# log pip and python locations
+echo $ENV_PREFIX
+which python
+which pip
+
+# print the version of crabs package (last number is the commit hash)
+echo "Git branch: $GIT_BRANCH"
+conda list crabs
+echo "-----"
+
+# ------------------------------------
+# GPU specs
+# ------------------------------------
+echo "Memory used per GPU before training"
+echo $(nvidia-smi --query-gpu=name,memory.total,memory.free,memory.used --format=csv) #noheader
+echo "-----"
+
+# -------------------
+# Run evaluation script
+# -------------------
+detect-and-track-video  \
+ --trained_model_path $TRAINED_MODEL_PATH \
+ --video_path $VIDEO_PATH \
+ --config_file $CONFIG_FILE \
+ --output_dir $OUTPUT_DIR \
+ --gt_path $GT_PATH

bash_scripts/run_tracking_all_escape_events.sh

@@ -0,0 +1,116 @@
+#!/bin/bash
+
+#SBATCH -p gpu # a100 # partition
+#SBATCH --gres=gpu:1
+#SBATCH -N 1   # number of nodes
+#SBATCH --ntasks-per-node 8 # 2 # max number of tasks per node
+#SBATCH --mem 64G # memory pool for all cores
+#SBATCH -t 3-00:00 # time (D-HH:MM)
+#SBATCH -o slurm.%A.%N.out
+#SBATCH -e slurm.%A.%N.err
+#SBATCH --mail-type=ALL
+#SBATCH [email protected]
+
+# ---------------------
+# Source bashrc
+# ----------------------
+# Otherwise `which python` points to the miniconda module's Python
+source ~/.bashrc
+
+# memory
+# see https://pytorch.org/docs/stable/notes/cuda.html#environment-variables
+PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+
+# -----------------------------
+# Error settings for bash
+# -----------------------------
+# see https://wizardzines.com/comics/bash-errors/
+set -e  # do not continue after errors
+set -u  # throw error if variable is unset
+set -o pipefail  # make the pipe fail if any part of it fails
+
+# ---------------------
+# Define variables
+# ----------------------
+
+# video and inference config
+VIDEO_DIR=/ceph/zoo/raw/CrabField/ramalhete_2023/Escapes
+CONFIG_FILE=/ceph/zoo/users/sminano/cluster_tracking_config.yaml
+
+# checkpoint
+TRAINED_MODEL_PATH=/ceph/zoo/users/sminano/ml-runs-all/ml_runs-nikkna-copy/243676951438603508/8dbe61069f17453a87c27b4f61f6e681/checkpoints/last.ckpt
+
+# output directory
+OUTPUT_DIR=/ceph/zoo/users/sminano/crabs_track_output
+
+# version of the codebase
+GIT_BRANCH=main
+
+# Check if the target is not a directory
+if [ ! -d "$VIDEO_DIR" ]; then
+  exit 1
+fi
+
+# -----------------------------
+# Create virtual environment
+# -----------------------------
+module load miniconda
+
+# Define a environment for each job in the
+# temporary directory of the compute node
+ENV_NAME=crabs-dev-$SLURM_JOB_ID
+ENV_PREFIX=$TMPDIR/$ENV_NAME
+
+# create environment
+conda create \
+    --prefix $ENV_PREFIX \
+    -y \
+    python=3.10
+
+# activate environment
+conda activate $ENV_PREFIX
+
+# install crabs package in virtual env
+python -m pip install git+https://github.com/SainsburyWellcomeCentre/crabs-exploration.git@$GIT_BRANCH
+
+# log pip and python locations
+echo $ENV_PREFIX
+which python
+which pip
+
+# print the version of crabs package (last number is the commit hash)
+echo "Git branch: $GIT_BRANCH"
+conda list crabs
+echo "-----"
+
+# ------------------------------------
+# GPU specs
+# ------------------------------------
+echo "Memory used per GPU before training"
+echo $(nvidia-smi --query-gpu=name,memory.total,memory.free,memory.used --format=csv) #noheader
+echo "-----"
+
+# -------------------
+# Run evaluation script for each .mov file in VIDEO_DIR
+# -------------------
+
+TIMESTAMP=$(date +"%Y%m%d_%H%M%S")
+PARENT_OUTPUT_DIR="${OUTPUT_DIR}_${TIMESTAMP}"
+mkdir -p "$PARENT_OUTPUT_DIR"
+
+for VIDEO_PATH in "$VIDEO_DIR"/*.mov; do
+  VIDEO_BASENAME=$(basename "$VIDEO_PATH" .mov)
+
+  echo "Processing video: $VIDEO_PATH"
+
+  VIDEO_OUTPUT_DIR="$PARENT_OUTPUT_DIR/$VIDEO_BASENAME"
+
+  mkdir -p "$VIDEO_OUTPUT_DIR"
+
+  detect-and-track-video  \
+    --trained_model_path "$TRAINED_MODEL_PATH" \
+    --video_path "$VIDEO_PATH" \
+    --config_file "$CONFIG_FILE" \
+    --output_dir "$VIDEO_OUTPUT_DIR" \
+    --run_on_video_dir
+done

crabs/tracker/track_video.py

@@ -63,6 +63,11 @@ def setup(self):
         """
         Load tracking config, trained model and input video path.
         """
+        # Check for CUDA availability
+        if self.device == "cuda" and not torch.cuda.is_available():
+            print("CUDA is not available. Falling back to CPU.")
+            self.device = "cpu"
+
         with open(self.config_file, "r") as f:
             self.config = yaml.safe_load(f)
 
@@ -87,7 +92,11 @@ def prep_outputs(self):
             self.csv_writer,
             self.csv_file,
             self.tracking_output_dir,
-        ) = prep_csv_writer(self.args.output_dir, self.video_file_root)
+        ) = prep_csv_writer(
+            self.args.output_dir,
+            self.video_file_root,
+            self.args.run_on_video_dir,
+        )
 
         if self.args.save_video:
             frame_width = int(self.video.get(cv2.CAP_PROP_FRAME_WIDTH))
@@ -299,6 +308,11 @@ def tracking_parse_args(args):
         action="store_true",
         help="Save frame to be used in correcting track labelling",
     )
+    parser.add_argument(
+        "--run_on_video_dir",
+        action="store_true",
+        help="option to run track video on directory instead of a video.",
+    )
     parser.add_argument(
         "--device",
         type=str,

crabs/tracker/utils/io.py

@@ -13,7 +13,9 @@
 )
 
 
-def prep_csv_writer(output_dir: str, video_file_root: str):
+def prep_csv_writer(
+    output_dir: str, video_file_root: str, run_on_video_dir: bool
+):
     """
     Prepare csv writer to output tracking results.
 
@@ -29,11 +31,15 @@ def prep_csv_writer(output_dir: str, video_file_root: str):
     Tuple
         A tuple containing the CSV writer, the CSV file object, and the tracking output directory path.
     """
-
-    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-    tracking_output_dir = Path(output_dir + f"_{timestamp}") / video_file_root
-    # Create the subdirectory for the specific video file root
-    tracking_output_dir.mkdir(parents=True, exist_ok=True)
+    if run_on_video_dir:
+        tracking_output_dir = Path(output_dir)
+    else:
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        tracking_output_dir = (
+            Path(output_dir + f"_{timestamp}") / video_file_root
+        )
+        # Create the subdirectory for the specific video file root
+        tracking_output_dir.mkdir(parents=True, exist_ok=True)
 
     csv_file = open(
         f"{str(tracking_output_dir)}/predicted_tracks.csv",