Spaces:

Yuxihenry
/

SpatialTrackerV2

Build error

App Files Files Community

xiaoyuxi commited on Jul 6, 2025

Commit

c8d9d42

0 Parent(s):

Cleaned history, reset to current state

Browse files

This view is limited to 50 files because it contains too many changes. See raw diff

Files changed (50) hide show

.gitattributes +36 -0
.gitignore +69 -0
README.md +14 -0
_viz/viz_template.html +1778 -0
app.py +1118 -0
app_3rd/README.md +12 -0
app_3rd/sam_utils/hf_sam_predictor.py +129 -0
app_3rd/sam_utils/inference.py +123 -0
app_3rd/spatrack_utils/infer_track.py +194 -0
config/__init__.py +0 -0
config/magic_infer_moge.yaml +48 -0
examples/backpack.mp4 +3 -0
examples/ball.mp4 +3 -0
examples/basketball.mp4 +3 -0
examples/biker.mp4 +3 -0
examples/cinema_0.mp4 +3 -0
examples/cinema_1.mp4 +3 -0
examples/drifting.mp4 +3 -0
examples/ego_kc1.mp4 +3 -0
examples/ego_teaser.mp4 +3 -0
examples/handwave.mp4 +3 -0
examples/hockey.mp4 +3 -0
examples/ken_block_0.mp4 +3 -0
examples/kiss.mp4 +3 -0
examples/kitchen.mp4 +3 -0
examples/kitchen_egocentric.mp4 +3 -0
examples/pillow.mp4 +3 -0
examples/protein.mp4 +3 -0
examples/pusht.mp4 +3 -0
examples/robot1.mp4 +3 -0
examples/robot2.mp4 +3 -0
examples/robot_3.mp4 +3 -0
examples/robot_unitree.mp4 +3 -0
examples/running.mp4 +3 -0
examples/teleop2.mp4 +3 -0
examples/vertical_place.mp4 +3 -0
models/SpaTrackV2/models/SpaTrack.py +759 -0
models/SpaTrackV2/models/__init__.py +0 -0
models/SpaTrackV2/models/blocks.py +519 -0
models/SpaTrackV2/models/camera_transform.py +248 -0
models/SpaTrackV2/models/depth_refiner/backbone.py +472 -0
models/SpaTrackV2/models/depth_refiner/decode_head.py +619 -0
models/SpaTrackV2/models/depth_refiner/depth_refiner.py +115 -0
models/SpaTrackV2/models/depth_refiner/network.py +429 -0
models/SpaTrackV2/models/depth_refiner/stablilization_attention.py +1187 -0
models/SpaTrackV2/models/depth_refiner/stablizer.py +342 -0
models/SpaTrackV2/models/predictor.py +153 -0
models/SpaTrackV2/models/tracker3D/TrackRefiner.py +1478 -0
models/SpaTrackV2/models/tracker3D/co_tracker/cotracker_base.py +418 -0
models/SpaTrackV2/models/tracker3D/co_tracker/utils.py +929 -0

.gitattributes ADDED Viewed

	@@ -0,0 +1,36 @@

+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.mp4 filter=lfs diff=lfs merge=lfs -text

.gitignore ADDED Viewed

	@@ -0,0 +1,69 @@

+# ignore the multi media
+checkpoints
+**/checkpoints/
+**/temp/
+temp
+assets_dev
+assets/example0/results
+assets/example0/snowboard.npz
+assets/example1/results
+assets/davis_eval
+assets/*/results
+*gradio*
+#
+models/monoD/zoeDepth/ckpts/*
+models/monoD/depth_anything/ckpts/*
+vis_results
+dist_encrypted
+# remove the dependencies
+deps
+# filter the __pycache__ files
+__pycache__/
+/**/**/__pycache__
+/**/__pycache__
+outputs
+scripts/lauch_exp/config
+scripts/lauch_exp/submit_job.log
+scripts/lauch_exp/hydra_output
+scripts/lauch_wulan
+scripts/custom_video
+# ignore the visualizer
+viser
+viser_result
+benchmark/results
+benchmark
+ossutil_output
+prev_version
+spat_ceres
+wandb
+*.log
+seg_target.py
+eval_davis.py
+eval_multiple_gpu.py
+eval_pose_scan.py
+eval_single_gpu.py
+infer_cam.py
+infer_stream.py
+*.egg-info/
+**/*.egg-info
+eval_kinectics.py
+models/SpaTrackV2/datasets
+scripts
+config/fix_2d.yaml
+models/SpaTrackV2/datasets
+scripts/
+models/**/build
+models/**/dist
+temp_local

README.md ADDED Viewed

	@@ -0,0 +1,14 @@

+---
+title: SpatialTrackerV2
+emoji: ⚡️
+colorFrom: yellow
+colorTo: red
+sdk: gradio
+sdk_version: 5.31.0
+app_file: app.py
+pinned: false
+license: mit
+short_description: Official Space for SpatialTrackerV2
+---
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

_viz/viz_template.html ADDED Viewed

	@@ -0,0 +1,1778 @@

