Use SuperAnimal detection fallback for Space

Author: lexiyutou

README.md

@@ -163,8 +163,8 @@ The `s1ckpt_inference.ckpt` checkpoint is downloaded automatically if missing.
 | | **Local** (`python app.py`) | **Hugging Face Space** |
 |--|--|--|
 | PRIMA device | GPU if available, else CPU | CPU only |
-| Detector | Detectron2 X-101-FPN | full-image crop fallback |
-| Default TTA iterations | 30 | 0 (PRIMA-only by default) |
+| Detector | Detectron2 X-101-FPN | DeepLabCut SuperAnimal detector |
+| Default TTA iterations | 30 | 30 |
 | Save `.obj` meshes | on | off |
 | Preload checkpoint at startup | off | on |
 
@@ -194,7 +194,7 @@ The script rsyncs only the Git-tracked files needed by the Space from the
 working tree (not `git archive`) so image files are materialized before
 `git add` turns them into LFS blobs.
 During deployment, `detectron2` is removed from the Space `requirements.txt`;
-the app uses its full-image crop fallback on the CPU Space.
+the app uses the DeepLabCut SuperAnimal detector fallback on the CPU Space.
 
 ---
 

app.py

@@ -163,9 +163,9 @@ def resolve_detectron_device(self) -> str:
     ),
     description=(
         "**Hugging Face Space (cpu-basic)** — lightweight demo: **CPU-only** PRIMA inference. "
-        "The Space build skips Detectron2 and uses a full-image crop fallback. TTA is optional "
-        "(30 iterations by default, matching the local demo; set to 0 to skip). Mesh `.obj` export "
-        "is off by default to save time and disk."
+        "The Space build skips Detectron2 and uses the DeepLabCut SuperAnimal detector for animal "
+        "crops. TTA is optional (30 iterations by default, matching the local demo; set to 0 to "
+        "skip). Mesh `.obj` export is off by default to save time and disk."
     ),
     interface_title="PRIMA on Hugging Face — lightweight CPU demo",
 )
@@ -298,7 +298,7 @@ def _build_detector(profile: Optional[DemoProfile] = None):
         import detectron2.engine
         from detectron2 import model_zoo
     except Exception as e:
-        print(f"[warn] Detectron2 unavailable ({type(e).__name__}: {e}); using full-image fallback bbox.")
+        print(f"[warn] Detectron2 unavailable ({type(e).__name__}: {e}); using SuperAnimal detector fallback.")
         return None
 
     if profile is None:
@@ -317,6 +317,84 @@ def _build_detector(profile: Optional[DemoProfile] = None):
     return detector
 
 