+<!DOCTYPE html>
+<html lang="en">
+<head>
+  <meta charset="UTF-8">
+  <meta name="viewport" content="width=device-width, initial-scale=1.0">
+  <title>3D Point Cloud Visualizer</title>
+  <style>
+    :root {
+      --primary: #9b59b6; /* Brighter purple for dark mode */
+      --primary-light: #3a2e4a;
+      --secondary: #a86add;
+      --accent: #ff6e6e;
+      --bg: #1a1a1a;
+      --surface: #2c2c2c;
+      --text: #e0e0e0;
+      --text-secondary: #a0a0a0;
+      --border: #444444;
+      --shadow: rgba(0, 0, 0, 0.2);
+      --shadow-hover: rgba(0, 0, 0, 0.3);
+      --space-sm: 16px;
+      --space-md: 24px;
+      --space-lg: 32px;
+    }
+    body {
+      margin: 0;
+      overflow: hidden;
+      background: var(--bg);
+      color: var(--text);
+      font-family: 'Inter', sans-serif;
+      -webkit-font-smoothing: antialiased;
+    }
+    #canvas-container {
+      position: absolute;
+      width: 100%;
+      height: 100%;
+    }
+    #ui-container {
+      position: absolute;
+      top: 0;
+      left: 0;
+      width: 100%;
+      height: 100%;
+      pointer-events: none;
+      z-index: 10;
+    }
+    #status-bar {
+      position: absolute;
+      top: 16px;
+      left: 16px;
+      background: rgba(30, 30, 30, 0.9);
+      padding: 8px 16px;
+      border-radius: 8px;
+      pointer-events: auto;
+      box-shadow: 0 4px 6px var(--shadow);
+      backdrop-filter: blur(4px);
+      border: 1px solid var(--border);
+      color: var(--text);
+      transition: opacity 0.5s ease, transform 0.5s ease;
+      font-weight: 500;
+    }
+    #status-bar.hidden {
+      opacity: 0;
+      transform: translateY(-20px);
+      pointer-events: none;
+    }
+    #control-panel {
+      position: absolute;
+      bottom: 16px;
+      left: 50%;
+      transform: translateX(-50%);
+      background: rgba(44, 44, 44, 0.95);
+      padding: 6px 8px;
+      border-radius: 6px;
+      display: flex;
+      gap: 8px;
+      align-items: center;
+      justify-content: space-between;
+      pointer-events: auto;
+      box-shadow: 0 4px 10px var(--shadow);
+      backdrop-filter: blur(4px);
+      border: 1px solid var(--border);
+    }
+    #timeline {
+      width: 150px;
+      height: 4px;
+      background: rgba(255, 255, 255, 0.1);
+      border-radius: 2px;
+      position: relative;
+      cursor: pointer;
+    }
+    #progress {
+      position: absolute;
+      height: 100%;
+      background: var(--primary);
+      border-radius: 2px;
+      width: 0%;
+    }
+    #playback-controls {
+      display: flex;
+      gap: 4px;
+      align-items: center;
+    }
+    button {
+      background: rgba(255, 255, 255, 0.08);
+      border: 1px solid var(--border);
+      color: var(--text);
+      padding: 4px 6px;
+      border-radius: 3px;
+      cursor: pointer;
+      display: flex;
+      align-items: center;
+      justify-content: center;
+      transition: background 0.2s, transform 0.2s;
+      font-family: 'Inter', sans-serif;
+      font-weight: 500;
+      font-size: 6px;
+    }
+    button:hover {
+      background: rgba(255, 255, 255, 0.15);
+      transform: translateY(-1px);
+    }
+    button.active {
+      background: var(--primary);
+      color: white;
+      box-shadow: 0 2px 8px rgba(155, 89, 182, 0.4);
+    }
+    select, input {
+      background: rgba(255, 255, 255, 0.08);
+      border: 1px solid var(--border);
+      color: var(--text);
+      padding: 4px 6px;
+      border-radius: 3px;
+      cursor: pointer;
+      font-family: 'Inter', sans-serif;
+      font-size: 6px;
+    }
+    .icon {
+      width: 10px;
+      height: 10px;
+      fill: currentColor;
+    }
+    .tooltip {
+      position: absolute;
+      bottom: 100%;
+      left: 50%;
+      transform: translateX(-50%);
+      background: var(--surface);
+      color: var(--text);
+      padding: 3px 6px;
+      border-radius: 3px;
+      font-size: 7px;
+      white-space: nowrap;
+      margin-bottom: 4px;
+      opacity: 0;
+      transition: opacity 0.2s;
+      pointer-events: none;
+      box-shadow: 0 2px 4px var(--shadow);
+      border: 1px solid var(--border);
+    }
+    button:hover .tooltip {
+      opacity: 1;
+    }
+    #settings-panel {
+      position: absolute;
+      top: 16px;
+      right: 16px;
+      background: rgba(44, 44, 44, 0.98);
+      padding: 10px;
+      border-radius: 6px;
+      width: 195px;
+      max-height: calc(100vh - 40px);
+      overflow-y: auto;
+      pointer-events: auto;
+      box-shadow: 0 4px 15px var(--shadow);
+      backdrop-filter: blur(4px);
+      border: 1px solid var(--border);
+      display: block;
+      opacity: 1;
+      scrollbar-width: thin;
+      scrollbar-color: var(--primary-light) transparent;
+      transition: transform 0.35s ease-in-out, opacity 0.3s ease-in-out;
+    }
+    #settings-panel.is-hidden {
+      transform: translateX(calc(100% + 20px));
+      opacity: 0;
+      pointer-events: none;
+    }
+    #settings-panel::-webkit-scrollbar {
+      width: 3px;
+    }
+    #settings-panel::-webkit-scrollbar-track {
+      background: transparent;
+    }
+    #settings-panel::-webkit-scrollbar-thumb {
+      background-color: var(--primary-light);
+      border-radius: 3px;
+    }
+    @media (max-height: 700px) {
+      #settings-panel {
+        max-height: calc(100vh - 40px);
+      }
+    }
+    @media (max-width: 768px) {
+      #control-panel {
+        width: 90%;
+        flex-wrap: wrap;
+        justify-content: center;
+      }
+      #timeline {
+        width: 100%;
+        order: 3;
+        margin-top: 10px;
+      }
+      #settings-panel {
+        width: 140px;
+        right: 10px;
+        top: 10px;
+        max-height: calc(100vh - 20px);
+      }
+    }
+    .settings-group {
+      margin-bottom: 8px;
+    }
+    .settings-group h3 {
+      margin: 0 0 6px 0;
+      font-size: 10px;
+      font-weight: 500;
+      color: var(--text-secondary);
+    }
+    .slider-container {
+      display: flex;
+      align-items: center;
+      gap: 6px;
+      width: 100%;
+    }
+    .slider-container label {
+      min-width: 60px;
+      font-size: 10px;
+      flex-shrink: 0;
+    }
+    input[type="range"] {
+      flex: 1;
+      height: 2px;
+      -webkit-appearance: none;
+      background: rgba(255, 255, 255, 0.1);
+      border-radius: 1px;
+      min-width: 0;
+    }
+    input[type="range"]::-webkit-slider-thumb {
+      -webkit-appearance: none;
+      width: 8px;
+      height: 8px;
+      border-radius: 50%;
+      background: var(--primary);
+      cursor: pointer;
+    }
+    .toggle-switch {
+      position: relative;
+      display: inline-block;
+      width: 20px;
+      height: 10px;
+    }
+    .toggle-switch input {
+      opacity: 0;
+      width: 0;
+      height: 0;
+    }
+    .toggle-slider {
+      position: absolute;
+      cursor: pointer;
+      top: 0;
+      left: 0;
+      right: 0;
+      bottom: 0;
+      background: rgba(255, 255, 255, 0.1);
+      transition: .4s;
+      border-radius: 10px;
+    }
+    .toggle-slider:before {
+      position: absolute;
+      content: "";
+      height: 8px;
+      width: 8px;
+      left: 1px;
+      bottom: 1px;
+      background: var(--surface);
+      border: 1px solid var(--border);
+      transition: .4s;
+      border-radius: 50%;
+    }
+    input:checked + .toggle-slider {
+      background: var(--primary);
+    }
+    input:checked + .toggle-slider:before {
+      transform: translateX(10px);
+    }
+    .checkbox-container {
+      display: flex;
+      align-items: center;
+      gap: 4px;
+      margin-bottom: 4px;
+    }
+    .checkbox-container label {
+      font-size: 10px;
+      cursor: pointer;
+    }
+    #loading-overlay {
+      position: absolute;
+      top: 0;
+      left: 0;
+      width: 100%;
+      height: 100%;
+      background: var(--bg);
+      display: flex;
+      flex-direction: column;
+      align-items: center;
+      justify-content: center;
+      z-index: 100;
+      transition: opacity 0.5s;
+    }
+    #loading-overlay.fade-out {
+      opacity: 0;
+      pointer-events: none;
+    }
+    .spinner {
+      width: 50px;
+      height: 50px;
+      border: 5px solid rgba(155, 89, 182, 0.2);
+      border-radius: 50%;
+      border-top-color: var(--primary);
+      animation: spin 1s ease-in-out infinite;
+      margin-bottom: 16px;
+    }
+    @keyframes spin {
+      to { transform: rotate(360deg); }
+    }
+    #loading-text {
+      margin-top: 16px;
+      font-size: 18px;
+      color: var(--text);
+      font-weight: 500;
+    }
+    #frame-counter {
+      color: var(--text-secondary);
+      font-size: 7px;
+      font-weight: 500;
+      min-width: 60px;
+      text-align: center;
+      padding: 0 4px;
+    }
+    .control-btn {
+      background: rgba(255, 255, 255, 0.08);
+      border: 1px solid var(--border);
+      padding: 4px 6px;
+      border-radius: 3px;
+      cursor: pointer;
+      display: flex;
+      align-items: center;
+      justify-content: center;
+      transition: all 0.2s ease;
+      font-size: 6px;
+    }
+    .control-btn:hover {
+      background: rgba(255, 255, 255, 0.15);
+      transform: translateY(-1px);
+    }
+    .control-btn.active {
+      background: var(--primary);
+      color: white;
+    }
+    .control-btn.active:hover {
+      background: var(--primary);
+      box-shadow: 0 2px 8px rgba(155, 89, 182, 0.4);
+    }
+    #settings-toggle-btn {
+      position: relative;
+      border-radius: 6px;
+      z-index: 20;
+    }
+    #settings-toggle-btn.active {
+      background: var(--primary);
+      color: white;
+    }
+    #status-bar,
+    #control-panel,
+    #settings-panel,
+    button,
+    input,
+    select,
+    .toggle-switch {
+      pointer-events: auto;
+    }
+    h2 {
+      font-size: 0.9rem;
+      font-weight: 600;
+      margin-top: 0;
+      margin-bottom: 12px;
+      color: var(--primary);
+      cursor: move;
+      user-select: none;
+      display: flex;
+      align-items: center;
+    }
+    .drag-handle {
+      font-size: 10px;
+      margin-right: 4px;
+      opacity: 0.6;
+    }
+    h2:hover .drag-handle {
+      opacity: 1;
+    }
+    .loading-subtitle {
+      font-size: 7px;
+      color: var(--text-secondary);
+      margin-top: 4px;
+    }
+    #reset-view-btn {
+      background: var(--primary-light);
+      color: var(--primary);
+      border: 1px solid rgba(155, 89, 182, 0.2);
+      font-weight: 600;
+      font-size: 9px;
+      padding: 4px 6px;
+      transition: all 0.2s;
+    }
+    #reset-view-btn:hover {
+      background: var(--primary);
+      color: white;
+      transform: translateY(-2px);
+      box-shadow: 0 4px 8px rgba(155, 89, 182, 0.3);
+    }
+    #show-settings-btn {
+      position: absolute;
+      top: 16px;
+      right: 16px;
+      z-index: 15;
+      display: none;
+    }
+    #settings-panel.visible {
+      display: block;
+      opacity: 1;
+      animation: slideIn 0.3s ease forwards;
+    }
+    @keyframes slideIn {
+      from {
+        transform: translateY(20px);
+        opacity: 0;
+      }
+      to {
+        transform: translateY(0);
+        opacity: 1;
+      }
+    }
+    .dragging {
+      opacity: 0.9;
+      box-shadow: 0 8px 20px rgba(0, 0, 0, 0.15) !important;
+      transition: none !important;
+    }
+    /* Tooltip for draggable element */
+    .tooltip-drag {
+      position: absolute;
+      left: 50%;
+      transform: translateX(-50%);
+      background: var(--primary);
+      color: white;
+      font-size: 9px;
+      padding: 2px 4px;
+      border-radius: 2px;
+      opacity: 0;
+      pointer-events: none;
+      transition: opacity 0.3s;
+      white-space: nowrap;
+      bottom: 100%;
+      margin-bottom: 4px;
+    }
+    h2:hover .tooltip-drag {
+      opacity: 1;
+    }
+    .btn-group {
+      display: flex;
+      margin-top: 8px;
+    }
+    #reset-settings-btn {
+      background: var(--primary-light);
+      color: var(--primary);
+      border: 1px solid rgba(155, 89, 182, 0.2);
+      font-weight: 600;
+      font-size: 9px;
+      padding: 4px 6px;
+      transition: all 0.2s;
+    }
+    #reset-settings-btn:hover {
+      background: var(--primary);
+      color: white;
+      transform: translateY(-2px);
+      box-shadow: 0 4px 8px rgba(155, 89, 182, 0.3);
+    }
+  </style>
+</head>
+<body>
+  <link rel="preconnect" href="https://fonts.googleapis.com">
+  <link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
+  <link href="https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&display=swap" rel="stylesheet">
+  <div id="canvas-container"></div>
+  <div id="ui-container">
+    <div id="status-bar">Initializing...</div>
+    <div id="control-panel">
+      <button id="play-pause-btn" class="control-btn">
+        <svg class="icon" viewBox="0 0 24 24">
+          <path id="play-icon" d="M8 5v14l11-7z"/>
+          <path id="pause-icon" d="M6 19h4V5H6v14zm8-14v14h4V5h-4z" style="display: none;"/>
+        </svg>
+        <span class="tooltip">Play/Pause</span>
+      </button>
+      <div id="timeline">
+        <div id="progress"></div>
+      </div>
+      <div id="frame-counter">Frame 0 / 0</div>
+      <div id="playback-controls">
+        <button id="speed-btn" class="control-btn">1x</button>
+      </div>
+    </div>
+    <div id="settings-panel">
+      <h2>
+        <span class="drag-handle">☰</span>
+        Visualization Settings
+        <button id="hide-settings-btn" class="control-btn" style="margin-left: auto; padding: 2px;" title="Hide Panel">
+          <svg class="icon" viewBox="0 0 24 24" style="width: 9px; height: 9px;">
+            <path d="M14.59 7.41L18.17 11H4v2h14.17l-3.58 3.59L16 18l6-6-6-6-1.41 1.41z"/>
+          </svg>
+        </button>
+      </h2>
+      <div class="settings-group">
+        <h3>Point Cloud</h3>
+        <div class="slider-container">
+          <label for="point-size">Size</label>
+          <input type="range" id="point-size" min="0.005" max="0.1" step="0.005" value="0.03">
+        </div>
+        <div class="slider-container">
+          <label for="point-opacity">Opacity</label>
+          <input type="range" id="point-opacity" min="0.1" max="1" step="0.05" value="1">
+        </div>
+        <div class="slider-container">
+          <label for="max-depth">Max Depth</label>
+          <input type="range" id="max-depth" min="0.1" max="10" step="0.2" value="100">
+        </div>
+      </div>
+      <div class="settings-group">
+        <h3>Trajectory</h3>
+        <div class="checkbox-container">
+          <label class="toggle-switch">
+            <input type="checkbox" id="show-trajectory" checked>
+            <span class="toggle-slider"></span>
+          </label>
+          <label for="show-trajectory">Show Trajectory</label>
+        </div>
+        <div class="checkbox-container">
+          <label class="toggle-switch">
+            <input type="checkbox" id="enable-rich-trail">
+            <span class="toggle-slider"></span>
+          </label>
+          <label for="enable-rich-trail">Visual-Rich Trail</label>
+        </div>
+        <div class="slider-container">
+          <label for="trajectory-line-width">Line Width</label>
+          <input type="range" id="trajectory-line-width" min="0.5" max="5" step="0.5" value="1.5">
+        </div>
+        <div class="slider-container">
+          <label for="trajectory-ball-size">Ball Size</label>
+          <input type="range" id="trajectory-ball-size" min="0.005" max="0.05" step="0.001" value="0.02">
+        </div>
+        <div class="slider-container">
+          <label for="trajectory-history">History Frames</label>
+          <input type="range" id="trajectory-history" min="1" max="500" step="1" value="30">
+        </div>
+        <div class="slider-container" id="tail-opacity-container" style="display: none;">
+          <label for="trajectory-fade">Tail Opacity</label>
+          <input type="range" id="trajectory-fade" min="0" max="1" step="0.05" value="0.0">
+        </div>
+      </div>
+      <div class="settings-group">
+        <h3>Camera</h3>
+        <div class="checkbox-container">
+          <label class="toggle-switch">
+            <input type="checkbox" id="show-camera-frustum" checked>
+            <span class="toggle-slider"></span>
+          </label>
+          <label for="show-camera-frustum">Show Camera Frustum</label>
+        </div>
+        <div class="slider-container">
+          <label for="frustum-size">Size</label>
+          <input type="range" id="frustum-size" min="0.02" max="0.5" step="0.01" value="0.2">
+        </div>
+      </div>
+      <div class="settings-group">
+        <div class="btn-group">
+          <button id="reset-view-btn" style="flex: 1; margin-right: 5px;">Reset View</button>
+          <button id="reset-settings-btn" style="flex: 1; margin-left: 5px;">Reset Settings</button>
+        </div>
+      </div>
+    </div>
+    <button id="show-settings-btn" class="control-btn" title="Show Settings">
+      <svg class="icon" viewBox="0 0 24 24">
+        <path d="M19.14,12.94c0.04-0.3,0.06-0.61,0.06-0.94c0-0.32-0.02-0.64-0.07-0.94l2.03-1.58c0.18-0.14,0.23-0.41,0.12-0.61 l-1.92-3.32c-0.12-0.22-0.37-0.29-0.59-0.22l-2.39,0.96c-0.5-0.38-1.03-0.7-1.62-0.94L14.4,2.81c-0.04-0.24-0.24-0.41-0.48-0.41 h-3.84c-0.24,0-0.43,0.17-0.47,0.41L9.25,5.35C8.66,5.59,8.12,5.92,7.63,6.29L5.24,5.33c-0.22-0.08-0.47,0-0.59,0.22L2.74,8.87 C2.62,9.08,2.66,9.34,2.86,9.48l2.03,1.58C4.84,11.36,4.8,11.69,4.8,12s0.02,0.64,0.07,0.94l-2.03,1.58 c-0.18,0.14-0.23,0.41-0.12,0.61l1.92,3.32c0.12,0.22,0.37,0.29,0.59,0.22l2.39-0.96c0.5,0.38,1.03,0.7,1.62,0.94l0.36,2.54 c0.04,0.24,0.24,0.41,0.48,0.41h3.84c0.24,0,0.44-0.17,0.47-0.41l0.36-2.54c0.59-0.24,1.13-0.56,1.62-0.94l2.39,0.96 c0.22,0.08,0.47,0,0.59-0.22l1.92-3.32c0.12-0.22,0.07-0.47-0.12-0.61L19.14,12.94z M12,15.6c-1.98,0-3.6-1.62-3.6-3.6 s1.62-3.6,3.6-3.6s3.6,1.62,3.6,3.6S13.98,15.6,12,15.6z"/>
+      </svg>
+    </button>
+  </div>
+  <div id="loading-overlay">
+    <!-- <div class="spinner"></div> -->
+    <div id="loading-text"></div>
+    <div class="loading-subtitle" style="font-size: medium;">Interactive Viewer of 3D Tracking</div>
+  </div>
+  <!-- Libraries -->
+  <script src="https://cdnjs.cloudflare.com/ajax/libs/pako/2.1.0/pako.min.js"></script>
+  <script src="https://cdn.jsdelivr.net/npm/three@0.132.2/build/three.min.js"></script>
+  <script src="https://cdn.jsdelivr.net/npm/three@0.132.2/examples/js/controls/OrbitControls.js"></script>
+  <script src="https://cdn.jsdelivr.net/npm/dat.gui@0.7.7/build/dat.gui.min.js"></script>
+  <script src="https://cdn.jsdelivr.net/npm/three@0.132.2/examples/js/lines/LineSegmentsGeometry.js"></script>
+  <script src="https://cdn.jsdelivr.net/npm/three@0.132.2/examples/js/lines/LineGeometry.js"></script>
+  <script src="https://cdn.jsdelivr.net/npm/three@0.132.2/examples/js/lines/LineMaterial.js"></script>
+  <script src="https://cdn.jsdelivr.net/npm/three@0.132.2/examples/js/lines/LineSegments2.js"></script>
+  <script src="https://cdn.jsdelivr.net/npm/three@0.132.2/examples/js/lines/Line2.js"></script>
+  <script>
+    class PointCloudVisualizer {
+      constructor() {
+        this.data = null;
+        this.config = {};
+        this.currentFrame = 0;
+        this.isPlaying = false;
+        this.playbackSpeed = 1;
+        this.lastFrameTime = 0;
+        this.defaultSettings = null;
+        this.ui = {
+          statusBar: document.getElementById('status-bar'),
+          playPauseBtn: document.getElementById('play-pause-btn'),
+          speedBtn: document.getElementById('speed-btn'),
+          timeline: document.getElementById('timeline'),
+          progress: document.getElementById('progress'),
+          settingsPanel: document.getElementById('settings-panel'),
+          loadingOverlay: document.getElementById('loading-overlay'),
+          loadingText: document.getElementById('loading-text'),
+          settingsToggleBtn: document.getElementById('settings-toggle-btn'),
+          frameCounter: document.getElementById('frame-counter'),
+          pointSize: document.getElementById('point-size'),
+          pointOpacity: document.getElementById('point-opacity'),
+          maxDepth: document.getElementById('max-depth'),
+          showTrajectory: document.getElementById('show-trajectory'),
+          enableRichTrail: document.getElementById('enable-rich-trail'),
+          trajectoryLineWidth: document.getElementById('trajectory-line-width'),
+          trajectoryBallSize: document.getElementById('trajectory-ball-size'),
+          trajectoryHistory: document.getElementById('trajectory-history'),
+          trajectoryFade: document.getElementById('trajectory-fade'),
+          tailOpacityContainer: document.getElementById('tail-opacity-container'),
+          resetViewBtn: document.getElementById('reset-view-btn'),
+          showCameraFrustum: document.getElementById('show-camera-frustum'),
+          frustumSize: document.getElementById('frustum-size'),
+          hideSettingsBtn: document.getElementById('hide-settings-btn'),
+          showSettingsBtn: document.getElementById('show-settings-btn')
+        };
+        this.scene = null;
+        this.camera = null;
+        this.renderer = null;
+        this.controls = null;
+        this.pointCloud = null;
+        this.trajectories = [];
+        this.cameraFrustum = null;
+        this.initThreeJS();
+        this.loadDefaultSettings().then(() => {
+          this.initEventListeners();
+          this.loadData();
+        });
+      }
+      async loadDefaultSettings() {
+        try {
+          const urlParams = new URLSearchParams(window.location.search);
+          const dataPath = urlParams.get('data') || '';
+          const defaultSettings = {
+            pointSize: 0.03,
+            pointOpacity: 1.0,
+            showTrajectory: true,
+            trajectoryLineWidth: 2.5,
+            trajectoryBallSize: 0.015,
+            trajectoryHistory: 0,
+            showCameraFrustum: true,
+            frustumSize: 0.2
+          };
+          if (!dataPath) {
+            this.defaultSettings = defaultSettings;
+            this.applyDefaultSettings();
+            return;
+          }
+          // Try to extract dataset and videoId from the data path
+          // Expected format: demos/datasetname/videoid.bin
+          const pathParts = dataPath.split('/');
+          if (pathParts.length < 3) {
+            this.defaultSettings = defaultSettings;
+            this.applyDefaultSettings();
+            return;
+          }
+          const datasetName = pathParts[pathParts.length - 2];
+          let videoId = pathParts[pathParts.length - 1].replace('.bin', '');
+          // Load settings from data.json
+          const response = await fetch('./data.json');
+          if (!response.ok) {
+            this.defaultSettings = defaultSettings;
+            this.applyDefaultSettings();
+            return;
+          }
+          const settingsData = await response.json();
+          // Check if this dataset and video exist
+          if (settingsData[datasetName] && settingsData[datasetName][videoId]) {
+            this.defaultSettings = settingsData[datasetName][videoId];
+          } else {
+            this.defaultSettings = defaultSettings;
+          }
+          this.applyDefaultSettings();
+        } catch (error) {
+          console.error("Error loading default settings:", error);
+          this.defaultSettings = {
+            pointSize: 0.03,
+            pointOpacity: 1.0,
+            showTrajectory: true,
+            trajectoryLineWidth: 2.5,
+            trajectoryBallSize: 0.015,
+            trajectoryHistory: 0,
+            showCameraFrustum: true,
+            frustumSize: 0.2
+          };
+          this.applyDefaultSettings();
+        }
+      }
+      applyDefaultSettings() {
+        if (!this.defaultSettings) return;
+        if (this.ui.pointSize) {
+          this.ui.pointSize.value = this.defaultSettings.pointSize;
+        }
+        if (this.ui.pointOpacity) {
+          this.ui.pointOpacity.value = this.defaultSettings.pointOpacity;
+        }
+        if (this.ui.maxDepth) {
+          this.ui.maxDepth.value = this.defaultSettings.maxDepth || 100.0;
+        }
+        if (this.ui.showTrajectory) {
+          this.ui.showTrajectory.checked = this.defaultSettings.showTrajectory;
+        }
+        if (this.ui.trajectoryLineWidth) {
+          this.ui.trajectoryLineWidth.value = this.defaultSettings.trajectoryLineWidth;
+        }
+        if (this.ui.trajectoryBallSize) {
+          this.ui.trajectoryBallSize.value = this.defaultSettings.trajectoryBallSize;
+        }
+        if (this.ui.trajectoryHistory) {
+          this.ui.trajectoryHistory.value = this.defaultSettings.trajectoryHistory;
+        }
+        if (this.ui.showCameraFrustum) {
+          this.ui.showCameraFrustum.checked = this.defaultSettings.showCameraFrustum;
+        }
+        if (this.ui.frustumSize) {
+          this.ui.frustumSize.value = this.defaultSettings.frustumSize;
+        }
+      }
+      initThreeJS() {
+        this.scene = new THREE.Scene();
+        this.scene.background = new THREE.Color(0x1a1a1a);
+        this.camera = new THREE.PerspectiveCamera(60, window.innerWidth / window.innerHeight, 0.1, 10000);
+        this.camera.position.set(0, 0, 0);
+        this.renderer = new THREE.WebGLRenderer({ antialias: true });
+        this.renderer.setPixelRatio(window.devicePixelRatio);
+        this.renderer.setSize(window.innerWidth, window.innerHeight);
+        document.getElementById('canvas-container').appendChild(this.renderer.domElement);
+        this.controls = new THREE.OrbitControls(this.camera, this.renderer.domElement);
+        this.controls.enableDamping = true;
+        this.controls.dampingFactor = 0.05;
+        this.controls.target.set(0, 0, 0);
+        this.controls.minDistance = 0.1;
+        this.controls.maxDistance = 1000;
+        this.controls.update();
+        const ambientLight = new THREE.AmbientLight(0xffffff, 0.5);
+        this.scene.add(ambientLight);
+        const directionalLight = new THREE.DirectionalLight(0xffffff, 0.8);
+        directionalLight.position.set(1, 1, 1);
+        this.scene.add(directionalLight);
+      }
+      initEventListeners() {
+        window.addEventListener('resize', () => this.onWindowResize());
+        this.ui.playPauseBtn.addEventListener('click', () => this.togglePlayback());
+        this.ui.timeline.addEventListener('click', (e) => {
+          const rect = this.ui.timeline.getBoundingClientRect();
+          const pos = (e.clientX - rect.left) / rect.width;
+          this.seekTo(pos);
+        });
+        this.ui.speedBtn.addEventListener('click', () => this.cyclePlaybackSpeed());
+        this.ui.pointSize.addEventListener('input', () => this.updatePointCloudSettings());
+        this.ui.pointOpacity.addEventListener('input', () => this.updatePointCloudSettings());
+        this.ui.maxDepth.addEventListener('input', () => this.updatePointCloudSettings());
+        this.ui.showTrajectory.addEventListener('change', () => {
+          this.trajectories.forEach(trajectory => {
+            trajectory.visible = this.ui.showTrajectory.checked;
+          });
+        });
+        this.ui.enableRichTrail.addEventListener('change', () => {
+            this.ui.tailOpacityContainer.style.display = this.ui.enableRichTrail.checked ? 'flex' : 'none';
+            this.updateTrajectories(this.currentFrame);
+        });
+        this.ui.trajectoryLineWidth.addEventListener('input', () => this.updateTrajectorySettings());
+        this.ui.trajectoryBallSize.addEventListener('input', () => this.updateTrajectorySettings());
+        this.ui.trajectoryHistory.addEventListener('input', () => {
+          this.updateTrajectories(this.currentFrame);
+        });
+        this.ui.trajectoryFade.addEventListener('input', () => {
+          this.updateTrajectories(this.currentFrame);
+        });
+        this.ui.resetViewBtn.addEventListener('click', () => this.resetView());
+        const resetSettingsBtn = document.getElementById('reset-settings-btn');
+        if (resetSettingsBtn) {
+          resetSettingsBtn.addEventListener('click', () => this.resetSettings());
+        }
+        document.addEventListener('keydown', (e) => {
+          if (e.key === 'Escape' && this.ui.settingsPanel.classList.contains('visible')) {
+            this.ui.settingsPanel.classList.remove('visible');
+            this.ui.settingsToggleBtn.classList.remove('active');
+          }
+        });
+        if (this.ui.settingsToggleBtn) {
+          this.ui.settingsToggleBtn.addEventListener('click', () => {
+            const isVisible = this.ui.settingsPanel.classList.toggle('visible');
+            this.ui.settingsToggleBtn.classList.toggle('active', isVisible);
+            if (isVisible) {
+              const panelRect = this.ui.settingsPanel.getBoundingClientRect();
+              const viewportHeight = window.innerHeight;
+              if (panelRect.bottom > viewportHeight) {
+                this.ui.settingsPanel.style.bottom = 'auto';
+                this.ui.settingsPanel.style.top = '80px';
+              }
+            }
+          });
+        }
+        if (this.ui.frustumSize) {
+          this.ui.frustumSize.addEventListener('input', () => this.updateFrustumDimensions());
+        }
+        if (this.ui.hideSettingsBtn && this.ui.showSettingsBtn && this.ui.settingsPanel) {
+          this.ui.hideSettingsBtn.addEventListener('click', () => {
+            this.ui.settingsPanel.classList.add('is-hidden');
+            this.ui.showSettingsBtn.style.display = 'flex';
+          });
+          this.ui.showSettingsBtn.addEventListener('click', () => {
+            this.ui.settingsPanel.classList.remove('is-hidden');
+            this.ui.showSettingsBtn.style.display = 'none';
+          });
+        }
+      }
+      makeElementDraggable(element) {
+        let pos1 = 0, pos2 = 0, pos3 = 0, pos4 = 0;
+        const dragHandle = element.querySelector('h2');
+        if (dragHandle) {
+          dragHandle.onmousedown = dragMouseDown;
+          dragHandle.title = "Drag to move panel";
+        } else {
+          element.onmousedown = dragMouseDown;
+        }
+        function dragMouseDown(e) {
+          e = e || window.event;
+          e.preventDefault();
+          pos3 = e.clientX;
+          pos4 = e.clientY;
+          document.onmouseup = closeDragElement;
+          document.onmousemove = elementDrag;
+          element.classList.add('dragging');
+        }
+        function elementDrag(e) {
+          e = e || window.event;
+          e.preventDefault();
+          pos1 = pos3 - e.clientX;
+          pos2 = pos4 - e.clientY;
+          pos3 = e.clientX;
+          pos4 = e.clientY;
+          const newTop = element.offsetTop - pos2;
+          const newLeft = element.offsetLeft - pos1;
+          const viewportWidth = window.innerWidth;
+          const viewportHeight = window.innerHeight;
+          const panelRect = element.getBoundingClientRect();
+          const maxTop = viewportHeight - 50;
+          const maxLeft = viewportWidth - 50;
+          element.style.top = Math.min(Math.max(newTop, 0), maxTop) + "px";
+          element.style.left = Math.min(Math.max(newLeft, 0), maxLeft) + "px";
+          // Remove bottom/right settings when dragging
+          element.style.bottom = 'auto';
+          element.style.right = 'auto';
+        }
+        function closeDragElement() {
+          document.onmouseup = null;
+          document.onmousemove = null;
+          element.classList.remove('dragging');
+        }
+      }
+      async loadData() {
+        try {
+          // this.ui.loadingText.textContent = "Loading binary data...";
+          let arrayBuffer;
+          if (window.embeddedBase64) {
+            // Base64 embedded path
+            const binaryString = atob(window.embeddedBase64);
+            const len = binaryString.length;
+            const bytes = new Uint8Array(len);
+            for (let i = 0; i < len; i++) {
+              bytes[i] = binaryString.charCodeAt(i);
+            }
+            arrayBuffer = bytes.buffer;
+          } else {
+            // Default fetch path (fallback)
+            const urlParams = new URLSearchParams(window.location.search);
+            const dataPath = urlParams.get('data') || 'data.bin';
+            const response = await fetch(dataPath);
+            if (!response.ok) throw new Error(`Failed to load ${dataPath}`);
+            arrayBuffer = await response.arrayBuffer();
+          }
+          const dataView = new DataView(arrayBuffer);
+          const headerLen = dataView.getUint32(0, true);
+          const headerText = new TextDecoder("utf-8").decode(arrayBuffer.slice(4, 4 + headerLen));
+          const header = JSON.parse(headerText);
+          const compressedBlob = new Uint8Array(arrayBuffer, 4 + headerLen);
+          const decompressed = pako.inflate(compressedBlob).buffer;
+          const arrays = {};
+          for (const key in header) {
+            if (key === "meta") continue;
+            const meta = header[key];
+            const { dtype, shape, offset, length } = meta;
+            const slice = decompressed.slice(offset, offset + length);
+            let typedArray;
+            switch (dtype) {
+              case "uint8": typedArray = new Uint8Array(slice); break;
+              case "uint16": typedArray = new Uint16Array(slice); break;
+              case "float32": typedArray = new Float32Array(slice); break;
+              case "float64": typedArray = new Float64Array(slice); break;
+              default: throw new Error(`Unknown dtype: ${dtype}`);
+            }
+            arrays[key] = { data: typedArray, shape: shape };
+          }
+          this.data = arrays;
+          this.config = header.meta;
+          this.initCameraWithCorrectFOV();
+          this.ui.loadingText.textContent = "Creating point cloud...";
+          this.initPointCloud();
+          this.initTrajectories();
+          setTimeout(() => {
+            this.ui.loadingOverlay.classList.add('fade-out');
+            this.ui.statusBar.classList.add('hidden');
+            this.startAnimation();
+          }, 500);
+        } catch (error) {
+          console.error("Error loading data:", error);
+          this.ui.statusBar.textContent = `Error: ${error.message}`;
+          // this.ui.loadingText.textContent = `Error loading data: ${error.message}`;
+        }
+      }
+      initPointCloud() {
+        const numPoints = this.config.resolution[0] * this.config.resolution[1];
+        const positions = new Float32Array(numPoints * 3);
+        const colors = new Float32Array(numPoints * 3);
+        const geometry = new THREE.BufferGeometry();
+        geometry.setAttribute('position', new THREE.BufferAttribute(positions, 3).setUsage(THREE.DynamicDrawUsage));
+        geometry.setAttribute('color', new THREE.BufferAttribute(colors, 3).setUsage(THREE.DynamicDrawUsage));
+        const pointSize = parseFloat(this.ui.pointSize.value) || this.defaultSettings.pointSize;
+        const pointOpacity = parseFloat(this.ui.pointOpacity.value) || this.defaultSettings.pointOpacity;
+        const material = new THREE.PointsMaterial({
+          size: pointSize,
+          vertexColors: true,
+          transparent: true,
+          opacity: pointOpacity,
+          sizeAttenuation: true
+        });
+        this.pointCloud = new THREE.Points(geometry, material);
+        this.scene.add(this.pointCloud);
+      }
+      initTrajectories() {
+        if (!this.data.trajectories) return;
+        this.trajectories.forEach(trajectory => {
+            if (trajectory.userData.lineSegments) {
+                trajectory.userData.lineSegments.forEach(segment => {
+                    segment.geometry.dispose();
+                    segment.material.dispose();
+                });
+            }
+            this.scene.remove(trajectory);
+        });
+        this.trajectories = [];
+        const shape = this.data.trajectories.shape;
+        if (!shape || shape.length < 2) return;
+        const [totalFrames, numTrajectories] = shape;
+        const palette = this.createColorPalette(numTrajectories);
+        const resolution = new THREE.Vector2(window.innerWidth, window.innerHeight);
+        const maxHistory = 500; // Max value of the history slider, for the object pool
+        for (let i = 0; i < numTrajectories; i++) {
+            const trajectoryGroup = new THREE.Group();
+            const ballSize = parseFloat(this.ui.trajectoryBallSize.value);
+            const sphereGeometry = new THREE.SphereGeometry(ballSize, 16, 16);
+            const sphereMaterial = new THREE.MeshBasicMaterial({ color: palette[i], transparent: true });
+            const positionMarker = new THREE.Mesh(sphereGeometry, sphereMaterial);
+            trajectoryGroup.add(positionMarker);
+            // High-Performance Line (default)
+            const simpleLineGeometry = new THREE.BufferGeometry();
+            const simpleLinePositions = new Float32Array(maxHistory * 3);
+            simpleLineGeometry.setAttribute('position', new THREE.BufferAttribute(simpleLinePositions, 3).setUsage(THREE.DynamicDrawUsage));
+            const simpleLine = new THREE.Line(simpleLineGeometry, new THREE.LineBasicMaterial({ color: palette[i] }));
+            simpleLine.frustumCulled = false;
+            trajectoryGroup.add(simpleLine);
+            // High-Quality Line Segments (for rich trail)
+            const lineSegments = [];
+            const lineWidth = parseFloat(this.ui.trajectoryLineWidth.value);
+            // Create a pool of line segment objects
+            for (let j = 0; j < maxHistory - 1; j++) {
+                const lineGeometry = new THREE.LineGeometry();
+                lineGeometry.setPositions([0, 0, 0, 0, 0, 0]);
+                const lineMaterial = new THREE.LineMaterial({
+                    color: palette[i],
+                    linewidth: lineWidth,
+                    resolution: resolution,
+                    transparent: true,
+                    depthWrite: false, // Correctly handle transparency
+                    opacity: 0
+                });
+                const segment = new THREE.Line2(lineGeometry, lineMaterial);
+                segment.frustumCulled = false;
+                segment.visible = false; // Start with all segments hidden
+                trajectoryGroup.add(segment);
+                lineSegments.push(segment);
+            }
+            trajectoryGroup.userData = {
+                marker: positionMarker,
+                simpleLine: simpleLine,
+                lineSegments: lineSegments,
+                color: palette[i]
+            };
+            this.scene.add(trajectoryGroup);
+            this.trajectories.push(trajectoryGroup);
+        }
+        const showTrajectory = this.ui.showTrajectory.checked;
+        this.trajectories.forEach(trajectory => trajectory.visible = showTrajectory);
+      }
+      createColorPalette(count) {
+        const colors = [];
+        const hueStep = 360 / count;
+        for (let i = 0; i < count; i++) {
+          const hue = (i * hueStep) % 360;
+          const color = new THREE.Color().setHSL(hue / 360, 0.8, 0.6);
+          colors.push(color);
+        }
+        return colors;
+      }
+      updatePointCloud(frameIndex) {
+        if (!this.data || !this.pointCloud) return;
+        const positions = this.pointCloud.geometry.attributes.position.array;
+        const colors = this.pointCloud.geometry.attributes.color.array;
+        const rgbVideo = this.data.rgb_video;
+        const depthsRgb = this.data.depths_rgb;
+        const intrinsics = this.data.intrinsics;
+        const invExtrinsics = this.data.inv_extrinsics;
+        const width = this.config.resolution[0];
+        const height = this.config.resolution[1];
+        const numPoints = width * height;
+        const K = this.get3x3Matrix(intrinsics.data, intrinsics.shape, frameIndex);
+        const fx = K[0][0], fy = K[1][1], cx = K[0][2], cy = K[1][2];
+        const invExtrMat = this.get4x4Matrix(invExtrinsics.data, invExtrinsics.shape, frameIndex);
+        const transform = this.getTransformElements(invExtrMat);
+        const rgbFrame = this.getFrame(rgbVideo.data, rgbVideo.shape, frameIndex);
+        const depthFrame = this.getFrame(depthsRgb.data, depthsRgb.shape, frameIndex);
+        const maxDepth = parseFloat(this.ui.maxDepth.value) || 10.0;
+        let validPointCount = 0;
+        for (let i = 0; i < numPoints; i++) {
+          const xPix = i % width;
+          const yPix = Math.floor(i / width);
+          const d0 = depthFrame[i * 3];
+          const d1 = depthFrame[i * 3 + 1];
+          const depthEncoded = d0 | (d1 << 8);
+          const depthValue = (depthEncoded / ((1 << 16) - 1)) *
+                           (this.config.depthRange[1] - this.config.depthRange[0]) +
+                           this.config.depthRange[0];
+          if (depthValue === 0 || depthValue > maxDepth) {
+            continue;
+          }
+          const X = ((xPix - cx) * depthValue) / fx;
+          const Y = ((yPix - cy) * depthValue) / fy;
+          const Z = depthValue;
+          const tx = transform.m11 * X + transform.m12 * Y + transform.m13 * Z + transform.m14;
+          const ty = transform.m21 * X + transform.m22 * Y + transform.m23 * Z + transform.m24;
+          const tz = transform.m31 * X + transform.m32 * Y + transform.m33 * Z + transform.m34;
+          const index = validPointCount * 3;
+          positions[index] = tx;
+          positions[index + 1] = -ty;
+          positions[index + 2] = -tz;
+          colors[index] = rgbFrame[i * 3] / 255;
+          colors[index + 1] = rgbFrame[i * 3 + 1] / 255;
+          colors[index + 2] = rgbFrame[i * 3 + 2] / 255;
+          validPointCount++;
+        }
+        this.pointCloud.geometry.setDrawRange(0, validPointCount);
+        this.pointCloud.geometry.attributes.position.needsUpdate = true;
+        this.pointCloud.geometry.attributes.color.needsUpdate = true;
+        this.pointCloud.geometry.computeBoundingSphere(); // Important for camera culling
+        this.updateTrajectories(frameIndex);
+        const progress = (frameIndex + 1) / this.config.totalFrames;
+        this.ui.progress.style.width = `${progress * 100}%`;
+        if (this.ui.frameCounter && this.config.totalFrames) {
+          this.ui.frameCounter.textContent = `Frame ${frameIndex} / ${this.config.totalFrames - 1}`;
+        }
+        this.updateCameraFrustum(frameIndex);
+      }
+      updateTrajectories(frameIndex) {
+        if (!this.data.trajectories || this.trajectories.length === 0) return;
+        const trajectoryData = this.data.trajectories.data;
+        const [totalFrames, numTrajectories] = this.data.trajectories.shape;
+        const historyFrames = parseInt(this.ui.trajectoryHistory.value);
+        const tailOpacity = parseFloat(this.ui.trajectoryFade.value);
+        const isRichMode = this.ui.enableRichTrail.checked;
+        for (let i = 0; i < numTrajectories; i++) {
+          const trajectoryGroup = this.trajectories[i];
+          const { marker, simpleLine, lineSegments } = trajectoryGroup.userData;
+          const currentPos = new THREE.Vector3();
+          const currentOffset = (frameIndex * numTrajectories + i) * 3;
+          currentPos.x = trajectoryData[currentOffset];
+          currentPos.y = -trajectoryData[currentOffset + 1];
+          currentPos.z = -trajectoryData[currentOffset + 2];
+          marker.position.copy(currentPos);
+          marker.material.opacity = 1.0;
+          const historyToShow = Math.min(historyFrames, frameIndex + 1);
+          if (isRichMode) {
+              // --- High-Quality Mode ---
+              simpleLine.visible = false;
+              for (let j = 0; j < lineSegments.length; j++) {
+                  const segment = lineSegments[j];
+                  if (j < historyToShow - 1) {
+                      const headFrame = frameIndex - j;
+                      const tailFrame = frameIndex - j - 1;
+                      const headOffset = (headFrame * numTrajectories + i) * 3;
+                      const tailOffset = (tailFrame * numTrajectories + i) * 3;
+                      const positions = [
+                          trajectoryData[headOffset], -trajectoryData[headOffset + 1], -trajectoryData[headOffset + 2],
+                          trajectoryData[tailOffset], -trajectoryData[tailOffset + 1], -trajectoryData[tailOffset + 2]
+                      ];
+                      segment.geometry.setPositions(positions);
+                      const headOpacity = 1.0;
+                      const normalizedAge = j / Math.max(1, historyToShow - 2);
+                      const alpha = headOpacity - (headOpacity - tailOpacity) * normalizedAge;
+                      segment.material.opacity = Math.max(0, alpha);
+                      segment.visible = true;
+                  } else {
+                      segment.visible = false;
+                  }
+              }
+          } else {
+              // --- Performance Mode ---
+              lineSegments.forEach(s => s.visible = false);
+              simpleLine.visible = true;
+              const positions = simpleLine.geometry.attributes.position.array;
+              for (let j = 0; j < historyToShow; j++) {
+                  const historyFrame = Math.max(0, frameIndex - j);
+                  const offset = (historyFrame * numTrajectories + i) * 3;
+                  positions[j * 3] = trajectoryData[offset];
+                  positions[j * 3 + 1] = -trajectoryData[offset + 1];
+                  positions[j * 3 + 2] = -trajectoryData[offset + 2];
+              }
+              simpleLine.geometry.setDrawRange(0, historyToShow);
+              simpleLine.geometry.attributes.position.needsUpdate = true;
+          }
+        }
+      }
+      updateTrajectorySettings() {
+        if (!this.trajectories || this.trajectories.length === 0) return;
+        const ballSize = parseFloat(this.ui.trajectoryBallSize.value);
+        const lineWidth = parseFloat(this.ui.trajectoryLineWidth.value);
+        this.trajectories.forEach(trajectoryGroup => {
+          const { marker, lineSegments } = trajectoryGroup.userData;
+          marker.geometry.dispose();
+          marker.geometry = new THREE.SphereGeometry(ballSize, 16, 16);
+          // Line width only affects rich mode
+          lineSegments.forEach(segment => {
+            if (segment.material) {
+              segment.material.linewidth = lineWidth;
+            }
+          });
+        });
+        this.updateTrajectories(this.currentFrame);
+      }
+      getDepthColor(normalizedDepth) {
+        const hue = (1 - normalizedDepth) * 240 / 360;
+        const color = new THREE.Color().setHSL(hue, 1.0, 0.5);
+        return color;
+      }
+      getFrame(typedArray, shape, frameIndex) {
+        const [T, H, W, C] = shape;
+        const frameSize = H * W * C;
+        const offset = frameIndex * frameSize;
+        return typedArray.subarray(offset, offset + frameSize);
+      }
+      get3x3Matrix(typedArray, shape, frameIndex) {
+        const frameSize = 9;
+        const offset = frameIndex * frameSize;
+        const K = [];
+        for (let i = 0; i < 3; i++) {
+          const row = [];
+          for (let j = 0; j < 3; j++) {
+            row.push(typedArray[offset + i * 3 + j]);
+          }
+          K.push(row);
+        }
+        return K;
+      }
+      get4x4Matrix(typedArray, shape, frameIndex) {
+        const frameSize = 16;
+        const offset = frameIndex * frameSize;
+        const M = [];
+        for (let i = 0; i < 4; i++) {
+          const row = [];
+          for (let j = 0; j < 4; j++) {
+            row.push(typedArray[offset + i * 4 + j]);
+          }
+          M.push(row);
+        }
+        return M;
+      }
+      getTransformElements(matrix) {
+        return {
+          m11: matrix[0][0], m12: matrix[0][1], m13: matrix[0][2], m14: matrix[0][3],
+          m21: matrix[1][0], m22: matrix[1][1], m23: matrix[1][2], m24: matrix[1][3],
+          m31: matrix[2][0], m32: matrix[2][1], m33: matrix[2][2], m34: matrix[2][3]
+        };
+      }
+      togglePlayback() {
+        this.isPlaying = !this.isPlaying;
+        const playIcon = document.getElementById('play-icon');
+        const pauseIcon = document.getElementById('pause-icon');
+        if (this.isPlaying) {
+          playIcon.style.display = 'none';
+          pauseIcon.style.display = 'block';
+          this.lastFrameTime = performance.now();
+        } else {
+          playIcon.style.display = 'block';
+          pauseIcon.style.display = 'none';
+        }
+      }
+      cyclePlaybackSpeed() {
+        const speeds = [0.5, 1, 2, 4, 8];
+        const speedRates = speeds.map(s => s * this.config.baseFrameRate);
+        let currentIndex = 0;
+        const normalizedSpeed = this.playbackSpeed / this.config.baseFrameRate;
+        for (let i = 0; i < speeds.length; i++) {
+          if (Math.abs(normalizedSpeed - speeds[i]) < Math.abs(normalizedSpeed - speeds[currentIndex])) {
+            currentIndex = i;
+          }
+        }
+        const nextIndex = (currentIndex + 1) % speeds.length;
+        this.playbackSpeed = speedRates[nextIndex];
+        this.ui.speedBtn.textContent = `${speeds[nextIndex]}x`;
+        if (speeds[nextIndex] === 1) {
+          this.ui.speedBtn.classList.remove('active');
+        } else {
+          this.ui.speedBtn.classList.add('active');
+        }
+      }
+      seekTo(position) {
+        const frameIndex = Math.floor(position * this.config.totalFrames);
+        this.currentFrame = Math.max(0, Math.min(frameIndex, this.config.totalFrames - 1));
+        this.updatePointCloud(this.currentFrame);
+      }
+      updatePointCloudSettings() {
+        if (!this.pointCloud) return;
+        const size = parseFloat(this.ui.pointSize.value);
+        const opacity = parseFloat(this.ui.pointOpacity.value);
+        this.pointCloud.material.size = size;
+        this.pointCloud.material.opacity = opacity;
+        this.pointCloud.material.needsUpdate = true;
+        this.updatePointCloud(this.currentFrame);
+      }
+      updateControls() {
+        if (!this.controls) return;
+        this.controls.update();
+      }
+      resetView() {
+        if (!this.camera || !this.controls) return;
+        // Reset camera position
+        this.camera.position.set(0, 0, this.config.cameraZ || 0);
+        // Reset controls
+        this.controls.reset();
+        // Set target slightly in front of camera
+        this.controls.target.set(0, 0, -1);
+        this.controls.update();
+        // Show status message
+        this.ui.statusBar.textContent = "View reset";
+        this.ui.statusBar.classList.remove('hidden');
+        // Hide status message after a few seconds
+        setTimeout(() => {
+          this.ui.statusBar.classList.add('hidden');
+        }, 3000);
+      }
+      onWindowResize() {
+        if (!this.camera || !this.renderer) return;
+        const windowAspect = window.innerWidth / window.innerHeight;
+        this.camera.aspect = windowAspect;
+        this.camera.updateProjectionMatrix();
+        this.renderer.setSize(window.innerWidth, window.innerHeight);
+        if (this.trajectories && this.trajectories.length > 0) {
+          const resolution = new THREE.Vector2(window.innerWidth, window.innerHeight);
+          this.trajectories.forEach(trajectory => {
+            const { lineSegments } = trajectory.userData;
+            if (lineSegments && lineSegments.length > 0) {
+              lineSegments.forEach(segment => {
+                if (segment.material && segment.material.resolution) {
+                  segment.material.resolution.copy(resolution);
+                }
+              });
+            }
+          });
+        }
+        if (this.cameraFrustum) {
+          const resolution = new THREE.Vector2(window.innerWidth, window.innerHeight);
+          this.cameraFrustum.children.forEach(line => {
+            if (line.material && line.material.resolution) {
+              line.material.resolution.copy(resolution);
+            }
+          });
+        }
+      }
+      startAnimation() {
+        this.isPlaying = true;
+        this.lastFrameTime = performance.now();
+        this.camera.position.set(0, 0, this.config.cameraZ || 0);
+        this.controls.target.set(0, 0, -1);
+        this.controls.update();
+        this.playbackSpeed = this.config.baseFrameRate;
+        document.getElementById('play-icon').style.display = 'none';
+        document.getElementById('pause-icon').style.display = 'block';
+        this.animate();
+      }
+      animate() {
+        requestAnimationFrame(() => this.animate());
+        if (this.controls) {
+          this.controls.update();
+        }
+        if (this.isPlaying && this.data) {
+          const now = performance.now();
+          const delta = (now - this.lastFrameTime) / 1000;
+          const framesToAdvance = Math.floor(delta * this.config.baseFrameRate * this.playbackSpeed);
+          if (framesToAdvance > 0) {
+            this.currentFrame = (this.currentFrame + framesToAdvance) % this.config.totalFrames;
+            this.lastFrameTime = now;
+            this.updatePointCloud(this.currentFrame);
+          }
+        }
+        if (this.renderer && this.scene && this.camera) {
+          this.renderer.render(this.scene, this.camera);
+        }
+      }
+      initCameraWithCorrectFOV() {
+        const fov = this.config.fov || 60;
+        const windowAspect = window.innerWidth / window.innerHeight;
+        this.camera = new THREE.PerspectiveCamera(
+          fov,
+          windowAspect,
+          0.1,
+          10000
+        );
+        this.controls.object = this.camera;
+        this.controls.update();
+        this.initCameraFrustum();
+      }
+      initCameraFrustum() {
+        this.cameraFrustum = new THREE.Group();
+        this.scene.add(this.cameraFrustum);
+        this.initCameraFrustumGeometry();
+        const showCameraFrustum = this.ui.showCameraFrustum ? this.ui.showCameraFrustum.checked : (this.defaultSettings ? this.defaultSettings.showCameraFrustum : false);
+        this.cameraFrustum.visible = showCameraFrustum;
+      }
+      initCameraFrustumGeometry() {
+        const fov = this.config.fov || 60;
+        const originalAspect = this.config.original_aspect_ratio || 1.33;
+        const size = parseFloat(this.ui.frustumSize.value) || this.defaultSettings.frustumSize;
+        const halfHeight = Math.tan(THREE.MathUtils.degToRad(fov / 2)) * size;
+        const halfWidth = halfHeight * originalAspect;
+        const vertices = [
+          new THREE.Vector3(0, 0, 0),
+          new THREE.Vector3(-halfWidth, -halfHeight, size),
+          new THREE.Vector3(halfWidth, -halfHeight, size),
+          new THREE.Vector3(halfWidth, halfHeight, size),
+          new THREE.Vector3(-halfWidth, halfHeight, size)
+        ];
+        const resolution = new THREE.Vector2(window.innerWidth, window.innerHeight);
+        const linePairs = [
+          [1, 2], [2, 3], [3, 4], [4, 1],
+          [0, 1], [0, 2], [0, 3], [0, 4]
+        ];
+        const colors = {
+          edge: new THREE.Color(0x3366ff),
+          ray: new THREE.Color(0x33cc66)
+        };
+        linePairs.forEach((pair, index) => {
+          const positions = [
+            vertices[pair[0]].x, vertices[pair[0]].y, vertices[pair[0]].z,
+            vertices[pair[1]].x, vertices[pair[1]].y, vertices[pair[1]].z
+          ];
+          const lineGeometry = new THREE.LineGeometry();
+          lineGeometry.setPositions(positions);
+          let color = index < 4 ? colors.edge : colors.ray;
+          const lineMaterial = new THREE.LineMaterial({
+            color: color,
+            linewidth: 2,
+            resolution: resolution,
+            dashed: false
+          });
+          const line = new THREE.Line2(lineGeometry, lineMaterial);
+          this.cameraFrustum.add(line);
+        });
+      }
+      updateCameraFrustum(frameIndex) {
+        if (!this.cameraFrustum || !this.data) return;
+        const invExtrinsics = this.data.inv_extrinsics;
+        if (!invExtrinsics) return;
+        const invExtrMat = this.get4x4Matrix(invExtrinsics.data, invExtrinsics.shape, frameIndex);
+        const matrix = new THREE.Matrix4();
+        matrix.set(
+          invExtrMat[0][0], invExtrMat[0][1], invExtrMat[0][2], invExtrMat[0][3],
+          invExtrMat[1][0], invExtrMat[1][1], invExtrMat[1][2], invExtrMat[1][3],
+          invExtrMat[2][0], invExtrMat[2][1], invExtrMat[2][2], invExtrMat[2][3],
+          invExtrMat[3][0], invExtrMat[3][1], invExtrMat[3][2], invExtrMat[3][3]
+        );
+        const position = new THREE.Vector3();
+        position.setFromMatrixPosition(matrix);
+        const rotMatrix = new THREE.Matrix4().extractRotation(matrix);
+        const coordinateCorrection = new THREE.Matrix4().makeRotationX(Math.PI);
+        const finalRotation = new THREE.Matrix4().multiplyMatrices(coordinateCorrection, rotMatrix);
+        const quaternion = new THREE.Quaternion();
+        quaternion.setFromRotationMatrix(finalRotation);
+        position.y = -position.y;
+        position.z = -position.z;
+        this.cameraFrustum.position.copy(position);
+        this.cameraFrustum.quaternion.copy(quaternion);
+        const showCameraFrustum = this.ui.showCameraFrustum ? this.ui.showCameraFrustum.checked : this.defaultSettings.showCameraFrustum;
+        if (this.cameraFrustum.visible !== showCameraFrustum) {
+          this.cameraFrustum.visible = showCameraFrustum;
+        }
+        const resolution = new THREE.Vector2(window.innerWidth, window.innerHeight);
+        this.cameraFrustum.children.forEach(line => {
+          if (line.material && line.material.resolution) {
+            line.material.resolution.copy(resolution);
+          }
+        });
+      }
+      updateFrustumDimensions() {
+        if (!this.cameraFrustum) return;
+        while(this.cameraFrustum.children.length > 0) {
+          const child = this.cameraFrustum.children[0];
+          if (child.geometry) child.geometry.dispose();
+          if (child.material) child.material.dispose();
+          this.cameraFrustum.remove(child);
+        }
+        this.initCameraFrustumGeometry();
+        this.updateCameraFrustum(this.currentFrame);
+      }
+      resetSettings() {
+        if (!this.defaultSettings) return;
+        this.applyDefaultSettings();
+        this.updatePointCloudSettings();
+        this.updateTrajectorySettings();
+        this.updateFrustumDimensions();
+        this.ui.statusBar.textContent = "Settings reset to defaults";
+        this.ui.statusBar.classList.remove('hidden');
+        setTimeout(() => {
+          this.ui.statusBar.classList.add('hidden');
+        }, 3000);
+      }
+    }
+    window.addEventListener('DOMContentLoaded', () => {
+      new PointCloudVisualizer();
+    });
+  </script>
+</body>
+</html>

app.py ADDED Viewed

	@@ -0,0 +1,1118 @@

+import gradio as gr
+import os
+import json
+import numpy as np
+import cv2
+import base64
+import time
+import tempfile
+import shutil
+import glob
+import threading
+import subprocess
+import struct
+import zlib
+from pathlib import Path
+from einops import rearrange
+from typing import List, Tuple, Union
+try:
+    import spaces
+except ImportError:
+    # Fallback for local development
+    def spaces(func):
+        return func
+import torch
+import logging
+from concurrent.futures import ThreadPoolExecutor
+import atexit
+import uuid
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+# Import custom modules with error handling
+try:
+    from app_3rd.sam_utils.inference import SamPredictor, get_sam_predictor, run_inference
+    from app_3rd.spatrack_utils.infer_track import get_tracker_predictor, run_tracker, get_points_on_a_grid
+except ImportError as e:
+    logger.error(f"Failed to import custom modules: {e}")
+    raise
+# Constants
+MAX_FRAMES = 80
+COLORS = [(0, 0, 255), (0, 255, 255)]  # BGR: Red for negative, Yellow for positive
+MARKERS = [1, 5]  # Cross for negative, Star for positive
+MARKER_SIZE = 8
+# Thread pool for delayed deletion
+thread_pool_executor = ThreadPoolExecutor(max_workers=2)
+def delete_later(path: Union[str, os.PathLike], delay: int = 600):
+    """Delete file or directory after specified delay (default 10 minutes)"""
+    def _delete():
+        try:
+            if os.path.isfile(path):
+                os.remove(path)
+            elif os.path.isdir(path):
+                shutil.rmtree(path)
+        except Exception as e:
+            logger.warning(f"Failed to delete {path}: {e}")
+    def _wait_and_delete():
+        time.sleep(delay)
+        _delete()
+    thread_pool_executor.submit(_wait_and_delete)
+    atexit.register(_delete)
+def create_user_temp_dir():
+    """Create a unique temporary directory for each user session"""
+    session_id = str(uuid.uuid4())[:8]  # Short unique ID
+    temp_dir = os.path.join("temp_local", f"session_{session_id}")
+    os.makedirs(temp_dir, exist_ok=True)
+    # Schedule deletion after 10 minutes
+    delete_later(temp_dir, delay=600)
+    return temp_dir
+from huggingface_hub import hf_hub_download
+# init the model
+os.environ["VGGT_DIR"] = hf_hub_download("Yuxihenry/SpatialTrackerCkpts", "spatrack_front.pth") #, force_download=True)
+if os.environ.get("VGGT_DIR", None) is not None:
+    from models.vggt.vggt.models.vggt_moe import VGGT_MoE
+    from models.vggt.vggt.utils.load_fn import preprocess_image
+    vggt_model = VGGT_MoE()
+    vggt_model.load_state_dict(torch.load(os.environ.get("VGGT_DIR")), strict=False)
+    vggt_model.eval()
+    vggt_model = vggt_model.to("cuda")
+# Global model initialization
+print("🚀 Initializing local models...")
+tracker_model, _ = get_tracker_predictor(".", vo_points=756)
+predictor = get_sam_predictor()
+print("✅ Models loaded successfully!")
+gr.set_static_paths(paths=[Path.cwd().absolute()/"_viz"])
+@spaces.GPU
+def gpu_run_inference(predictor_arg, image, points, boxes):
+    """GPU-accelerated SAM inference"""
+    if predictor_arg is None:
+        print("Initializing SAM predictor inside GPU function...")
+        predictor_arg = get_sam_predictor(predictor=predictor)
+    # Ensure predictor is on GPU
+    try:
+        if hasattr(predictor_arg, 'model'):
+            predictor_arg.model = predictor_arg.model.cuda()
+        elif hasattr(predictor_arg, 'sam'):
+            predictor_arg.sam = predictor_arg.sam.cuda()
+        elif hasattr(predictor_arg, 'to'):
+            predictor_arg = predictor_arg.to('cuda')
+        if hasattr(image, 'cuda'):
+            image = image.cuda()
+    except Exception as e:
+        print(f"Warning: Could not move predictor to GPU: {e}")
+    return run_inference(predictor_arg, image, points, boxes)
+@spaces.GPU
+def gpu_run_tracker(tracker_model_arg, tracker_viser_arg, temp_dir, video_name, grid_size, vo_points, fps, mode="offline"):
+    """GPU-accelerated tracking"""
+    import torchvision.transforms as T
+    import decord
+    if tracker_model_arg is None or tracker_viser_arg is None:
+        print("Initializing tracker models inside GPU function...")
+        out_dir = os.path.join(temp_dir, "results")
+        os.makedirs(out_dir, exist_ok=True)
+        tracker_model_arg, tracker_viser_arg = get_tracker_predictor(out_dir, vo_points=vo_points, tracker_model=tracker_model)
+    # Setup paths
+    video_path = os.path.join(temp_dir, f"{video_name}.mp4")
+    mask_path = os.path.join(temp_dir, f"{video_name}.png")
+    out_dir = os.path.join(temp_dir, "results")
+    os.makedirs(out_dir, exist_ok=True)
+    # Load video using decord
+    video_reader = decord.VideoReader(video_path)
+    video_tensor = torch.from_numpy(video_reader.get_batch(range(len(video_reader))).asnumpy()).permute(0, 3, 1, 2)
+    # Resize to ensure minimum side is 336
+    h, w = video_tensor.shape[2:]
+    scale = max(224 / h, 224 / w)
+    if scale < 1:
+        new_h, new_w = int(h * scale), int(w * scale)
+        video_tensor = T.Resize((new_h, new_w))(video_tensor)
+    video_tensor = video_tensor[::fps].float()[:MAX_FRAMES]
+    # Move to GPU
+    video_tensor = video_tensor.cuda()
+    print(f"Video tensor shape: {video_tensor.shape}, device: {video_tensor.device}")
+    depth_tensor = None
+    intrs = None
+    extrs = None
+    data_npz_load = {}
+    # run vggt
+    if os.environ.get("VGGT_DIR", None) is not None:
+        # process the image tensor
+        video_tensor = preprocess_image(video_tensor)[None]
+        with torch.no_grad():
+            with torch.cuda.amp.autocast(dtype=torch.bfloat16):
+                # Predict attributes including cameras, depth maps, and point maps.
+                predictions = vggt_model(video_tensor.cuda()/255)
+                extrinsic, intrinsic = predictions["poses_pred"], predictions["intrs"]
+                depth_map, depth_conf = predictions["points_map"][..., 2], predictions["unc_metric"]
+        depth_tensor = depth_map.squeeze().cpu().numpy()
+        extrs = np.eye(4)[None].repeat(len(depth_tensor), axis=0)
+        extrs = extrinsic.squeeze().cpu().numpy()
+        intrs = intrinsic.squeeze().cpu().numpy()
+        video_tensor = video_tensor.squeeze()
+        #NOTE: 20% of the depth is not reliable
+        # threshold = depth_conf.squeeze()[0].view(-1).quantile(0.6).item()
+        unc_metric = depth_conf.squeeze().cpu().numpy() > 0.5
+    # Load and process mask
+    if os.path.exists(mask_path):
+        mask = cv2.imread(mask_path)
+        mask = cv2.resize(mask, (video_tensor.shape[3], video_tensor.shape[2]))
+        mask = mask.sum(axis=-1)>0
+    else:
+        mask = np.ones_like(video_tensor[0,0].cpu().numpy())>0
+        grid_size = 10
+    # Get frame dimensions and create grid points
+    frame_H, frame_W = video_tensor.shape[2:]
+    grid_pts = get_points_on_a_grid(grid_size, (frame_H, frame_W), device="cuda")
+    # Sample mask values at grid points and filter
+    if os.path.exists(mask_path):
+        grid_pts_int = grid_pts[0].long()
+        mask_values = mask[grid_pts_int.cpu()[...,1], grid_pts_int.cpu()[...,0]]
+        grid_pts = grid_pts[:, mask_values]
+    query_xyt = torch.cat([torch.zeros_like(grid_pts[:, :, :1]), grid_pts], dim=2)[0].cpu().numpy()
+    print(f"Query points shape: {query_xyt.shape}")
+    # Run model inference
+    with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16):
+        (
+            c2w_traj, intrs, point_map, conf_depth,
+            track3d_pred, track2d_pred, vis_pred, conf_pred, video
+        ) = tracker_model_arg.forward(video_tensor, depth=depth_tensor,
+                            intrs=intrs, extrs=extrs,
+                            queries=query_xyt,
+                            fps=1, full_point=False, iters_track=4,
+                            query_no_BA=True, fixed_cam=False, stage=1, unc_metric=unc_metric,
+                            support_frame=len(video_tensor)-1, replace_ratio=0.2)
+        # Resize results to avoid large I/O
+        max_size = 224
+        h, w = video.shape[2:]
+        scale = min(max_size / h, max_size / w)
+        if scale < 1:
+            new_h, new_w = int(h * scale), int(w * scale)
+            video = T.Resize((new_h, new_w))(video)
+            video_tensor = T.Resize((new_h, new_w))(video_tensor)
+            point_map = T.Resize((new_h, new_w))(point_map)
+            track2d_pred[...,:2] = track2d_pred[...,:2] * scale
+            intrs[:,:2,:] = intrs[:,:2,:] * scale
+            conf_depth = T.Resize((new_h, new_w))(conf_depth)
+        # Visualize tracks
+        tracker_viser_arg.visualize(video=video[None],
+                        tracks=track2d_pred[None][...,:2],
+                        visibility=vis_pred[None],filename="test")
+        # Save in tapip3d format
+        data_npz_load["coords"] = (torch.einsum("tij,tnj->tni", c2w_traj[:,:3,:3], track3d_pred[:,:,:3].cpu()) + c2w_traj[:,:3,3][:,None,:]).numpy()
+        data_npz_load["extrinsics"] = torch.inverse(c2w_traj).cpu().numpy()
+        data_npz_load["intrinsics"] = intrs.cpu().numpy()
+        data_npz_load["depths"] = point_map[:,2,...].cpu().numpy()
+        data_npz_load["video"] = (video_tensor).cpu().numpy()/255
+        data_npz_load["visibs"] = vis_pred.cpu().numpy()
+        data_npz_load["confs"] = conf_pred.cpu().numpy()
+        data_npz_load["confs_depth"] = conf_depth.cpu().numpy()
+        np.savez(os.path.join(out_dir, f'result.npz'), **data_npz_load)
+    return None
+def compress_and_write(filename, header, blob):
+    header_bytes = json.dumps(header).encode("utf-8")
+    header_len = struct.pack("<I", len(header_bytes))
+    with open(filename, "wb") as f:
+        f.write(header_len)
+        f.write(header_bytes)
+        f.write(blob)
+def process_point_cloud_data(npz_file, width=256, height=192, fps=4):
+    fixed_size = (width, height)
+    data = np.load(npz_file)
+    extrinsics = data["extrinsics"]
+    intrinsics = data["intrinsics"]
+    trajs = data["coords"]
+    T, C, H, W = data["video"].shape
+    fx = intrinsics[0, 0, 0]
+    fy = intrinsics[0, 1, 1]
+    fov_y = 2 * np.arctan(H / (2 * fy)) * (180 / np.pi)
+    fov_x = 2 * np.arctan(W / (2 * fx)) * (180 / np.pi)
+    original_aspect_ratio = (W / fx) / (H / fy)
+    rgb_video = (rearrange(data["video"], "T C H W -> T H W C") * 255).astype(np.uint8)
+    rgb_video = np.stack([cv2.resize(frame, fixed_size, interpolation=cv2.INTER_AREA)
+                          for frame in rgb_video])
+    depth_video = data["depths"].astype(np.float32)
+    if "confs_depth" in data.keys():
+        confs = (data["confs_depth"].astype(np.float32) > 0.5).astype(np.float32)
+        depth_video = depth_video * confs
+    depth_video = np.stack([cv2.resize(frame, fixed_size, interpolation=cv2.INTER_NEAREST)
+                            for frame in depth_video])
+    scale_x = fixed_size[0] / W
+    scale_y = fixed_size[1] / H
+    intrinsics = intrinsics.copy()
+    intrinsics[:, 0, :] *= scale_x
+    intrinsics[:, 1, :] *= scale_y
+    min_depth = float(depth_video.min()) * 0.8
+    max_depth = float(depth_video.max()) * 1.5
+    depth_normalized = (depth_video - min_depth) / (max_depth - min_depth)
+    depth_int = (depth_normalized * ((1 << 16) - 1)).astype(np.uint16)
+    depths_rgb = np.zeros((T, fixed_size[1], fixed_size[0], 3), dtype=np.uint8)
+    depths_rgb[:, :, :, 0] = (depth_int & 0xFF).astype(np.uint8)
+    depths_rgb[:, :, :, 1] = ((depth_int >> 8) & 0xFF).astype(np.uint8)
+    first_frame_inv = np.linalg.inv(extrinsics[0])
+    normalized_extrinsics = np.array([first_frame_inv @ ext for ext in extrinsics])
+    normalized_trajs = np.zeros_like(trajs)
+    for t in range(T):
+        homogeneous_trajs = np.concatenate([trajs[t], np.ones((trajs.shape[1], 1))], axis=1)
+        transformed_trajs = (first_frame_inv @ homogeneous_trajs.T).T
+        normalized_trajs[t] = transformed_trajs[:, :3]
+    arrays = {
+        "rgb_video": rgb_video,
+        "depths_rgb": depths_rgb,
+        "intrinsics": intrinsics,
+        "extrinsics": normalized_extrinsics,
+        "inv_extrinsics": np.linalg.inv(normalized_extrinsics),
+        "trajectories": normalized_trajs.astype(np.float32),
+        "cameraZ": 0.0
+    }
+    header = {}
+    blob_parts = []
+    offset = 0
+    for key, arr in arrays.items():
+        arr = np.ascontiguousarray(arr)
+        arr_bytes = arr.tobytes()
+        header[key] = {
+            "dtype": str(arr.dtype),
+            "shape": arr.shape,
+            "offset": offset,
+            "length": len(arr_bytes)
+        }
+        blob_parts.append(arr_bytes)
+        offset += len(arr_bytes)
+    raw_blob = b"".join(blob_parts)
+    compressed_blob = zlib.compress(raw_blob, level=9)
+    header["meta"] = {
+        "depthRange": [min_depth, max_depth],
+        "totalFrames": int(T),
+        "resolution": fixed_size,
+        "baseFrameRate": fps,
+        "numTrajectoryPoints": normalized_trajs.shape[1],
+        "fov": float(fov_y),
+        "fov_x": float(fov_x),
+        "original_aspect_ratio": float(original_aspect_ratio),
+        "fixed_aspect_ratio": float(fixed_size[0]/fixed_size[1])
+    }
+    compress_and_write('./_viz/data.bin', header, compressed_blob)
+    with open('./_viz/data.bin', "rb") as f:
+        encoded_blob = base64.b64encode(f.read()).decode("ascii")
+    os.unlink('./_viz/data.bin')
+    random_path = f'./_viz/_{time.time()}.html'
+    with open('./_viz/viz_template.html') as f:
+        html_template = f.read()
+    html_out = html_template.replace(
+        "<head>",
+        f"<head>\n<script>window.embeddedBase64 = `{encoded_blob}`;</script>"
+    )
+    with open(random_path,'w') as f:
+        f.write(html_out)
+    return random_path
+def numpy_to_base64(arr):
+    """Convert numpy array to base64 string"""
+    return base64.b64encode(arr.tobytes()).decode('utf-8')
+def base64_to_numpy(b64_str, shape, dtype):
+    """Convert base64 string back to numpy array"""
+    return np.frombuffer(base64.b64decode(b64_str), dtype=dtype).reshape(shape)
+def get_video_name(video_path):
+    """Extract video name without extension"""
+    return os.path.splitext(os.path.basename(video_path))[0]
+def extract_first_frame(video_path):
+    """Extract first frame from video file"""
+    try:
+        cap = cv2.VideoCapture(video_path)
+        ret, frame = cap.read()
+        cap.release()
+        if ret:
+            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            return frame_rgb
+        else:
+            return None
+    except Exception as e:
+        print(f"Error extracting first frame: {e}")
+        return None
+def handle_video_upload(video):
+    """Handle video upload and extract first frame"""
+    if video is None:
+        return (None, None, [],
+                gr.update(value=50),
+                gr.update(value=756),
+                gr.update(value=3))
+    # Create user-specific temporary directory
+    user_temp_dir = create_user_temp_dir()
+    # Get original video name and copy to temp directory
+    if isinstance(video, str):
+        video_name = get_video_name(video)
+        video_path = os.path.join(user_temp_dir, f"{video_name}.mp4")
+        shutil.copy(video, video_path)
+    else:
+        video_name = get_video_name(video.name)
+        video_path = os.path.join(user_temp_dir, f"{video_name}.mp4")
+        with open(video_path, 'wb') as f:
+            f.write(video.read())
+    print(f"📁 Video saved to: {video_path}")
+    # Extract first frame
+    frame = extract_first_frame(video_path)
+    if frame is None:
+        return (None, None, [],
+                gr.update(value=50),
+                gr.update(value=756),
+                gr.update(value=3))
+    # Resize frame to have minimum side length of 336
+    h, w = frame.shape[:2]
+    scale = 336 / min(h, w)
+    new_h, new_w = int(h * scale)//2*2, int(w * scale)//2*2
+    frame = cv2.resize(frame, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
+    # Store frame data with temp directory info
+    frame_data = {
+        'data': numpy_to_base64(frame),
+        'shape': frame.shape,
+        'dtype': str(frame.dtype),
+        'temp_dir': user_temp_dir,
+        'video_name': video_name,
+        'video_path': video_path
+    }
+    # Get video-specific settings
+    print(f"🎬 Video path: '{video}' -> Video name: '{video_name}'")
+    grid_size_val, vo_points_val, fps_val = get_video_settings(video_name)
+    print(f"🎬 Video settings for '{video_name}': grid_size={grid_size_val}, vo_points={vo_points_val}, fps={fps_val}")
+    return (json.dumps(frame_data), frame, [],
+            gr.update(value=grid_size_val),
+            gr.update(value=vo_points_val),
+            gr.update(value=fps_val))
+def save_masks(o_masks, video_name, temp_dir):
+    """Save binary masks to files in user-specific temp directory"""
+    o_files = []
+    for mask, _ in o_masks:
+        o_mask = np.uint8(mask.squeeze() * 255)
+        o_file = os.path.join(temp_dir, f"{video_name}.png")
+        cv2.imwrite(o_file, o_mask)
+        o_files.append(o_file)
+    return o_files
+def select_point(original_img: str, sel_pix: list, point_type: str, evt: gr.SelectData):
+    """Handle point selection for SAM"""
+    if original_img is None:
+        return None, []
+    try:
+        # Convert stored image data back to numpy array
+        frame_data = json.loads(original_img)
+        original_img_array = base64_to_numpy(frame_data['data'], frame_data['shape'], frame_data['dtype'])
+        temp_dir = frame_data.get('temp_dir', 'temp_local')
+        video_name = frame_data.get('video_name', 'video')
+        # Create a display image for visualization
+        display_img = original_img_array.copy()
+        new_sel_pix = sel_pix.copy() if sel_pix else []
+        new_sel_pix.append((evt.index, 1 if point_type == 'positive_point' else 0))
+        print(f"🎯 Running SAM inference for point: {evt.index}, type: {point_type}")
+        # Run SAM inference
+        o_masks = gpu_run_inference(None, original_img_array, new_sel_pix, [])
+        # Draw points on display image
+        for point, label in new_sel_pix:
+            cv2.drawMarker(display_img, point, COLORS[label], markerType=MARKERS[label], markerSize=MARKER_SIZE, thickness=2)
+        # Draw mask overlay on display image
+        if o_masks:
+            mask = o_masks[0][0]
+            overlay = display_img.copy()
+            overlay[mask.squeeze()!=0] = [20, 60, 200]  # Light blue
+            display_img = cv2.addWeighted(overlay, 0.6, display_img, 0.4, 0)
+            # Save mask for tracking
+            save_masks(o_masks, video_name, temp_dir)
+            print(f"✅ Mask saved for video: {video_name}")
+        return display_img, new_sel_pix
+    except Exception as e:
+        print(f"❌ Error in select_point: {e}")
+        return None, []
+def reset_points(original_img: str, sel_pix):
+    """Reset all points and clear the mask"""
+    if original_img is None:
+        return None, []
+    try:
+        # Convert stored image data back to numpy array
+        frame_data = json.loads(original_img)
+        original_img_array = base64_to_numpy(frame_data['data'], frame_data['shape'], frame_data['dtype'])
+        temp_dir = frame_data.get('temp_dir', 'temp_local')
+        # Create a display image (just the original image)
+        display_img = original_img_array.copy()
+        # Clear all points
+        new_sel_pix = []
+        # Clear any existing masks
+        for mask_file in glob.glob(os.path.join(temp_dir, "*.png")):
+            try:
+                os.remove(mask_file)
+            except Exception as e:
+                logger.warning(f"Failed to remove mask file {mask_file}: {e}")
+        print("🔄 Points and masks reset")
+        return display_img, new_sel_pix
+    except Exception as e:
+        print(f"❌ Error in reset_points: {e}")
+        return None, []
+def launch_viz(grid_size, vo_points, fps, original_image_state, mode="offline"):
+    """Launch visualization with user-specific temp directory"""
+    if original_image_state is None:
+        return None, None, None
+    try:
+        # Get user's temp directory from stored frame data
+        frame_data = json.loads(original_image_state)
+        temp_dir = frame_data.get('temp_dir', 'temp_local')
+        video_name = frame_data.get('video_name', 'video')
+        print(f"🚀 Starting tracking for video: {video_name}")
+        print(f"📊 Parameters: grid_size={grid_size}, vo_points={vo_points}, fps={fps}")
+        # Check for mask files
+        mask_files = glob.glob(os.path.join(temp_dir, "*.png"))
+        video_files = glob.glob(os.path.join(temp_dir, "*.mp4"))
+        if not video_files:
+            print("❌ No video file found")
+            return "❌ Error: No video file found", None, None
+        video_path = video_files[0]
+        mask_path = mask_files[0] if mask_files else None
+        # Run tracker
+        print("🎯 Running tracker...")
+        out_dir = os.path.join(temp_dir, "results")
+        os.makedirs(out_dir, exist_ok=True)
+        gpu_run_tracker(None, None, temp_dir, video_name, grid_size, vo_points, fps, mode=mode)
+        # Process results
+        npz_path = os.path.join(out_dir, "result.npz")
+        track2d_video = os.path.join(out_dir, "test_pred_track.mp4")
+        if os.path.exists(npz_path):
+            print("📊 Processing 3D visualization...")
+            html_path = process_point_cloud_data(npz_path)
+            # Schedule deletion of generated files
+            delete_later(html_path, delay=600)
+            if os.path.exists(track2d_video):
+                delete_later(track2d_video, delay=600)
+            delete_later(npz_path, delay=600)
+            # Create iframe HTML
+            iframe_html = f"""
+            <div style='border: 3px solid #667eea; border-radius: 10px;
+                        background: #f8f9ff; height: 650px; width: 100%;
+                        box-shadow: 0 8px 32px rgba(102, 126, 234, 0.3);
+                        margin: 0; padding: 0; box-sizing: border-box; overflow: hidden;'>
+                <iframe id="viz_iframe" src="/gradio_api/file={html_path}"
+                        width="100%" height="650" frameborder="0"
+                        style="border: none; display: block; width: 100%; height: 650px;
+                               margin: 0; padding: 0; border-radius: 7px;">
+                </iframe>
+            </div>
+            """
+            print("✅ Tracking completed successfully!")
+            return iframe_html, track2d_video if os.path.exists(track2d_video) else None, html_path
+        else:
+            print("❌ Tracking failed - no results generated")
+            return "❌ Error: Tracking failed to generate results", None, None
+    except Exception as e:
+        print(f"❌ Error in launch_viz: {e}")
+        return f"❌ Error: {str(e)}", None, None
+def clear_all():
+    """Clear all buffers and temporary files"""
+    return (None, None, [],
+            gr.update(value=50),
+            gr.update(value=756),
+            gr.update(value=3))
+def clear_all_with_download():
+    """Clear all buffers including both download components"""
+    return (None, None, [],
+            gr.update(value=50),
+            gr.update(value=756),
+            gr.update(value=3),
+            None,  # tracking_video_download
+            None)  # HTML download component
+def get_video_settings(video_name):
+    """Get video-specific settings based on video name"""
+    video_settings = {
+        "running": (50, 512, 2),
+        "backpack": (40, 600, 2),
+        "kitchen": (60, 800, 3),
+        "pillow": (35, 500, 2),
+        "handwave": (35, 500, 8),
+        "hockey": (45, 700, 2),
+        "drifting": (35, 1000, 6),
+        "basketball": (45, 1500, 5),
+        "ken_block_0": (45, 700, 2),
+        "ego_kc1": (45, 500, 4),
+        "vertical_place": (45, 500, 3),
+        "ego_teaser": (45, 1200, 10),
+        "robot_unitree": (45, 500, 4),
+        "robot_3": (35, 400, 5),
+        "teleop2": (45, 256, 7),
+        "pusht": (45, 256, 10),
+        "cinema_0": (45, 356, 5),
+        "cinema_1": (45, 756, 3),
+        "robot1": (45, 600, 2),
+        "robot2": (45, 600, 2),
+        "protein": (45, 600, 2),
+        "kitchen_egocentric": (45, 600, 2),
+    }
+    return video_settings.get(video_name, (50, 756, 3))
+# Create the Gradio interface
+print("🎨 Creating Gradio interface...")
+with gr.Blocks(
+    theme=gr.themes.Soft(),
+    title="🎯 [SpatialTracker V2](https://github.com/henry123-boy/SpaTrackerV2)",
+    css="""
+    .gradio-container {
+        max-width: 1200px !important;
+        margin: auto !important;
+    }
+    .gr-button {
+        margin: 5px;
+    }
+    .gr-form {
+        background: white;
+        border-radius: 10px;
+        padding: 20px;
+        box-shadow: 0 2px 10px rgba(0,0,0,0.1);
+    }
+    /* 移除 gr.Group 的默认灰色背景 */
+    .gr-form {
+        background: transparent !important;
+        border: none !important;
+        box-shadow: none !important;
+        padding: 0 !important;
+    }
+    /* 固定3D可视化器尺寸 */
+    #viz_container {
+        height: 650px !important;
+        min-height: 650px !important;
+        max-height: 650px !important;
+        width: 100% !important;
+        margin: 0 !important;
+        padding: 0 !important;
+        overflow: hidden !important;
+    }
+    #viz_container > div {
+        height: 650px !important;
+        min-height: 650px !important;
+        max-height: 650px !important;
+        width: 100% !important;
+        margin: 0 !important;
+        padding: 0 !important;
+        box-sizing: border-box !important;
+    }
+    #viz_container iframe {
+        height: 650px !important;
+        min-height: 650px !important;
+        max-height: 650px !important;
+        width: 100% !important;
+        border: none !important;
+        display: block !important;
+        margin: 0 !important;
+        padding: 0 !important;
+        box-sizing: border-box !important;
+    }
+    /* 固定视频上传组件高度 */
+    .gr-video {
+        height: 300px !important;
+        min-height: 300px !important;
+        max-height: 300px !important;
+    }
+    .gr-video video {
+        height: 260px !important;
+        max-height: 260px !important;
+        object-fit: contain !important;
+        background: #f8f9fa;
+    }
+    .gr-video .gr-video-player {
+        height: 260px !important;
+        max-height: 260px !important;
+    }
+    /* 强力移除examples的灰色背景 - 使用更通用的选择器 */
+    .horizontal-examples,
+    .horizontal-examples > *,
+    .horizontal-examples * {
+        background: transparent !important;
+        background-color: transparent !important;
+        border: none !important;
+    }
+    /* Examples组件水平滚动样式 */
+    .horizontal-examples [data-testid="examples"] {
+        background: transparent !important;
+        background-color: transparent !important;
+    }
+    .horizontal-examples [data-testid="examples"] > div {
+        background: transparent !important;
+        background-color: transparent !important;
+        overflow-x: auto !important;
+        overflow-y: hidden !important;
+        scrollbar-width: thin;
+        scrollbar-color: #667eea transparent;
+        padding: 0 !important;
+        margin-top: 10px;
+        border: none !important;
+    }
+    .horizontal-examples [data-testid="examples"] table {
+        display: flex !important;
+        flex-wrap: nowrap !important;
+        min-width: max-content !important;
+        gap: 15px !important;
+        padding: 10px 0;
+        background: transparent !important;
+        border: none !important;
+    }
+    .horizontal-examples [data-testid="examples"] tbody {
+        display: flex !important;
+        flex-direction: row !important;
+        flex-wrap: nowrap !important;
+        gap: 15px !important;
+        background: transparent !important;
+    }
+    .horizontal-examples [data-testid="examples"] tr {
+        display: flex !important;
+        flex-direction: column !important;
+        min-width: 160px !important;
+        max-width: 160px !important;
+        margin: 0 !important;
+        background: white !important;
+        border-radius: 12px;
+        box-shadow: 0 3px 12px rgba(0,0,0,0.12);
+        transition: all 0.3s ease;
+        cursor: pointer;
+        overflow: hidden;
+        border: none !important;
+    }
+    .horizontal-examples [data-testid="examples"] tr:hover {
+        transform: translateY(-4px);
+        box-shadow: 0 8px 20px rgba(102, 126, 234, 0.25);
+    }
+    .horizontal-examples [data-testid="examples"] td {
+        text-align: center !important;
+        padding: 0 !important;
+        border: none !important;
+        background: transparent !important;
+    }
+    .horizontal-examples [data-testid="examples"] td:first-child {
+        padding: 0 !important;
+        background: transparent !important;
+    }
+    .horizontal-examples [data-testid="examples"] video {
+        border-radius: 8px 8px 0 0 !important;
+        width: 100% !important;
+        height: 90px !important;
+        object-fit: cover !important;
+        background: #f8f9fa !important;
+    }
+    .horizontal-examples [data-testid="examples"] td:last-child {
+        font-size: 11px !important;
+        font-weight: 600 !important;
+        color: #333 !important;
+        padding: 8px 12px !important;
+        background: linear-gradient(135deg, #f8f9ff 0%, #e6f3ff 100%) !important;
+        border-radius: 0 0 8px 8px;
+    }
+    /* 滚动条样式 */
+    .horizontal-examples [data-testid="examples"] > div::-webkit-scrollbar {
+        height: 8px;
+    }
+    .horizontal-examples [data-testid="examples"] > div::-webkit-scrollbar-track {
+        background: transparent;
+        border-radius: 4px;
+    }
+    .horizontal-examples [data-testid="examples"] > div::-webkit-scrollbar-thumb {
+        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+        border-radius: 4px;
+    }
+    .horizontal-examples [data-testid="examples"] > div::-webkit-scrollbar-thumb:hover {
+        background: linear-gradient(135deg, #5a6fd8 0%, #6a4190 100%);
+    }
+    """
+) as demo:
+    # Add prominent main title
+    gr.Markdown("""
+    # ✨ SpatialTrackerV2
+    Welcome to [SpatialTracker V2](https://github.com/henry123-boy/SpaTrackerV2)! This interface allows you to track any pixels in 3D using our model.
+    **⚡ Quick Start:** Upload video → Click "Start Tracking Now!"
+    **🔬 Advanced Usage with SAM:**
+    1. Upload a video file or select from examples below
+    2. Expand "Manual Point Selection" to click on specific objects for SAM-guided tracking
+    3. Adjust tracking parameters for optimal performance
+    4. Click "Start Tracking Now!" to begin 3D tracking with SAM guidance
+    """)
+    # Status indicator
+    gr.Markdown("**Status:** 🟢 Local Processing Mode")
+    # Main content area - video upload left, 3D visualization right
+    with gr.Row():
+        with gr.Column(scale=1):
+            # Video upload section
+            gr.Markdown("### 📂 Select Video")
+            # Define video_input here so it can be referenced in examples
+            video_input = gr.Video(
+                label="Upload Video or Select Example",
+                format="mp4",
+                height=250  # Matched height with 3D viz
+            )
+            # Traditional examples but with horizontal scroll styling
+            gr.Markdown("🎨**Examples:** (scroll horizontally to see all videos)")
+            with gr.Row(elem_classes=["horizontal-examples"]):
+                # Horizontal video examples with slider
+                # gr.HTML("<div style='margin-top: 5px;'></div>")
+                gr.Examples(
+                    examples=[
+                        ["./examples/robot1.mp4"],
+                        ["./examples/robot2.mp4"],
+                        ["./examples/protein.mp4"],
+                        ["./examples/kitchen_egocentric.mp4"],
+                        ["./examples/hockey.mp4"],
+                        ["./examples/running.mp4"],
+                        ["./examples/robot_3.mp4"],
+                        ["./examples/backpack.mp4"],
+                        ["./examples/kitchen.mp4"],
+                        ["./examples/pillow.mp4"],
+                        ["./examples/handwave.mp4"],
+                        ["./examples/drifting.mp4"],
+                        ["./examples/basketball.mp4"],
+                        ["./examples/ken_block_0.mp4"],
+                        ["./examples/ego_kc1.mp4"],
+                        ["./examples/vertical_place.mp4"],
+                        ["./examples/ego_teaser.mp4"],
+                        ["./examples/robot_unitree.mp4"],
+                        ["./examples/teleop2.mp4"],
+                        ["./examples/pusht.mp4"],
+                        ["./examples/cinema_0.mp4"],
+                        ["./examples/cinema_1.mp4"],
+                    ],
+                    inputs=[video_input],
+                    outputs=[video_input],
+                    fn=None,
+                    cache_examples=False,
+                    label="",
+                    examples_per_page=6  # Show 6 examples per page so they can wrap to multiple rows
+                )
+        with gr.Column(scale=2):
+            # 3D Visualization - wider and taller to match left side
+            with gr.Group():
+                gr.Markdown("### 🌐 3D Trajectory Visualization")
+                viz_html = gr.HTML(
+                    label="3D Trajectory Visualization",
+                    value="""
+                    <div style='border: 3px solid #667eea; border-radius: 10px;
+                                background: linear-gradient(135deg, #f8f9ff 0%, #e6f3ff 100%);
+                                text-align: center; height: 650px; display: flex;
+                                flex-direction: column; justify-content: center; align-items: center;
+                                box-shadow: 0 4px 16px rgba(102, 126, 234, 0.15);
+                                margin: 0; padding: 20px; box-sizing: border-box;'>
+                        <div style='font-size: 56px; margin-bottom: 25px;'>🌐</div>
+                        <h3 style='color: #667eea; margin-bottom: 18px; font-size: 28px; font-weight: 600;'>
+                            3D Trajectory Visualization
+                        </h3>
+                        <p style='color: #666; font-size: 18px; line-height: 1.6; max-width: 550px; margin-bottom: 30px;'>
+                            Track any pixels in 3D space with camera motion
+                        </p>
+                        <div style='background: rgba(102, 126, 234, 0.1); border-radius: 30px;
+                                    padding: 15px 30px; border: 1px solid rgba(102, 126, 234, 0.2);'>
+                            <span style='color: #667eea; font-weight: 600; font-size: 16px;'>
+                                ⚡ Powered by SpatialTracker V2
+                            </span>
+                        </div>
+                    </div>
+                    """,
+                    elem_id="viz_container"
+                )
+    # Start button section - below video area
+    with gr.Row():
+        with gr.Column(scale=3):
+            launch_btn = gr.Button("🚀 Start Tracking Now!", variant="primary", size="lg")
+        with gr.Column(scale=1):
+            clear_all_btn = gr.Button("🗑️ Clear All", variant="secondary", size="sm")
+    # Tracking parameters section
+    with gr.Row():
+        gr.Markdown("### ⚙️ Tracking Parameters")
+    with gr.Row():
+        grid_size = gr.Slider(
+            minimum=10, maximum=100, step=10, value=50,
+            label="Grid Size", info="Tracking detail level"
+        )
+        vo_points = gr.Slider(
+            minimum=100, maximum=2000, step=50, value=756,
+            label="VO Points", info="Motion accuracy"
+        )
+        fps = gr.Slider(
+            minimum=1, maximum=20, step=1, value=3,
+            label="FPS", info="Processing speed"
+        )
+    # Advanced Point Selection with SAM - Collapsed by default
+    with gr.Row():
+        gr.Markdown("### 🎯 Advanced: Manual Point Selection with SAM")
+    with gr.Accordion("🔬 SAM Point Selection Controls", open=False):
+        gr.HTML("""
+        <div style='margin-bottom: 15px;'>
+            <ul style='color: #4a5568; font-size: 14px; line-height: 1.6; margin: 0; padding-left: 20px;'>
+                <li>Click on target objects in the image for SAM-guided segmentation</li>
+                <li>Positive points: include these areas | Negative points: exclude these areas</li>
+                <li>Get more accurate 3D tracking results with SAM's powerful segmentation</li>
+            </ul>
+        </div>
+        """)
+        with gr.Row():
+            with gr.Column():
+                interactive_frame = gr.Image(
+                    label="Click to select tracking points with SAM guidance",
+                    type="numpy",
+                    interactive=True,
+                    height=300
+                )
+                with gr.Row():
+                    point_type = gr.Radio(
+                        choices=["positive_point", "negative_point"],
+                        value="positive_point",
+                        label="Point Type",
+                        info="Positive: track these areas | Negative: avoid these areas"
+                    )
+                with gr.Row():
+                    reset_points_btn = gr.Button("🔄 Reset Points", variant="secondary", size="sm")
+    # Downloads section - hidden but still functional for local processing
+    with gr.Row(visible=False):
+        with gr.Column(scale=1):
+            tracking_video_download = gr.File(
+                label="📹 Download 2D Tracking Video",
+                interactive=False,
+                visible=False
+            )
+        with gr.Column(scale=1):
+            html_download = gr.File(
+                label="📄 Download 3D Visualization HTML",
+                interactive=False,
+                visible=False
+            )
+    # GitHub Star Section
+    gr.HTML("""
+    <div style='background: linear-gradient(135deg, #e8eaff 0%, #f0f2ff 100%);
+                border-radius: 8px; padding: 20px; margin: 15px 0;
+                box-shadow: 0 2px 8px rgba(102, 126, 234, 0.1);
+                border: 1px solid rgba(102, 126, 234, 0.15);'>
+        <div style='text-align: center;'>
+            <h3 style='color: #4a5568; margin: 0 0 10px 0; font-size: 18px; font-weight: 600;'>
+                ⭐ Love SpatialTracker? Give us a Star! ⭐
+            </h3>
+            <p style='color: #666; margin: 0 0 15px 0; font-size: 14px; line-height: 1.5;'>
+                Help us grow by starring our repository on GitHub! Your support means a lot to the community. 🚀
+            </p>
+            <a href="https://github.com/henry123-boy/SpaTrackerV2" target="_blank"
+               style='display: inline-flex; align-items: center; gap: 8px;
+                      background: rgba(102, 126, 234, 0.1); color: #4a5568;
+                      padding: 10px 20px; border-radius: 25px; text-decoration: none;
+                      font-weight: bold; font-size: 14px; border: 1px solid rgba(102, 126, 234, 0.2);
+                      transition: all 0.3s ease;'
+               onmouseover="this.style.background='rgba(102, 126, 234, 0.15)'; this.style.transform='translateY(-2px)'"
+               onmouseout="this.style.background='rgba(102, 126, 234, 0.1)'; this.style.transform='translateY(0)'">
+                <span style='font-size: 16px;'>⭐</span>
+                Star SpatialTracker V2 on GitHub
+            </a>
+        </div>
+    </div>
+    """)
+    # Acknowledgments Section
+    gr.HTML("""
+    <div style='background: linear-gradient(135deg, #fff8e1 0%, #fffbf0 100%);
+                border-radius: 8px; padding: 20px; margin: 15px 0;
+                box-shadow: 0 2px 8px rgba(255, 193, 7, 0.1);
+                border: 1px solid rgba(255, 193, 7, 0.2);'>
+        <div style='text-align: center;'>
+            <h3 style='color: #5d4037; margin: 0 0 10px 0; font-size: 18px; font-weight: 600;'>
+                📚 Acknowledgments
+            </h3>
+            <p style='color: #5d4037; margin: 0 0 15px 0; font-size: 14px; line-height: 1.5;'>
+                Our 3D visualizer is adapted from <strong>TAPIP3D</strong>. We thank the authors for their excellent work and contribution to the computer vision community!
+            </p>
+            <a href="https://github.com/zbw001/TAPIP3D" target="_blank"
+               style='display: inline-flex; align-items: center; gap: 8px;
+                      background: rgba(255, 193, 7, 0.15); color: #5d4037;
+                      padding: 10px 20px; border-radius: 25px; text-decoration: none;
+                      font-weight: bold; font-size: 14px; border: 1px solid rgba(255, 193, 7, 0.3);
+                      transition: all 0.3s ease;'
+               onmouseover="this.style.background='rgba(255, 193, 7, 0.25)'; this.style.transform='translateY(-2px)'"
+               onmouseout="this.style.background='rgba(255, 193, 7, 0.15)'; this.style.transform='translateY(0)'">
+                📚 Visit TAPIP3D Repository
+            </a>
+        </div>
+    </div>
+    """)
+    # Footer
+    gr.HTML("""
+    <div style='text-align: center; margin: 20px 0 10px 0;'>
+        <span style='font-size: 12px; color: #888; font-style: italic;'>
+            Powered by SpatialTracker V2 | Built with ❤️ for the Computer Vision Community
+        </span>
+    </div>
+    """)
+    # Hidden state variables
+    original_image_state = gr.State(None)
+    selected_points = gr.State([])
+    # Event handlers
+    video_input.change(
+        fn=handle_video_upload,
+        inputs=[video_input],
+        outputs=[original_image_state, interactive_frame, selected_points, grid_size, vo_points, fps]
+    )
+    interactive_frame.select(
+        fn=select_point,
+        inputs=[original_image_state, selected_points, point_type],
+        outputs=[interactive_frame, selected_points]
+    )
+    reset_points_btn.click(
+        fn=reset_points,
+        inputs=[original_image_state, selected_points],
+        outputs=[interactive_frame, selected_points]
+    )
+    clear_all_btn.click(
+        fn=clear_all_with_download,
+        outputs=[video_input, interactive_frame, selected_points, grid_size, vo_points, fps, tracking_video_download, html_download]
+    )
+    launch_btn.click(
+        fn=launch_viz,
+        inputs=[grid_size, vo_points, fps, original_image_state],
+        outputs=[viz_html, tracking_video_download, html_download]
+    )
+# Launch the interface
+if __name__ == "__main__":
+    print("🌟 Launching SpatialTracker V2 Local Version...")
+    print("🔗 Running in Local Processing Mode")
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=True,
+        debug=True,
+        show_error=True
+    )