+def _filter_superanimal_boxes(
+    payload: Dict[str, Any],
+    det_thresh: float,
+    img_shape: Tuple[int, int],
+) -> Optional[np.ndarray]:
+    boxes = payload.get("bboxes")
+    if boxes is None:
+        return None
+
+    boxes = np.asarray(boxes, dtype=np.float32).reshape(-1, 4)
+    if len(boxes) == 0:
+        return None
+
+    scores = payload.get("bbox_scores")
+    if scores is None:
+        scores = np.ones((len(boxes),), dtype=np.float32)
+    scores = np.asarray(scores, dtype=np.float32).reshape(-1)
+    if len(scores) != len(boxes):
+        scores = np.ones((len(boxes),), dtype=np.float32)
+
+    h, w = img_shape
+    valid = (
+        (scores > float(det_thresh))
+        & np.isfinite(boxes).all(axis=1)
+        & (boxes[:, 2] > 0.0)
+        & (boxes[:, 3] > 0.0)
+    )
+    if not np.any(valid):
+        return None
+
+    xywh = boxes[valid]
+    scores = scores[valid]
+    boxes = xywh.copy()
+    boxes[:, 2] = xywh[:, 0] + xywh[:, 2]
+    boxes[:, 3] = xywh[:, 1] + xywh[:, 3]
+    boxes[:, [0, 2]] = np.clip(boxes[:, [0, 2]], 0.0, float(max(1, w - 1)))
+    boxes[:, [1, 3]] = np.clip(boxes[:, [1, 3]], 0.0, float(max(1, h - 1)))
+    valid_size = (boxes[:, 2] > boxes[:, 0]) & (boxes[:, 3] > boxes[:, 1])
+    boxes = boxes[valid_size]
+    scores = scores[valid_size]
+    if len(boxes) == 0:
+        return None
+    order = np.argsort(scores)[::-1]
+    return boxes[order].astype(np.float32, copy=False)
+
+
+def _detect_superanimal_boxes(img_rgb: np.ndarray, det_thresh: float) -> Optional[np.ndarray]:
+    try:
+        from deeplabcut.pose_estimation_pytorch.apis import superanimal_analyze_images
+    except Exception as e:
+        print(f"[warn] DeepLabCut SuperAnimal unavailable ({type(e).__name__}: {e}); no fallback bbox.")
+        return None
+
+    with tempfile.TemporaryDirectory(prefix="sa_detect_") as tmp_dir:
+        img_path = os.path.join(tmp_dir, "image.png")
+        cv2.imwrite(img_path, cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR))
+
+        dlc_device = "cuda" if torch.cuda.is_available() else "cpu"
+        preds = superanimal_analyze_images(
+            superanimal_name=SUPER_ANIMAL_ARGS.superanimal_name,
+            model_name=SUPER_ANIMAL_ARGS.superanimal_model_name,
+            detector_name=SUPER_ANIMAL_ARGS.superanimal_detector_name,
+            images=img_path,
+            max_individuals=SUPER_ANIMAL_ARGS.superanimal_max_individuals,
+            out_folder=tmp_dir,
+            progress_bar=False,
+            device=dlc_device,
+            pose_threshold=0.1,
+            bbox_threshold=float(det_thresh),
+            plot_skeleton=False,
+        )
+
+        payload = preds.get(img_path)
+        if payload is None:
+            return None
+        return _filter_superanimal_boxes(payload, det_thresh, img_rgb.shape[:2])
+
+
 def _load_model_and_detector_for_demo(checkpoint_path: str, profile: DemoProfile):
     """Load PRIMA and Detectron2 once for the Gradio session (main thread only)."""
     model, model_cfg, renderer, cam_crop_to_full_fn, device = _load_prima_model(checkpoint_path)
@@ -331,8 +409,8 @@ def _detect_animal_boxes(
 ) -> Optional[np.ndarray]:
     """Return Nx4 XYXY boxes or None if no animal detections."""
     if detector is None:
-        h, w = img_bgr.shape[:2]
-        return np.array([[0.0, 0.0, float(max(1, w - 1)), float(max(1, h - 1))]], dtype=np.float32)
+        img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
+        return _detect_superanimal_boxes(img_rgb, det_thresh)
 
     det_out = detector(img_bgr)
     det_instances = det_out["instances"]
@@ -403,7 +481,7 @@ def report(message: str) -> None:
     img_bgr = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR)
     if boxes is None:
         if detector is None:
-            report("Detectron2 unavailable; using full-image crop...")
+            report("Detectron2 unavailable; detecting animals with SuperAnimal...")
         else:
             report("Detecting animals with Detectron2...")
         boxes = _detect_animal_boxes(detector, img_bgr, det_thresh)

scripts/deploy_hf_space.sh

@@ -70,7 +70,7 @@ done
 README_FILE="${TMP}/README.md"
 REQ_FILE="${TMP}/requirements.txt"
 
-echo "[deploy] Removing Detectron2 from Space requirements (app falls back to full-image crops) ..."
+echo "[deploy] Removing Detectron2 from Space requirements (app falls back to SuperAnimal detection) ..."
 grep -vE '^[[:space:]]*detectron2([[:space:]]|@|$)' "$REQ_FILE" > "${REQ_FILE}.tmp"
 mv "${REQ_FILE}.tmp" "$REQ_FILE"
 