app_3rd/README.md ADDED Viewed

	@@ -0,0 +1,12 @@

+# 🌟 SpatialTrackerV2 Integrated with SAM 🌟
+SAM receives a point prompt and generates a mask for the target object, facilitating easy interaction to obtain the object's 3D trajectories with SpaTrack2.
+## Installation
+```
+python -m pip install git+https://github.com/facebookresearch/segment-anything.git
+cd app_3rd/sam_utils
+mkdir checkpoints
+cd checkpoints
+wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth
+```

app_3rd/sam_utils/hf_sam_predictor.py ADDED Viewed

	@@ -0,0 +1,129 @@

+import gc
+import numpy as np
+import torch
+from typing import Optional, Tuple, List, Union
+import warnings
+import cv2
+try:
+    from transformers import SamModel, SamProcessor
+    from huggingface_hub import hf_hub_download
+    HF_AVAILABLE = True
+except ImportError:
+    HF_AVAILABLE = False
+    warnings.warn("transformers or huggingface_hub not available. HF SAM models will not work.")
+# Hugging Face model mapping
+HF_MODELS = {
+    'vit_b': 'facebook/sam-vit-base',
+    'vit_l': 'facebook/sam-vit-large',
+    'vit_h': 'facebook/sam-vit-huge'
+}
+class HFSamPredictor:
+    """
+    Hugging Face version of SamPredictor that wraps the transformers SAM models.
+    This class provides the same interface as the original SamPredictor for seamless integration.
+    """
+    def __init__(self, model: SamModel, processor: SamProcessor, device: Optional[str] = None):
+        """
+        Initialize the HF SAM predictor.
+        Args:
+            model: The SAM model from transformers
+            processor: The SAM processor from transformers
+            device: Device to run the model on ('cuda', 'cpu', etc.)
+        """
+        self.model = model
+        self.processor = processor
+        self.device = device or ('cuda' if torch.cuda.is_available() else 'cpu')
+        self.model.to(self.device)
+        self.model.eval()
+        # Store the current image and its features
+        self.original_size = None
+        self.input_size = None
+        self.features = None
+        self.image = None
+    @classmethod
+    def from_pretrained(cls, model_name: str, device: Optional[str] = None) -> 'HFSamPredictor':
+        """
+        Load a SAM model from Hugging Face Hub.
+        Args:
+            model_name: Model name from HF_MODELS or direct HF model path
+            device: Device to load the model on
+        Returns:
+            HFSamPredictor instance
+        """
+        if not HF_AVAILABLE:
+            raise ImportError("transformers and huggingface_hub are required for HF SAM models")
+        # Map model type to HF model name if needed
+        if model_name in HF_MODELS:
+            model_name = HF_MODELS[model_name]
+        print(f"Loading SAM model from Hugging Face: {model_name}")
+        # Load model and processor
+        model = SamModel.from_pretrained(model_name)
+        processor = SamProcessor.from_pretrained(model_name)
+        return cls(model, processor, device)
+    def preprocess(self, image: np.ndarray,
+                         input_points: List[List[float]], input_labels: List[int]) -> None:
+        """
+        Set the image for prediction. This preprocesses the image and extracts features.
+        Args:
+            image: Input image as numpy array (H, W, C) in RGB format
+        """
+        if image.dtype != np.uint8:
+            image = (image * 255).astype(np.uint8)
+        self.image = image
+        self.original_size = image.shape[:2]
+        # Use dummy point to ensure processor returns original_sizes & reshaped_input_sizes
+        inputs = self.processor(
+            images=image,
+            input_points=input_points,
+            input_labels=input_labels,
+            return_tensors="pt"
+        )
+        inputs = {k: v.to(self.device) for k, v in inputs.items()}
+        self.input_size = inputs['pixel_values'].shape[-2:]
+        self.features = inputs
+        return inputs
+def get_hf_sam_predictor(model_type: str = 'vit_h', device: Optional[str] = None,
+                        image: Optional[np.ndarray] = None) -> HFSamPredictor:
+    """
+    Get a Hugging Face SAM predictor with the same interface as the original get_sam_predictor.
+    Args:
+        model_type: Model type ('vit_b', 'vit_l', 'vit_h')
+        device: Device to run the model on
+        image: Optional image to set immediately
+    Returns:
+        HFSamPredictor instance
+    """
+    if not HF_AVAILABLE:
+        raise ImportError("transformers and huggingface_hub are required for HF SAM models")
+    if device is None:
+        device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    # Load the predictor
+    predictor = HFSamPredictor.from_pretrained(model_type, device)
+    # Set image if provided
+    if image is not None:
+        predictor.set_image(image)
+    return predictor

app_3rd/sam_utils/inference.py ADDED Viewed

	@@ -0,0 +1,123 @@

+import gc
+import numpy as np
+import torch
+from segment_anything import SamPredictor, sam_model_registry
+# Try to import HF SAM support
+try:
+    from app_3rd.sam_utils.hf_sam_predictor import get_hf_sam_predictor, HFSamPredictor
+    HF_AVAILABLE = True
+except ImportError:
+    HF_AVAILABLE = False
+models = {
+  'vit_b': 'app_3rd/sam_utils/checkpoints/sam_vit_b_01ec64.pth',
+  'vit_l': 'app_3rd/sam_utils/checkpoints/sam_vit_l_0b3195.pth',
+  'vit_h': 'app_3rd/sam_utils/checkpoints/sam_vit_h_4b8939.pth'
+}
+def get_sam_predictor(model_type='vit_b', device=None, image=None, use_hf=True, predictor=None):
+  """
+  Get SAM predictor with option to use HuggingFace version
+  Args:
+      model_type: Model type ('vit_b', 'vit_l', 'vit_h')
+      device: Device to run on
+      image: Optional image to set immediately
+      use_hf: Whether to use HuggingFace SAM instead of original SAM
+  """
+  if predictor is not None:
+    return predictor
+  if use_hf:
+    if not HF_AVAILABLE:
+      raise ImportError("HuggingFace SAM not available. Install transformers and huggingface_hub.")
+    return get_hf_sam_predictor(model_type, device, image)
+  # Original SAM logic
+  if device is None and torch.cuda.is_available():
+    device = 'cuda'
+  elif device is None:
+    device = 'cpu'
+  # sam model
+  sam = sam_model_registry[model_type](checkpoint=models[model_type])
+  sam = sam.to(device)
+  predictor = SamPredictor(sam)
+  if image is not None:
+    predictor.set_image(image)
+  return predictor
+def run_inference(predictor, input_x, selected_points, multi_object: bool = False):
+  """
+  Run inference with either original SAM or HF SAM predictor
+  Args:
+      predictor: SamPredictor or HFSamPredictor instance
+      input_x: Input image
+      selected_points: List of (point, label) tuples
+      multi_object: Whether to handle multiple objects
+  """
+  if len(selected_points) == 0:
+    return []
+  # Check if using HF SAM
+  if isinstance(predictor, HFSamPredictor):
+    return _run_hf_inference(predictor, input_x, selected_points, multi_object)
+  else:
+    return _run_original_inference(predictor, input_x, selected_points, multi_object)
+def _run_original_inference(predictor: SamPredictor, input_x, selected_points, multi_object: bool = False):
+  """Run inference with original SAM"""
+  points = torch.Tensor(
+      [p for p, _ in selected_points]
+  ).to(predictor.device).unsqueeze(1)
+  labels = torch.Tensor(
+      [int(l) for _, l in selected_points]
+  ).to(predictor.device).unsqueeze(1)
+  transformed_points = predictor.transform.apply_coords_torch(
+      points, input_x.shape[:2])
+  masks, scores, logits = predictor.predict_torch(
+    point_coords=transformed_points[:,0][None],
+    point_labels=labels[:,0][None],
+    multimask_output=False,
+  )
+  masks = masks[0].cpu().numpy()  # N 1 H W   N is the number of points
+  gc.collect()
+  torch.cuda.empty_cache()
+  return [(masks, 'final_mask')]
+def _run_hf_inference(predictor: HFSamPredictor, input_x, selected_points, multi_object: bool = False):
+  """Run inference with HF SAM"""
+  # Prepare points and labels for HF SAM
+  select_pts = [[list(p) for p, _ in selected_points]]
+  select_lbls = [[int(l) for _, l in selected_points]]
+  # Preprocess inputs
+  inputs = predictor.preprocess(input_x, select_pts, select_lbls)
+  # Run inference
+  with torch.no_grad():
+    outputs = predictor.model(**inputs)
+  # Post-process masks
+  masks = predictor.processor.image_processor.post_process_masks(
+    outputs.pred_masks.cpu(),
+    inputs["original_sizes"].cpu(),
+    inputs["reshaped_input_sizes"].cpu(),
+  )
+  masks = masks[0][:,:1,...].cpu().numpy()
+  gc.collect()
+  torch.cuda.empty_cache()
+  return [(masks, 'final_mask')]

app_3rd/spatrack_utils/infer_track.py ADDED Viewed

	@@ -0,0 +1,194 @@

+from models.SpaTrackV2.models.predictor import Predictor
+import yaml
+import easydict
+import os
+import numpy as np
+import cv2
+import torch
+import torchvision.transforms as T
+from PIL import Image
+import io
+import moviepy.editor as mp
+from models.SpaTrackV2.utils.visualizer import Visualizer
+import tqdm
+from models.SpaTrackV2.models.utils import get_points_on_a_grid
+import glob
+from rich import print
+import argparse
+import decord
+from huggingface_hub import hf_hub_download
+config = {
+    "ckpt_dir": "Yuxihenry/SpatialTrackerCkpts",  # HuggingFace repo ID
+    "cfg_dir": "config/magic_infer_moge.yaml",
+}
+def get_tracker_predictor(output_dir: str, vo_points: int = 756, tracker_model=None):
+    """
+    Initialize and return the tracker predictor and visualizer
+    Args:
+        output_dir: Directory to save visualization results
+        vo_points: Number of points for visual odometry
+    Returns:
+        Tuple of (tracker_predictor, visualizer)
+    """
+    viz = True
+    os.makedirs(output_dir, exist_ok=True)
+    with open(config["cfg_dir"], "r") as f:
+        cfg = yaml.load(f, Loader=yaml.FullLoader)
+    cfg = easydict.EasyDict(cfg)
+    cfg.out_dir = output_dir
+    cfg.model.track_num = vo_points
+    # Check if it's a local path or HuggingFace repo
+    if tracker_model is not None:
+        model = tracker_model
+        model.spatrack.track_num = vo_points
+    else:
+        if os.path.exists(config["ckpt_dir"]):
+            # Local file
+            model = Predictor.from_pretrained(config["ckpt_dir"], model_cfg=cfg["model"])
+        else:
+            # HuggingFace repo - download the model
+            print(f"Downloading model from HuggingFace: {config['ckpt_dir']}")
+            checkpoint_path = hf_hub_download(
+                repo_id=config["ckpt_dir"],
+                repo_type="model",
+                filename="SpaTrack3_offline.pth"
+            )
+            model = Predictor.from_pretrained(checkpoint_path, model_cfg=cfg["model"])
+        model.eval()
+        model.to("cuda")
+    viser = Visualizer(save_dir=cfg.out_dir, grayscale=True,
+                     fps=10, pad_value=0, tracks_leave_trace=5)
+    return model, viser
+def run_tracker(model, viser, temp_dir, video_name, grid_size, vo_points, fps=3):
+    """
+    Run tracking on a video sequence
+    Args:
+        model: Tracker predictor instance
+        viser: Visualizer instance
+        temp_dir: Directory containing temporary files
+        video_name: Name of the video file (without extension)
+        grid_size: Size of the tracking grid
+        vo_points: Number of points for visual odometry
+        fps: Frames per second for visualization
+    """
+    # Setup paths
+    video_path = os.path.join(temp_dir, f"{video_name}.mp4")
+    mask_path = os.path.join(temp_dir, f"{video_name}.png")
+    out_dir = os.path.join(temp_dir, "results")
+    os.makedirs(out_dir, exist_ok=True)
+    # Load video using decord
+    video_reader = decord.VideoReader(video_path)
+    video_tensor = torch.from_numpy(video_reader.get_batch(range(len(video_reader))).asnumpy()).permute(0, 3, 1, 2)  # Convert to tensor and permute to (N, C, H, W)
+    # resize make sure the shortest side is 336
+    h, w = video_tensor.shape[2:]
+    scale = max(336 / h, 336 / w)
+    if scale < 1:
+        new_h, new_w = int(h * scale), int(w * scale)
+        video_tensor = T.Resize((new_h, new_w))(video_tensor)
+    video_tensor = video_tensor[::fps].float()
+    depth_tensor = None
+    intrs = None
+    extrs = None
+    data_npz_load = {}
+    # Load and process mask
+    if os.path.exists(mask_path):
+        mask = cv2.imread(mask_path)
+        mask = cv2.resize(mask, (video_tensor.shape[3], video_tensor.shape[2]))
+        mask = mask.sum(axis=-1)>0
+    else:
+        mask = np.ones_like(video_tensor[0,0].numpy())>0
+    # Get frame dimensions and create grid points
+    frame_H, frame_W = video_tensor.shape[2:]
+    grid_pts = get_points_on_a_grid(grid_size, (frame_H, frame_W), device="cpu")
+    # Sample mask values at grid points and filter out points where mask=0
+    if os.path.exists(mask_path):
+        grid_pts_int = grid_pts[0].long()
+        mask_values = mask[grid_pts_int[...,1], grid_pts_int[...,0]]
+        grid_pts = grid_pts[:, mask_values]
+    query_xyt = torch.cat([torch.zeros_like(grid_pts[:, :, :1]), grid_pts], dim=2)[0].numpy()
+    # run vggt
+    if os.environ.get("VGGT_DIR", None) is not None:
+        vggt_model = VGGT()
+        vggt_model.load_state_dict(torch.load(VGGT_DIR))
+        vggt_model.eval()
+        vggt_model = vggt_model.to("cuda")
+        # process the image tensor
+        video_tensor = preprocess_image(video_tensor)[None]
+        with torch.cuda.amp.autocast(dtype=torch.bfloat16):
+            # Predict attributes including cameras, depth maps, and point maps.
+            aggregated_tokens_list, ps_idx = vggt_model.aggregator(video_tensor.cuda()/255)
+            pose_enc = vggt_model.camera_head(aggregated_tokens_list)[-1]
+            # Extrinsic and intrinsic matrices, following OpenCV convention (camera from world)
+            extrinsic, intrinsic = pose_encoding_to_extri_intri(pose_enc, video_tensor.shape[-2:])
+            # Predict Depth Maps
+            depth_map, depth_conf = vggt_model.depth_head(aggregated_tokens_list, video_tensor.cuda()/255, ps_idx)
+            # clear the cache
+            del vggt_model, aggregated_tokens_list, ps_idx, pose_enc
+            torch.cuda.empty_cache()
+        depth_tensor = depth_map.squeeze().cpu().numpy()
+        extrs = np.eye(4)[None].repeat(len(depth_tensor), axis=0)
+        extrs[:, :3, :4] = extrinsic.squeeze().cpu().numpy()
+        intrs = intrinsic.squeeze().cpu().numpy()
+        video_tensor = video_tensor.squeeze()
+        #NOTE: 20% of the depth is not reliable
+        # threshold = depth_conf.squeeze().view(-1).quantile(0.5)
+        unc_metric = depth_conf.squeeze().cpu().numpy() > 0.5
+    # Run model inference
+    with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16):
+        (
+            c2w_traj, intrs, point_map, conf_depth,
+            track3d_pred, track2d_pred, vis_pred, conf_pred, video
+        ) = model.forward(video_tensor, depth=depth_tensor,
+                            intrs=intrs, extrs=extrs,
+                            queries=query_xyt,
+                            fps=1, full_point=False, iters_track=4,
+                            query_no_BA=True, fixed_cam=False, stage=1,
+                            support_frame=len(video_tensor)-1, replace_ratio=0.2)
+        # Resize results to avoid too large I/O Burden
+        max_size = 336
+        h, w = video.shape[2:]
+        scale = min(max_size / h, max_size / w)
+        if scale < 1:
+            new_h, new_w = int(h * scale), int(w * scale)
+            video = T.Resize((new_h, new_w))(video)
+            video_tensor = T.Resize((new_h, new_w))(video_tensor)
+            point_map = T.Resize((new_h, new_w))(point_map)
+            track2d_pred[...,:2] = track2d_pred[...,:2] * scale
+            intrs[:,:2,:] = intrs[:,:2,:] * scale
+            if depth_tensor is not None:
+                depth_tensor = T.Resize((new_h, new_w))(depth_tensor)
+            conf_depth = T.Resize((new_h, new_w))(conf_depth)
+        # Visualize tracks
+        viser.visualize(video=video[None],
+                        tracks=track2d_pred[None][...,:2],
+                        visibility=vis_pred[None],filename="test")
+        # Save in tapip3d format
+        data_npz_load["coords"] = (torch.einsum("tij,tnj->tni", c2w_traj[:,:3,:3], track3d_pred[:,:,:3].cpu()) + c2w_traj[:,:3,3][:,None,:]).numpy()
+        data_npz_load["extrinsics"] = torch.inverse(c2w_traj).cpu().numpy()
+        data_npz_load["intrinsics"] = intrs.cpu().numpy()
+        data_npz_load["depths"] = point_map[:,2,...].cpu().numpy()
+        data_npz_load["video"] = (video_tensor).cpu().numpy()/255
+        data_npz_load["visibs"] = vis_pred.cpu().numpy()
+        data_npz_load["confs"] = conf_pred.cpu().numpy()
+        data_npz_load["confs_depth"] = conf_depth.cpu().numpy()
+        np.savez(os.path.join(out_dir, f'result.npz'), **data_npz_load)
+        print(f"Results saved to {out_dir}.\nTo visualize them with tapip3d, run: [bold yellow]python tapip3d_viz.py {out_dir}/result.npz[/bold yellow]")

config/__init__.py ADDED Viewed

File without changes

config/magic_infer_moge.yaml ADDED Viewed

	@@ -0,0 +1,48 @@

+seed: 0
+# config the hydra logger, only in hydra `$` can be decoded as cite
+data: ./assets/room
+vis_track: false
+hydra:
+  run:
+    dir: .
+  output_subdir: null
+  job_logging: {}
+  hydra_logging: {}
+mixed_precision: bf16
+visdom:
+  viz_ip: "localhost"
+  port: 6666
+relax_load: false
+res_all: 336
+# config the ckpt path
+# ckpts: "/mnt/bn/xyxdata/home/codes/my_projs/SpaTrack2/checkpoints/new_base.pth"
+ckpts: "Yuxihenry/SpatialTracker_Files"
+batch_size: 1
+input:
+  type: image
+fps: 1
+model_wind_size: 32
+model:
+  backbone_cfg:
+    ckpt_dir: "checkpoints/model.pt"
+  chunk_size: 24        # downsample factor for patchified features
+  ckpt_fwd: true
+  ft_cfg:
+    mode: "fix"
+    paras_name:  []
+  resolution: 336
+  max_len: 512
+  Track_cfg:
+    base_ckpt: "checkpoints/scaled_offline.pth"
+    base:
+      stride: 4
+      corr_radius: 3
+      window_len: 60
+    stablizer: True
+    mode: "online"
+    s_wind: 200
+    overlap: 4
+  track_num: 0
+dist_train:
+  num_nodes: 1

examples/backpack.mp4 ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:4b5ac6b2285ffb48e3a740e419e38c781df9c963589a5fd894e5b4e13dd6a8b8
+size 1208738

examples/ball.mp4 ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:31f6e3bf875a85284b376c05170b4c08b546b7d5e95106848b1e3818a9d0db91
+size 3030268

examples/basketball.mp4 ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:0df3b429d5fd64c298f2d79b2d818a4044e7341a71d70b957f60b24e313c3760
+size 2487837

examples/biker.mp4 ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:fba880c24bdb8fa3b84b1b491d52f2c1f426fb09e34c3013603e5a549cf3b22b
+size 249196

examples/cinema_0.mp4 ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:a68a5643c14f61c05d48e25a98ddf5cf0344d3ffcda08ad4a0adc989d49d7a9c
+size 1774022

examples/cinema_1.mp4 ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:99624e2d0fb2e9f994e46aefb904e884de37a6d78e7f6b6670e286eaa397e515
+size 2370749

examples/drifting.mp4 ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:4f3937871117d3cc5d7da3ef31d1edf5626fc8372126b73590f75f05713fe97c
+size 4695804

examples/ego_kc1.mp4 ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:22fe64e458e329e8b3c3e20b3725ffd85c3a2e725fd03909cf883d3fd02c80b3
+size 1365980

examples/ego_teaser.mp4 ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:8780b291b48046b1c7dea90712c1c3f59d60c03216df1c489f6f03e8d61fae5c
+size 7365665

examples/handwave.mp4 ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:e6dde7cf4ffa7c66b6861bb5abdedc49dfc4b5b4dd9dd46ee8415dd4953935b6
+size 2099369

examples/hockey.mp4 ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:8c3be095777b442dc401e7d1f489b749611ffade3563a01e4e3d1e511311bd86
+size 1795810

examples/ken_block_0.mp4 ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:7b788faeb4d3206fa604d622a05268f1321ad6a229178fe12319d20c9438deb1
+size 196343

examples/kiss.mp4 ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:f78fffc5108d95d4e5837d7607226f3dd9796615ea3481f2629c69ccd2ccb12f
+size 1073570

examples/kitchen.mp4 ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:3120e942a9b3d7b300928e43113b000fb5ccc209012a2c560ec26b8a04c2d5f9
+size 543970

examples/kitchen_egocentric.mp4 ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:5468ab10d8d39b68b51fa616adc3d099dab7543e38dd221a0a7a20a2401824a2
+size 2176685

examples/pillow.mp4 ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:8f05818f586d7b0796fcd4714ea4be489c93701598cadc86ce7973fc24655fee
+size 1407147

examples/protein.mp4 ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:b2dc9cfceb0984b61ebc62fda4c826135ebe916c8c966a8123dcc3315d43b73f
+size 2002300

examples/pusht.mp4 ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:996d1923e36811a1069e4d6b5e8c0338d9068c0870ea09c4c04e13e9fbcd207a
+size 5256495

examples/robot1.mp4 ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:7a3b9e4449572129fdd96a751938e211241cdd86bcc56ffd33bfd23fc4d6e9c0
+size 1178671

examples/robot2.mp4 ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:188b2d8824ce345c86a603bff210639a6158d72cf6119cc1d3f79d409ac68bb3
+size 867261

examples/robot_3.mp4 ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:784a0f9c36a316d0da5745075dbc8cefd9ce60c25b067d3d80a1d52830df8a37
+size 1153015

examples/robot_unitree.mp4 ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:99bc274f7613a665c6135085fe01691ebfaa9033101319071f37c550ab21d1ea
+size 1964268

examples/running.mp4 ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:9ceb96b287fefb1c090dcd2f5db7634f808d2079413500beeb7b33023dfae51b
+size 7307897

examples/teleop2.mp4 ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:59ea006a18227da8cf5db1fa50cd48e71ec7eb66fef48ea2158c325088bd9fee
+size 1077267

examples/vertical_place.mp4 ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:6c8061ae449f986113c2ecb17aefc2c13f737aecbcd41d6c057c88e6d41ac3ee
+size 719810

models/SpaTrackV2/models/SpaTrack.py ADDED Viewed

	@@ -0,0 +1,759 @@