@@ -109,10 +109,37 @@ if [[ "$PUSH_URL" == https://huggingface.co/* && -z "${HF_TOKEN:-}" && -f "${HF_
   HF_TOKEN="$(<"${HF_HOME:-$HOME/.cache/huggingface}/token")"
 fi
 if [[ "$PUSH_URL" == https://huggingface.co/* && -n "${HF_TOKEN:-}" ]]; then
-  PUSH_URL="${PUSH_URL/https:\/\/huggingface.co/https:\/\/hf_user:${HF_TOKEN}@huggingface.co}"
+  export HF_TOKEN
+  git config credential.helper "store --file=${TMP}/git-credentials"
+  printf 'protocol=https\nhost=huggingface.co\nusername=hf_user\npassword=%s\n\n' "$HF_TOKEN" | git credential approve
+  ASKPASS="${TMP}/git-askpass.sh"
+  cat > "$ASKPASS" <<'SH'
+#!/usr/bin/env bash
+case "$1" in
+  *Username*) printf '%s\n' 'hf_user' ;;
+  *Password*) printf '%s\n' "${HF_TOKEN}" ;;
+  *) printf '%s\n' "${HF_TOKEN}" ;;
+esac
+SH
+  chmod 700 "$ASKPASS"
+  export GIT_ASKPASS="$ASKPASS"
 fi
 
 git remote add hf "$PUSH_URL"
+echo "[deploy] Uploading LFS objects to Hugging Face Space ..."
+mapfile -t LFS_OIDS < <(git lfs ls-files -l | awk '{print $1}')
+if [[ "${#LFS_OIDS[@]}" -gt 0 ]]; then
+  if ! GIT_TERMINAL_PROMPT=0 git lfs push --object-id hf "${LFS_OIDS[@]}"; then
+    echo "[deploy] ERROR: LFS upload failed. Ensure HF_TOKEN has write access to ${SPACE_URL}." >&2
+    exit 1
+  fi
+else
+  echo "[deploy] No LFS objects found in this snapshot."
+fi
+
 echo "[deploy] Force-pushing to Hugging Face Space ..."
-GIT_TERMINAL_PROMPT=0 git -c credential.helper= push hf HEAD:main --force
+# This deploy repo is freshly initialized, so older git-lfs pre-push hooks can
+# fail when they try to inspect the remote's previous main commit. LFS objects
+# are uploaded explicitly above; skip the hook for the Git ref update.
+GIT_TERMINAL_PROMPT=0 git push hf HEAD:main --force --no-verify
 echo "[deploy] Done."

scripts/local_infer.py

@@ -71,9 +71,9 @@ def main() -> int:
     model, model_cfg, renderer, cam_crop_to_full_fn, device = app._load_prima_model()
     print(f"[local_infer] device={device}")
 
-    print("[local_infer] Building detector (Detectron2 if installed, else fallback) ...")
+    print("[local_infer] Building detector (Detectron2 if installed, else SuperAnimal fallback) ...")
     detector = app._build_detector()
-    print(f"[local_infer] detector={'detectron2' if detector is not None else 'fallback'}")
+    print(f"[local_infer] detector={'detectron2' if detector is not None else 'superanimal fallback'}")
 
     print("[local_infer] Running inference ...")
     before, after, kpts, mesh_before, mesh_after = app._collect_animal_results(

scripts/run_local_demo_once.py

@@ -37,9 +37,9 @@ def main() -> int:
     model, cfg, renderer, device = app._load_prima_model()
     print(f"      device={device}")
 
-    print("[2/4] Building detector (Detectron2 if installed, else full-image bbox) …")
+    print("[2/4] Building detector (Detectron2 if installed, else SuperAnimal detector) …")
     det = app._build_detector()
-    print(f"      detector={'detectron2' if det is not None else 'full-image fallback'}")
+    print(f"      detector={'detectron2' if det is not None else 'superanimal fallback'}")
 
     img = cv2.cvtColor(cv2.imread(str(img_path)), cv2.COLOR_BGR2RGB)
     print(f"[3/4] Running inference on {img_path.name} (TTA iterations=0) …")