+#python
+"""
+SpaTrackerV2, which is an unified model to estimate 'intrinsic',
+'video depth', 'extrinsic' and '3D Tracking' from casual video frames.
+Contact: DM yuxixiao@zju.edu.cn
+"""
+import os
+import numpy as np
+from typing import Literal, Union, List, Tuple, Dict
+import cv2
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+# from depth anything v2
+from huggingface_hub import PyTorchModelHubMixin  # used for model hub
+from einops import rearrange
+from models.monoD.depth_anything_v2.dpt import DepthAnythingV2
+from models.moge.model.v1 import MoGeModel
+import copy
+from functools import partial
+from models.SpaTrackV2.models.tracker3D.TrackRefiner import TrackRefiner3D
+import kornia
+from models.SpaTrackV2.utils.model_utils import sample_features5d
+import utils3d
+from models.SpaTrackV2.models.tracker3D.spatrack_modules.utils import depth_to_points_colmap, get_nth_visible_time_index
+from models.SpaTrackV2.models.utils import pose_enc2mat, matrix_to_quaternion, get_track_points, normalize_rgb
+import random
+class SpaTrack2(nn.Module, PyTorchModelHubMixin):
+    def __init__(
+        self,
+        loggers: list,   # include [ viz, logger_tf, logger]
+        backbone_cfg,
+        Track_cfg=None,
+        chunk_size=24,
+        ckpt_fwd: bool = False,
+        ft_cfg=None,
+        resolution=518,
+        max_len=600,  # the maximum video length we can preprocess,
+        track_num=768,
+    ):
+        self.chunk_size = chunk_size
+        self.max_len = max_len
+        self.resolution = resolution
+        # config the T-Lora Dinov2
+        #NOTE: initial the base model
+        base_cfg = copy.deepcopy(backbone_cfg)
+        backbone_ckpt_dir = base_cfg.pop('ckpt_dir', None)
+        super(SpaTrack2, self).__init__()
+        if os.path.exists(backbone_ckpt_dir)==False:
+            base_model = MoGeModel.from_pretrained('Ruicheng/moge-vitl')
+        else:
+            checkpoint = torch.load(backbone_ckpt_dir, map_location='cpu', weights_only=True)
+            base_model = MoGeModel(**checkpoint["model_config"])
+            base_model.load_state_dict(checkpoint['model'])
+        # avoid the base_model is a member of SpaTrack2
+        object.__setattr__(self, 'base_model', base_model)
+        # Tracker model
+        self.Track3D = TrackRefiner3D(Track_cfg)
+        track_base_ckpt_dir = Track_cfg.base_ckpt
+        if os.path.exists(track_base_ckpt_dir):
+            track_pretrain = torch.load(track_base_ckpt_dir)
+            self.Track3D.load_state_dict(track_pretrain, strict=False)
+        # wrap the function of make lora trainable
+        self.make_paras_trainable = partial(self.make_paras_trainable,
+                                            mode=ft_cfg.mode,
+                                            paras_name=ft_cfg.paras_name)
+        self.track_num = track_num
+    def make_paras_trainable(self, mode: str = 'fix', paras_name: List[str] = []):
+        # gradient required for the lora_experts and gate
+        for name, param in self.named_parameters():
+            if any(x in name for x in paras_name):
+                if mode == 'fix':
+                    param.requires_grad = False
+                else:
+                    param.requires_grad = True
+            else:
+                if mode == 'fix':
+                    param.requires_grad = True
+                else:
+                    param.requires_grad = False
+        total_params = sum(p.numel() for p in self.parameters())
+        trainable_params = sum(p.numel() for p in self.parameters() if p.requires_grad)
+        print(f"Total parameters: {total_params}")
+        print(f"Trainable parameters: {trainable_params/total_params*100:.2f}%")
+    def ProcVid(self,
+                        x: torch.Tensor):
+        """
+        split the video into several overlapped windows.
+        args:
+            x: the input video frames.   [B, T, C, H, W]
+        outputs:
+            patch_size: the patch size of the video features
+        raises:
+            ValueError: if the input video is longer than `max_len`.
+        """
+        # normalize the input images
+        num_types = x.dtype
+        x = normalize_rgb(x, input_size=self.resolution)
+        x = x.to(num_types)
+        # get the video features
+        B, T, C, H, W = x.size()
+        if T > self.max_len:
+            raise ValueError(f"the video length should no more than {self.max_len}.")
+        # get the video features
+        patch_h, patch_w = H // 14, W // 14
+        patch_size = (patch_h,  patch_w)
+        # resize and get the video features
+        x = x.view(B * T, C, H, W)
+        # operate the temporal encoding
+        return patch_size, x
+    def forward_stream(
+            self,
+            video: torch.Tensor,
+            queries: torch.Tensor = None,
+            T_org: int = None,
+            depth: torch.Tensor|np.ndarray|str=None,
+            unc_metric_in: torch.Tensor|np.ndarray|str=None,
+            intrs: torch.Tensor|np.ndarray|str=None,
+            extrs: torch.Tensor|np.ndarray|str=None,
+            queries_3d: torch.Tensor = None,
+            window_len: int = 16,
+            overlap_len: int = 4,
+            full_point: bool = False,
+            track2d_gt: torch.Tensor = None,
+            fixed_cam: bool = False,
+            query_no_BA: bool = False,
+            stage: int = 0,
+            support_frame: int = 0,
+            replace_ratio: float = 0.6,
+            annots_train: Dict = None,
+            iters_track=4,
+            **kwargs,
+    ):
+        # step 1 allocate the query points on the grid
+        T, C, H, W = video.shape
+        if annots_train is not None:
+            vis_gt = annots_train["vis"]
+            _, _, N = vis_gt.shape
+            number_visible = vis_gt.sum(dim=1)
+            ratio_rand = torch.rand(1, N, device=vis_gt.device)
+            first_positive_inds = get_nth_visible_time_index(vis_gt, (number_visible*ratio_rand).long().clamp(min=1, max=T))
+            assert (torch.gather(vis_gt, 1, first_positive_inds[:, None, :].repeat(1, T, 1)) < 0).sum() == 0
+            first_positive_inds = first_positive_inds.long()
+            gather = torch.gather(
+                annots_train["traj_3d"][...,:2], 1, first_positive_inds[:, :, None, None].repeat(1, 1, N, 2)
+                )
+            xys = torch.diagonal(gather, dim1=1, dim2=2).permute(0, 2, 1)
+            queries = torch.cat([first_positive_inds[:, :, None], xys], dim=-1)[0].cpu().numpy()
+        # Unfold video into segments of window_len with overlap_len
+        step_slide = window_len - overlap_len
+        if T < window_len:
+            video_unf = video.unsqueeze(0)
+            if depth is not None:
+                depth_unf = depth.unsqueeze(0)
+            else:
+                depth_unf = None
+            if unc_metric_in is not None:
+                unc_metric_unf = unc_metric_in.unsqueeze(0)
+            else:
+                unc_metric_unf = None
+            if intrs is not None:
+                intrs_unf = intrs.unsqueeze(0)
+            else:
+                intrs_unf = None
+            if extrs is not None:
+                extrs_unf = extrs.unsqueeze(0)
+            else:
+                extrs_unf = None
+        else:
+            video_unf = video.unfold(0, window_len, step_slide).permute(0, 4, 1, 2, 3)  # [B, S, C, H, W]
+            if depth is not None:
+                depth_unf = depth.unfold(0, window_len, step_slide).permute(0, 3, 1, 2)
+                intrs_unf = intrs.unfold(0, window_len, step_slide).permute(0, 3, 1, 2)
+            else:
+                depth_unf = None
+                intrs_unf = None
+            if extrs is not None:
+                extrs_unf = extrs.unfold(0, window_len, step_slide).permute(0, 3, 1, 2)
+            else:
+                extrs_unf = None
+            if unc_metric_in is not None:
+                unc_metric_unf = unc_metric_in.unfold(0, window_len, step_slide).permute(0, 3, 1, 2)
+            else:
+                unc_metric_unf = None
+        # parallel
+        # Get number of segments
+        B = video_unf.shape[0]
+        #TODO: Process each segment in parallel using torch.nn.DataParallel
+        c2w_traj = torch.eye(4, 4)[None].repeat(T, 1, 1)
+        intrs_out = torch.eye(3, 3)[None].repeat(T, 1, 1)
+        point_map = torch.zeros(T, 3, H, W).cuda()
+        unc_metric = torch.zeros(T, H, W).cuda()
+        # set the queries
+        N, _ = queries.shape
+        track3d_pred = torch.zeros(T, N, 6).cuda()
+        track2d_pred = torch.zeros(T, N, 3).cuda()
+        vis_pred = torch.zeros(T, N, 1).cuda()
+        conf_pred = torch.zeros(T, N, 1).cuda()
+        dyn_preds = torch.zeros(T, N, 1).cuda()
+        # sort the queries by time
+        sorted_indices = np.argsort(queries[...,0])
+        sorted_inv_indices = np.argsort(sorted_indices)
+        sort_query = queries[sorted_indices]
+        sort_query = torch.from_numpy(sort_query).cuda()
+        if queries_3d is not None:
+            sort_query_3d = queries_3d[sorted_indices]
+            sort_query_3d = torch.from_numpy(sort_query_3d).cuda()
+        queries_len = 0
+        overlap_d = None
+        cache = None
+        loss = 0.0
+        for i in range(B):
+            segment = video_unf[i:i+1].cuda()
+            # Forward pass through model
+            # detect the key points for each frames
+            queries_new_mask = (sort_query[...,0] < i * step_slide + window_len) * (sort_query[...,0] >= (i * step_slide + overlap_len if i > 0 else 0))
+            if queries_3d is not None:
+                queries_new_3d = sort_query_3d[queries_new_mask]
+                queries_new_3d = queries_new_3d.float()
+            else:
+                queries_new_3d = None
+            queries_new = sort_query[queries_new_mask.bool()]
+            queries_new = queries_new.float()
+            if i > 0:
+                overlap2d = track2d_pred[i*step_slide:(i+1)*step_slide, :queries_len, :]
+                overlapvis = vis_pred[i*step_slide:(i+1)*step_slide, :queries_len, :]
+                overlapconf = conf_pred[i*step_slide:(i+1)*step_slide, :queries_len, :]
+                overlap_query = (overlapvis * overlapconf).max(dim=0)[1][None, ...]
+                overlap_xy = torch.gather(overlap2d, 0, overlap_query.repeat(1,1,2))
+                overlap_d = torch.gather(overlap2d, 0, overlap_query.repeat(1,1,3))[...,2].detach()
+                overlap_query = torch.cat([overlap_query[...,:1], overlap_xy], dim=-1)[0]
+                queries_new[...,0] -= i*step_slide
+                queries_new = torch.cat([overlap_query, queries_new], dim=0).detach()
+            if annots_train is None:
+                annots = {}
+            else:
+                annots = copy.deepcopy(annots_train)
+                annots["traj_3d"] = annots["traj_3d"][:, i*step_slide:i*step_slide+window_len, sorted_indices,:][...,:len(queries_new),:]
+                annots["vis"] = annots["vis"][:, i*step_slide:i*step_slide+window_len, sorted_indices][...,:len(queries_new)]
+                annots["poses_gt"] =  annots["poses_gt"][:, i*step_slide:i*step_slide+window_len]
+                annots["depth_gt"] = annots["depth_gt"][:, i*step_slide:i*step_slide+window_len]
+                annots["intrs"] = annots["intrs"][:, i*step_slide:i*step_slide+window_len]
+                annots["traj_mat"] = annots["traj_mat"][:,i*step_slide:i*step_slide+window_len]
+            if depth is not None:
+                annots["depth_gt"] = depth_unf[i:i+1].to(segment.device).to(segment.dtype)
+            if unc_metric_in is not None:
+                annots["unc_metric"] = unc_metric_unf[i:i+1].to(segment.device).to(segment.dtype)
+            if intrs is not None:
+                intr_seg = intrs_unf[i:i+1].to(segment.device).to(segment.dtype)[0].clone()
+                focal = (intr_seg[:,0,0] / segment.shape[-1] + intr_seg[:,1,1]/segment.shape[-2]) / 2
+                pose_fake = torch.zeros(1, 8).to(depth.device).to(depth.dtype).repeat(segment.shape[1], 1)
+                pose_fake[:, -1] = focal
+                pose_fake[:,3]=1
+                annots["intrs_gt"] = intr_seg
+            if extrs is not None:
+                extrs_unf_norm = extrs_unf[i:i+1][0].clone()
+                extrs_unf_norm = torch.inverse(extrs_unf_norm[:1,...]) @ extrs_unf[i:i+1][0]
+                rot_vec = matrix_to_quaternion(extrs_unf_norm[:,:3,:3])
+                annots["poses_gt"] = torch.zeros(1, rot_vec.shape[0], 7).to(segment.device).to(segment.dtype)
+                annots["poses_gt"][:, :, 3:7] = rot_vec.to(segment.device).to(segment.dtype)[None]
+                annots["poses_gt"][:, :, :3] = extrs_unf_norm[:,:3,3].to(segment.device).to(segment.dtype)[None]
+                annots["use_extr"] = True
+            kwargs.update({"stage": stage})
+            #TODO: DEBUG
+            out = self.forward(segment, pts_q=queries_new,
+                                pts_q_3d=queries_new_3d, overlap_d=overlap_d,
+                                full_point=full_point,
+                                fixed_cam=fixed_cam, query_no_BA=query_no_BA,
+                                support_frame=segment.shape[1]-1,
+                                cache=cache, replace_ratio=replace_ratio,
+                                iters_track=iters_track,
+                                **kwargs, annots=annots)
+            if self.training:
+                loss += out["loss"].squeeze()
+            # from models.SpaTrackV2.utils.visualizer import Visualizer
+            # vis_track = Visualizer(grayscale=False,
+            #                     fps=10, pad_value=50, tracks_leave_trace=0)
+            # vis_track.visualize(video=segment,
+            #                         tracks=out["traj_est"][...,:2],
+            #                         visibility=out["vis_est"],
+            #                         save_video=True)
+            # # visualize 4d
+            # import os, json
+            # import os.path as osp
+            # viser4d_dir = os.path.join("viser_4d_results")
+            # os.makedirs(viser4d_dir, exist_ok=True)
+            # depth_est = annots["depth_gt"][0]
+            # unc_metric = out["unc_metric"]
+            # mask = (unc_metric > 0.5).squeeze(1)
+            # # pose_est = out["poses_pred"].squeeze(0)
+            # pose_est = annots["traj_mat"][0]
+            # rgb_tracks = out["rgb_tracks"].squeeze(0)
+            # intrinsics = out["intrs"].squeeze(0)
+            # for i_k in range(out["depth"].shape[0]):
+            #     img_i = out["imgs_raw"][0][i_k].permute(1, 2, 0).cpu().numpy()
+            #     img_i = cv2.cvtColor(img_i, cv2.COLOR_BGR2RGB)
+            #     cv2.imwrite(osp.join(viser4d_dir, f'frame_{i_k:04d}.png'), img_i)
+            #     if stage == 1:
+            #         depth = depth_est[i_k].squeeze().cpu().numpy()
+            #         np.save(osp.join(viser4d_dir, f'frame_{i_k:04d}.npy'), depth)
+            #     else:
+            #         point_map_vis = out["points_map"][i_k].cpu().numpy()
+            #         np.save(osp.join(viser4d_dir, f'point_{i_k:04d}.npy'), point_map_vis)
+            # np.save(os.path.join(viser4d_dir, f'intrinsics.npy'), intrinsics.cpu().numpy())
+            # np.save(os.path.join(viser4d_dir, f'extrinsics.npy'), pose_est.cpu().numpy())
+            # np.save(os.path.join(viser4d_dir, f'conf.npy'), mask.float().cpu().numpy())
+            # np.save(os.path.join(viser4d_dir, f'colored_track3d.npy'), rgb_tracks.cpu().numpy())
+            queries_len = len(queries_new)
+            # update the track3d and track2d
+            left_len = len(track3d_pred[i*step_slide:i*step_slide+window_len, :queries_len, :])
+            track3d_pred[i*step_slide:i*step_slide+window_len, :queries_len, :] = out["rgb_tracks"][0,:left_len,:queries_len,:]
+            track2d_pred[i*step_slide:i*step_slide+window_len, :queries_len, :] = out["traj_est"][0,:left_len,:queries_len,:3]
+            vis_pred[i*step_slide:i*step_slide+window_len, :queries_len, :] = out["vis_est"][0,:left_len,:queries_len,None]
+            conf_pred[i*step_slide:i*step_slide+window_len, :queries_len, :] = out["conf_pred"][0,:left_len,:queries_len,None]
+            dyn_preds[i*step_slide:i*step_slide+window_len, :queries_len, :] = out["dyn_preds"][0,:left_len,:queries_len,None]
+            # process the output for each segment
+            seg_c2w = out["poses_pred"][0]
+            seg_intrs = out["intrs"][0]
+            seg_point_map = out["points_map"]
+            seg_conf_depth = out["unc_metric"]
+            # cache management
+            cache = out["cache"]
+            for k in cache.keys():
+                if "_pyramid" in k:
+                    for j in range(len(cache[k])):
+                        if len(cache[k][j].shape) == 5:
+                            cache[k][j] = cache[k][j][:,:,:,:queries_len,:]
+                        elif len(cache[k][j].shape) == 4:
+                            cache[k][j] = cache[k][j][:,:1,:queries_len,:]
+                elif "_pred_cache" in k:
+                    cache[k] = cache[k][-overlap_len:,:queries_len,:]
+                else:
+                    cache[k] = cache[k][-overlap_len:]
+            # update the results
+            idx_glob = i * step_slide
+            # refine part
+            # mask_update = sort_query[..., 0] < i * step_slide + window_len
+            # sort_query_pick = sort_query[mask_update]
+            intrs_out[idx_glob:idx_glob+window_len] = seg_intrs
+            point_map[idx_glob:idx_glob+window_len] = seg_point_map
+            unc_metric[idx_glob:idx_glob+window_len] = seg_conf_depth
+            # update the camera poses
+            # if using the ground truth pose
+            # if extrs_unf is not None:
+            #     c2w_traj[idx_glob:idx_glob+window_len] = extrs_unf[i:i+1][0].to(c2w_traj.device).to(c2w_traj.dtype)
+            # else:
+            prev_c2w = c2w_traj[idx_glob:idx_glob+window_len][:1]
+            c2w_traj[idx_glob:idx_glob+window_len] = prev_c2w@seg_c2w.to(c2w_traj.device).to(c2w_traj.dtype)
+        track2d_pred = track2d_pred[:T_org,sorted_inv_indices,:]
+        track3d_pred = track3d_pred[:T_org,sorted_inv_indices,:]
+        vis_pred = vis_pred[:T_org,sorted_inv_indices,:]
+        conf_pred = conf_pred[:T_org,sorted_inv_indices,:]
+        dyn_preds = dyn_preds[:T_org,sorted_inv_indices,:]
+        unc_metric = unc_metric[:T_org,:]
+        point_map = point_map[:T_org,:]
+        intrs_out = intrs_out[:T_org,:]
+        c2w_traj = c2w_traj[:T_org,:]
+        if self.training:
+            ret = {
+                "loss": loss,
+                "depth_loss": 0.0,
+                "ab_loss": 0.0,
+                "vis_loss": out["vis_loss"],
+                "track_loss": out["track_loss"],
+                "conf_loss": out["conf_loss"],
+                "dyn_loss": out["dyn_loss"],
+                "sync_loss": out["sync_loss"],
+                "poses_pred": c2w_traj[None],
+                "intrs": intrs_out[None],
+                "points_map": point_map,
+                "track3d_pred": track3d_pred[None],
+                "rgb_tracks": track3d_pred[None],
+                "track2d_pred": track2d_pred[None],
+                "traj_est": track2d_pred[None],
+                "vis_est": vis_pred[None], "conf_pred": conf_pred[None],
+                "dyn_preds": dyn_preds[None],
+                "imgs_raw": video[None],
+                "unc_metric": unc_metric,
+                }
+            return ret
+        else:
+            return c2w_traj, intrs_out, point_map, unc_metric, track3d_pred, track2d_pred, vis_pred, conf_pred
+    def forward(self,
+                 x: torch.Tensor,
+                 annots: Dict = {},
+                 pts_q: torch.Tensor = None,
+                 pts_q_3d: torch.Tensor = None,
+                 overlap_d: torch.Tensor = None,
+                 full_point = False,
+                 fixed_cam = False,
+                 support_frame = 0,
+                 query_no_BA = False,
+                 cache = None,
+                 replace_ratio = 0.6,
+                 iters_track=4,
+                 **kwargs):
+        """
+        forward the video camera model, which predict (
+            `intr` `camera poses` `video depth`
+            )
+        args:
+            x: the input video frames. [B, T, C, H, W]
+            annots: the annotations for video frames.
+                    {
+                        "poses_gt": the pose encoding for the video frames. [B, T, 7]
+                        "depth_gt": the ground truth depth for the video frames. [B, T, 1, H, W],
+                        "metric": bool, whether to calculate the metric for the video frames.
+                    }
+        """
+        self.support_frame = support_frame
+        #TODO: to adjust a little bit
+        track_loss=ab_loss=vis_loss=track_loss=conf_loss=dyn_loss=0.0
+        B, T, _, H, W = x.shape
+        imgs_raw = x.clone()
+        # get the video split and features for each segment
+        patch_size, x_resize = self.ProcVid(x)
+        x_resize = rearrange(x_resize, "(b t) c h w -> b t c h w", b=B)
+        H_resize, W_resize = x_resize.shape[-2:]
+        prec_fx = W / W_resize
+        prec_fy = H / H_resize
+        # get patch size
+        P_H, P_W = patch_size
+        # get the depth, pointmap and mask
+        #TODO: Release DepthAnything Version
+        points_map_gt = None
+        with torch.no_grad():
+            if_gt_depth = (("depth_gt" in annots.keys())) and (kwargs.get('stage', 0)==1 or kwargs.get('stage', 0)==3)
+            if if_gt_depth==False:
+                if cache is not None:
+                    T_cache = cache["points_map"].shape[0]
+                    T_new = T - T_cache
+                    x_resize_new = x_resize[:, T_cache:]
+                else:
+                    T_new = T
+                    x_resize_new = x_resize
+                # infer with chunk
+                chunk_size = self.chunk_size
+                metric_depth = []
+                intrs = []
+                unc_metric = []
+                mask = []
+                points_map = []
+                normals = []
+                normals_mask = []
+                for i in range(0, B*T_new, chunk_size):
+                    output = self.base_model.infer(x_resize_new.view(B*T_new, -1, H_resize, W_resize)[i:i+chunk_size])
+                    metric_depth.append(output['depth'])
+                    intrs.append(output['intrinsics'])
+                    unc_metric.append(output['mask_prob'])
+                    mask.append(output['mask'])
+                    points_map.append(output['points'])
+                    normals_i, normals_mask_i = utils3d.torch.points_to_normals(output['points'], mask=output['mask'])
+                    normals.append(normals_i)
+                    normals_mask.append(normals_mask_i)
+                metric_depth = torch.cat(metric_depth, dim=0).view(B*T_new, 1, H_resize, W_resize).to(x.dtype)
+                intrs = torch.cat(intrs, dim=0).view(B, T_new, 3, 3).to(x.dtype)
+                intrs[:,:,0,:] *= W_resize
+                intrs[:,:,1,:] *= H_resize
+                # points_map = torch.cat(points_map, dim=0)
+                mask = torch.cat(mask, dim=0).view(B*T_new, 1, H_resize, W_resize).to(x.dtype)
+                # cat the normals
+                normals = torch.cat(normals, dim=0)
+                normals_mask = torch.cat(normals_mask, dim=0)
+                metric_depth = metric_depth.clone()
+                metric_depth[metric_depth == torch.inf] = 0
+                _depths = metric_depth[metric_depth > 0].reshape(-1)
+                q25 = torch.kthvalue(_depths, int(0.25 * len(_depths))).values
+                q75 = torch.kthvalue(_depths, int(0.75 * len(_depths))).values
+                iqr = q75 - q25
+                upper_bound = (q75 + 0.8*iqr).clamp(min=1e-6, max=10*q25)
+                _depth_roi = torch.tensor(
+                    [1e-1, upper_bound.item()],
+                    dtype=metric_depth.dtype,
+                    device=metric_depth.device
+                )
+                mask_roi = (metric_depth > _depth_roi[0]) & (metric_depth < _depth_roi[1])
+                mask = mask * mask_roi
+                mask = mask * (~(utils3d.torch.depth_edge(metric_depth, rtol=0.03, mask=mask.bool()))) * normals_mask[:,None,...]
+                points_map = depth_to_points_colmap(metric_depth.squeeze(1), intrs.view(B*T_new, 3, 3))
+                unc_metric = torch.cat(unc_metric, dim=0).view(B*T_new, 1, H_resize, W_resize).to(x.dtype)
+                unc_metric *= mask
+                if full_point:
+                    unc_metric = (~(utils3d.torch.depth_edge(metric_depth, rtol=0.1, mask=torch.ones_like(metric_depth).bool()))).float() * (metric_depth != 0)
+                if cache is not None:
+                    assert B==1, "only support batch size 1 right now."
+                    unc_metric = torch.cat([cache["unc_metric"], unc_metric], dim=0)
+                    intrs = torch.cat([cache["intrs"][None], intrs], dim=1)
+                    points_map = torch.cat([cache["points_map"].permute(0,2,3,1), points_map], dim=0)
+                    metric_depth = torch.cat([cache["metric_depth"], metric_depth], dim=0)
+            if "poses_gt" in annots.keys():
+                intrs, c2w_traj_gt = pose_enc2mat(annots["poses_gt"],
+                                            H_resize, W_resize, self.resolution)
+            else:
+                c2w_traj_gt = None
+            if "intrs_gt" in annots.keys():
+                intrs = annots["intrs_gt"].view(B, T, 3, 3)
+                fx_factor = W_resize / W
+                fy_factor = H_resize / H
+                intrs[:,:,0,:] *= fx_factor
+                intrs[:,:,1,:] *= fy_factor
+            if "depth_gt" in annots.keys():
+                metric_depth_gt = annots['depth_gt'].view(B*T, 1, H, W)
+                metric_depth_gt = F.interpolate(metric_depth_gt,
+                                size=(H_resize, W_resize), mode='nearest')
+                _depths = metric_depth_gt[metric_depth_gt > 0].reshape(-1)
+                q25 = torch.kthvalue(_depths, int(0.25 * len(_depths))).values
+                q75 = torch.kthvalue(_depths, int(0.75 * len(_depths))).values
+                iqr = q75 - q25
+                upper_bound = (q75 + 0.8*iqr).clamp(min=1e-6, max=10*q25)
+                _depth_roi = torch.tensor(
+                    [1e-1, upper_bound.item()],
+                    dtype=metric_depth_gt.dtype,
+                    device=metric_depth_gt.device
+                )
+                mask_roi = (metric_depth_gt > _depth_roi[0]) & (metric_depth_gt < _depth_roi[1])
+                # if (upper_bound > 200).any():
+                #     import pdb; pdb.set_trace()
+                if (kwargs.get('stage', 0) == 2):
+                    unc_metric = ((metric_depth_gt > 0)*(mask_roi) * (unc_metric > 0.5)).float()
+                    metric_depth_gt[metric_depth_gt > 10*q25] = 0
+                else:
+                    unc_metric = ((metric_depth_gt > 0)*(mask_roi)).float()
+                    unc_metric *= (~(utils3d.torch.depth_edge(metric_depth_gt, rtol=0.03, mask=mask_roi.bool()))).float()
+                    # filter the sky
+                    metric_depth_gt[metric_depth_gt > 10*q25] = 0
+                if "unc_metric" in annots.keys():
+                    unc_metric_ = F.interpolate(annots["unc_metric"].permute(1,0,2,3),
+                                size=(H_resize, W_resize), mode='nearest')
+                    unc_metric = unc_metric * unc_metric_
+                if if_gt_depth:
+                    points_map = depth_to_points_colmap(metric_depth_gt.squeeze(1), intrs.view(B*T, 3, 3))
+                    metric_depth = metric_depth_gt
+                    points_map_gt = points_map
+                else:
+                    points_map_gt = depth_to_points_colmap(metric_depth_gt.squeeze(1), intrs.view(B*T, 3, 3))
+        # track the 3d points
+        ret_track = None
+        regular_track = True
+        dyn_preds, final_tracks = None, None
+        if "use_extr" in annots.keys():
+            init_pose = True
+        else:
+            init_pose = False
+        # set the custom vid and valid only
+        custom_vid = annots.get("custom_vid", False)
+        valid_only = annots.get("data_dir", [""])[0] == "stereo4d"
+        if self.training:
+            if (annots["vis"].sum() > 0) and (kwargs.get('stage', 0)==1 or kwargs.get('stage', 0)==3):
+                traj3d = annots['traj_3d']
+                if (kwargs.get('stage', 0)==1) and (annots.get("custom_vid", False)==False):
+                    support_pts_q = get_track_points(H_resize, W_resize,
+                                                    T, x.device, query_size=self.track_num // 2,
+                                                    support_frame=self.support_frame, unc_metric=unc_metric, mode="incremental")[None]
+                else:
+                    support_pts_q = get_track_points(H_resize, W_resize,
+                                                    T, x.device, query_size=random.randint(32, 256),
+                                                    support_frame=self.support_frame, unc_metric=unc_metric, mode="incremental")[None]
+                if pts_q is not None:
+                    pts_q = pts_q[None,None]
+                    ret_track, dyn_preds, final_tracks, rgb_tracks, intrs_org, point_map_org_refined, cache = self.Track3D(imgs_raw,
+                                                        metric_depth,
+                                                        unc_metric.detach(), points_map, pts_q,
+                                                        intrs=intrs.clone(), cache=cache,
+                                                        prec_fx=prec_fx, prec_fy=prec_fy, overlap_d=overlap_d,
+                                                        vis_gt=annots['vis'], traj3d_gt=traj3d, iters=iters_track,
+                                                        cam_gt=c2w_traj_gt, support_pts_q=support_pts_q, custom_vid=custom_vid,
+                                                        init_pose=init_pose, fixed_cam=fixed_cam, stage=kwargs.get('stage', 0),
+                                                        points_map_gt=points_map_gt, valid_only=valid_only, replace_ratio=replace_ratio)
+                else:
+                    ret_track, dyn_preds, final_tracks, rgb_tracks, intrs_org, point_map_org_refined, cache = self.Track3D(imgs_raw,
+                                                        metric_depth,
+                                                        unc_metric.detach(), points_map, traj3d[..., :2],
+                                                        intrs=intrs.clone(), cache=cache,
+                                                        prec_fx=prec_fx, prec_fy=prec_fy, overlap_d=overlap_d,
+                                                        vis_gt=annots['vis'], traj3d_gt=traj3d, iters=iters_track,
+                                                        cam_gt=c2w_traj_gt, support_pts_q=support_pts_q, custom_vid=custom_vid,
+                                                        init_pose=init_pose, fixed_cam=fixed_cam, stage=kwargs.get('stage', 0),
+                                                        points_map_gt=points_map_gt, valid_only=valid_only, replace_ratio=replace_ratio)
+                regular_track = False
+        if regular_track:
+            if pts_q is None:
+                pts_q = get_track_points(H_resize, W_resize,
+                                            T, x.device, query_size=self.track_num,
+                                            support_frame=self.support_frame, unc_metric=unc_metric, mode="incremental" if self.training else "incremental")[None]
+                support_pts_q = None
+            else:
+                pts_q = pts_q[None,None]
+                # resize the query points
+                pts_q[...,1] *= W_resize / W
+                pts_q[...,2] *= H_resize / H
+                if pts_q_3d is not None:
+                    pts_q_3d = pts_q_3d[None,None]
+                    # resize the query points
+                    pts_q_3d[...,1] *= W_resize / W
+                    pts_q_3d[...,2] *= H_resize / H
+                else:
+                    # adjust the query with uncertainty
+                    if (full_point==False) and (overlap_d is None):
+                        pts_q_unc = sample_features5d(unc_metric[None], pts_q).squeeze()
+                        pts_q = pts_q[:,:,pts_q_unc>0.5,:]
+                        if (pts_q_unc<0.5).sum() > 0:
+                            # pad the query points
+                            pad_num = pts_q_unc.shape[0] - pts_q.shape[2]
+                            # pick the random indices
+                            indices = torch.randint(0, pts_q.shape[2], (pad_num,), device=pts_q.device)
+                            pad_pts = indices
+                            pts_q = torch.cat([pts_q, pts_q[:,:,pad_pts,:]], dim=-2)
+                support_pts_q = get_track_points(H_resize, W_resize,
+                                            T, x.device, query_size=self.track_num,
+                                            support_frame=self.support_frame,
+                                            unc_metric=unc_metric, mode="mixed")[None]
+            points_map[points_map>1e3] = 0
+            points_map = depth_to_points_colmap(metric_depth.squeeze(1), intrs.view(B*T, 3, 3))
+            ret_track, dyn_preds, final_tracks, rgb_tracks, intrs_org, point_map_org_refined, cache = self.Track3D(imgs_raw,
+                                                    metric_depth,
+                                                    unc_metric.detach(), points_map, pts_q,
+                                                    pts_q_3d=pts_q_3d, intrs=intrs.clone(),cache=cache,
+                                                    overlap_d=overlap_d, cam_gt=c2w_traj_gt if kwargs.get('stage', 0)==1 else None,
+                                                    prec_fx=prec_fx, prec_fy=prec_fy, support_pts_q=support_pts_q, custom_vid=custom_vid, valid_only=valid_only,
+                                                    fixed_cam=fixed_cam, query_no_BA=query_no_BA, init_pose=init_pose, iters=iters_track,
+                                                    stage=kwargs.get('stage', 0), points_map_gt=points_map_gt, replace_ratio=replace_ratio)
+        intrs = intrs_org
+        points_map = point_map_org_refined
+        c2w_traj = ret_track["cam_pred"]
+        if ret_track is not None:
+            if ret_track["loss"] is not None:
+                track_loss, conf_loss, dyn_loss, vis_loss, point_map_loss, scale_loss, shift_loss, sync_loss= ret_track["loss"]
+        # update the cache
+        cache.update({"metric_depth": metric_depth, "unc_metric": unc_metric, "points_map": points_map, "intrs": intrs[0]})
+        # output
+        depth = F.interpolate(metric_depth,
+                            size=(H, W), mode='bilinear', align_corners=True).squeeze(1)
+        points_map = F.interpolate(points_map,
+                            size=(H, W), mode='bilinear', align_corners=True).squeeze(1)
+        unc_metric = F.interpolate(unc_metric,
+                            size=(H, W), mode='bilinear', align_corners=True).squeeze(1)
+        if self.training:
+            loss = track_loss + conf_loss + dyn_loss + sync_loss + vis_loss + point_map_loss + (scale_loss + shift_loss)*50
+            ret = {"loss": loss,
+                    "depth_loss": point_map_loss,
+                    "ab_loss": (scale_loss + shift_loss)*50,
+                    "vis_loss": vis_loss, "track_loss": track_loss,
+                    "poses_pred": c2w_traj, "dyn_preds": dyn_preds, "traj_est": final_tracks, "conf_loss": conf_loss,
+                    "imgs_raw": imgs_raw, "rgb_tracks": rgb_tracks, "vis_est": ret_track['vis_pred'],
+                    "depth": depth, "points_map": points_map, "unc_metric": unc_metric, "intrs": intrs, "dyn_loss": dyn_loss,
+                    "sync_loss": sync_loss, "conf_pred": ret_track['conf_pred'], "cache": cache,
+                    }
+        else:
+            if ret_track is not None:
+                traj_est = ret_track['preds']
+                traj_est[..., 0] *= W / W_resize
+                traj_est[..., 1] *= H / H_resize
+                vis_est = ret_track['vis_pred']
+            else:
+                traj_est = torch.zeros(B, self.track_num // 2, 3).to(x.device)
+                vis_est = torch.zeros(B, self.track_num // 2).to(x.device)
+            if intrs is not None:
+                intrs[..., 0, :] *= W / W_resize
+                intrs[..., 1, :] *= H / H_resize
+            ret = {"poses_pred": c2w_traj, "dyn_preds": dyn_preds,
+                    "depth": depth, "traj_est": traj_est, "vis_est": vis_est, "imgs_raw": imgs_raw,
+                    "rgb_tracks": rgb_tracks, "intrs": intrs, "unc_metric": unc_metric, "points_map": points_map,
+                    "conf_pred": ret_track['conf_pred'], "cache": cache,
+                    }
+        return ret
+# three stages of training
+# stage 1:
+# gt depth and intrinsics synthetic (includes Dynamic Replica, Kubric, Pointodyssey, Vkitti, TartanAir and Indoor() ) Motion Patern (tapvid3d)
+# Tracking and Pose as well -> based on gt depth and intrinsics
+# (Finished) -> (megasam + base model) vs. tapip3d.     (use depth from megasam or pose, which keep the same setting as tapip3d.)
+# stage 2: fixed 3D tracking
+# Joint depth refiner
+# input depth from whatever + rgb -> temporal module + scale and shift token -> coarse alignment -> scale and shift
+# estimate the 3D tracks -> 3D tracks combine with pointmap -> update for pointmap (iteratively) -> residual map B T 3 H W
+# ongoing two days
+# stage 3: train multi windows by propagation
+# 4 frames overlapped -> train on 64 -> fozen image encoder and finetuning the transformer  (learnable parameters pretty small)
+# types of scenarioes:
+# 1. auto driving (waymo open dataset)
+# 2. robot
+# 3. internet ego video
+# Iterative Transformer -- Solver -- General Neural MegaSAM + Tracks
+# Update Variables:
+# 1. 3D tracks B T N 3  xyz.
+# 2. 2D tracks B T N 2  x y.
+# 3. Dynamic Mask B T H W.
+# 4. Camera Pose B T 4 4.
+# 5. Video Depth.
+# (RGB, RGBD, RGBD+Pose) x (Static, Dynamic)
+# Campatiablity by product.

models/SpaTrackV2/models/__init__.py ADDED Viewed

File without changes

models/SpaTrackV2/models/blocks.py ADDED Viewed

	@@ -0,0 +1,519 @@

+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.cuda.amp import autocast
+from einops import rearrange
+import collections
+from functools import partial
+from itertools import repeat
+import torchvision.models as tvm
+from torch.utils.checkpoint import checkpoint
+from models.monoD.depth_anything.dpt import DPTHeadEnc, DPTHead
+from typing import Union, Tuple
+from torch import Tensor
+# From PyTorch internals
+def _ntuple(n):
+    def parse(x):
+        if isinstance(x, collections.abc.Iterable) and not isinstance(x, str):
+            return tuple(x)
+        return tuple(repeat(x, n))
+    return parse
+def exists(val):
+    return val is not None
+def default(val, d):
+    return val if exists(val) else d
+to_2tuple = _ntuple(2)
+class LayerScale(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        init_values: Union[float, Tensor] = 1e-5,
+        inplace: bool = False,
+    ) -> None:
+        super().__init__()
+        self.inplace = inplace
+        self.gamma = nn.Parameter(init_values * torch.ones(dim))
+    def forward(self, x: Tensor) -> Tensor:
+        return x.mul_(self.gamma) if self.inplace else x * self.gamma
+class Mlp(nn.Module):
+    """MLP as used in Vision Transformer, MLP-Mixer and related networks"""
+    def __init__(
+        self,
+        in_features,
+        hidden_features=None,
+        out_features=None,
+        act_layer=nn.GELU,
+        norm_layer=None,
+        bias=True,
+        drop=0.0,
+        use_conv=False,
+    ):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        bias = to_2tuple(bias)
+        drop_probs = to_2tuple(drop)
+        linear_layer = partial(nn.Conv2d, kernel_size=1) if use_conv else nn.Linear
+        self.fc1 = linear_layer(in_features, hidden_features, bias=bias[0])
+        self.act = act_layer()
+        self.drop1 = nn.Dropout(drop_probs[0])
+        self.norm = norm_layer(hidden_features) if norm_layer is not None else nn.Identity()
+        self.fc2 = linear_layer(hidden_features, out_features, bias=bias[1])
+        self.drop2 = nn.Dropout(drop_probs[1])
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop1(x)
+        x = self.fc2(x)
+        x = self.drop2(x)
+        return x
+class Attention(nn.Module):
+    def __init__(self, query_dim, context_dim=None,
+                  num_heads=8, dim_head=48, qkv_bias=False, flash=False):
+        super().__init__()
+        inner_dim = self.inner_dim = dim_head * num_heads
+        context_dim = default(context_dim, query_dim)
+        self.scale = dim_head**-0.5
+        self.heads = num_heads
+        self.flash = flash
+        self.to_q = nn.Linear(query_dim, inner_dim, bias=qkv_bias)
+        self.to_kv = nn.Linear(context_dim, inner_dim * 2, bias=qkv_bias)
+        self.to_out = nn.Linear(inner_dim, query_dim)
+    def forward(self, x, context=None, attn_bias=None):
+        B, N1, _ = x.shape
+        C = self.inner_dim
+        h = self.heads
+        q = self.to_q(x).reshape(B, N1, h, C // h).permute(0, 2, 1, 3)
+        context = default(context, x)
+        k, v = self.to_kv(context).chunk(2, dim=-1)
+        N2 = context.shape[1]
+        k = k.reshape(B, N2, h, C // h).permute(0, 2, 1, 3)
+        v = v.reshape(B, N2, h, C // h).permute(0, 2, 1, 3)
+        with torch.autocast("cuda", enabled=True, dtype=torch.bfloat16):
+            if self.flash==False:
+                sim = (q @ k.transpose(-2, -1)) * self.scale
+                if attn_bias is not None:
+                    sim = sim + attn_bias
+                if sim.abs().max()>1e2:
+                    import pdb; pdb.set_trace()
+                attn = sim.softmax(dim=-1)
+                x = (attn @ v).transpose(1, 2).reshape(B, N1, C)
+            else:
+                input_args = [x.contiguous() for x in [q, k, v]]
+                x = F.scaled_dot_product_attention(*input_args).permute(0,2,1,3).reshape(B,N1,-1)  # type: ignore
+            if self.to_out.bias.dtype != x.dtype:
+                x = x.to(self.to_out.bias.dtype)
+        return self.to_out(x)
+class VGG19(nn.Module):
+    def __init__(self, pretrained=False, amp = False, amp_dtype = torch.float16) -> None:
+        super().__init__()
+        self.layers = nn.ModuleList(tvm.vgg19_bn(pretrained=pretrained).features[:40])
+        self.amp = amp
+        self.amp_dtype = amp_dtype
+    def forward(self, x, **kwargs):
+        with torch.autocast("cuda", enabled=self.amp, dtype = self.amp_dtype):
+            feats = {}
+            scale = 1
+            for layer in self.layers:
+                if isinstance(layer, nn.MaxPool2d):
+                    feats[scale] = x
+                    scale = scale*2
+                x = layer(x)
+            return feats
+class CNNandDinov2(nn.Module):
+    def __init__(self, cnn_kwargs = None, amp = True, amp_dtype = torch.float16):
+        super().__init__()
+        # in case the Internet connection is not stable, please load the DINOv2 locally
+        self.dinov2_vitl14 = torch.hub.load('models/torchhub/facebookresearch_dinov2_main',
+                                          'dinov2_{:}14'.format("vitl"), source='local', pretrained=False)
+        state_dict = torch.load("models/monoD/zoeDepth/ckpts/dinov2_vitl14_pretrain.pth")
+        self.dinov2_vitl14.load_state_dict(state_dict, strict=True)
+        cnn_kwargs = cnn_kwargs if cnn_kwargs is not None else {}
+        self.cnn = VGG19(**cnn_kwargs)
+        self.amp = amp
+        self.amp_dtype = amp_dtype
+        if self.amp:
+            dinov2_vitl14 = dinov2_vitl14.to(self.amp_dtype)
+        self.dinov2_vitl14 = [dinov2_vitl14] # ugly hack to not show parameters to DDP
+    def train(self, mode: bool = True):
+        return self.cnn.train(mode)
+    def forward(self, x, upsample = False):
+        B,C,H,W = x.shape
+        feature_pyramid = self.cnn(x)
+        if not upsample:
+            with torch.no_grad():
+                if self.dinov2_vitl14[0].device != x.device:
+                    self.dinov2_vitl14[0] = self.dinov2_vitl14[0].to(x.device).to(self.amp_dtype)
+                dinov2_features_16 = self.dinov2_vitl14[0].forward_features(x.to(self.amp_dtype))
+                features_16 = dinov2_features_16['x_norm_patchtokens'].permute(0,2,1).reshape(B,1024,H//14, W//14)
+                del dinov2_features_16
+                feature_pyramid[16] = features_16
+        return feature_pyramid
+class Dinov2(nn.Module):
+    def __init__(self, amp = True, amp_dtype = torch.float16):
+        super().__init__()
+        # in case the Internet connection is not stable, please load the DINOv2 locally
+        self.dinov2_vitl14 = torch.hub.load('models/torchhub/facebookresearch_dinov2_main',
+                                          'dinov2_{:}14'.format("vitl"), source='local', pretrained=False)
+        state_dict = torch.load("models/monoD/zoeDepth/ckpts/dinov2_vitl14_pretrain.pth")
+        self.dinov2_vitl14.load_state_dict(state_dict, strict=True)
+        self.amp = amp
+        self.amp_dtype = amp_dtype
+        if self.amp:
+            self.dinov2_vitl14 = self.dinov2_vitl14.to(self.amp_dtype)
+    def forward(self, x, upsample = False):
+        B,C,H,W = x.shape
+        mean_ = torch.tensor([0.485, 0.456, 0.406],
+                             device=x.device).view(1, 3, 1, 1)
+        std_ = torch.tensor([0.229, 0.224, 0.225],
+                            device=x.device).view(1, 3, 1, 1)
+        x = (x+1)/2
+        x = (x - mean_)/std_
+        h_re, w_re = 560, 560
+        x_resize = F.interpolate(x, size=(h_re, w_re),
+                                  mode='bilinear', align_corners=True)
+        if not upsample:
+            with torch.no_grad():
+                dinov2_features_16 = self.dinov2_vitl14.forward_features(x_resize.to(self.amp_dtype))
+                features_16 = dinov2_features_16['x_norm_patchtokens'].permute(0,2,1).reshape(B,1024,h_re//14, w_re//14)
+                del dinov2_features_16
+        features_16 = F.interpolate(features_16, size=(H//8, W//8), mode="bilinear", align_corners=True)
+        return features_16
+class AttnBlock(nn.Module):
+    """
+    A DiT block with adaptive layer norm zero (adaLN-Zero) conditioning.
+    """
+    def __init__(self, hidden_size, num_heads, mlp_ratio=4.0,
+                  flash=False, ckpt_fwd=False, debug=False, **block_kwargs):
+        super().__init__()
+        self.debug=debug
+        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.flash=flash
+        self.attn = Attention(
+            hidden_size, num_heads=num_heads, qkv_bias=True, flash=flash,
+            **block_kwargs
+        )
+        self.ls = LayerScale(hidden_size, init_values=0.005)
+        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        approx_gelu = lambda: nn.GELU(approximate="tanh")
+        self.mlp = Mlp(
+            in_features=hidden_size,
+            hidden_features=mlp_hidden_dim,
+            act_layer=approx_gelu,
+        )
+        self.ckpt_fwd = ckpt_fwd
+    def forward(self, x):
+        if self.debug:
+            print(x.max(), x.min(), x.mean())
+        if self.ckpt_fwd:
+            x = x + checkpoint(self.attn, self.norm1(x), use_reentrant=False)
+        else:
+            x = x + self.attn(self.norm1(x))
+        x = x + self.ls(self.mlp(self.norm2(x)))
+        return x
+class CrossAttnBlock(nn.Module):
+    def __init__(self, hidden_size, context_dim, num_heads=1, mlp_ratio=4.0, head_dim=48,
+                 flash=False, ckpt_fwd=False, **block_kwargs):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.norm_context = nn.LayerNorm(hidden_size)
+        self.cross_attn = Attention(
+            hidden_size, context_dim=context_dim, dim_head=head_dim,
+            num_heads=num_heads, qkv_bias=True, **block_kwargs, flash=flash,
+        )
+        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        approx_gelu = lambda: nn.GELU(approximate="tanh")
+        self.mlp = Mlp(
+            in_features=hidden_size,
+            hidden_features=mlp_hidden_dim,
+            act_layer=approx_gelu,
+            drop=0,
+        )
+        self.ckpt_fwd = ckpt_fwd
+    def forward(self, x, context):
+        if self.ckpt_fwd:
+            with autocast():
+                x = x + checkpoint(self.cross_attn,
+                                    self.norm1(x), self.norm_context(context), use_reentrant=False)
+        else:
+            with autocast():
+                x = x + self.cross_attn(
+                    self.norm1(x), self.norm_context(context)
+                )
+        x = x + self.mlp(self.norm2(x))
+        return x
+def bilinear_sampler(img, coords, mode="bilinear", mask=False):
+    """Wrapper for grid_sample, uses pixel coordinates"""
+    H, W = img.shape[-2:]
+    xgrid, ygrid = coords.split([1, 1], dim=-1)
+    # go to 0,1 then 0,2 then -1,1
+    xgrid = 2 * xgrid / (W - 1) - 1
+    ygrid = 2 * ygrid / (H - 1) - 1
+    grid = torch.cat([xgrid, ygrid], dim=-1)
+    img = F.grid_sample(img, grid, align_corners=True, mode=mode)
+    if mask:
+        mask = (xgrid > -1) & (ygrid > -1) & (xgrid < 1) & (ygrid < 1)
+        return img, mask.float()
+    return img
+class CorrBlock:
+    def __init__(self, fmaps, num_levels=4, radius=4, depths_dnG=None):
+        B, S, C, H_prev, W_prev = fmaps.shape
+        self.S, self.C, self.H, self.W = S, C, H_prev, W_prev
+        self.num_levels = num_levels
+        self.radius = radius
+        self.fmaps_pyramid = []
+        self.depth_pyramid = []
+        self.fmaps_pyramid.append(fmaps)
+        if depths_dnG is not None:
+           self.depth_pyramid.append(depths_dnG)
+        for i in range(self.num_levels - 1):
+            if depths_dnG is not None:
+                depths_dnG_ = depths_dnG.reshape(B * S, 1, H_prev, W_prev)
+                depths_dnG_ = F.avg_pool2d(depths_dnG_, 2, stride=2)
+                _, _, H, W = depths_dnG_.shape
+                depths_dnG = depths_dnG_.reshape(B, S, 1, H, W)
+                self.depth_pyramid.append(depths_dnG)
+            fmaps_ = fmaps.reshape(B * S, C, H_prev, W_prev)
+            fmaps_ = F.avg_pool2d(fmaps_, 2, stride=2)
+            _, _, H, W = fmaps_.shape
+            fmaps = fmaps_.reshape(B, S, C, H, W)
+            H_prev = H
+            W_prev = W
+            self.fmaps_pyramid.append(fmaps)
+    def sample(self, coords):
+        r = self.radius
+        B, S, N, D = coords.shape
+        assert D == 2
+        H, W = self.H, self.W
+        out_pyramid = []
+        for i in range(self.num_levels):
+            corrs = self.corrs_pyramid[i]  # B, S, N, H, W
+            _, _, _, H, W = corrs.shape
+            dx = torch.linspace(-r, r, 2 * r + 1)
+            dy = torch.linspace(-r, r, 2 * r + 1)
+            delta = torch.stack(torch.meshgrid(dy, dx, indexing="ij"), axis=-1).to(
+                coords.device
+            )
+            centroid_lvl = coords.reshape(B * S * N, 1, 1, 2) / 2 ** i
+            delta_lvl = delta.view(1, 2 * r + 1, 2 * r + 1, 2)
+            coords_lvl = centroid_lvl + delta_lvl
+            corrs = bilinear_sampler(corrs.reshape(B * S * N, 1, H, W), coords_lvl)
+            corrs = corrs.view(B, S, N, -1)
+            out_pyramid.append(corrs)
+        out = torch.cat(out_pyramid, dim=-1)  # B, S, N, LRR*2
+        return out.contiguous().float()
+    def corr(self, targets):
+        B, S, N, C = targets.shape
+        assert C == self.C
+        assert S == self.S
+        fmap1 = targets
+        self.corrs_pyramid = []
+        for fmaps in self.fmaps_pyramid:
+            _, _, _, H, W = fmaps.shape
+            fmap2s = fmaps.view(B, S, C, H * W)
+            corrs = torch.matmul(fmap1, fmap2s)
+            corrs = corrs.view(B, S, N, H, W)
+            corrs = corrs / torch.sqrt(torch.tensor(C).float())
+            self.corrs_pyramid.append(corrs)
+    def corr_sample(self, targets, coords, coords_dp=None):
+        B, S, N, C = targets.shape
+        r = self.radius
+        Dim_c = (2*r+1)**2
+        assert C == self.C
+        assert S == self.S
+        out_pyramid = []
+        out_pyramid_dp = []
+        for i in range(self.num_levels):
+            dx = torch.linspace(-r, r, 2 * r + 1)
+            dy = torch.linspace(-r, r, 2 * r + 1)
+            delta = torch.stack(torch.meshgrid(dy, dx, indexing="ij"), axis=-1).to(
+                coords.device
+            )
+            centroid_lvl = coords.reshape(B * S * N, 1, 1, 2) / 2 ** i
+            delta_lvl = delta.view(1, 2 * r + 1, 2 * r + 1, 2)
+            coords_lvl = centroid_lvl + delta_lvl
+            fmaps = self.fmaps_pyramid[i]
+            _, _, _, H, W = fmaps.shape
+            fmap2s = fmaps.view(B*S, C, H, W)
+            if len(self.depth_pyramid)>0:
+                depths_dnG_i = self.depth_pyramid[i]
+                depths_dnG_i = depths_dnG_i.view(B*S, 1, H, W)
+                dnG_sample = bilinear_sampler(depths_dnG_i, coords_lvl.view(B*S,1,N*Dim_c,2))
+                dp_corrs = (dnG_sample.view(B*S,N,-1) - coords_dp[0]).abs()/coords_dp[0]
+                out_pyramid_dp.append(dp_corrs)
+            fmap2s_sample = bilinear_sampler(fmap2s, coords_lvl.view(B*S,1,N*Dim_c,2))
+            fmap2s_sample = fmap2s_sample.permute(0, 3, 1, 2) # B*S, N*Dim_c, C, -1
+            corrs = torch.matmul(targets.reshape(B*S*N, 1, -1), fmap2s_sample.reshape(B*S*N, Dim_c, -1).permute(0, 2, 1))
+            corrs = corrs / torch.sqrt(torch.tensor(C).float())
+            corrs = corrs.view(B, S, N, -1)
+            out_pyramid.append(corrs)
+        out = torch.cat(out_pyramid, dim=-1)  # B, S, N, LRR*2
+        if len(self.depth_pyramid)>0:
+            out_dp = torch.cat(out_pyramid_dp, dim=-1)
+            self.fcorrD = out_dp.contiguous().float()
+        else:
+            self.fcorrD = torch.zeros_like(out).contiguous().float()
+        return out.contiguous().float()
+class EUpdateFormer(nn.Module):
+    """
+    Transformer model that updates track estimates.
+    """
+    def __init__(
+        self,
+        space_depth=12,
+        time_depth=12,
+        input_dim=320,
+        hidden_size=384,
+        num_heads=8,
+        output_dim=130,
+        mlp_ratio=4.0,
+        vq_depth=3,
+        add_space_attn=True,
+        add_time_attn=True,
+        flash=True
+    ):
+        super().__init__()
+        self.out_channels = 2
+        self.num_heads = num_heads
+        self.hidden_size = hidden_size
+        self.add_space_attn = add_space_attn
+        self.input_transform = torch.nn.Linear(input_dim, hidden_size, bias=True)
+        self.flash = flash
+        self.flow_head = nn.Sequential(
+            nn.Linear(hidden_size, output_dim, bias=True),
+            nn.ReLU(inplace=True),
+            nn.Linear(output_dim, output_dim, bias=True),
+            nn.ReLU(inplace=True),
+            nn.Linear(output_dim, output_dim, bias=True)
+        )
+        self.norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        cfg = xLSTMBlockStackConfig(
+            mlstm_block=mLSTMBlockConfig(
+                mlstm=mLSTMLayerConfig(
+                    conv1d_kernel_size=4, qkv_proj_blocksize=4, num_heads=4
+                )
+            ),
+            slstm_block=sLSTMBlockConfig(
+                slstm=sLSTMLayerConfig(
+                    backend="cuda",
+                    num_heads=4,
+                    conv1d_kernel_size=4,
+                    bias_init="powerlaw_blockdependent",
+                ),
+                feedforward=FeedForwardConfig(proj_factor=1.3, act_fn="gelu"),
+            ),
+            context_length=50,
+            num_blocks=7,
+            embedding_dim=384,
+            slstm_at=[1],
+        )
+        self.xlstm_fwd = xLSTMBlockStack(cfg)
+        self.xlstm_bwd = xLSTMBlockStack(cfg)
+        self.initialize_weights()
+    def initialize_weights(self):
+        def _basic_init(module):
+            if isinstance(module, nn.Linear):
+                torch.nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+        self.apply(_basic_init)
+    def forward(self,
+                input_tensor,
+                track_mask=None):
+        """ Updating with Transformer
+        Args:
+            input_tensor: B, N, T, C
+            arap_embed: B, N, T, C
+        """
+        B, N, T, C = input_tensor.shape
+        x = self.input_transform(input_tensor)
+        track_mask = track_mask.permute(0,2,1,3).float()
+        fwd_x = x*track_mask
+        bwd_x = x.flip(2)*track_mask.flip(2)
+        feat_fwd = self.xlstm_fwd(self.norm(fwd_x.view(B*N, T, -1)))
+        feat_bwd = self.xlstm_bwd(self.norm(bwd_x.view(B*N, T, -1)))
+        feat = (feat_bwd.flip(1) + feat_fwd).view(B, N, T, -1)
+        flow = self.flow_head(feat)
+        return flow[..., :2], flow[..., 2:]

models/SpaTrackV2/models/camera_transform.py ADDED Viewed

	@@ -0,0 +1,248 @@

+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# Adapted from https://github.com/amyxlase/relpose-plus-plus
+import torch
+import numpy as np
+import math
+def bbox_xyxy_to_xywh(xyxy):
+    wh = xyxy[2:] - xyxy[:2]
+    xywh = np.concatenate([xyxy[:2], wh])
+    return xywh
+def adjust_camera_to_bbox_crop_(fl, pp, image_size_wh: torch.Tensor, clamp_bbox_xywh: torch.Tensor):
+    focal_length_px, principal_point_px = _convert_ndc_to_pixels(fl, pp, image_size_wh)
+    principal_point_px_cropped = principal_point_px - clamp_bbox_xywh[:2]
+    focal_length, principal_point_cropped = _convert_pixels_to_ndc(
+        focal_length_px, principal_point_px_cropped, clamp_bbox_xywh[2:]
+    )
+    return focal_length, principal_point_cropped
+def adjust_camera_to_image_scale_(fl, pp, original_size_wh: torch.Tensor, new_size_wh: torch.LongTensor):
+    focal_length_px, principal_point_px = _convert_ndc_to_pixels(fl, pp, original_size_wh)
+    # now scale and convert from pixels to NDC
+    image_size_wh_output = new_size_wh.float()
+    scale = (image_size_wh_output / original_size_wh).min(dim=-1, keepdim=True).values
+    focal_length_px_scaled = focal_length_px * scale
+    principal_point_px_scaled = principal_point_px * scale
+    focal_length_scaled, principal_point_scaled = _convert_pixels_to_ndc(
+        focal_length_px_scaled, principal_point_px_scaled, image_size_wh_output
+    )
+    return focal_length_scaled, principal_point_scaled
+def _convert_ndc_to_pixels(focal_length: torch.Tensor, principal_point: torch.Tensor, image_size_wh: torch.Tensor):
+    half_image_size = image_size_wh / 2
+    rescale = half_image_size.min()
+    principal_point_px = half_image_size - principal_point * rescale
+    focal_length_px = focal_length * rescale
+    return focal_length_px, principal_point_px
+def _convert_pixels_to_ndc(
+    focal_length_px: torch.Tensor, principal_point_px: torch.Tensor, image_size_wh: torch.Tensor
+):
+    half_image_size = image_size_wh / 2
+    rescale = half_image_size.min()
+    principal_point = (half_image_size - principal_point_px) / rescale
+    focal_length = focal_length_px / rescale
+    return focal_length, principal_point
+def normalize_cameras(
+    cameras, compute_optical=True, first_camera=True, normalize_trans=True, scale=1.0, points=None, max_norm=False,
+    pose_mode="C2W"
+):
+    """
+    Normalizes cameras such that
+    (1) the optical axes point to the origin and the average distance to the origin is 1
+    (2) the first camera is the origin
+    (3) the translation vector is normalized
+    TODO: some transforms overlap with others. no need to do so many transforms
+    Args:
+        cameras (List[camera]).
+    """
+    # Let distance from first camera to origin be unit
+    new_cameras = cameras.clone()
+    scale = 1.0
+    if compute_optical:
+        new_cameras, points = compute_optical_transform(new_cameras, points=points)
+    if first_camera:
+        new_cameras, points = first_camera_transform(new_cameras, points=points, pose_mode=pose_mode)
+    if normalize_trans:
+        new_cameras, points, scale = normalize_translation(new_cameras,
+                                                            points=points, max_norm=max_norm)
+    return new_cameras, points, scale
+def compute_optical_transform(new_cameras, points=None):
+    """
+    adapted from https://github.com/amyxlase/relpose-plus-plus
+    """
+    new_transform = new_cameras.get_world_to_view_transform()
+    p_intersect, dist, p_line_intersect, pp, r = compute_optical_axis_intersection(new_cameras)
+    t = Translate(p_intersect)
+    scale = dist.squeeze()[0]
+    if points is not None:
+        points = t.inverse().transform_points(points)
+        points = points / scale
+    # Degenerate case
+    if scale == 0:
+        scale = torch.norm(new_cameras.T, dim=(0, 1))
+        scale = torch.sqrt(scale)
+        new_cameras.T = new_cameras.T / scale
+    else:
+        new_matrix = t.compose(new_transform).get_matrix()
+        new_cameras.R = new_matrix[:, :3, :3]
+        new_cameras.T = new_matrix[:, 3, :3] / scale
+    return new_cameras, points
+def compute_optical_axis_intersection(cameras):
+    centers = cameras.get_camera_center()
+    principal_points = cameras.principal_point
+    one_vec = torch.ones((len(cameras), 1))
+    optical_axis = torch.cat((principal_points, one_vec), -1)
+    pp = cameras.unproject_points(optical_axis, from_ndc=True, world_coordinates=True)
+    pp2 = pp[torch.arange(pp.shape[0]), torch.arange(pp.shape[0])]
+    directions = pp2 - centers
+    centers = centers.unsqueeze(0).unsqueeze(0)
+    directions = directions.unsqueeze(0).unsqueeze(0)
+    p_intersect, p_line_intersect, _, r = intersect_skew_line_groups(p=centers, r=directions, mask=None)
+    p_intersect = p_intersect.squeeze().unsqueeze(0)
+    dist = (p_intersect - centers).norm(dim=-1)
+    return p_intersect, dist, p_line_intersect, pp2, r
+def intersect_skew_line_groups(p, r, mask):
+    # p, r both of shape (B, N, n_intersected_lines, 3)
+    # mask of shape (B, N, n_intersected_lines)
+    p_intersect, r = intersect_skew_lines_high_dim(p, r, mask=mask)
+    _, p_line_intersect = _point_line_distance(p, r, p_intersect[..., None, :].expand_as(p))
+    intersect_dist_squared = ((p_line_intersect - p_intersect[..., None, :]) ** 2).sum(dim=-1)
+    return p_intersect, p_line_intersect, intersect_dist_squared, r
+def intersect_skew_lines_high_dim(p, r, mask=None):
+    # Implements https://en.wikipedia.org/wiki/Skew_lines In more than two dimensions
+    dim = p.shape[-1]
+    # make sure the heading vectors are l2-normed
+    if mask is None:
+        mask = torch.ones_like(p[..., 0])
+    r = torch.nn.functional.normalize(r, dim=-1)
+    eye = torch.eye(dim, device=p.device, dtype=p.dtype)[None, None]
+    I_min_cov = (eye - (r[..., None] * r[..., None, :])) * mask[..., None, None]
+    sum_proj = I_min_cov.matmul(p[..., None]).sum(dim=-3)
+    p_intersect = torch.linalg.lstsq(I_min_cov.sum(dim=-3), sum_proj).solution[..., 0]
+    if torch.any(torch.isnan(p_intersect)):
+        print(p_intersect)
+        raise ValueError(f"p_intersect is NaN")
+    return p_intersect, r
+def _point_line_distance(p1, r1, p2):
+    df = p2 - p1
+    proj_vector = df - ((df * r1).sum(dim=-1, keepdim=True) * r1)
+    line_pt_nearest = p2 - proj_vector
+    d = (proj_vector).norm(dim=-1)
+    return d, line_pt_nearest
+def first_camera_transform(cameras, rotation_only=False,
+                            points=None, pose_mode="C2W"):
+    """
+    Transform so that the first camera is the origin
+    """
+    new_cameras = cameras.clone()
+    # new_transform = new_cameras.get_world_to_view_transform()
+    R = cameras.R
+    T = cameras.T
+    Tran_M = torch.cat([R, T.unsqueeze(-1)], dim=-1) # [B, 3, 4]
+    Tran_M = torch.cat([Tran_M,
+                        torch.tensor([[[0, 0, 0, 1]]], device=Tran_M.device).expand(Tran_M.shape[0], -1, -1)], dim=1)
+    if pose_mode == "C2W":
+        Tran_M_new = (Tran_M[:1,...].inverse())@Tran_M
+    elif pose_mode == "W2C":
+        Tran_M_new = Tran_M@(Tran_M[:1,...].inverse())
+    if False:
+        tR = Rotate(new_cameras.R[0].unsqueeze(0))
+        if rotation_only:
+            t = tR.inverse()
+        else:
+            tT = Translate(new_cameras.T[0].unsqueeze(0))
+            t = tR.compose(tT).inverse()
+        if points is not None:
+            points = t.inverse().transform_points(points)
+        if pose_mode == "C2W":
+            new_matrix = new_transform.compose(t).get_matrix()
+        else:
+            import ipdb; ipdb.set_trace()
+            new_matrix = t.compose(new_transform).get_matrix()
+    new_cameras.R = Tran_M_new[:, :3, :3]
+    new_cameras.T = Tran_M_new[:, :3, 3]
+    return new_cameras, points
+def normalize_translation(new_cameras, points=None, max_norm=False):
+    t_gt = new_cameras.T.clone()
+    t_gt = t_gt[1:, :]
+    if max_norm:
+        t_gt_norm = torch.norm(t_gt, dim=(-1))
+        t_gt_scale = t_gt_norm.max()
+        if t_gt_norm.max() < 0.001:
+            t_gt_scale = torch.ones_like(t_gt_scale)
+        t_gt_scale = t_gt_scale.clamp(min=0.01, max=1e5)
+    else:
+        t_gt_norm = torch.norm(t_gt, dim=(0, 1))
+        t_gt_scale = t_gt_norm / math.sqrt(len(t_gt))
+        t_gt_scale = t_gt_scale / 2
+        if t_gt_norm.max()  < 0.001:
+            t_gt_scale = torch.ones_like(t_gt_scale)
+        t_gt_scale = t_gt_scale.clamp(min=0.01, max=1e5)
+    new_cameras.T = new_cameras.T / t_gt_scale
+    if points is not None:
+        points = points / t_gt_scale
+    return new_cameras, points, t_gt_scale

models/SpaTrackV2/models/depth_refiner/backbone.py ADDED Viewed

	@@ -0,0 +1,472 @@

+# ---------------------------------------------------------------
+# Copyright (c) 2021, NVIDIA Corporation. All rights reserved.
+#
+# This work is licensed under the NVIDIA Source Code License
+# ---------------------------------------------------------------
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from functools import partial
+from timm.layers import DropPath, to_2tuple, trunc_normal_
+from timm.models import register_model
+from timm.models.vision_transformer import _cfg
+import math
+class Mlp(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.dwconv = DWConv(hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+        self.apply(self._init_weights)
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+        elif isinstance(m, nn.Conv2d):
+            fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+            fan_out //= m.groups
+            m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
+            if m.bias is not None:
+                m.bias.data.zero_()
+    def forward(self, x, H, W):
+        x = self.fc1(x)
+        x = self.dwconv(x, H, W)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+class Attention(nn.Module):
+    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0., sr_ratio=1):
+        super().__init__()
+        assert dim % num_heads == 0, f"dim {dim} should be divided by num_heads {num_heads}."
+        self.dim = dim
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+        self.q = nn.Linear(dim, dim, bias=qkv_bias)
+        self.kv = nn.Linear(dim, dim * 2, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+        self.sr_ratio = sr_ratio
+        if sr_ratio > 1:
+            self.sr = nn.Conv2d(dim, dim, kernel_size=sr_ratio, stride=sr_ratio)
+            self.norm = nn.LayerNorm(dim)
+        self.apply(self._init_weights)
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+        elif isinstance(m, nn.Conv2d):
+            fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+            fan_out //= m.groups
+            m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
+            if m.bias is not None:
+                m.bias.data.zero_()
+    def forward(self, x, H, W):
+        B, N, C = x.shape
+        q = self.q(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
+        if self.sr_ratio > 1:
+            x_ = x.permute(0, 2, 1).reshape(B, C, H, W)
+            x_ = self.sr(x_).reshape(B, C, -1).permute(0, 2, 1)
+            x_ = self.norm(x_)
+            kv = self.kv(x_).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        else:
+            kv = self.kv(x).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        k, v = kv[0], kv[1]
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+class Block(nn.Module):
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, sr_ratio=1):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim,
+            num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,
+            attn_drop=attn_drop, proj_drop=drop, sr_ratio=sr_ratio)
+        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+        self.apply(self._init_weights)
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+        elif isinstance(m, nn.Conv2d):
+            fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+            fan_out //= m.groups
+            m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
+            if m.bias is not None:
+                m.bias.data.zero_()
+    def forward(self, x, H, W):
+        x = x + self.drop_path(self.attn(self.norm1(x), H, W))
+        x = x + self.drop_path(self.mlp(self.norm2(x), H, W))
+        return x
+class OverlapPatchEmbed(nn.Module):
+    """ Image to Patch Embedding
+    """
+    def __init__(self, img_size=224, patch_size=7, stride=4, in_chans=3, embed_dim=768):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.H, self.W = img_size[0] // patch_size[0], img_size[1] // patch_size[1]
+        self.num_patches = self.H * self.W
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride,
+                              padding=(patch_size[0] // 2, patch_size[1] // 2))
+        self.norm = nn.LayerNorm(embed_dim)
+        self.apply(self._init_weights)
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+        elif isinstance(m, nn.Conv2d):
+            fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+            fan_out //= m.groups
+            m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
+            if m.bias is not None:
+                m.bias.data.zero_()
+    def forward(self, x):
+        x = self.proj(x)
+        _, _, H, W = x.shape
+        x = x.flatten(2).transpose(1, 2)
+        x = self.norm(x)
+        return x, H, W
+class OverlapPatchEmbed43(nn.Module):
+    """ Image to Patch Embedding
+    """
+    def __init__(self, img_size=224, patch_size=7, stride=4, in_chans=3, embed_dim=768):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.H, self.W = img_size[0] // patch_size[0], img_size[1] // patch_size[1]
+        self.num_patches = self.H * self.W
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride,
+                              padding=(patch_size[0] // 2, patch_size[1] // 2))
+        self.norm = nn.LayerNorm(embed_dim)
+        self.apply(self._init_weights)
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+        elif isinstance(m, nn.Conv2d):
+            fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+            fan_out //= m.groups
+            m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
+            if m.bias is not None:
+                m.bias.data.zero_()
+    def forward(self, x):
+        if x.shape[1]==4:
+            x = self.proj_4c(x)
+        else:
+            x = self.proj(x)
+        _, _, H, W = x.shape
+        x = x.flatten(2).transpose(1, 2)
+        x = self.norm(x)
+        return x, H, W
+class MixVisionTransformer(nn.Module):
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dims=[64, 128, 256, 512],
+                 num_heads=[1, 2, 4, 8], mlp_ratios=[4, 4, 4, 4], qkv_bias=False, qk_scale=None, drop_rate=0.,
+                 attn_drop_rate=0., drop_path_rate=0., norm_layer=nn.LayerNorm,
+                 depths=[3, 4, 6, 3], sr_ratios=[8, 4, 2, 1]):
+        super().__init__()
+        self.num_classes = num_classes
+        self.depths = depths
+        # patch_embed    43
+        self.patch_embed1 = OverlapPatchEmbed(img_size=img_size, patch_size=7, stride=4, in_chans=in_chans,
+                                              embed_dim=embed_dims[0])
+        self.patch_embed2 = OverlapPatchEmbed(img_size=img_size // 4, patch_size=3, stride=2, in_chans=embed_dims[0],
+                                              embed_dim=embed_dims[1])
+        self.patch_embed3 = OverlapPatchEmbed(img_size=img_size // 8, patch_size=3, stride=2, in_chans=embed_dims[1],
+                                              embed_dim=embed_dims[2])
+        self.patch_embed4 = OverlapPatchEmbed(img_size=img_size // 16, patch_size=3, stride=2, in_chans=embed_dims[2],
+                                              embed_dim=embed_dims[3])
+        # transformer encoder
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]  # stochastic depth decay rule
+        cur = 0
+        self.block1 = nn.ModuleList([Block(
+            dim=embed_dims[0], num_heads=num_heads[0], mlp_ratio=mlp_ratios[0], qkv_bias=qkv_bias, qk_scale=qk_scale,
+            drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer,
+            sr_ratio=sr_ratios[0])
+            for i in range(depths[0])])
+        self.norm1 = norm_layer(embed_dims[0])
+        cur += depths[0]
+        self.block2 = nn.ModuleList([Block(
+            dim=embed_dims[1], num_heads=num_heads[1], mlp_ratio=mlp_ratios[1], qkv_bias=qkv_bias, qk_scale=qk_scale,
+            drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer,
+            sr_ratio=sr_ratios[1])
+            for i in range(depths[1])])
+        self.norm2 = norm_layer(embed_dims[1])
+        cur += depths[1]
+        self.block3 = nn.ModuleList([Block(
+            dim=embed_dims[2], num_heads=num_heads[2], mlp_ratio=mlp_ratios[2], qkv_bias=qkv_bias, qk_scale=qk_scale,
+            drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer,
+            sr_ratio=sr_ratios[2])
+            for i in range(depths[2])])
+        self.norm3 = norm_layer(embed_dims[2])
+        cur += depths[2]
+        self.block4 = nn.ModuleList([Block(
+            dim=embed_dims[3], num_heads=num_heads[3], mlp_ratio=mlp_ratios[3], qkv_bias=qkv_bias, qk_scale=qk_scale,
+            drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer,
+            sr_ratio=sr_ratios[3])
+            for i in range(depths[3])])
+        self.norm4 = norm_layer(embed_dims[3])
+        # classification head
+        # self.head = nn.Linear(embed_dims[3], num_classes) if num_classes > 0 else nn.Identity()
+        self.apply(self._init_weights)
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+        elif isinstance(m, nn.Conv2d):
+            fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+            fan_out //= m.groups
+            m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
+            if m.bias is not None:
+                m.bias.data.zero_()
+    def init_weights(self, pretrained=None):
+        if isinstance(pretrained, str):
+            logger = get_root_logger()
+            load_checkpoint(self, pretrained, map_location='cpu', strict=False, logger=logger)
+    def reset_drop_path(self, drop_path_rate):
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(self.depths))]
+        cur = 0
+        for i in range(self.depths[0]):
+            self.block1[i].drop_path.drop_prob = dpr[cur + i]
+        cur += self.depths[0]
+        for i in range(self.depths[1]):
+            self.block2[i].drop_path.drop_prob = dpr[cur + i]
+        cur += self.depths[1]
+        for i in range(self.depths[2]):
+            self.block3[i].drop_path.drop_prob = dpr[cur + i]
+        cur += self.depths[2]
+        for i in range(self.depths[3]):
+            self.block4[i].drop_path.drop_prob = dpr[cur + i]
+    def freeze_patch_emb(self):
+        self.patch_embed1.requires_grad = False
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'pos_embed1', 'pos_embed2', 'pos_embed3', 'pos_embed4', 'cls_token'}  # has pos_embed may be better
+    def get_classifier(self):
+        return self.head
+    def reset_classifier(self, num_classes, global_pool=''):
+        self.num_classes = num_classes
+        self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+    def forward_features(self, x):
+        B = x.shape[0]
+        outs = []
+        # stage 1
+        x, H, W = self.patch_embed1(x)
+        for i, blk in enumerate(self.block1):
+            x = blk(x, H, W)
+        x = self.norm1(x)
+        x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()
+        outs.append(x)
+        # stage 2
+        x, H, W = self.patch_embed2(x)
+        for i, blk in enumerate(self.block2):
+            x = blk(x, H, W)
+        x = self.norm2(x)
+        x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()
+        outs.append(x)
+        # stage 3
+        x, H, W = self.patch_embed3(x)
+        for i, blk in enumerate(self.block3):
+            x = blk(x, H, W)
+        x = self.norm3(x)
+        x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()
+        outs.append(x)
+        # stage 4
+        x, H, W = self.patch_embed4(x)
+        for i, blk in enumerate(self.block4):
+            x = blk(x, H, W)
+        x = self.norm4(x)
+        x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()
+        outs.append(x)
+        return outs
+    def forward(self, x):
+        if x.dim() == 5:
+            x = x.reshape(x.shape[0]*x.shape[1],x.shape[2],x.shape[3],x.shape[4])
+        x = self.forward_features(x)
+        # x = self.head(x)
+        return x
+class DWConv(nn.Module):
+    def __init__(self, dim=768):
+        super(DWConv, self).__init__()
+        self.dwconv = nn.Conv2d(dim, dim, 3, 1, 1, bias=True, groups=dim)
+    def forward(self, x, H, W):
+        B, N, C = x.shape
+        x = x.transpose(1, 2).view(B, C, H, W)
+        x = self.dwconv(x)
+        x = x.flatten(2).transpose(1, 2)
+        return x
+#@BACKBONES.register_module()
+class mit_b0(MixVisionTransformer):
+    def __init__(self, **kwargs):
+        super(mit_b0, self).__init__(
+            patch_size=4, embed_dims=[32, 64, 160, 256], num_heads=[1, 2, 5, 8], mlp_ratios=[4, 4, 4, 4],
+            qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[2, 2, 2, 2], sr_ratios=[8, 4, 2, 1],
+            drop_rate=0.0, drop_path_rate=0.1)
+#@BACKBONES.register_module()
+class mit_b1(MixVisionTransformer):
+    def __init__(self, **kwargs):
+        super(mit_b1, self).__init__(
+            patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[4, 4, 4, 4],
+            qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[2, 2, 2, 2], sr_ratios=[8, 4, 2, 1],
+            drop_rate=0.0, drop_path_rate=0.1)
+#@BACKBONES.register_module()
+class mit_b2(MixVisionTransformer):
+    def __init__(self, **kwargs):
+        super(mit_b2, self).__init__(
+            patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[4, 4, 4, 4],
+            qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 4, 6, 3], sr_ratios=[8, 4, 2, 1],
+            drop_rate=0.0, drop_path_rate=0.1)
+#@BACKBONES.register_module()
+class mit_b3(MixVisionTransformer):
+    def __init__(self, **kwargs):
+        super(mit_b3, self).__init__(
+            patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[4, 4, 4, 4],
+            qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 4, 18, 3], sr_ratios=[8, 4, 2, 1],
+            drop_rate=0.0, drop_path_rate=0.1)
+#@BACKBONES.register_module()
+class mit_b4(MixVisionTransformer):
+    def __init__(self, **kwargs):
+        super(mit_b4, self).__init__(
+            patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[4, 4, 4, 4],
+            qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 8, 27, 3], sr_ratios=[8, 4, 2, 1],
+            drop_rate=0.0, drop_path_rate=0.1)
+#@BACKBONES.register_module()
+class mit_b5(MixVisionTransformer):
+    def __init__(self, **kwargs):
+        super(mit_b5, self).__init__(
+            patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[4, 4, 4, 4],
+            qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 6, 40, 3], sr_ratios=[8, 4, 2, 1],
+            drop_rate=0.0, drop_path_rate=0.1)

models/SpaTrackV2/models/depth_refiner/decode_head.py ADDED Viewed

	@@ -0,0 +1,619 @@

+from abc import ABCMeta, abstractmethod
+import torch
+import torch.nn as nn
+# from mmcv.cnn import normal_init
+# from mmcv.runner import auto_fp16, force_fp32
+# from mmseg.core import build_pixel_sampler
+# from mmseg.ops import resize
+class BaseDecodeHead(nn.Module, metaclass=ABCMeta):
+    """Base class for BaseDecodeHead.
+    Args:
+        in_channels (int|Sequence[int]): Input channels.
+        channels (int): Channels after modules, before conv_seg.
+        num_classes (int): Number of classes.
+        dropout_ratio (float): Ratio of dropout layer. Default: 0.1.
+        conv_cfg (dict|None): Config of conv layers. Default: None.
+        norm_cfg (dict|None): Config of norm layers. Default: None.
+        act_cfg (dict): Config of activation layers.
+            Default: dict(type='ReLU')
+        in_index (int|Sequence[int]): Input feature index. Default: -1
+        input_transform (str|None): Transformation type of input features.
+            Options: 'resize_concat', 'multiple_select', None.
+            'resize_concat': Multiple feature maps will be resize to the
+                same size as first one and than concat together.
+                Usually used in FCN head of HRNet.
+            'multiple_select': Multiple feature maps will be bundle into
+                a list and passed into decode head.
+            None: Only one select feature map is allowed.
+            Default: None.
+        loss_decode (dict): Config of decode loss.
+            Default: dict(type='CrossEntropyLoss').
+        ignore_index (int | None): The label index to be ignored. When using
+            masked BCE loss, ignore_index should be set to None. Default: 255
+        sampler (dict|None): The config of segmentation map sampler.
+            Default: None.
+        align_corners (bool): align_corners argument of F.interpolate.
+            Default: False.
+    """
+    def __init__(self,
+                 in_channels,
+                 channels,
+                 *,
+                 num_classes,
+                 dropout_ratio=0.1,
+                 conv_cfg=None,
+                 norm_cfg=None,
+                 act_cfg=dict(type='ReLU'),
+                 in_index=-1,
+                 input_transform=None,
+                 loss_decode=dict(
+                     type='CrossEntropyLoss',
+                     use_sigmoid=False,
+                     loss_weight=1.0),
+                 decoder_params=None,
+                 ignore_index=255,
+                 sampler=None,
+                 align_corners=False):
+        super(BaseDecodeHead, self).__init__()
+        self._init_inputs(in_channels, in_index, input_transform)
+        self.channels = channels
+        self.num_classes = num_classes
+        self.dropout_ratio = dropout_ratio
+        self.conv_cfg = conv_cfg
+        self.norm_cfg = norm_cfg
+        self.act_cfg = act_cfg
+        self.in_index = in_index
+        self.ignore_index = ignore_index
+        self.align_corners = align_corners
+        if sampler is not None:
+            self.sampler = build_pixel_sampler(sampler, context=self)
+        else:
+            self.sampler = None
+        self.conv_seg = nn.Conv2d(channels, num_classes, kernel_size=1)
+        if dropout_ratio > 0:
+            self.dropout = nn.Dropout2d(dropout_ratio)
+        else:
+            self.dropout = None
+        self.fp16_enabled = False
+    def extra_repr(self):
+        """Extra repr."""
+        s = f'input_transform={self.input_transform}, ' \
+            f'ignore_index={self.ignore_index}, ' \
+            f'align_corners={self.align_corners}'
+        return s
+    def _init_inputs(self, in_channels, in_index, input_transform):
+        """Check and initialize input transforms.
+        The in_channels, in_index and input_transform must match.
+        Specifically, when input_transform is None, only single feature map
+        will be selected. So in_channels and in_index must be of type int.
+        When input_transform
+        Args:
+            in_channels (int|Sequence[int]): Input channels.
+            in_index (int|Sequence[int]): Input feature index.
+            input_transform (str|None): Transformation type of input features.
+                Options: 'resize_concat', 'multiple_select', None.
+                'resize_concat': Multiple feature maps will be resize to the
+                    same size as first one and than concat together.
+                    Usually used in FCN head of HRNet.
+                'multiple_select': Multiple feature maps will be bundle into
+                    a list and passed into decode head.
+                None: Only one select feature map is allowed.
+        """
+        if input_transform is not None:
+            assert input_transform in ['resize_concat', 'multiple_select']
+        self.input_transform = input_transform
+        self.in_index = in_index
+        if input_transform is not None:
+            assert isinstance(in_channels, (list, tuple))
+            assert isinstance(in_index, (list, tuple))
+            assert len(in_channels) == len(in_index)
+            if input_transform == 'resize_concat':
+                self.in_channels = sum(in_channels)
+            else:
+                self.in_channels = in_channels
+        else:
+            assert isinstance(in_channels, int)
+            assert isinstance(in_index, int)
+            self.in_channels = in_channels
+    def init_weights(self):
+        """Initialize weights of classification layer."""
+        normal_init(self.conv_seg, mean=0, std=0.01)
+    def _transform_inputs(self, inputs):
+        """Transform inputs for decoder.
+        Args:
+            inputs (list[Tensor]): List of multi-level img features.
+        Returns:
+            Tensor: The transformed inputs
+        """
+        if self.input_transform == 'resize_concat':
+            inputs = [inputs[i] for i in self.in_index]
+            upsampled_inputs = [
+                resize(
+                    input=x,
+                    size=inputs[0].shape[2:],
+                    mode='bilinear',
+                    align_corners=self.align_corners) for x in inputs
+            ]
+            inputs = torch.cat(upsampled_inputs, dim=1)
+        elif self.input_transform == 'multiple_select':
+            inputs = [inputs[i] for i in self.in_index]
+        else:
+            inputs = inputs[self.in_index]
+        return inputs
+    # @auto_fp16()
+    @abstractmethod
+    def forward(self, inputs):
+        """Placeholder of forward function."""
+        pass
+    def forward_train(self, inputs, img_metas, gt_semantic_seg, train_cfg):
+        """Forward function for training.
+        Args:
+            inputs (list[Tensor]): List of multi-level img features.
+            img_metas (list[dict]): List of image info dict where each dict
+                has: 'img_shape', 'scale_factor', 'flip', and may also contain
+                'filename', 'ori_shape', 'pad_shape', and 'img_norm_cfg'.
+                For details on the values of these keys see
+                `mmseg/datasets/pipelines/formatting.py:Collect`.
+            gt_semantic_seg (Tensor): Semantic segmentation masks
+                used if the architecture supports semantic segmentation task.
+            train_cfg (dict): The training config.
+        Returns:
+            dict[str, Tensor]: a dictionary of loss components
+        """
+        seg_logits = self.forward(inputs)
+        losses = self.losses(seg_logits, gt_semantic_seg)
+        return losses
+    def forward_test(self, inputs, img_metas, test_cfg):
+        """Forward function for testing.
+        Args:
+            inputs (list[Tensor]): List of multi-level img features.
+            img_metas (list[dict]): List of image info dict where each dict
+                has: 'img_shape', 'scale_factor', 'flip', and may also contain
+                'filename', 'ori_shape', 'pad_shape', and 'img_norm_cfg'.
+                For details on the values of these keys see
+                `mmseg/datasets/pipelines/formatting.py:Collect`.
+            test_cfg (dict): The testing config.
+        Returns:
+            Tensor: Output segmentation map.
+        """
+        return self.forward(inputs)
+    def cls_seg(self, feat):
+        """Classify each pixel."""
+        if self.dropout is not None:
+            feat = self.dropout(feat)
+        output = self.conv_seg(feat)
+        return output
+class BaseDecodeHead_clips(nn.Module, metaclass=ABCMeta):
+    """Base class for BaseDecodeHead_clips.
+    Args:
+        in_channels (int|Sequence[int]): Input channels.
+        channels (int): Channels after modules, before conv_seg.
+        num_classes (int): Number of classes.
+        dropout_ratio (float): Ratio of dropout layer. Default: 0.1.
+        conv_cfg (dict|None): Config of conv layers. Default: None.
+        norm_cfg (dict|None): Config of norm layers. Default: None.
+        act_cfg (dict): Config of activation layers.
+            Default: dict(type='ReLU')
+        in_index (int|Sequence[int]): Input feature index. Default: -1
+        input_transform (str|None): Transformation type of input features.
+            Options: 'resize_concat', 'multiple_select', None.
+            'resize_concat': Multiple feature maps will be resize to the
+                same size as first one and than concat together.
+                Usually used in FCN head of HRNet.
+            'multiple_select': Multiple feature maps will be bundle into
+                a list and passed into decode head.
+            None: Only one select feature map is allowed.
+            Default: None.
+        loss_decode (dict): Config of decode loss.
+            Default: dict(type='CrossEntropyLoss').
+        ignore_index (int | None): The label index to be ignored. When using
+            masked BCE loss, ignore_index should be set to None. Default: 255
+        sampler (dict|None): The config of segmentation map sampler.
+            Default: None.
+        align_corners (bool): align_corners argument of F.interpolate.
+            Default: False.
+    """
+    def __init__(self,
+                 in_channels,
+                 channels,
+                 *,
+                 num_classes,
+                 dropout_ratio=0.1,
+                 conv_cfg=None,
+                 norm_cfg=None,
+                 act_cfg=dict(type='ReLU'),
+                 in_index=-1,
+                 input_transform=None,
+                 loss_decode=dict(
+                     type='CrossEntropyLoss',
+                     use_sigmoid=False,
+                     loss_weight=1.0),
+                 decoder_params=None,
+                 ignore_index=255,
+                 sampler=None,
+                 align_corners=False,
+                 num_clips=5):
+        super(BaseDecodeHead_clips, self).__init__()
+        self._init_inputs(in_channels, in_index, input_transform)
+        self.channels = channels
+        self.num_classes = num_classes
+        self.dropout_ratio = dropout_ratio
+        self.conv_cfg = conv_cfg
+        self.norm_cfg = norm_cfg
+        self.act_cfg = act_cfg
+        self.in_index = in_index
+        self.ignore_index = ignore_index
+        self.align_corners = align_corners
+        self.num_clips=num_clips
+        if sampler is not None:
+            self.sampler = build_pixel_sampler(sampler, context=self)
+        else:
+            self.sampler = None
+        self.conv_seg = nn.Conv2d(channels, num_classes, kernel_size=1)
+        if dropout_ratio > 0:
+            self.dropout = nn.Dropout2d(dropout_ratio)
+        else:
+            self.dropout = None
+        self.fp16_enabled = False
+    def extra_repr(self):
+        """Extra repr."""
+        s = f'input_transform={self.input_transform}, ' \
+            f'ignore_index={self.ignore_index}, ' \
+            f'align_corners={self.align_corners}'
+        return s
+    def _init_inputs(self, in_channels, in_index, input_transform):
+        """Check and initialize input transforms.
+        The in_channels, in_index and input_transform must match.
+        Specifically, when input_transform is None, only single feature map
+        will be selected. So in_channels and in_index must be of type int.
+        When input_transform
+        Args:
+            in_channels (int|Sequence[int]): Input channels.
+            in_index (int|Sequence[int]): Input feature index.
+            input_transform (str|None): Transformation type of input features.
+                Options: 'resize_concat', 'multiple_select', None.
+                'resize_concat': Multiple feature maps will be resize to the
+                    same size as first one and than concat together.
+                    Usually used in FCN head of HRNet.
+                'multiple_select': Multiple feature maps will be bundle into
+                    a list and passed into decode head.
+                None: Only one select feature map is allowed.
+        """
+        if input_transform is not None:
+            assert input_transform in ['resize_concat', 'multiple_select']
+        self.input_transform = input_transform
+        self.in_index = in_index
+        if input_transform is not None:
+            assert isinstance(in_channels, (list, tuple))
+            assert isinstance(in_index, (list, tuple))
+            assert len(in_channels) == len(in_index)
+            if input_transform == 'resize_concat':
+                self.in_channels = sum(in_channels)
+            else:
+                self.in_channels = in_channels
+        else:
+            assert isinstance(in_channels, int)
+            assert isinstance(in_index, int)
+            self.in_channels = in_channels
+    def init_weights(self):
+        """Initialize weights of classification layer."""
+        normal_init(self.conv_seg, mean=0, std=0.01)
+    def _transform_inputs(self, inputs):
+        """Transform inputs for decoder.
+        Args:
+            inputs (list[Tensor]): List of multi-level img features.
+        Returns:
+            Tensor: The transformed inputs
+        """
+        if self.input_transform == 'resize_concat':
+            inputs = [inputs[i] for i in self.in_index]
+            upsampled_inputs = [
+                resize(
+                    input=x,
+                    size=inputs[0].shape[2:],
+                    mode='bilinear',
+                    align_corners=self.align_corners) for x in inputs
+            ]
+            inputs = torch.cat(upsampled_inputs, dim=1)
+        elif self.input_transform == 'multiple_select':
+            inputs = [inputs[i] for i in self.in_index]
+        else:
+            inputs = inputs[self.in_index]
+        return inputs
+    # @auto_fp16()
+    @abstractmethod
+    def forward(self, inputs):
+        """Placeholder of forward function."""
+        pass
+    def forward_train(self, inputs, img_metas, gt_semantic_seg, train_cfg,batch_size, num_clips):
+        """Forward function for training.
+        Args:
+            inputs (list[Tensor]): List of multi-level img features.
+            img_metas (list[dict]): List of image info dict where each dict
+                has: 'img_shape', 'scale_factor', 'flip', and may also contain
+                'filename', 'ori_shape', 'pad_shape', and 'img_norm_cfg'.
+                For details on the values of these keys see
+                `mmseg/datasets/pipelines/formatting.py:Collect`.
+            gt_semantic_seg (Tensor): Semantic segmentation masks
+                used if the architecture supports semantic segmentation task.
+            train_cfg (dict): The training config.
+        Returns:
+            dict[str, Tensor]: a dictionary of loss components
+        """
+        seg_logits = self.forward(inputs,batch_size, num_clips)
+        losses = self.losses(seg_logits, gt_semantic_seg)
+        return losses
+    def forward_test(self, inputs, img_metas, test_cfg, batch_size, num_clips):
+        """Forward function for testing.
+        Args:
+            inputs (list[Tensor]): List of multi-level img features.
+            img_metas (list[dict]): List of image info dict where each dict
+                has: 'img_shape', 'scale_factor', 'flip', and may also contain
+                'filename', 'ori_shape', 'pad_shape', and 'img_norm_cfg'.
+                For details on the values of these keys see
+                `mmseg/datasets/pipelines/formatting.py:Collect`.
+            test_cfg (dict): The testing config.
+        Returns:
+            Tensor: Output segmentation map.
+        """
+        return self.forward(inputs, batch_size, num_clips)
+    def cls_seg(self, feat):
+        """Classify each pixel."""
+        if self.dropout is not None:
+            feat = self.dropout(feat)
+        output = self.conv_seg(feat)
+        return output
+class BaseDecodeHead_clips_flow(nn.Module, metaclass=ABCMeta):
+    """Base class for BaseDecodeHead_clips_flow.
+    Args:
+        in_channels (int|Sequence[int]): Input channels.
+        channels (int): Channels after modules, before conv_seg.
+        num_classes (int): Number of classes.
+        dropout_ratio (float): Ratio of dropout layer. Default: 0.1.
+        conv_cfg (dict|None): Config of conv layers. Default: None.
+        norm_cfg (dict|None): Config of norm layers. Default: None.
+        act_cfg (dict): Config of activation layers.
+            Default: dict(type='ReLU')
+        in_index (int|Sequence[int]): Input feature index. Default: -1
+        input_transform (str|None): Transformation type of input features.
+            Options: 'resize_concat', 'multiple_select', None.
+            'resize_concat': Multiple feature maps will be resize to the
+                same size as first one and than concat together.
+                Usually used in FCN head of HRNet.
+            'multiple_select': Multiple feature maps will be bundle into
+                a list and passed into decode head.
+            None: Only one select feature map is allowed.
+            Default: None.
+        loss_decode (dict): Config of decode loss.
+            Default: dict(type='CrossEntropyLoss').
+        ignore_index (int | None): The label index to be ignored. When using
+            masked BCE loss, ignore_index should be set to None. Default: 255
+        sampler (dict|None): The config of segmentation map sampler.
+            Default: None.
+        align_corners (bool): align_corners argument of F.interpolate.
+            Default: False.
+    """
+    def __init__(self,
+                 in_channels,
+                 channels,
+                 *,
+                 num_classes,
+                 dropout_ratio=0.1,
+                 conv_cfg=None,
+                 norm_cfg=None,
+                 act_cfg=dict(type='ReLU'),
+                 in_index=-1,
+                 input_transform=None,
+                 loss_decode=dict(
+                     type='CrossEntropyLoss',
+                     use_sigmoid=False,
+                     loss_weight=1.0),
+                 decoder_params=None,
+                 ignore_index=255,
+                 sampler=None,
+                 align_corners=False,
+                 num_clips=5):
+        super(BaseDecodeHead_clips_flow, self).__init__()
+        self._init_inputs(in_channels, in_index, input_transform)
+        self.channels = channels
+        self.num_classes = num_classes
+        self.dropout_ratio = dropout_ratio
+        self.conv_cfg = conv_cfg
+        self.norm_cfg = norm_cfg
+        self.act_cfg = act_cfg
+        self.in_index = in_index
+        self.ignore_index = ignore_index
+        self.align_corners = align_corners
+        self.num_clips=num_clips
+        if sampler is not None:
+            self.sampler = build_pixel_sampler(sampler, context=self)
+        else:
+            self.sampler = None
+        self.conv_seg = nn.Conv2d(channels, num_classes, kernel_size=1)
+        if dropout_ratio > 0:
+            self.dropout = nn.Dropout2d(dropout_ratio)
+        else:
+            self.dropout = None
+        self.fp16_enabled = False
+    def extra_repr(self):
+        """Extra repr."""
+        s = f'input_transform={self.input_transform}, ' \
+            f'ignore_index={self.ignore_index}, ' \
+            f'align_corners={self.align_corners}'
+        return s
+    def _init_inputs(self, in_channels, in_index, input_transform):
+        """Check and initialize input transforms.
+        The in_channels, in_index and input_transform must match.
+        Specifically, when input_transform is None, only single feature map
+        will be selected. So in_channels and in_index must be of type int.
+        When input_transform
+        Args:
+            in_channels (int|Sequence[int]): Input channels.
+            in_index (int|Sequence[int]): Input feature index.
+            input_transform (str|None): Transformation type of input features.
+                Options: 'resize_concat', 'multiple_select', None.
+                'resize_concat': Multiple feature maps will be resize to the
+                    same size as first one and than concat together.
+                    Usually used in FCN head of HRNet.
+                'multiple_select': Multiple feature maps will be bundle into
+                    a list and passed into decode head.
+                None: Only one select feature map is allowed.
+        """
+        if input_transform is not None:
+            assert input_transform in ['resize_concat', 'multiple_select']
+        self.input_transform = input_transform
+        self.in_index = in_index
+        if input_transform is not None:
+            assert isinstance(in_channels, (list, tuple))
+            assert isinstance(in_index, (list, tuple))
+            assert len(in_channels) == len(in_index)
+            if input_transform == 'resize_concat':
+                self.in_channels = sum(in_channels)
+            else:
+                self.in_channels = in_channels
+        else:
+            assert isinstance(in_channels, int)
+            assert isinstance(in_index, int)
+            self.in_channels = in_channels
+    def init_weights(self):
+        """Initialize weights of classification layer."""
+        normal_init(self.conv_seg, mean=0, std=0.01)
+    def _transform_inputs(self, inputs):
+        """Transform inputs for decoder.
+        Args:
+            inputs (list[Tensor]): List of multi-level img features.
+        Returns:
+            Tensor: The transformed inputs
+        """
+        if self.input_transform == 'resize_concat':
+            inputs = [inputs[i] for i in self.in_index]
+            upsampled_inputs = [
+                resize(
+                    input=x,
+                    size=inputs[0].shape[2:],
+                    mode='bilinear',
+                    align_corners=self.align_corners) for x in inputs
+            ]
+            inputs = torch.cat(upsampled_inputs, dim=1)
+        elif self.input_transform == 'multiple_select':
+            inputs = [inputs[i] for i in self.in_index]
+        else:
+            inputs = inputs[self.in_index]
+        return inputs
+    # @auto_fp16()
+    @abstractmethod
+    def forward(self, inputs):
+        """Placeholder of forward function."""
+        pass
+    def forward_train(self, inputs, img_metas, gt_semantic_seg, train_cfg,batch_size, num_clips,img=None):
+        """Forward function for training.
+        Args:
+            inputs (list[Tensor]): List of multi-level img features.
+            img_metas (list[dict]): List of image info dict where each dict
+                has: 'img_shape', 'scale_factor', 'flip', and may also contain
+                'filename', 'ori_shape', 'pad_shape', and 'img_norm_cfg'.
+                For details on the values of these keys see
+                `mmseg/datasets/pipelines/formatting.py:Collect`.
+            gt_semantic_seg (Tensor): Semantic segmentation masks
+                used if the architecture supports semantic segmentation task.
+            train_cfg (dict): The training config.
+        Returns:
+            dict[str, Tensor]: a dictionary of loss components
+        """
+        seg_logits = self.forward(inputs,batch_size, num_clips,img)
+        losses = self.losses(seg_logits, gt_semantic_seg)
+        return losses
+    def forward_test(self, inputs, img_metas, test_cfg, batch_size=None, num_clips=None, img=None):
+        """Forward function for testing.
+        Args:
+            inputs (list[Tensor]): List of multi-level img features.
+            img_metas (list[dict]): List of image info dict where each dict
+                has: 'img_shape', 'scale_factor', 'flip', and may also contain
+                'filename', 'ori_shape', 'pad_shape', and 'img_norm_cfg'.
+                For details on the values of these keys see
+                `mmseg/datasets/pipelines/formatting.py:Collect`.
+            test_cfg (dict): The testing config.
+        Returns:
+            Tensor: Output segmentation map.
+        """
+        return self.forward(inputs, batch_size, num_clips,img)
+    def cls_seg(self, feat):
+        """Classify each pixel."""
+        if self.dropout is not None:
+            feat = self.dropout(feat)
+        output = self.conv_seg(feat)
+        return output

models/SpaTrackV2/models/depth_refiner/depth_refiner.py ADDED Viewed

	@@ -0,0 +1,115 @@

+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from models.monoD.depth_anything_v2.dinov2_layers.patch_embed import PatchEmbed
+from models.SpaTrackV2.models.depth_refiner.backbone import mit_b3
+from models.SpaTrackV2.models.depth_refiner.stablizer import Stabilization_Network_Cross_Attention
+from einops import rearrange
+class TrackStablizer(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.backbone = mit_b3()
+        old_conv = self.backbone.patch_embed1.proj
+        new_conv = nn.Conv2d(old_conv.in_channels + 4, old_conv.out_channels, kernel_size=old_conv.kernel_size, stride=old_conv.stride, padding=old_conv.padding)
+        new_conv.weight[:, :3, :, :].data.copy_(old_conv.weight.clone())
+        self.backbone.patch_embed1.proj = new_conv
+        self.Track_Stabilizer = Stabilization_Network_Cross_Attention(in_channels=[64, 128, 320, 512],
+                    in_index=[0, 1, 2, 3],
+                    feature_strides=[4, 8, 16, 32],
+                    channels=128,
+                    dropout_ratio=0.1,
+                    num_classes=1,
+                    align_corners=False,
+                    decoder_params=dict(embed_dim=256, depths=4),
+                    num_clips=16,
+                    norm_cfg = dict(type='SyncBN', requires_grad=True))
+        self.edge_conv = nn.Sequential(nn.Conv2d(in_channels=4, out_channels=64, kernel_size=3, padding=1, stride=1, bias=True),\
+                                  nn.ReLU(inplace=True))
+        self.edge_conv1 = nn.Sequential(nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, padding=1, stride=2, bias=True),\
+                                  nn.ReLU(inplace=True))
+        self.success = False
+        self.x = None
+    def buffer_forward(self, inputs, num_clips=16):
+        """
+            buffer forward for getting the pointmap and image features
+        """
+        B, T, C, H, W = inputs.shape
+        self.x = self.backbone(inputs)
+        scale, shift = self.Track_Stabilizer.buffer_forward(self.x, num_clips=num_clips)
+        self.success = True
+        return scale, shift
+    def forward(self, inputs, tracks, tracks_uvd, num_clips=16, imgs=None, vis_track=None):
+        """
+        Args:
+            inputs: [B, T, C, H, W], RGB + PointMap + Mask
+            tracks: [B, T, N, 4], 3D tracks in camera coordinate + visibility
+            num_clips: int, number of clips to use
+        """
+        B, T, C, H, W = inputs.shape
+        edge_feat = self.edge_conv(inputs.view(B*T,4,H,W))
+        edge_feat1 = self.edge_conv1(edge_feat)
+        if not self.success:
+            scale, shift = self.Track_Stabilizer.buffer_forward(self.x,num_clips=num_clips)
+            self.success = True
+            update = self.Track_Stabilizer(self.x,edge_feat,edge_feat1,tracks,tracks_uvd,num_clips=num_clips, imgs=imgs, vis_track=vis_track)
+        else:
+            update = self.Track_Stabilizer(self.x,edge_feat,edge_feat1,tracks,tracks_uvd,num_clips=num_clips, imgs=imgs, vis_track=vis_track)
+        return update
+    def reset_success(self):
+        self.success = False
+        self.x = None
+        self.Track_Stabilizer.reset_success()
+if __name__ == "__main__":
+    # Create test input tensors
+    batch_size = 1
+    seq_len = 16
+    channels = 7  # 3 for RGB + 3 for PointMap + 1 for Mask
+    height = 384
+    width = 512
+    # Create random input tensor with shape [B, T, C, H, W]
+    inputs = torch.randn(batch_size, seq_len, channels, height, width)
+    # Create random tracks
+    tracks = torch.randn(batch_size, seq_len, 1024, 4)
+    # Create random test images
+    test_imgs = torch.randn(batch_size, seq_len, 3, height, width)
+    # Initialize model and move to GPU
+    model = TrackStablizer().cuda()
+    # Move inputs to GPU and run forward pass
+    inputs = inputs.cuda()
+    tracks = tracks.cuda()
+    outputs = model.buffer_forward(inputs, num_clips=seq_len)
+    import time
+    start_time = time.time()
+    outputs = model(inputs, tracks, num_clips=seq_len)
+    end_time = time.time()
+    print(f"Time taken: {end_time - start_time} seconds")
+    import pdb; pdb.set_trace()
+    # # Print shapes for verification
+    # print(f"Input shape: {inputs.shape}")
+    # print(f"Output shape: {outputs.shape}")
+    # # Basic tests
+    # assert outputs.shape[0] == batch_size, "Batch size mismatch"
+    # assert len(outputs.shape) == 4, "Output should be 4D: [B,C,H,W]"
+    # assert torch.all(outputs >= 0), "Output should be non-negative after ReLU"
+    # print("All tests passed!")

models/SpaTrackV2/models/depth_refiner/network.py ADDED Viewed

	@@ -0,0 +1,429 @@

+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+'''
+Author: Ke Xian
+Email: kexian@hust.edu.cn
+Date: 2020/07/20
+'''
+import torch
+import torch.nn as nn
+import torch.nn.init as init
+# ==============================================================================================================
+class FTB(nn.Module):
+    def __init__(self, inchannels, midchannels=512):
+        super(FTB, self).__init__()
+        self.in1 = inchannels
+        self.mid = midchannels
+        self.conv1 = nn.Conv2d(in_channels=self.in1, out_channels=self.mid, kernel_size=3, padding=1, stride=1, bias=True)
+        self.conv_branch = nn.Sequential(nn.ReLU(inplace=True),\
+                                         nn.Conv2d(in_channels=self.mid, out_channels=self.mid, kernel_size=3, padding=1, stride=1, bias=True),\
+                                         #nn.BatchNorm2d(num_features=self.mid),\
+                                         nn.ReLU(inplace=True),\
+                                         nn.Conv2d(in_channels=self.mid, out_channels= self.mid, kernel_size=3, padding=1, stride=1, bias=True))
+        self.relu = nn.ReLU(inplace=True)
+        self.init_params()
+    def forward(self, x):
+        x = self.conv1(x)
+        x = x + self.conv_branch(x)
+        x = self.relu(x)
+        return x
+    def init_params(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                #init.kaiming_normal_(m.weight, mode='fan_out')
+                init.normal_(m.weight, std=0.01)
+                # init.xavier_normal_(m.weight)
+                if m.bias is not None:
+                    init.constant_(m.bias, 0)
+            elif isinstance(m, nn.ConvTranspose2d):
+                #init.kaiming_normal_(m.weight, mode='fan_out')
+                init.normal_(m.weight, std=0.01)
+                # init.xavier_normal_(m.weight)
+                if m.bias is not None:
+                    init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm2d):  #nn.BatchNorm2d
+                init.constant_(m.weight, 1)
+                init.constant_(m.bias, 0)
+            elif isinstance(m, nn.Linear):
+                init.normal_(m.weight, std=0.01)
+                if m.bias is not None:
+                    init.constant_(m.bias, 0)
+class ATA(nn.Module):
+    def __init__(self, inchannels, reduction = 8):
+        super(ATA, self).__init__()
+        self.inchannels = inchannels
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.fc = nn.Sequential(nn.Linear(self.inchannels*2, self.inchannels // reduction),
+                                nn.ReLU(inplace=True),
+                                nn.Linear(self.inchannels // reduction, self.inchannels),
+                                nn.Sigmoid())
+        self.init_params()
+    def forward(self, low_x, high_x):
+        n, c, _, _ = low_x.size()
+        x = torch.cat([low_x, high_x], 1)
+        x = self.avg_pool(x)
+        x = x.view(n, -1)
+        x = self.fc(x).view(n,c,1,1)
+        x = low_x * x + high_x
+        return x
+    def init_params(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                #init.kaiming_normal_(m.weight, mode='fan_out')
+                #init.normal(m.weight, std=0.01)
+                init.xavier_normal_(m.weight)
+                if m.bias is not None:
+                    init.constant_(m.bias, 0)
+            elif isinstance(m, nn.ConvTranspose2d):
+                #init.kaiming_normal_(m.weight, mode='fan_out')
+                #init.normal_(m.weight, std=0.01)
+                init.xavier_normal_(m.weight)
+                if m.bias is not None:
+                    init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm2d): #nn.BatchNorm2d
+                init.constant_(m.weight, 1)
+                init.constant_(m.bias, 0)
+            elif isinstance(m, nn.Linear):
+                init.normal_(m.weight, std=0.01)
+                if m.bias is not None:
+                    init.constant_(m.bias, 0)
+class FFM(nn.Module):
+    def __init__(self, inchannels, midchannels, outchannels, upfactor=2):
+        super(FFM, self).__init__()
+        self.inchannels = inchannels
+        self.midchannels = midchannels
+        self.outchannels = outchannels
+        self.upfactor = upfactor
+        self.ftb1 = FTB(inchannels=self.inchannels, midchannels=self.midchannels)
+        self.ftb2 = FTB(inchannels=self.midchannels, midchannels=self.outchannels)
+        self.upsample = nn.Upsample(scale_factor=self.upfactor, mode='bilinear', align_corners=True)
+        self.init_params()
+        #self.p1 = nn.Conv2d(512, 256, kernel_size=1, padding=0, bias=False)
+        #self.p2 = nn.Conv2d(512, 256, kernel_size=1, padding=0, bias=False)
+        #self.p3 = nn.Conv2d(512, 256, kernel_size=1, padding=0, bias=False)
+    def forward(self, low_x, high_x):
+        x = self.ftb1(low_x)
+        '''
+        x = torch.cat((x,high_x),1)
+        if x.shape[2] == 12:
+            x = self.p1(x)
+        elif x.shape[2] == 24:
+            x = self.p2(x)
+        elif x.shape[2] == 48:
+            x = self.p3(x)
+        '''
+        x = x + high_x            ###high_x
+        x = self.ftb2(x)
+        x = self.upsample(x)
+        return x
+    def init_params(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                #init.kaiming_normal_(m.weight, mode='fan_out')
+                init.normal_(m.weight, std=0.01)
+                #init.xavier_normal_(m.weight)
+                if m.bias is not None:
+                    init.constant_(m.bias, 0)
+            elif isinstance(m, nn.ConvTranspose2d):
+                #init.kaiming_normal_(m.weight, mode='fan_out')
+                init.normal_(m.weight, std=0.01)
+                #init.xavier_normal_(m.weight)
+                if m.bias is not None:
+                    init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm2d): #nn.Batchnorm2d
+                init.constant_(m.weight, 1)
+                init.constant_(m.bias, 0)
+            elif isinstance(m, nn.Linear):
+                init.normal_(m.weight, std=0.01)
+                if m.bias is not None:
+                    init.constant_(m.bias, 0)
+class noFFM(nn.Module):
+    def __init__(self, inchannels, midchannels, outchannels, upfactor=2):
+        super(noFFM, self).__init__()
+        self.inchannels = inchannels
+        self.midchannels = midchannels
+        self.outchannels = outchannels
+        self.upfactor = upfactor
+        self.ftb2 = FTB(inchannels=self.midchannels, midchannels=self.outchannels)
+        self.upsample = nn.Upsample(scale_factor=self.upfactor, mode='bilinear', align_corners=True)
+        self.init_params()
+        #self.p1 = nn.Conv2d(512, 256, kernel_size=1, padding=0, bias=False)
+        #self.p2 = nn.Conv2d(512, 256, kernel_size=1, padding=0, bias=False)
+        #self.p3 = nn.Conv2d(512, 256, kernel_size=1, padding=0, bias=False)
+    def forward(self, low_x, high_x):
+        #x = self.ftb1(low_x)
+        x = high_x            ###high_x
+        x = self.ftb2(x)
+        x = self.upsample(x)
+        return x
+    def init_params(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                #init.kaiming_normal_(m.weight, mode='fan_out')
+                init.normal_(m.weight, std=0.01)
+                #init.xavier_normal_(m.weight)
+                if m.bias is not None:
+                    init.constant_(m.bias, 0)
+            elif isinstance(m, nn.ConvTranspose2d):
+                #init.kaiming_normal_(m.weight, mode='fan_out')
+                init.normal_(m.weight, std=0.01)
+                #init.xavier_normal_(m.weight)
+                if m.bias is not None:
+                    init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm2d): #nn.Batchnorm2d
+                init.constant_(m.weight, 1)
+                init.constant_(m.bias, 0)
+            elif isinstance(m, nn.Linear):
+                init.normal_(m.weight, std=0.01)
+                if m.bias is not None:
+                    init.constant_(m.bias, 0)
+class AO(nn.Module):
+    # Adaptive output module
+    def __init__(self, inchannels, outchannels, upfactor=2):
+        super(AO, self).__init__()
+        self.inchannels = inchannels
+        self.outchannels = outchannels
+        self.upfactor = upfactor
+        """
+        self.adapt_conv = nn.Sequential(nn.Conv2d(in_channels=self.inchannels, out_channels=self.inchannels//2, kernel_size=3, padding=1, stride=1, bias=True),\
+                                  nn.BatchNorm2d(num_features=self.inchannels//2),\
+                                  nn.ReLU(inplace=True),\
+                                  nn.Conv2d(in_channels=self.inchannels//2, out_channels=self.outchannels, kernel_size=3, padding=1, stride=1, bias=True),\
+                                  nn.Upsample(scale_factor=self.upfactor, mode='bilinear', align_corners=True) )#,\
+                                  #nn.ReLU(inplace=True))  ## get positive values
+        """
+        self.adapt_conv = nn.Sequential(nn.Conv2d(in_channels=self.inchannels, out_channels=self.inchannels//2, kernel_size=3, padding=1, stride=1, bias=True),\
+                                  #nn.BatchNorm2d(num_features=self.inchannels//2),\
+                                  nn.ReLU(inplace=True),\
+                                  nn.Upsample(scale_factor=self.upfactor, mode='bilinear', align_corners=True), \
+                                  nn.Conv2d(in_channels=self.inchannels//2, out_channels=self.outchannels, kernel_size=1, padding=0, stride=1))
+                                  #nn.ReLU(inplace=True))  ## get positive values
+        self.init_params()
+    def forward(self, x):
+        x = self.adapt_conv(x)
+        return x
+    def init_params(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                #init.kaiming_normal_(m.weight, mode='fan_out')
+                init.normal_(m.weight, std=0.01)
+                #init.xavier_normal_(m.weight)
+                if m.bias is not None:
+                    init.constant_(m.bias, 0)
+            elif isinstance(m, nn.ConvTranspose2d):
+                #init.kaiming_normal_(m.weight, mode='fan_out')
+                init.normal_(m.weight, std=0.01)
+                #init.xavier_normal_(m.weight)
+                if m.bias is not None:
+                    init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm2d): #nn.Batchnorm2d
+                init.constant_(m.weight, 1)
+                init.constant_(m.bias, 0)
+            elif isinstance(m, nn.Linear):
+                init.normal_(m.weight, std=0.01)
+                if m.bias is not None:
+                    init.constant_(m.bias, 0)
+class ASPP(nn.Module):
+    def __init__(self, inchannels=256, planes=128, rates = [1, 6, 12, 18]):
+        super(ASPP, self).__init__()
+        self.inchannels = inchannels
+        self.planes = planes
+        self.rates = rates
+        self.kernel_sizes = []
+        self.paddings = []
+        for rate in self.rates:
+            if rate == 1:
+                self.kernel_sizes.append(1)
+                self.paddings.append(0)
+            else:
+                self.kernel_sizes.append(3)
+                self.paddings.append(rate)
+        self.atrous_0 = nn.Sequential(nn.Conv2d(in_channels=self.inchannels, out_channels=self.planes, kernel_size=self.kernel_sizes[0],
+                                                     stride=1, padding=self.paddings[0], dilation=self.rates[0], bias=True),
+                                      nn.ReLU(inplace=True),
+                                      nn.BatchNorm2d(num_features=self.planes)
+                                      )
+        self.atrous_1 = nn.Sequential(nn.Conv2d(in_channels=self.inchannels, out_channels=self.planes, kernel_size=self.kernel_sizes[1],
+                                                     stride=1, padding=self.paddings[1], dilation=self.rates[1], bias=True),
+                                      nn.ReLU(inplace=True),
+                                      nn.BatchNorm2d(num_features=self.planes),
+                                      )
+        self.atrous_2 = nn.Sequential(nn.Conv2d(in_channels=self.inchannels, out_channels=self.planes, kernel_size=self.kernel_sizes[2],
+                                                     stride=1, padding=self.paddings[2], dilation=self.rates[2], bias=True),
+                                      nn.ReLU(inplace=True),
+                                      nn.BatchNorm2d(num_features=self.planes),
+                                      )
+        self.atrous_3 = nn.Sequential(nn.Conv2d(in_channels=self.inchannels, out_channels=self.planes, kernel_size=self.kernel_sizes[3],
+                                                     stride=1, padding=self.paddings[3], dilation=self.rates[3], bias=True),
+                                      nn.ReLU(inplace=True),
+                                      nn.BatchNorm2d(num_features=self.planes),
+                                      )
+        #self.conv = nn.Conv2d(in_channels=self.planes * 4, out_channels=self.inchannels, kernel_size=3, padding=1, stride=1, bias=True)
+    def forward(self, x):
+        x = torch.cat([self.atrous_0(x), self.atrous_1(x), self.atrous_2(x), self.atrous_3(x)],1)
+        #x = self.conv(x)
+        return x
+# ==============================================================================================================
+class ResidualConv(nn.Module):
+    def __init__(self, inchannels):
+        super(ResidualConv, self).__init__()
+        #nn.BatchNorm2d
+        self.conv = nn.Sequential(
+                                  #nn.BatchNorm2d(num_features=inchannels),
+                                  nn.ReLU(inplace=False),
+                                  #nn.Conv2d(in_channels=inchannels, out_channels=inchannels, kernel_size=3, padding=1, stride=1, groups=inchannels,bias=True),
+                                  #nn.Conv2d(in_channels=inchannels, out_channels=inchannels, kernel_size=1, padding=0, stride=1, groups=1,bias=True)
+                                  nn.Conv2d(in_channels=inchannels, out_channels=inchannels//2, kernel_size=3, padding=1, stride=1, bias=False),
+                                  nn.BatchNorm2d(num_features=inchannels//2),
+                                  nn.ReLU(inplace=False),
+                                  nn.Conv2d(in_channels=inchannels//2, out_channels=inchannels, kernel_size=3, padding=1, stride=1, bias=False)
+                                  )
+        self.init_params()
+    def forward(self, x):
+        x = self.conv(x)+x
+        return x
+    def init_params(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                #init.kaiming_normal_(m.weight, mode='fan_out')
+                init.normal_(m.weight, std=0.01)
+                #init.xavier_normal_(m.weight)
+                if m.bias is not None:
+                    init.constant_(m.bias, 0)
+            elif isinstance(m, nn.ConvTranspose2d):
+                #init.kaiming_normal_(m.weight, mode='fan_out')
+                init.normal_(m.weight, std=0.01)
+                #init.xavier_normal_(m.weight)
+                if m.bias is not None:
+                    init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm2d): #nn.BatchNorm2d
+                init.constant_(m.weight, 1)
+                init.constant_(m.bias, 0)
+            elif isinstance(m, nn.Linear):
+                init.normal_(m.weight, std=0.01)
+                if m.bias is not None:
+                    init.constant_(m.bias, 0)
+class FeatureFusion(nn.Module):
+    def __init__(self, inchannels, outchannels):
+        super(FeatureFusion, self).__init__()
+        self.conv = ResidualConv(inchannels=inchannels)
+        #nn.BatchNorm2d
+        self.up = nn.Sequential(ResidualConv(inchannels=inchannels),
+                                nn.ConvTranspose2d(in_channels=inchannels, out_channels=outchannels, kernel_size=3,stride=2, padding=1, output_padding=1),
+                                nn.BatchNorm2d(num_features=outchannels),
+                                nn.ReLU(inplace=True))
+    def forward(self, lowfeat, highfeat):
+        return self.up(highfeat + self.conv(lowfeat))
+    def init_params(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                #init.kaiming_normal_(m.weight, mode='fan_out')
+                init.normal_(m.weight, std=0.01)
+                #init.xavier_normal_(m.weight)
+                if m.bias is not None:
+                    init.constant_(m.bias, 0)
+            elif isinstance(m, nn.ConvTranspose2d):
+                #init.kaiming_normal_(m.weight, mode='fan_out')
+                init.normal_(m.weight, std=0.01)
+                #init.xavier_normal_(m.weight)
+                if m.bias is not None:
+                    init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm2d): #nn.BatchNorm2d
+                init.constant_(m.weight, 1)
+                init.constant_(m.bias, 0)
+            elif isinstance(m, nn.Linear):
+                init.normal_(m.weight, std=0.01)
+                if m.bias is not None:
+                    init.constant_(m.bias, 0)
+class SenceUnderstand(nn.Module):
+    def __init__(self, channels):
+        super(SenceUnderstand, self).__init__()
+        self.channels = channels
+        self.conv1 = nn.Sequential(nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=1),
+                                   nn.ReLU(inplace = True))
+        self.pool = nn.AdaptiveAvgPool2d(8)
+        self.fc = nn.Sequential(nn.Linear(512*8*8, self.channels),
+                                nn.ReLU(inplace = True))
+        self.conv2 = nn.Sequential(nn.Conv2d(in_channels=self.channels, out_channels=self.channels, kernel_size=1, padding=0),
+                                   nn.ReLU(inplace=True))
+        self.initial_params()
+    def forward(self, x):
+        n,c,h,w = x.size()
+        x = self.conv1(x)
+        x = self.pool(x)
+        x = x.view(n,-1)
+        x = self.fc(x)
+        x = x.view(n, self.channels, 1, 1)
+        x = self.conv2(x)
+        x = x.repeat(1,1,h,w)
+        return x
+    def initial_params(self, dev=0.01):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                #print torch.sum(m.weight)
+                m.weight.data.normal_(0, dev)
+                if m.bias is not None:
+                    m.bias.data.fill_(0)
+            elif isinstance(m, nn.ConvTranspose2d):
+                #print torch.sum(m.weight)
+                m.weight.data.normal_(0, dev)
+                if m.bias is not None:
+                    m.bias.data.fill_(0)
+            elif isinstance(m, nn.Linear):
+                m.weight.data.normal_(0, dev)

models/SpaTrackV2/models/depth_refiner/stablilization_attention.py ADDED Viewed

	@@ -0,0 +1,1187 @@