Upload EMS-superquadric_fitting_inference

.gitattributes

@@ -33,3 +33,13 @@ saved_model/**/* 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
+EMS-superquadric_fitting_inference/src/EMS/__pycache__/EMS_recovery.CostFunc-307.py312.1.nbc filter=lfs diff=lfs merge=lfs -text
+EMS-superquadric_fitting_inference/src/EMS/__pycache__/EMS_recovery.EigenAnalysis-272.py311.1.nbc filter=lfs diff=lfs merge=lfs -text
+EMS-superquadric_fitting_inference/src/EMS/__pycache__/EMS_recovery.Distance-286.py312.1.nbc filter=lfs diff=lfs merge=lfs -text
+EMS-superquadric_fitting_inference/src/EMS/__pycache__/EMS_recovery.SwitchCost-265.py311.1.nbc filter=lfs diff=lfs merge=lfs -text
+EMS-superquadric_fitting_inference/src/EMS/__pycache__/EMS_recovery.SimilarityCandidates-138.py311.1.nbc filter=lfs diff=lfs merge=lfs -text
+EMS-superquadric_fitting_inference/src/EMS/__pycache__/EMS_recovery.EigenAnalysis-272.py312.1.nbc filter=lfs diff=lfs merge=lfs -text
+EMS-superquadric_fitting_inference/src/EMS/__pycache__/EMS_recovery.Distance-286.py311.1.nbc filter=lfs diff=lfs merge=lfs -text
+EMS-superquadric_fitting_inference/src/EMS/__pycache__/EMS_recovery.SwitchCost-265.py312.1.nbc filter=lfs diff=lfs merge=lfs -text
+EMS-superquadric_fitting_inference/src/EMS/__pycache__/EMS_recovery.SimilarityCandidates-138.py312.1.nbc filter=lfs diff=lfs merge=lfs -text
+EMS-superquadric_fitting_inference/src/EMS/__pycache__/EMS_recovery.CostFunc-307.py311.1.nbc filter=lfs diff=lfs merge=lfs -text

EMS-superquadric_fitting_inference/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2022 Weixiao Liu
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

EMS-superquadric_fitting_inference/README.md

@@ -0,0 +1,37 @@
+## Guidelines for Python Implementation
+
+This is the guideline and structual explanation of the Python implementation of the EMS algorithm.
+
+### Dependency
+
+The code is tested under Python 3.8.8, but should have little compatibility concerns.
+The following packages are required to run the EMS algorithm:
+
+1. numpy 1.19.2
+2. scipy 1.5.2
+3. numba 0.53.1 -- for acceleration based on JIT (Just-In-Time compiler)
+
+
+For demo, the following packages are needed:
+
+1. plyfile -- for loading `.ply` point cloud files
+2. mayavi -- for visualization of meshes and point clouds
+
+### Installation
+
+We recommend to install the EMS package with `pip`.
+
+1. Change directory to `/Python`
+2. Install package: `pip install .`
+    
+
+### Run Demo
+
+The demo script is `/Python/tests/test_script.py`.
+The demo reads a `.ply` point cloud and returns the parameters of the recovered superquadric, runtime, and visualization as required.
+
+For example, in terminal type in
+
+     python test_script.py *.ply file location* --result --runtime --visualize
+
+Note the first run of the code takes longer, since the JIT will translate the Python and NumPy code into fast machine code (and will be cached for futher calls).

EMS-superquadric_fitting_inference/README_MOVI.md

@@ -0,0 +1,188 @@
+# MOVi-A Dataset for Superquadric Fitting
+
+This directory contains scripts to download and process the MOVi-A dataset for superquadric fitting experiments.
+
+● 完成了MOVi-A数据集下载脚本！
+
+  我在 EMS-superquadric_fitting_inference
+  目录下创建了：
+
+  1. download_movi_a.py - 主下载脚本
+
+  - 交互式选择下载数量（10/100/全部样本）
+  - 自动提取点云（从深度图）
+  - 保存为numpy格式，方便后续使用
+  - 分离每个物体的点云
+
+  2. 数据组织结构
+
+  data/
+  └── movi_a_128x128_train/
+      ├── sample_00000/
+      │   ├── metadata.json      #
+  物体形状、材质等信息
+      │   ├── rgb/              # RGB图像
+      │   ├── depth/            # 深度图
+      │   ├── segmentation/    # 实例分割
+      │   ├── point_clouds/     # 提取的3D点云
+      │   └── trajectories.npz # 运动轨迹
+      └── ...
+
+  3. 加载器脚本 (自动生成)
+
+  下载完成后会生成
+  load_movi_a.py，无需TensorFlow即可加载数据：
+  loader = MOViALoader('data/movi_a_128x128_train')
+  frame = loader.load_frame(0, 0)  # 
+  加载第0个样本的第0帧
+
+## Dataset Overview
+
+MOVi-A is a synthetic dataset with:
+- Simple geometric shapes (cube, sphere, cylinder) - perfect for superquadric fitting
+- 3-10 objects per scene with physics simulation
+- 128x128 resolution (smaller, faster to process)
+- Rich annotations including depth, segmentation, 3D trajectories, physics properties, and collision events
+
+## Quick Start
+
+1. **Install minimal dependencies** (in WaveGen environment):
+```bash
+conda create -n movi python=3.9
+conda activate movi
+pip install tensorflow-cpu tensorflow-datasets opencv-python tqdm
+```
+
+2. **Download the dataset**:
+```bash
+python download_movi_a.py
+```
+
+Press Enter to download ALL samples, or specify a number for testing.
+
+3. **After download, you can uninstall TensorFlow** if desired:
+```bash
+pip uninstall tensorflow tensorflow-datasets
+```
+
+## Data Structure
+
+After downloading, the data is organized as:
+
+```
+data/
+└── movi_a_128x128_train/
+    ├── dataset_info.json          # Overall dataset metadata
+    ├── sample_00000/              # Each sample
+    │   ├── metadata.json          # Sample metadata + physics properties
+    │   ├── rgb/                   # RGB frames (PNG)
+    │   │   ├── frame_000.png
+    │   │   └── ...
+    │   ├── depth/                 # Depth maps (NPY)
+    │   │   ├── frame_000.npy
+    │   │   └── ...
+    │   ├── segmentation/          # Instance masks (NPY)
+    │   │   ├── frame_000.npy
+    │   │   └── ...
+    │   ├── normal/                # Surface normals (NPY)
+    │   │   ├── frame_000.npy
+    │   │   └── ...
+    │   ├── object_coordinates/    # Object-relative coordinates (NPY)
+    │   │   ├── frame_000.npy
+    │   │   └── ...
+    │   ├── point_clouds/          # Extracted 3D points
+    │   │   ├── frame_000_full.npy
+    │   │   ├── frame_000_instance_1.npy
+    │   │   └── ...
+    │   ├── trajectories.npz      # Full object motion data
+    │   ├── camera_trajectory.npz # Camera positions and rotations
+    │   └── collisions.npz        # Collision events data
+    └── ...
+```
+
+## Loading Data (No TensorFlow Required!)
+
+```python
+from load_movi_a import MOViALoader
+
+# Initialize loader
+loader = MOViALoader('data/movi_a_128x128_train')
+
+# Load sample metadata
+sample = loader.load_sample(0)
+print(f"Objects: {sample['metadata']['num_instances']}")
+
+# Load frame data
+frame = loader.load_frame(sample_idx=0, frame_idx=0)
+rgb = frame['rgb']           # (128, 128, 3)
+depth = frame['depth']       # (128, 128, 1)
+points = frame['point_cloud'] # (N, 3)
+
+# Load instance-specific point cloud
+instance_pc = loader.load_instance_point_cloud(0, 0, instance_id=1)
+```
+
+## Using with Superquadric Fitting
+
+```python
+from EMS.EMS_recovery import EMS_recovery
+from load_movi_a import MOViALoader
+
+loader = MOViALoader('data/movi_a_128x128_train')
+
+# Fit superquadric to first object in first frame
+instance_pc = loader.load_instance_point_cloud(0, 0, 1)
+if instance_pc is not None and len(instance_pc) > 100:
+    sq, p = EMS_recovery(instance_pc, OutlierRatio=0.2)
+    print(f"Shape parameters: {sq.shape}")
+```
+
+## Object Properties
+
+Each object has (stored in `metadata.json`):
+- **Shape**: cube, sphere, or cylinder
+- **Size**: small or large
+- **Material**: metal or rubber
+- **Color**: 8 different colors + RGB values
+- **Physics properties**:
+  - mass: Mass of the object
+  - friction: Friction coefficient (metal=0.4, rubber=0.8)
+  - restitution: Bounciness (metal=0.3, rubber=0.7)
+
+## Additional Data
+
+### Trajectories (`trajectories.npz`)
+- positions: (num_objects, 24, 3) - 3D positions
+- quaternions: (num_objects, 24, 4) - Rotations
+- velocities: (num_objects, 24, 3) - Linear velocities
+- angular_velocities: (num_objects, 24, 3) - Angular velocities
+- visibility: (num_objects, 24) - Pixel count visibility
+- bboxes_3d: (num_objects, 24, 8, 3) - 3D bounding box corners
+- image_positions: (num_objects, 24, 2) - 2D center of mass
+
+### Camera (`camera_trajectory.npz`)
+- positions: (24, 3) - Camera positions (static in MOVi-A)
+- quaternions: (24, 4) - Camera rotations
+
+### Collisions (`collisions.npz`)
+- instances: (N, 2) - Pairs of colliding objects
+- frame: (N,) - Frame of collision
+- force: (N,) - Collision force
+- position: (N, 3) - 3D collision position
+- image_position: (N, 2) - 2D collision position
+- contact_normal: (N, 3) - Collision normal vector
+
+## Tips
+
+1. Start with 10 samples to test your pipeline
+2. The point clouds are already extracted from depth maps
+3. Instance segmentation helps separate objects
+4. Use trajectories.npz for temporal consistency
+
+## Storage Requirements (with all data)
+
+- 10 samples: ~250 MB
+- 100 samples: ~2.5 GB  
+- Full training set (9750 samples): ~25 GB
+- Validation set (250 samples): ~650 MB
+- **Total (train + validation): ~26 GB**

EMS-superquadric_fitting_inference/download_movi_a.py

@@ -0,0 +1,462 @@
+#!/usr/bin/env python3
+"""
+Simple MOVi-A dataset loader using the Kubric example code.
+This version works with the MOVi datasets hosted on Google Cloud.
+"""
+
+import os
+import numpy as np
+import tensorflow as tf
+import tensorflow_datasets as tfds
+from pathlib import Path
+import json
+import cv2
+from tqdm import tqdm
+
+# Reduce TF logging
+os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
+
+
+def extract_point_cloud_from_depth(depth, camera_K, segmentation=None, instance_id=None):
+    """
+    Convert depth image to 3D point cloud using camera intrinsics.
+    """
+    H, W = depth.shape[:2]
+    
+    # Create pixel coordinates
+    xx, yy = np.meshgrid(np.arange(W), np.arange(H))
+    
+    # Get valid depth values
+    if instance_id is not None and segmentation is not None:
+        mask = (segmentation == instance_id) & (depth > 0)
+    else:
+        mask = depth > 0
+    
+    # Extract valid coordinates
+    valid_x = xx[mask]
+    valid_y = yy[mask]
+    valid_z = depth[mask]
+    
+    # Unproject to 3D using camera intrinsics
+    # K = [[fx, 0, cx], [0, fy, cy], [0, 0, 1]]
+    fx, fy = camera_K[0, 0], camera_K[1, 1]
+    cx, cy = camera_K[0, 2], camera_K[1, 2]
+    
+    x_3d = (valid_x - cx) * valid_z / fx
+    y_3d = (valid_y - cy) * valid_z / fy
+    z_3d = valid_z
+    
+    return np.stack([x_3d, y_3d, z_3d], axis=-1)
+
+
+def process_and_save_sample(sample, output_dir, sample_idx):
+    """Process a single MOVi sample and save to disk."""
+    sample_dir = Path(output_dir) / f"sample_{sample_idx:05d}"
+    
+    # Check if sample already fully processed
+    if sample_dir.exists():
+        # Check if key files exist to determine if download was complete
+        required_files = [
+            sample_dir / "metadata.json",
+            sample_dir / "trajectories.npz",
+            sample_dir / "camera_trajectory.npz"
+        ]
+        if all(f.exists() for f in required_files):
+            # Also check if all frames are downloaded
+            num_frames = 24  # MOVi-A has 24 frames
+            frame_files_exist = all(
+                (sample_dir / "rgb" / f"frame_{i:03d}.png").exists() 
+                for i in range(num_frames)
+            )
+            if frame_files_exist:
+                print(f"  Sample {sample_idx:05d} already downloaded, skipping...")
+                return None
+    
+    sample_dir.mkdir(parents=True, exist_ok=True)
+    
+    # Decode depth values
+    minv, maxv = sample["metadata"]["depth_range"]
+    depth = sample["depth"] / 65535 * (maxv - minv) + minv
+    
+    # Get camera info
+    focal_length = float(sample["camera"]["focal_length"])
+    sensor_width = float(sample["camera"]["sensor_width"])
+    field_of_view = float(sample["camera"]["field_of_view"])
+    resolution = sample["video"].shape[1]  # Assuming square
+    
+    # Compute camera intrinsics
+    fx = fy = focal_length * resolution / sensor_width
+    cx = cy = resolution / 2
+    K = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]])
+    
+    # Extract metadata
+    metadata = {
+        "num_frames": int(sample["metadata"]["num_frames"]),
+        "num_instances": int(sample["metadata"]["num_instances"]),
+        "resolution": resolution,
+        "depth_range": [float(minv), float(maxv)],
+        "camera": {
+            "focal_length": focal_length,
+            "sensor_width": sensor_width,
+            "field_of_view": field_of_view,
+            "K": K.tolist()
+        },
+        "instances": []
+    }
+    
+    # Process instance information
+    for i in range(metadata["num_instances"]):
+        # Handle both string and integer labels
+        shape_label = sample["instances"]["shape_label"][i]
+        size_label = sample["instances"]["size_label"][i]
+        color_label = sample["instances"]["color_label"][i]
+        material_label = sample["instances"]["material_label"][i]
+        
+        # Decode if bytes, otherwise convert to string
+        if hasattr(shape_label, 'decode'):
+            shape = shape_label.decode()
+        else:
+            # Map integer labels to names
+            shape_names = ["cube", "cylinder", "sphere"]
+            shape = shape_names[int(shape_label)] if int(shape_label) < len(shape_names) else str(shape_label)
+            
+        if hasattr(size_label, 'decode'):
+            size = size_label.decode()
+        else:
+            size_names = ["small", "large"]
+            size = size_names[int(size_label)] if int(size_label) < len(size_names) else str(size_label)
+            
+        if hasattr(color_label, 'decode'):
+            color = color_label.decode()
+        else:
+            color_names = ["blue", "brown", "cyan", "gray", "green", "purple", "red", "yellow"]
+            color = color_names[int(color_label)] if int(color_label) < len(color_names) else str(color_label)
+            
+        if hasattr(material_label, 'decode'):
+            material = material_label.decode()
+        else:
+            material_names = ["metal", "rubber"]
+            material = material_names[int(material_label)] if int(material_label) < len(material_names) else str(material_label)
+        
+        # Extract physics properties
+        mass = float(sample["instances"]["mass"][i])
+        friction = float(sample["instances"]["friction"][i])
+        restitution = float(sample["instances"]["restitution"][i])
+        
+        # Extract color RGB values
+        color_rgb = sample["instances"]["color"][i].tolist()
+        
+        instance_info = {
+            "id": i + 1,  # 1-indexed in segmentation
+            "shape": shape,
+            "size": size,
+            "color": color,
+            "color_rgb": color_rgb,
+            "material": material,
+            "mass": mass,
+            "friction": friction,
+            "restitution": restitution
+        }
+        metadata["instances"].append(instance_info)
+    
+    # Save metadata
+    with open(sample_dir / "metadata.json", 'w') as f:
+        json.dump(metadata, f, indent=2)
+    
+    # Save trajectories with all motion data
+    np.savez_compressed(
+        sample_dir / "trajectories.npz",
+        positions=sample["instances"]["positions"],
+        quaternions=sample["instances"]["quaternions"],
+        velocities=sample["instances"]["velocities"],
+        angular_velocities=sample["instances"]["angular_velocities"],
+        visibility=sample["instances"]["visibility"],
+        bboxes_3d=sample["instances"]["bboxes_3d"],
+        image_positions=sample["instances"]["image_positions"]
+    )
+    
+    # Save camera trajectory
+    np.savez_compressed(
+        sample_dir / "camera_trajectory.npz",
+        positions=sample["camera"]["positions"],
+        quaternions=sample["camera"]["quaternions"]
+    )
+    
+    # Save collision events
+    if "events" in sample and "collisions" in sample["events"]:
+        collisions = sample["events"]["collisions"]
+        collision_data = {
+            "instances": collisions["instances"],
+            "frame": collisions["frame"],
+            "force": collisions["force"],
+            "position": collisions["position"],
+            "image_position": collisions["image_position"],
+            "contact_normal": collisions["contact_normal"]
+        }
+        np.savez_compressed(sample_dir / "collisions.npz", **collision_data)
+    
+    # Process and save frames
+    (sample_dir / "rgb").mkdir(exist_ok=True)
+    (sample_dir / "depth").mkdir(exist_ok=True)
+    (sample_dir / "segmentation").mkdir(exist_ok=True)
+    (sample_dir / "normal").mkdir(exist_ok=True)
+    (sample_dir / "object_coordinates").mkdir(exist_ok=True)
+    (sample_dir / "point_clouds").mkdir(exist_ok=True)
+    
+    for t in range(metadata["num_frames"]):
+        # Save RGB
+        rgb = sample["video"][t]
+        cv2.imwrite(str(sample_dir / "rgb" / f"frame_{t:03d}.png"), 
+                    cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR))
+        
+        # Save depth
+        np.save(sample_dir / "depth" / f"frame_{t:03d}.npy", depth[t])
+        
+        # Save segmentation
+        seg = sample["segmentations"][t, :, :, 0]
+        np.save(sample_dir / "segmentation" / f"frame_{t:03d}.npy", seg)
+        
+        # Save normal
+        normal = sample["normal"][t]
+        np.save(sample_dir / "normal" / f"frame_{t:03d}.npy", normal)
+        
+        # Save object coordinates
+        obj_coords = sample["object_coordinates"][t]
+        np.save(sample_dir / "object_coordinates" / f"frame_{t:03d}.npy", obj_coords)
+        
+        # Extract and save point clouds
+        # Full scene
+        pc_full = extract_point_cloud_from_depth(depth[t, :, :, 0], K)
+        np.save(sample_dir / "point_clouds" / f"frame_{t:03d}_full.npy", pc_full)
+        
+        # Per-instance
+        for i in range(metadata["num_instances"]):
+            instance_id = i + 1
+            pc_instance = extract_point_cloud_from_depth(
+                depth[t, :, :, 0], K, seg, instance_id
+            )
+            if len(pc_instance) > 0:
+                np.save(sample_dir / "point_clouds" / f"frame_{t:03d}_obj{instance_id}.npy", 
+                        pc_instance)
+    
+    return metadata
+
+
+def main():
+    print("MOVi-A Dataset Downloader (Simple Version)")
+    print("=" * 50)
+    
+    # Output directory
+    script_dir = Path(__file__).parent
+    output_base_dir = script_dir / ".." / "data" / "movi_a_128x128"
+    output_base_dir.mkdir(parents=True, exist_ok=True)
+    
+    print(f"Output directory: {output_base_dir}")
+    print("\nAttempting to load MOVi-A from Google Cloud Storage...")
+    
+    try:
+        # Download both train and validation splits
+        for split_name in ["train", "validation"]:
+            print(f"\n{'='*50}")
+            print(f"Processing {split_name.upper()} split...")
+            print(f"{'='*50}")
+            
+            output_dir = output_base_dir / split_name
+            output_dir.mkdir(exist_ok=True)
+            
+            # Check existing samples
+            existing_samples = len(list(output_dir.glob("sample_*")))
+            if existing_samples > 0:
+                print(f"Found {existing_samples} existing samples in {split_name} directory")
+            
+            # Load split
+            ds = tfds.load(
+                "movi_a/128x128",
+                split=split_name,
+                data_dir="gs://kubric-public/tfds",
+                with_info=False
+            )
+            
+            print(f"Successfully connected to MOVi-A {split_name} dataset!")
+            print(f"Processing ALL {split_name} samples (will skip existing)...")
+            
+            # Process all samples
+            total_processed = 0
+            total_skipped = 0
+            for idx, sample in enumerate(tqdm(tfds.as_numpy(ds), desc=f"Processing {split_name}")):
+                metadata = process_and_save_sample(sample, output_dir, idx)
+                if metadata is None:
+                    total_skipped += 1
+                else:
+                    total_processed += 1
+            
+            print(f"\nProcessed {split_name} split:")
+            print(f"  - Downloaded: {total_processed} samples")
+            print(f"  - Skipped (already exist): {total_skipped} samples")
+            print(f"  - Total: {total_processed + total_skipped} samples")
+            
+            # Save split info
+            dataset_info = {
+                "dataset": "movi_a",
+                "split": split_name,
+                "resolution": "128x128",
+                "num_samples": total_processed,
+                "fps": 12,
+                "num_frames_per_sample": 24
+            }
+            with open(output_dir / "dataset_info.json", 'w') as f:
+                json.dump(dataset_info, f, indent=2)
+        
+        print(f"\n{'='*50}")
+        print(f"All downloads complete!")
+        print(f"Data saved to: {output_base_dir}")
+        
+        # Count final samples
+        train_count = len(list((output_base_dir / "train").glob("sample_*")))
+        val_count = len(list((output_base_dir / "validation").glob("sample_*"))) if (output_base_dir / "validation").exists() else 0
+        
+        print(f"\nFinal dataset size:")
+        print(f"  - Train samples: {train_count} in {output_base_dir}/train")
+        print(f"  - Validation samples: {val_count} in {output_base_dir}/validation")
+        
+        # Create simple loader
+        create_loader_script(output_base_dir)
+        
+    except KeyboardInterrupt:
+        print("\n\nDownload interrupted by user. You can run the script again to resume.")
+        return
+    except Exception as e:
+        print(f"\nError: {e}")
+        import traceback
+        traceback.print_exc()
+        print("\nNote: You can run the script again to resume downloading.")
+        print("\nAlternative: Download manually using gsutil")
+        print("1. Install: conda install -c conda-forge google-cloud-sdk")
+        print("2. Download a few samples manually:")
+        print("   gsutil -m cp -r gs://kubric-public/tfds/movi_a/128x128/1.0.0/movi_a-train.tfrecord-00000-of-00256 ./")
+        print("\nOr try the original Kubric repository:")
+        print("   https://github.com/google-research/kubric")
+
+
+def create_loader_script(output_dir):
+    """Create a simple loader script."""
+    script = '''#!/usr/bin/env python3
+"""Simple MOVi-A data loader - no TensorFlow required!"""
+
+import numpy as np
+import json
+import cv2
+from pathlib import Path
+
+
+class MOViLoader:
+    def __init__(self, data_dir, split="train"):
+        self.data_dir = Path(data_dir)
+        self.split = split
+        self.split_dir = self.data_dir / split
+        self.samples = sorted(list(self.split_dir.glob("sample_*")))
+        
+    def load_sample(self, idx):
+        """Load metadata and trajectories for a sample."""
+        sample_dir = self.samples[idx]
+        
+        with open(sample_dir / "metadata.json", 'r') as f:
+            metadata = json.load(f)
+            
+        trajectories = np.load(sample_dir / "trajectories.npz")
+        
+        return {
+            "metadata": metadata,
+            "trajectories": trajectories,
+            "sample_dir": sample_dir,
+            "camera": None,
+            "collisions": None
+        }
+        
+        # Load camera trajectory if exists
+        camera_path = sample_dir / "camera_trajectory.npz"
+        if camera_path.exists():
+            data["camera"] = np.load(camera_path)
+            
+        # Load collision data if exists
+        collision_path = sample_dir / "collisions.npz"
+        if collision_path.exists():
+            data["collisions"] = np.load(collision_path)
+            
+        return data
+    
+    def load_frame(self, sample_idx, frame_idx):
+        """Load all data for a specific frame."""
+        sample_dir = self.samples[sample_idx]
+        
+        # Load RGB
+        rgb = cv2.imread(str(sample_dir / "rgb" / f"frame_{frame_idx:03d}.png"))
+        rgb = cv2.cvtColor(rgb, cv2.COLOR_BGR2RGB)
+        
+        # Load depth
+        depth = np.load(sample_dir / "depth" / f"frame_{frame_idx:03d}.npy")
+        
+        # Load segmentation
+        seg = np.load(sample_dir / "segmentation" / f"frame_{frame_idx:03d}.npy")
+        
+        # Load full point cloud
+        pc = np.load(sample_dir / "point_clouds" / f"frame_{frame_idx:03d}_full.npy")
+        
+        return {
+            "rgb": rgb,
+            "depth": depth,
+            "segmentation": seg,
+            "point_cloud": pc,
+            "normal": None,
+            "object_coordinates": None
+        }
+        
+        # Load normal if exists
+        normal_path = sample_dir / "normal" / f"frame_{frame_idx:03d}.npy"
+        if normal_path.exists():
+            frame_data["normal"] = np.load(normal_path)
+            
+        # Load object coordinates if exists
+        obj_coord_path = sample_dir / "object_coordinates" / f"frame_{frame_idx:03d}.npy"
+        if obj_coord_path.exists():
+            frame_data["object_coordinates"] = np.load(obj_coord_path)
+            
+        return frame_data
+    
+    def load_object_points(self, sample_idx, frame_idx, object_id):
+        """Load point cloud for a specific object."""
+        sample_dir = self.samples[sample_idx]
+        pc_file = sample_dir / "point_clouds" / f"frame_{frame_idx:03d}_obj{object_id}.npy"
+        
+        if pc_file.exists():
+            return np.load(pc_file)
+        return None
+
+
+if __name__ == "__main__":
+    # Example usage
+    loader_train = MOViLoader(".", split="train")
+    loader_val = MOViLoader(".", split="validation")
+    print(f"Found {len(loader_train.samples)} training samples")
+    print(f"Found {len(loader_val.samples)} validation samples")
+    
+    if len(loader_train.samples) > 0:
+        # Load first sample
+        sample = loader_train.load_sample(0)
+        print(f"\\nTrain sample 0: {sample['metadata']['num_instances']} objects")
+        
+        # Load first frame
+        frame = loader_train.load_frame(0, 0)
+        print(f"Point cloud shape: {frame['point_cloud'].shape}")
+'''
+    
+    loader_path = output_dir.parent / "load_movi.py"
+    with open(loader_path, 'w') as f:
+        f.write(script)
+    
+    print(f"\nCreated loader script: {loader_path}")
+
+
+if __name__ == "__main__":
+    print("\nTip: This script supports resuming downloads. If interrupted, just run it again!\n")
+    main()

EMS-superquadric_fitting_inference/process_movi_train.py

@@ -0,0 +1,886 @@
+#!/usr/bin/env python3
+"""
+Process MOVi-A train set with hierarchical multi-superquadric fitting
+Converts depth maps to normalized point clouds for superquadric fitting
+"""
+
+import numpy as np
+import sys
+import os
+import time
+import viser
+import json
+import cv2
+from pathlib import Path
+from sklearn.cluster import DBSCAN
+
+# Add the src directory to Python path
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src'))
+
+from EMS.EMS_recovery import EMS_recovery
+
+
+def depth_to_normalized_pointcloud(depth, segmentation, camera_K, camera_position=None, camera_quaternion=None, resolution=128, convert_to_zdepth=True):
+    """
+    Convert depth map to normalized point cloud in range [-10, 10]
+    
+    Args:
+        depth: (H, W, 1) depth array (euclidean distance from camera center)
+        segmentation: (H, W) instance segmentation mask
+        camera_K: 3x3 camera intrinsic matrix
+        camera_position: camera position in world coordinates
+        camera_quaternion: camera quaternion (x,y,z,w) in world coordinates
+        resolution: image resolution (assuming square)
+        convert_to_zdepth: bool, convert euclidean depth to z-depth before processing
+    
+    Returns:
+        dict: instance_id -> normalized point cloud in world coordinates
+    """
+    H, W = depth.shape[:2]
+    
+    # Get camera parameters
+    fx = camera_K[0, 0]
+    fy = camera_K[1, 1]
+    cx = camera_K[0, 2]
+    cy = camera_K[1, 2]
+    
+    # Create pixel grid
+    xx, yy = np.meshgrid(np.arange(W), np.arange(H))
+    
+    # Convert to normalized camera coordinates
+    x_norm = (xx - cx) / fx
+    y_norm = (yy - cy) / fy
+    
+    if convert_to_zdepth:
+        # MOVi uses euclidean distance, convert to z-depth (planar depth)
+        # For each pixel, we have: euclidean_dist^2 = x^2 + y^2 + z^2
+        # Where x = x_norm * z, y = y_norm * z
+        # So: euclidean_dist^2 = (x_norm^2 + y_norm^2 + 1) * z^2
+        z = depth[:, :, 0] / np.sqrt(x_norm**2 + y_norm**2 + 1)
+    else:
+        # Use depth as-is (assume it's already z-depth)
+        z = depth[:, :, 0]
+    
+    # Get 3D points
+    x = x_norm * z
+    y = y_norm * z
+    
+    # Stack to get point cloud (in camera coordinates)
+    points_3d_camera = np.stack([x, y, z], axis=-1)
+    
+    # Transform from camera to world coordinates if camera pose is provided
+    if camera_position is not None and camera_quaternion is not None:
+        from scipy.spatial.transform import Rotation
+        
+        # Convert quaternion to rotation matrix
+        # MOVi uses [x, y, z, w] format
+        cam_rot = Rotation.from_quat(camera_quaternion)
+        cam_rot_matrix = cam_rot.as_matrix()
+        
+        # Transform points: World = R * Camera + T
+        points_3d_flat = points_3d_camera.reshape(-1, 3)
+        points_3d_world = points_3d_flat @ cam_rot_matrix.T + camera_position
+        points_3d = points_3d_world.reshape(points_3d_camera.shape)
+    else:
+        points_3d = points_3d_camera
+    
+    
+    # Normalize entire scene to [-10, 10] range
+    # Find scene bounds (only valid depth points)
+    valid_mask = z > 0
+    valid_points = points_3d[valid_mask]
+    
+    if len(valid_points) > 0:
+        # Find scene extent
+        scene_min = np.min(valid_points, axis=0)
+        scene_max = np.max(valid_points, axis=0)
+        scene_center = (scene_min + scene_max) / 2
+        scene_extent = np.max(scene_max - scene_min)
+        
+        # Scale to [-10, 10]
+        if scene_extent > 0:
+            scale_factor = 20.0 / scene_extent  # 20 because we want -10 to 10
+            points_3d_normalized = (points_3d - scene_center) * scale_factor
+        else:
+            points_3d_normalized = points_3d - scene_center
+    else:
+        points_3d_normalized = points_3d
+    
+    # Get unique instance IDs (excluding background=0)
+    instance_ids = np.unique(segmentation)
+    instance_ids = instance_ids[instance_ids > 0]
+    
+    instance_pointclouds = {}
+    
+    for inst_id in instance_ids:
+        # Get mask for this instance
+        mask = segmentation == inst_id
+        
+        # Extract points for this instance (already normalized with scene)
+        instance_points = points_3d_normalized[mask]
+        
+        if len(instance_points) < 50:  # Skip if too few points
+            continue
+        
+        instance_pointclouds[int(inst_id)] = instance_points
+    
+    # Also return the full scene point cloud and segmentation for visualization
+    return instance_pointclouds, points_3d_normalized, segmentation, scene_center if 'scene_center' in locals() else np.zeros(3), scene_extent if 'scene_extent' in locals() else 1.0
+
+
+def hierarchical_ems(
+    point,
+    OutlierRatio=0.5,
+    MaxIterationEM=20,
+    ToleranceEM=1e-3,
+    RelativeToleranceEM=2e-1,
+    MaxOptiIterations=2,
+    Sigma=0.3,
+    MaxiSwitch=2,
+    AdaptiveUpperBound=True,
+    Rescale=False,
+    MaxLayer=3,
+    Eps=1.0,  # Adjusted for normalized [-10, 10] point clouds
+    MinPoints=50,
+):
+    """
+    Hierarchical EMS for extracting multiple superquadrics from a point cloud
+    """
+    point_seg = {key: [] for key in list(range(0, MaxLayer+1))}
+    point_outlier = {key: [] for key in list(range(0, MaxLayer+1))}
+    point_seg[0] = [point]
+    list_quadrics = []
+    quadric_info = []
+    
+    for h in range(MaxLayer):
+        if len(point_seg[h]) == 0:
+            break
+            
+        for c in range(len(point_seg[h])):
+            current_points = point_seg[h][c]
+            if len(current_points) < MinPoints * 2:
+                continue
+                
+            try:
+                # Fit superquadric
+                x_raw, p_raw = EMS_recovery(
+                    current_points,
+                    OutlierRatio,
+                    MaxIterationEM,
+                    ToleranceEM,
+                    RelativeToleranceEM,
+                    MaxOptiIterations,
+                    Sigma,
+                    MaxiSwitch,
+                    AdaptiveUpperBound,
+                    Rescale,
+                )
+                
+                # Calculate fitting quality
+                inlier_mask = p_raw > 0.5
+                inlier_ratio = np.sum(inlier_mask) / len(p_raw)
+                
+                if inlier_ratio > 0.3:  # Accept if at least 30% inliers
+                    list_quadrics.append(x_raw)
+                    quadric_info.append({
+                        'layer': h,
+                        'segment': c,
+                        'inlier_ratio': inlier_ratio,
+                        'num_points': len(current_points),
+                        'inlier_points': current_points[inlier_mask]
+                    })
+                
+                # Separate outliers for next layer
+                outlier_mask = p_raw < 0.1
+                outlier = current_points[outlier_mask]
+                
+                # If many outliers and not last layer, try clustering
+                if len(outlier) > MinPoints * 2 and h < MaxLayer - 1:
+                    clustering = DBSCAN(eps=Eps, min_samples=MinPoints).fit(outlier)
+                    labels = list(set(clustering.labels_))
+                    labels = [item for item in labels if item >= 0]
+                    
+                    if len(labels) >= 1:
+                        for i in range(len(labels)):
+                            cluster_points = outlier[clustering.labels_ == labels[i]]
+                            if len(cluster_points) > MinPoints:
+                                point_seg[h + 1].append(cluster_points)
+                
+            except Exception as e:
+                continue
+    
+    return list_quadrics, quadric_info
+
+
+def generate_superquadric_mesh(sq, num_samples=25):
+    """Generate mesh vertices and faces for superquadric surface"""
+    eta = np.linspace(-np.pi/2, np.pi/2, num_samples)
+    omega = np.linspace(-np.pi, np.pi, num_samples)
+    
+    vertices = []
+    faces = []
+    
+    # Generate vertices
+    for i, e in enumerate(eta):
+        for j, w in enumerate(omega):
+            # Superquadric parametric equations
+            cos_eta = np.sign(np.cos(e)) * np.abs(np.cos(e))**sq.shape[0]
+            sin_eta = np.sign(np.sin(e)) * np.abs(np.sin(e))**sq.shape[0]
+            cos_omega = np.sign(np.cos(w)) * np.abs(np.cos(w))**sq.shape[1]
+            sin_omega = np.sign(np.sin(w)) * np.abs(np.sin(w))**sq.shape[1]
+            
+            # Local coordinates
+            x_local = sq.scale[0] * cos_eta * cos_omega
+            y_local = sq.scale[1] * cos_eta * sin_omega
+            z_local = sq.scale[2] * sin_eta
+            
+            # Apply rotation and translation
+            point_local = np.array([x_local, y_local, z_local])
+            point_global = sq.RotM @ point_local + sq.translation
+            
+            vertices.append(point_global)
+    
+    vertices = np.array(vertices)
+    
+    # Generate faces (triangles)
+    for i in range(num_samples - 1):
+        for j in range(num_samples - 1):
+            # Current vertex indices
+            idx1 = i * num_samples + j
+            idx2 = i * num_samples + (j + 1) % num_samples
+            idx3 = (i + 1) * num_samples + j
+            idx4 = (i + 1) * num_samples + (j + 1) % num_samples
+            
+            # Two triangles per quad
+            faces.append([idx1, idx2, idx3])
+            faces.append([idx2, idx4, idx3])
+    
+    return vertices, np.array(faces)
+
+
+def preprocess_all_frames(samples_info):
+    """Preprocess all frames for all samples"""
+    print("\n" + "="*60)
+    print("PREPROCESSING ALL FRAMES")
+    print("="*60)
+    
+    all_results = {}
+    
+    for sample_idx, sample in enumerate(samples_info):
+        print(f"\nProcessing {sample['name']} ({sample_idx + 1}/{len(samples_info)})")
+        sample_results = {}
+        
+        for frame_idx in range(sample['num_frames']):
+            print(f"  Frame {frame_idx}/{sample['num_frames']-1}", end='', flush=True)
+            
+            try:
+                # Load depth, segmentation and RGB
+                depth = np.load(sample['dir'] / "depth" / f"frame_{frame_idx:03d}.npy")
+                segmentation = np.load(sample['dir'] / "segmentation" / f"frame_{frame_idx:03d}.npy")
+                
+                # Load RGB image
+                rgb_path = sample['dir'] / "rgb" / f"frame_{frame_idx:03d}.png"
+                rgb_image = cv2.imread(str(rgb_path))
+                rgb_image = cv2.cvtColor(rgb_image, cv2.COLOR_BGR2RGB)
+                
+                # Get camera intrinsics
+                camera_K = np.array(sample['metadata']['camera']['K'])
+                
+                # Load camera trajectory
+                camera_traj = np.load(sample['dir'] / "camera_trajectory.npz")
+                camera_position = camera_traj['positions'][frame_idx]
+                camera_quaternion = camera_traj['quaternions'][frame_idx]
+                
+                # Convert to normalized point clouds
+                instance_pointclouds, scene_points, scene_seg, scene_center, scene_extent = depth_to_normalized_pointcloud(
+                    depth, segmentation, camera_K, 
+                    camera_position=camera_position,
+                    camera_quaternion=camera_quaternion,
+                    convert_to_zdepth=True
+                )
+                
+                # Process each instance
+                instances = []
+                for inst_id, points in instance_pointclouds.items():
+                    inst_info = sample['metadata']['instances'][inst_id - 1]
+                    
+                    try:
+                        # Fit superquadric
+                        sq, p = EMS_recovery(
+                            points,
+                            OutlierRatio=0.13,
+                            MaxIterationEM=20,
+                            AdaptiveUpperBound=True,
+                            Rescale=False
+                        )
+                        
+                        inlier_ratio = np.sum(p > 0.5) / len(p)
+                        
+                        instances.append({
+                            'id': inst_id,
+                            'info': inst_info,
+                            'points': points,
+                            'quadric': sq,
+                            'inlier_ratio': inlier_ratio,
+                            'inlier_points': points[p > 0.5]
+                        })
+                    except Exception as e:
+                        print(f" [Failed instance {inst_id}: {str(e)[:30]}...]", end='')
+                
+                # Store frame result
+                sample_results[frame_idx] = {
+                    'metadata': sample['metadata'],
+                    'instances': instances,
+                    'scene_points': scene_points,
+                    'scene_seg': scene_seg,
+                    'rgb_image': rgb_image,
+                    'camera_position': camera_position,
+                    'camera_quaternion': camera_quaternion,
+                    'scene_scale': 20.0 / scene_extent if scene_extent > 0 else 1.0,
+                    'scene_center': scene_center
+                }
+                print(" ✓", end='', flush=True)
+                
+            except Exception as e:
+                print(f" [Error: {str(e)}]", end='')
+                sample_results[frame_idx] = None
+        
+        all_results[sample['name']] = sample_results
+        print()
+    
+    print(f"\nPreprocessing complete! Processed {len(all_results)} samples")
+    return all_results
+
+
+def main():
+    # Load MOVi-A train data
+    data_dir = Path("/research/cbim/vast/sf895/code/WaveGen/WaveGen_v33_使用超二次元函数_Transformer/data/movi_a_128x128/train")
+    
+    if not data_dir.exists():
+        print(f"Error: train data not found at {data_dir}")
+        print("Please run download_movi_simple.py first to download the MOVi-A dataset")
+        return
+    
+    # Get all train samples
+    sample_dirs = sorted(list(data_dir.glob("sample_*")))
+    print(f"Found {len(sample_dirs)} train samples")
+    
+    if len(sample_dirs) == 0:
+        print("No train samples found!")
+        return
+    
+    # Pre-load sample metadata
+    samples_info = []
+    print("\nLoading sample metadata...")
+    for sample_dir in sample_dirs[:10]:  # Process first 10 samples
+        with open(sample_dir / "metadata.json", 'r') as f:
+            metadata = json.load(f)
+        samples_info.append({
+            'dir': sample_dir,
+            'name': sample_dir.name,
+            'metadata': metadata,
+            'num_frames': metadata['num_frames']
+        })
+    
+    print(f"Loaded metadata for {len(samples_info)} samples")
+    
+    # Preprocess all frames for all samples
+    all_preprocessed_results = preprocess_all_frames(samples_info)
+    
+    # Start viser visualization
+    server = viser.ViserServer(port=8080)
+    print(f"\n{'='*60}")
+    print(f"Viser server started at http://localhost:8080")
+    print("Open this URL in your browser to view the 3D visualization")
+    print("Press Ctrl+C to stop the server")
+    print('='*60)
+    
+    # Colors for different objects
+    object_colors = {
+        'cube': (255, 0, 0),      # Red
+        'sphere': (0, 255, 0),    # Green
+        'cylinder': (0, 0, 255),  # Blue
+    }
+    
+    # Colors for instances
+    instance_colors = [
+        (255, 0, 0),    # Red
+        (0, 255, 0),    # Green
+        (0, 0, 255),    # Blue
+        (255, 255, 0),  # Yellow
+        (255, 0, 255),  # Magenta
+        (0, 255, 255),  # Cyan
+        (255, 128, 0),  # Orange
+        (128, 0, 255),  # Purple
+    ]
+    
+    # Create GUI
+    with server.gui.add_folder("Controls"):
+        # Sample selector
+        sample_names = [s['name'] for s in samples_info]
+        current_sample = server.gui.add_dropdown(
+            "Select Sample",
+            options=sample_names,
+            initial_value=sample_names[0] if sample_names else None
+        )
+        
+        # Frame selector
+        frame_slider = server.gui.add_slider(
+            "Frame",
+            min=0,
+            max=23,  # MOVi-A has 24 frames
+            step=1,
+            initial_value=0
+        )
+        
+        # Playback controls
+        with server.gui.add_folder("Playback Controls"):
+            play_button = server.gui.add_button("Play ▶")
+            pause_button = server.gui.add_button("Pause ⏸")
+            fps_slider = server.gui.add_slider(
+                "Playback FPS",
+                min=1,
+                max=24,
+                step=1,
+                initial_value=12
+            )
+        
+        # Status display
+        status_display = server.gui.add_markdown("**Status:** Ready")
+        
+        # Instance selector will be updated dynamically
+        instance_folder = server.gui.add_folder("Instances")
+        
+        # Visibility controls
+        show_scene = server.gui.add_checkbox("Show Background Points", initial_value=True)
+        show_points = server.gui.add_checkbox("Highlight only the identified instance points", initial_value=True)
+        show_quadrics = server.gui.add_checkbox("Show Superquadrics", initial_value=True)
+        show_labels = server.gui.add_checkbox("Show Labels", initial_value=False)
+        show_camera = server.gui.add_checkbox("Show Camera", initial_value=False)
+        use_rgb_colors = server.gui.add_checkbox("Show Point Colors", initial_value=True)
+        
+        # Camera view button
+        match_camera_view = server.gui.add_button("Match Frame Camera View")
+        
+        # Visual parameters
+        point_size = server.gui.add_slider(
+            "Point Size",
+            min=0.001,
+            max=0.05,
+            step=0.001,
+            initial_value=0.01
+        )
+        
+        mesh_opacity = server.gui.add_slider(
+            "Mesh Opacity",
+            min=0.0,
+            max=1.0,
+            step=0.1,
+            initial_value=0.7
+        )
+        
+        # Info display
+        info_display = server.gui.add_markdown("**Sample Info:**\n\nSelect a sample to view")
+    
+    # Store current visualization handles and results
+    current_viz = {
+        'scene_cloud': None,
+        'points': {},
+        'meshes': {},
+        'labels': {},
+        'camera_frustum': None,
+        'camera_label': None,
+        'instance_toggles': [],
+        'current_result': None,
+        'all_results': all_preprocessed_results,  # Store preprocessed results
+        'is_playing': False
+    }
+    
+    def load_frame():
+        """Load the current frame from preprocessed results"""
+        sample_name = current_sample.value
+        frame_idx = int(frame_slider.value)
+        
+        if sample_name not in current_viz['all_results']:
+            status_display.value = f"**Status:** Sample {sample_name} not found in preprocessed results"
+            return
+        
+        sample_results = current_viz['all_results'][sample_name]
+        if frame_idx not in sample_results or sample_results[frame_idx] is None:
+            status_display.value = f"**Status:** Frame {frame_idx} not available"
+            return
+        
+        # Get preprocessed result
+        current_viz['current_result'] = sample_results[frame_idx]
+        current_viz['current_result']['name'] = sample_name
+        current_viz['current_result']['frame'] = frame_idx
+        
+        # Update visualization
+        num_instances = len(current_viz['current_result']['instances'])
+        status_display.value = f"**Status:** Loaded frame {frame_idx} - {num_instances} instances"
+        update_scene()
+    
+    def update_scene():
+        """Update the 3D scene based on current result"""
+        # Clear existing visualization
+        if current_viz['scene_cloud'] is not None:
+            current_viz['scene_cloud'].remove()
+            current_viz['scene_cloud'] = None
+        
+        if current_viz['camera_frustum'] is not None:
+            current_viz['camera_frustum'].remove()
+            current_viz['camera_frustum'] = None
+        
+        if current_viz['camera_label'] is not None:
+            current_viz['camera_label'].remove()
+            current_viz['camera_label'] = None
+            
+        for pc in current_viz['points'].values():
+            pc.remove()
+        current_viz['points'] = {}
+        
+        for mesh in current_viz['meshes'].values():
+            mesh.remove()
+        current_viz['meshes'] = {}
+        
+        for label in current_viz['labels'].values():
+            label.remove()
+        current_viz['labels'] = {}
+        
+        # Clear instance toggles
+        for toggle in current_viz['instance_toggles']:
+            toggle.remove()
+        current_viz['instance_toggles'] = []
+        
+        # Get current result
+        selected = current_viz['current_result']
+        
+        if selected is None:
+            info_display.value = "**Sample Info:**\n\nClick 'Process Current Frame' to start"
+            return
+        
+        # Update info
+        info_text = f"**{selected['name']} - Frame {selected['frame']}**\n\n"
+        info_text += f"Total instances: {len(selected['instances'])}\n"
+        
+        # Show full scene point cloud if requested
+        if show_scene.value:
+            scene_points_flat = selected['scene_points'].reshape(-1, 3)
+            scene_seg_flat = selected['scene_seg'].reshape(-1)
+            
+            # Filter out invalid points
+            valid_mask = ~np.isnan(scene_points_flat).any(axis=1)
+            scene_points_valid = scene_points_flat[valid_mask]
+            scene_seg_valid = scene_seg_flat[valid_mask]
+            
+            if use_rgb_colors.value and 'rgb_image' in selected:
+                # Use RGB colors from image
+                rgb_flat = selected['rgb_image'].reshape(-1, 3)
+                rgb_valid = rgb_flat[valid_mask]
+                colors = rgb_valid.astype(np.uint8)
+            else:
+                # Use segmentation colors
+                colors = np.zeros((len(scene_points_valid), 3), dtype=np.uint8)
+                for i, seg_id in enumerate(scene_seg_valid):
+                    if seg_id == 0:
+                        colors[i] = [128, 128, 128]  # Gray for background
+                    else:
+                        colors[i] = instance_colors[(seg_id - 1) % len(instance_colors)]
+            
+            current_viz['scene_cloud'] = server.scene.add_point_cloud(
+                "/scene_points",
+                points=scene_points_valid,
+                colors=colors,
+                point_size=point_size.value,
+            )
+            info_text += f"Scene points shown: {len(scene_points_valid)}\n"
+        
+        info_text += "\n"
+        
+        # Show camera if requested
+        if show_camera.value and 'camera_position' in selected:
+            # Transform camera position to normalized scene coordinates
+            cam_pos = selected['camera_position']
+            scale = selected.get('scene_scale', 1.0)
+            center = selected.get('scene_center', np.zeros(3))
+            cam_pos_normalized = (cam_pos - center) * scale
+            
+            # Get camera parameters from metadata
+            focal_length = selected['metadata']['camera']['focal_length']
+            sensor_width = selected['metadata']['camera']['sensor_width']
+            resolution = selected['metadata']['resolution']
+            
+            # Calculate field of view
+            fov = 2 * np.arctan(sensor_width / (2 * focal_length))
+            
+            # Get camera orientation
+            cam_quat = selected['camera_quaternion']
+            
+            # Convert quaternion to wxyz format (viser uses w first)
+            # MOVi quaternion is [x, y, z, w], viser needs [w, x, y, z]
+            wxyz = np.array([cam_quat[3], cam_quat[0], cam_quat[1], cam_quat[2]])
+            
+            # Get the RGB image for the camera frustum
+            if 'rgb_image' in selected:
+                # Use original resolution for camera frustum display
+                # MOVi-A is 128x128, which should display clearly
+                small_rgb = selected['rgb_image']
+                
+                # Add camera frustum with image
+                current_viz['camera_frustum'] = server.scene.add_camera_frustum(
+                    "/camera_frustum",
+                    fov=fov,
+                    aspect=1.0,  # Square aspect ratio for MOVi
+                    scale=2.0,   # Size of frustum visualization
+                    wxyz=wxyz,
+                    position=cam_pos_normalized,
+                    image=small_rgb,
+                )
+            else:
+                # Add camera frustum without image
+                current_viz['camera_frustum'] = server.scene.add_camera_frustum(
+                    "/camera_frustum",
+                    fov=fov,
+                    aspect=1.0,
+                    scale=2.0,
+                    wxyz=wxyz,
+                    position=cam_pos_normalized,
+                    color=(255, 255, 0),
+                )
+            
+            # Add camera label
+            if show_labels.value:
+                current_viz['camera_label'] = server.scene.add_label(
+                    "/camera_label",
+                    text=f"Camera Frame {selected['frame']}",
+                    position=cam_pos_normalized + np.array([0, 0.5, 0]),
+                )
+        
+        # Create instance toggles
+        with instance_folder:
+            for inst in selected['instances']:
+                inst_info = inst['info']
+                toggle = server.gui.add_checkbox(
+                    f"Instance {inst['id']}: {inst_info['shape']} ({inst_info['color']})",
+                    initial_value=True
+                )
+                current_viz['instance_toggles'].append(toggle)
+        
+        # Add instances
+        for i, inst in enumerate(selected['instances']):
+            inst_id = inst['id']
+            inst_info = inst['info']
+            
+            # Check if this instance should be shown
+            show_this = i < len(current_viz['instance_toggles']) and current_viz['instance_toggles'][i].value
+            
+            if not show_this:
+                continue
+            
+            # Get color
+            shape_name = inst_info['shape']
+            color = object_colors.get(shape_name, (128, 128, 128))
+            
+            # Add point cloud
+            if show_points.value:
+                pc = server.scene.add_point_cloud(
+                    f"/instance_{inst_id}/points",
+                    points=inst['points'],
+                    colors=np.array([color] * len(inst['points']), dtype=np.uint8),
+                    point_size=point_size.value,
+                )
+                current_viz['points'][inst_id] = pc
+            
+            # Add superquadric
+            if show_quadrics.value:
+                try:
+                    vertices, faces = generate_superquadric_mesh(inst['quadric'], num_samples=20)
+                    
+                    mesh = server.scene.add_mesh_simple(
+                        f"/instance_{inst_id}/mesh",
+                        vertices=vertices,
+                        faces=faces,
+                        color=color,
+                        opacity=mesh_opacity.value,
+                    )
+                    current_viz['meshes'][inst_id] = mesh
+                    
+                    if show_labels.value:
+                        sq = inst['quadric']
+                        label_text = f"{inst_info['shape']}\n"
+                        label_text += f"ε₁={sq.shape[0]:.2f}, ε₂={sq.shape[1]:.2f}\n"
+                        label_text += f"Inliers: {inst['inlier_ratio']:.1%}\n"
+                        label_text += f"Outliers: {(1 - inst['inlier_ratio']):.1%}"
+                        
+                        label = server.scene.add_label(
+                            f"/instance_{inst_id}/label",
+                            text=label_text,
+                            position=sq.translation,
+                        )
+                        current_viz['labels'][inst_id] = label
+                    
+                except Exception as e:
+                    print(f"Error visualizing instance {inst_id}: {e}")
+            
+            # Update info
+            info_text += f"\n**Instance {inst_id}:**\n"
+            info_text += f"- Shape: {inst_info['shape']}\n"
+            info_text += f"- Size: {inst_info['size']}\n"
+            info_text += f"- Color: {inst_info['color']}\n"
+            info_text += f"- Points: {len(inst['points'])}\n"
+            info_text += f"- ε₁={inst['quadric'].shape[0]:.3f}, ε₂={inst['quadric'].shape[1]:.3f}\n"
+            info_text += f"- Inliers: {inst['inlier_ratio']:.1%}\n"
+        
+        info_display.value = info_text
+    
+    # Set up callbacks
+    @current_sample.on_update
+    def _(_):
+        # Update frame slider max value based on selected sample
+        for s in samples_info:
+            if s['name'] == current_sample.value:
+                frame_slider.max = s['num_frames'] - 1
+                frame_slider.value = 0  # Reset to first frame
+                break
+        load_frame()  # Automatically load when sample changes
+    
+    @frame_slider.on_update
+    def _(_):
+        if not current_viz['is_playing']:  # Only load if not playing (playback will handle it)
+            load_frame()
+    
+    # Playback functions
+    import threading
+    playback_thread = None
+    
+    def playback_loop():
+        """Playback loop in separate thread"""
+        while current_viz['is_playing']:
+            # Move to next frame
+            current_frame = int(frame_slider.value)
+            next_frame = (current_frame + 1) % (frame_slider.max + 1)
+            frame_slider.value = next_frame
+            load_frame()
+            
+            # Sleep based on FPS
+            time.sleep(1.0 / fps_slider.value)
+    
+    @play_button.on_click
+    def _(_):
+        if not current_viz['is_playing']:
+            current_viz['is_playing'] = True
+            play_button.disabled = True
+            pause_button.disabled = False
+            # Start playback thread
+            playback_thread = threading.Thread(target=playback_loop)
+            playback_thread.start()
+            status_display.value = "**Status:** Playing..."
+    
+    @pause_button.on_click
+    def _(_):
+        if current_viz['is_playing']:
+            current_viz['is_playing'] = False
+            play_button.disabled = False
+            pause_button.disabled = True
+            status_display.value = "**Status:** Paused"
+    
+    @show_scene.on_update
+    def _(_):
+        update_scene()
+    
+    @use_rgb_colors.on_update
+    def _(_):
+        if show_scene.value:
+            update_scene()
+    
+    @show_points.on_update
+    def _(_):
+        for pc in current_viz['points'].values():
+            pc.visible = show_points.value
+    
+    @show_quadrics.on_update
+    def _(_):
+        for mesh in current_viz['meshes'].values():
+            mesh.visible = show_quadrics.value
+    
+    @show_labels.on_update
+    def _(_):
+        for label in current_viz['labels'].values():
+            label.visible = show_labels.value
+    
+    @point_size.on_update
+    def _(event):
+        if current_viz['scene_cloud'] is not None:
+            current_viz['scene_cloud'].point_size = event.target.value
+        for pc in current_viz['points'].values():
+            pc.point_size = event.target.value
+    
+    @mesh_opacity.on_update
+    def _(event):
+        for mesh in current_viz['meshes'].values():
+            mesh.opacity = event.target.value
+    
+    @match_camera_view.on_click
+    def _(event: viser.GuiEvent):
+        """Set the viewer camera to match the current frame's camera"""
+        if current_viz['current_result'] is None:
+            status_display.value = "**Status:** Please process a frame first"
+            return
+            
+        result = current_viz['current_result']
+        if 'camera_position' not in result:
+            status_display.value = "**Status:** No camera data available"
+            return
+        
+        # Get normalized camera position and orientation
+        cam_pos = result['camera_position']
+        cam_quat = result['camera_quaternion']  # [x, y, z, w] format
+        
+        # Apply scene normalization to camera position
+        scale = result.get('scene_scale', 1.0)
+        center = result.get('scene_center', np.zeros(3))
+        cam_pos_normalized = (cam_pos - center) * scale
+        
+        # Convert quaternion from xyzw to wxyz format for viser
+        wxyz = np.array([cam_quat[3], cam_quat[0], cam_quat[1], cam_quat[2]])
+        
+        # Set camera position and orientation
+        client = event.client
+        if client is not None:
+            client.camera.position = cam_pos_normalized
+            client.camera.wxyz = wxyz
+            
+            # Also update the up direction based on camera orientation
+            from scipy.spatial.transform import Rotation
+            rot = Rotation.from_quat(cam_quat)  # xyzw format
+            # In MOVi, +Y is up in camera space, but viser might need -Y
+            # Try negative Y to fix the upside-down issue
+            camera_up = rot.apply([0, -1, 0])
+            client.camera.up_direction = camera_up
+            
+            status_display.value = f"**Status:** Matched camera view for frame {result['frame']}"
+    
+    # Initial setup
+    pause_button.disabled = True  # Initially disabled
+    
+    # Load first frame of first sample
+    if len(samples_info) > 0:
+        load_frame()
+    else:
+        info_display.value = "**Sample Info:**\n\nNo samples available"
+    
+    # Set initial viewer camera to a reasonable position
+    # Look at the origin from a distance
+    server.scene.set_up_direction("+y")
+    
+    # Keep server running
+    try:
+        while True:
+            time.sleep(0.1)
+    except KeyboardInterrupt:
+        print("\nShutting down server...")
+        server.stop()
+
+
+if __name__ == "__main__":
+    main()

EMS-superquadric_fitting_inference/process_movi_validation.py

@@ -0,0 +1,886 @@
+#!/usr/bin/env python3
+"""
+Process MOVi-A validation set with hierarchical multi-superquadric fitting
+Converts depth maps to normalized point clouds for superquadric fitting
+"""
+
+import numpy as np
+import sys
+import os
+import time
+import viser
+import json
+import cv2
+from pathlib import Path
+from sklearn.cluster import DBSCAN
+
+# Add the src directory to Python path
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src'))
+
+from EMS.EMS_recovery import EMS_recovery
+
+
+def depth_to_normalized_pointcloud(depth, segmentation, camera_K, camera_position=None, camera_quaternion=None, resolution=128, convert_to_zdepth=True):
+    """
+    Convert depth map to normalized point cloud in range [-10, 10]
+    
+    Args:
+        depth: (H, W, 1) depth array (euclidean distance from camera center)
+        segmentation: (H, W) instance segmentation mask
+        camera_K: 3x3 camera intrinsic matrix
+        camera_position: camera position in world coordinates
+        camera_quaternion: camera quaternion (x,y,z,w) in world coordinates
+        resolution: image resolution (assuming square)
+        convert_to_zdepth: bool, convert euclidean depth to z-depth before processing
+    
+    Returns:
+        dict: instance_id -> normalized point cloud in world coordinates
+    """
+    H, W = depth.shape[:2]
+    
+    # Get camera parameters
+    fx = camera_K[0, 0]
+    fy = camera_K[1, 1]
+    cx = camera_K[0, 2]
+    cy = camera_K[1, 2]
+    
+    # Create pixel grid
+    xx, yy = np.meshgrid(np.arange(W), np.arange(H))
+    
+    # Convert to normalized camera coordinates
+    x_norm = (xx - cx) / fx
+    y_norm = (yy - cy) / fy
+    
+    if convert_to_zdepth:
+        # MOVi uses euclidean distance, convert to z-depth (planar depth)
+        # For each pixel, we have: euclidean_dist^2 = x^2 + y^2 + z^2
+        # Where x = x_norm * z, y = y_norm * z
+        # So: euclidean_dist^2 = (x_norm^2 + y_norm^2 + 1) * z^2
+        z = depth[:, :, 0] / np.sqrt(x_norm**2 + y_norm**2 + 1)
+    else:
+        # Use depth as-is (assume it's already z-depth)
+        z = depth[:, :, 0]
+    
+    # Get 3D points
+    x = x_norm * z
+    y = y_norm * z
+    
+    # Stack to get point cloud (in camera coordinates)
+    points_3d_camera = np.stack([x, y, z], axis=-1)
+    
+    # Transform from camera to world coordinates if camera pose is provided
+    if camera_position is not None and camera_quaternion is not None:
+        from scipy.spatial.transform import Rotation
+        
+        # Convert quaternion to rotation matrix
+        # MOVi uses [x, y, z, w] format
+        cam_rot = Rotation.from_quat(camera_quaternion)
+        cam_rot_matrix = cam_rot.as_matrix()
+        
+        # Transform points: World = R * Camera + T
+        points_3d_flat = points_3d_camera.reshape(-1, 3)
+        points_3d_world = points_3d_flat @ cam_rot_matrix.T + camera_position
+        points_3d = points_3d_world.reshape(points_3d_camera.shape)
+    else:
+        points_3d = points_3d_camera
+    
+    
+    # Normalize entire scene to [-10, 10] range
+    # Find scene bounds (only valid depth points)
+    valid_mask = z > 0
+    valid_points = points_3d[valid_mask]
+    
+    if len(valid_points) > 0:
+        # Find scene extent
+        scene_min = np.min(valid_points, axis=0)
+        scene_max = np.max(valid_points, axis=0)
+        scene_center = (scene_min + scene_max) / 2
+        scene_extent = np.max(scene_max - scene_min)
+        
+        # Scale to [-10, 10]
+        if scene_extent > 0:
+            scale_factor = 20.0 / scene_extent  # 20 because we want -10 to 10
+            points_3d_normalized = (points_3d - scene_center) * scale_factor
+        else:
+            points_3d_normalized = points_3d - scene_center
+    else:
+        points_3d_normalized = points_3d
+    
+    # Get unique instance IDs (excluding background=0)
+    instance_ids = np.unique(segmentation)
+    instance_ids = instance_ids[instance_ids > 0]
+    
+    instance_pointclouds = {}
+    
+    for inst_id in instance_ids:
+        # Get mask for this instance
+        mask = segmentation == inst_id
+        
+        # Extract points for this instance (already normalized with scene)
+        instance_points = points_3d_normalized[mask]
+        
+        if len(instance_points) < 50:  # Skip if too few points
+            continue
+        
+        instance_pointclouds[int(inst_id)] = instance_points
+    
+    # Also return the full scene point cloud and segmentation for visualization
+    return instance_pointclouds, points_3d_normalized, segmentation, scene_center if 'scene_center' in locals() else np.zeros(3), scene_extent if 'scene_extent' in locals() else 1.0
+
+
+def hierarchical_ems(
+    point,
+    OutlierRatio=0.5,
+    MaxIterationEM=20,
+    ToleranceEM=1e-3,
+    RelativeToleranceEM=2e-1,
+    MaxOptiIterations=2,
+    Sigma=0.3,
+    MaxiSwitch=2,
+    AdaptiveUpperBound=True,
+    Rescale=False,
+    MaxLayer=3,
+    Eps=1.0,  # Adjusted for normalized [-10, 10] point clouds
+    MinPoints=50,
+):
+    """
+    Hierarchical EMS for extracting multiple superquadrics from a point cloud
+    """
+    point_seg = {key: [] for key in list(range(0, MaxLayer+1))}
+    point_outlier = {key: [] for key in list(range(0, MaxLayer+1))}
+    point_seg[0] = [point]
+    list_quadrics = []
+    quadric_info = []
+    
+    for h in range(MaxLayer):
+        if len(point_seg[h]) == 0:
+            break
+            
+        for c in range(len(point_seg[h])):
+            current_points = point_seg[h][c]
+            if len(current_points) < MinPoints * 2:
+                continue
+                
+            try:
+                # Fit superquadric
+                x_raw, p_raw = EMS_recovery(
+                    current_points,
+                    OutlierRatio,
+                    MaxIterationEM,
+                    ToleranceEM,
+                    RelativeToleranceEM,
+                    MaxOptiIterations,
+                    Sigma,
+                    MaxiSwitch,
+                    AdaptiveUpperBound,
+                    Rescale,
+                )
+                
+                # Calculate fitting quality
+                inlier_mask = p_raw > 0.5
+                inlier_ratio = np.sum(inlier_mask) / len(p_raw)
+                
+                if inlier_ratio > 0.3:  # Accept if at least 30% inliers
+                    list_quadrics.append(x_raw)
+                    quadric_info.append({
+                        'layer': h,
+                        'segment': c,
+                        'inlier_ratio': inlier_ratio,
+                        'num_points': len(current_points),
+                        'inlier_points': current_points[inlier_mask]
+                    })
+                
+                # Separate outliers for next layer
+                outlier_mask = p_raw < 0.1
+                outlier = current_points[outlier_mask]
+                
+                # If many outliers and not last layer, try clustering
+                if len(outlier) > MinPoints * 2 and h < MaxLayer - 1:
+                    clustering = DBSCAN(eps=Eps, min_samples=MinPoints).fit(outlier)
+                    labels = list(set(clustering.labels_))
+                    labels = [item for item in labels if item >= 0]
+                    
+                    if len(labels) >= 1:
+                        for i in range(len(labels)):
+                            cluster_points = outlier[clustering.labels_ == labels[i]]
+                            if len(cluster_points) > MinPoints:
+                                point_seg[h + 1].append(cluster_points)
+                
+            except Exception as e:
+                continue
+    
+    return list_quadrics, quadric_info
+
+
+def generate_superquadric_mesh(sq, num_samples=25):
+    """Generate mesh vertices and faces for superquadric surface"""
+    eta = np.linspace(-np.pi/2, np.pi/2, num_samples)
+    omega = np.linspace(-np.pi, np.pi, num_samples)
+    
+    vertices = []
+    faces = []
+    
+    # Generate vertices
+    for i, e in enumerate(eta):
+        for j, w in enumerate(omega):
+            # Superquadric parametric equations
+            cos_eta = np.sign(np.cos(e)) * np.abs(np.cos(e))**sq.shape[0]
+            sin_eta = np.sign(np.sin(e)) * np.abs(np.sin(e))**sq.shape[0]
+            cos_omega = np.sign(np.cos(w)) * np.abs(np.cos(w))**sq.shape[1]
+            sin_omega = np.sign(np.sin(w)) * np.abs(np.sin(w))**sq.shape[1]
+            
+            # Local coordinates
+            x_local = sq.scale[0] * cos_eta * cos_omega
+            y_local = sq.scale[1] * cos_eta * sin_omega
+            z_local = sq.scale[2] * sin_eta
+            
+            # Apply rotation and translation
+            point_local = np.array([x_local, y_local, z_local])
+            point_global = sq.RotM @ point_local + sq.translation
+            
+            vertices.append(point_global)
+    
+    vertices = np.array(vertices)
+    
+    # Generate faces (triangles)
+    for i in range(num_samples - 1):
+        for j in range(num_samples - 1):
+            # Current vertex indices
+            idx1 = i * num_samples + j
+            idx2 = i * num_samples + (j + 1) % num_samples
+            idx3 = (i + 1) * num_samples + j
+            idx4 = (i + 1) * num_samples + (j + 1) % num_samples
+            
+            # Two triangles per quad
+            faces.append([idx1, idx2, idx3])
+            faces.append([idx2, idx4, idx3])
+    
+    return vertices, np.array(faces)
+
+
+def preprocess_all_frames(samples_info):
+    """Preprocess all frames for all samples"""
+    print("\n" + "="*60)
+    print("PREPROCESSING ALL FRAMES")
+    print("="*60)
+    
+    all_results = {}
+    
+    for sample_idx, sample in enumerate(samples_info):
+        print(f"\nProcessing {sample['name']} ({sample_idx + 1}/{len(samples_info)})")
+        sample_results = {}
+        
+        for frame_idx in range(sample['num_frames']):
+            print(f"  Frame {frame_idx}/{sample['num_frames']-1}", end='', flush=True)
+            
+            try:
+                # Load depth, segmentation and RGB
+                depth = np.load(sample['dir'] / "depth" / f"frame_{frame_idx:03d}.npy")
+                segmentation = np.load(sample['dir'] / "segmentation" / f"frame_{frame_idx:03d}.npy")
+                
+                # Load RGB image
+                rgb_path = sample['dir'] / "rgb" / f"frame_{frame_idx:03d}.png"
+                rgb_image = cv2.imread(str(rgb_path))
+                rgb_image = cv2.cvtColor(rgb_image, cv2.COLOR_BGR2RGB)
+                
+                # Get camera intrinsics
+                camera_K = np.array(sample['metadata']['camera']['K'])
+                
+                # Load camera trajectory
+                camera_traj = np.load(sample['dir'] / "camera_trajectory.npz")
+                camera_position = camera_traj['positions'][frame_idx]
+                camera_quaternion = camera_traj['quaternions'][frame_idx]
+                
+                # Convert to normalized point clouds
+                instance_pointclouds, scene_points, scene_seg, scene_center, scene_extent = depth_to_normalized_pointcloud(
+                    depth, segmentation, camera_K, 
+                    camera_position=camera_position,
+                    camera_quaternion=camera_quaternion,
+                    convert_to_zdepth=True
+                )
+                
+                # Process each instance
+                instances = []
+                for inst_id, points in instance_pointclouds.items():
+                    inst_info = sample['metadata']['instances'][inst_id - 1]
+                    
+                    try:
+                        # Fit superquadric
+                        sq, p = EMS_recovery(
+                            points,
+                            OutlierRatio=0.13,
+                            MaxIterationEM=20,
+                            AdaptiveUpperBound=True,
+                            Rescale=False
+                        )
+                        
+                        inlier_ratio = np.sum(p > 0.5) / len(p)
+                        
+                        instances.append({
+                            'id': inst_id,
+                            'info': inst_info,
+                            'points': points,
+                            'quadric': sq,
+                            'inlier_ratio': inlier_ratio,
+                            'inlier_points': points[p > 0.5]
+                        })
+                    except Exception as e:
+                        print(f" [Failed instance {inst_id}: {str(e)[:30]}...]", end='')
+                
+                # Store frame result
+                sample_results[frame_idx] = {
+                    'metadata': sample['metadata'],
+                    'instances': instances,
+                    'scene_points': scene_points,
+                    'scene_seg': scene_seg,
+                    'rgb_image': rgb_image,
+                    'camera_position': camera_position,
+                    'camera_quaternion': camera_quaternion,
+                    'scene_scale': 20.0 / scene_extent if scene_extent > 0 else 1.0,
+                    'scene_center': scene_center
+                }
+                print(" ✓", end='', flush=True)
+                
+            except Exception as e:
+                print(f" [Error: {str(e)}]", end='')
+                sample_results[frame_idx] = None
+        
+        all_results[sample['name']] = sample_results
+        print()
+    
+    print(f"\nPreprocessing complete! Processed {len(all_results)} samples")
+    return all_results
+
+
+def main():
+    # Load MOVi-A validation data
+    data_dir = Path("/research/cbim/vast/sf895/code/WaveGen/WaveGen_v33_使用超二次元函数_Transformer/data/movi_a_128x128/validation")
+    
+    if not data_dir.exists():
+        print(f"Error: Validation data not found at {data_dir}")
+        print("Please run download_movi_simple.py first to download the MOVi-A dataset")
+        return
+    
+    # Get all validation samples
+    sample_dirs = sorted(list(data_dir.glob("sample_*")))
+    print(f"Found {len(sample_dirs)} validation samples")
+    
+    if len(sample_dirs) == 0:
+        print("No validation samples found!")
+        return
+    
+    # Pre-load sample metadata
+    samples_info = []
+    print("\nLoading sample metadata...")
+    for sample_dir in sample_dirs[:10]:  # Process first 10 samples
+        with open(sample_dir / "metadata.json", 'r') as f:
+            metadata = json.load(f)
+        samples_info.append({
+            'dir': sample_dir,
+            'name': sample_dir.name,
+            'metadata': metadata,
+            'num_frames': metadata['num_frames']
+        })
+    
+    print(f"Loaded metadata for {len(samples_info)} samples")
+    
+    # Preprocess all frames for all samples
+    all_preprocessed_results = preprocess_all_frames(samples_info)
+    
+    # Start viser visualization
+    server = viser.ViserServer(port=8080)
+    print(f"\n{'='*60}")
+    print(f"Viser server started at http://localhost:8080")
+    print("Open this URL in your browser to view the 3D visualization")
+    print("Press Ctrl+C to stop the server")
+    print('='*60)
+    
+    # Colors for different objects
+    object_colors = {
+        'cube': (255, 0, 0),      # Red
+        'sphere': (0, 255, 0),    # Green
+        'cylinder': (0, 0, 255),  # Blue
+    }
+    
+    # Colors for instances
+    instance_colors = [
+        (255, 0, 0),    # Red
+        (0, 255, 0),    # Green
+        (0, 0, 255),    # Blue
+        (255, 255, 0),  # Yellow
+        (255, 0, 255),  # Magenta
+        (0, 255, 255),  # Cyan
+        (255, 128, 0),  # Orange
+        (128, 0, 255),  # Purple
+    ]
+    
+    # Create GUI
+    with server.gui.add_folder("Controls"):
+        # Sample selector
+        sample_names = [s['name'] for s in samples_info]
+        current_sample = server.gui.add_dropdown(
+            "Select Sample",
+            options=sample_names,
+            initial_value=sample_names[0] if sample_names else None
+        )
+        
+        # Frame selector
+        frame_slider = server.gui.add_slider(
+            "Frame",
+            min=0,
+            max=23,  # MOVi-A has 24 frames
+            step=1,
+            initial_value=0
+        )
+        
+        # Playback controls
+        with server.gui.add_folder("Playback Controls"):
+            play_button = server.gui.add_button("Play ▶")
+            pause_button = server.gui.add_button("Pause ⏸")
+            fps_slider = server.gui.add_slider(
+                "Playback FPS",
+                min=1,
+                max=24,
+                step=1,
+                initial_value=12
+            )
+        
+        # Status display
+        status_display = server.gui.add_markdown("**Status:** Ready")
+        
+        # Instance selector will be updated dynamically
+        instance_folder = server.gui.add_folder("Instances")
+        
+        # Visibility controls
+        show_scene = server.gui.add_checkbox("Show Background Points", initial_value=True)
+        show_points = server.gui.add_checkbox("Highlight only the identified instance points", initial_value=True)
+        show_quadrics = server.gui.add_checkbox("Show Superquadrics", initial_value=True)
+        show_labels = server.gui.add_checkbox("Show Labels", initial_value=False)
+        show_camera = server.gui.add_checkbox("Show Camera", initial_value=False)
+        use_rgb_colors = server.gui.add_checkbox("Show Point Colors", initial_value=True)
+        
+        # Camera view button
+        match_camera_view = server.gui.add_button("Match Frame Camera View")
+        
+        # Visual parameters
+        point_size = server.gui.add_slider(
+            "Point Size",
+            min=0.001,
+            max=0.05,
+            step=0.001,
+            initial_value=0.01
+        )
+        
+        mesh_opacity = server.gui.add_slider(
+            "Mesh Opacity",
+            min=0.0,
+            max=1.0,
+            step=0.1,
+            initial_value=0.7
+        )
+        
+        # Info display
+        info_display = server.gui.add_markdown("**Sample Info:**\n\nSelect a sample to view")
+    
+    # Store current visualization handles and results
+    current_viz = {
+        'scene_cloud': None,
+        'points': {},
+        'meshes': {},
+        'labels': {},
+        'camera_frustum': None,
+        'camera_label': None,
+        'instance_toggles': [],
+        'current_result': None,
+        'all_results': all_preprocessed_results,  # Store preprocessed results
+        'is_playing': False
+    }
+    
+    def load_frame():
+        """Load the current frame from preprocessed results"""
+        sample_name = current_sample.value
+        frame_idx = int(frame_slider.value)
+        
+        if sample_name not in current_viz['all_results']:
+            status_display.value = f"**Status:** Sample {sample_name} not found in preprocessed results"
+            return
+        
+        sample_results = current_viz['all_results'][sample_name]
+        if frame_idx not in sample_results or sample_results[frame_idx] is None:
+            status_display.value = f"**Status:** Frame {frame_idx} not available"
+            return
+        
+        # Get preprocessed result
+        current_viz['current_result'] = sample_results[frame_idx]
+        current_viz['current_result']['name'] = sample_name
+        current_viz['current_result']['frame'] = frame_idx
+        
+        # Update visualization
+        num_instances = len(current_viz['current_result']['instances'])
+        status_display.value = f"**Status:** Loaded frame {frame_idx} - {num_instances} instances"
+        update_scene()
+    
+    def update_scene():
+        """Update the 3D scene based on current result"""
+        # Clear existing visualization
+        if current_viz['scene_cloud'] is not None:
+            current_viz['scene_cloud'].remove()
+            current_viz['scene_cloud'] = None
+        
+        if current_viz['camera_frustum'] is not None:
+            current_viz['camera_frustum'].remove()
+            current_viz['camera_frustum'] = None
+        
+        if current_viz['camera_label'] is not None:
+            current_viz['camera_label'].remove()
+            current_viz['camera_label'] = None
+            
+        for pc in current_viz['points'].values():
+            pc.remove()
+        current_viz['points'] = {}
+        
+        for mesh in current_viz['meshes'].values():
+            mesh.remove()
+        current_viz['meshes'] = {}
+        
+        for label in current_viz['labels'].values():
+            label.remove()
+        current_viz['labels'] = {}
+        
+        # Clear instance toggles
+        for toggle in current_viz['instance_toggles']:
+            toggle.remove()
+        current_viz['instance_toggles'] = []
+        
+        # Get current result
+        selected = current_viz['current_result']
+        
+        if selected is None:
+            info_display.value = "**Sample Info:**\n\nClick 'Process Current Frame' to start"
+            return
+        
+        # Update info
+        info_text = f"**{selected['name']} - Frame {selected['frame']}**\n\n"
+        info_text += f"Total instances: {len(selected['instances'])}\n"
+        
+        # Show full scene point cloud if requested
+        if show_scene.value:
+            scene_points_flat = selected['scene_points'].reshape(-1, 3)
+            scene_seg_flat = selected['scene_seg'].reshape(-1)
+            
+            # Filter out invalid points
+            valid_mask = ~np.isnan(scene_points_flat).any(axis=1)
+            scene_points_valid = scene_points_flat[valid_mask]
+            scene_seg_valid = scene_seg_flat[valid_mask]
+            
+            if use_rgb_colors.value and 'rgb_image' in selected:
+                # Use RGB colors from image
+                rgb_flat = selected['rgb_image'].reshape(-1, 3)
+                rgb_valid = rgb_flat[valid_mask]
+                colors = rgb_valid.astype(np.uint8)
+            else:
+                # Use segmentation colors
+                colors = np.zeros((len(scene_points_valid), 3), dtype=np.uint8)
+                for i, seg_id in enumerate(scene_seg_valid):
+                    if seg_id == 0:
+                        colors[i] = [128, 128, 128]  # Gray for background
+                    else:
+                        colors[i] = instance_colors[(seg_id - 1) % len(instance_colors)]
+            
+            current_viz['scene_cloud'] = server.scene.add_point_cloud(
+                "/scene_points",
+                points=scene_points_valid,
+                colors=colors,
+                point_size=point_size.value,
+            )
+            info_text += f"Scene points shown: {len(scene_points_valid)}\n"
+        
+        info_text += "\n"
+        
+        # Show camera if requested
+        if show_camera.value and 'camera_position' in selected:
+            # Transform camera position to normalized scene coordinates
+            cam_pos = selected['camera_position']
+            scale = selected.get('scene_scale', 1.0)
+            center = selected.get('scene_center', np.zeros(3))
+            cam_pos_normalized = (cam_pos - center) * scale
+            
+            # Get camera parameters from metadata
+            focal_length = selected['metadata']['camera']['focal_length']
+            sensor_width = selected['metadata']['camera']['sensor_width']
+            resolution = selected['metadata']['resolution']
+            
+            # Calculate field of view
+            fov = 2 * np.arctan(sensor_width / (2 * focal_length))
+            
+            # Get camera orientation
+            cam_quat = selected['camera_quaternion']
+            
+            # Convert quaternion to wxyz format (viser uses w first)
+            # MOVi quaternion is [x, y, z, w], viser needs [w, x, y, z]
+            wxyz = np.array([cam_quat[3], cam_quat[0], cam_quat[1], cam_quat[2]])
+            
+            # Get the RGB image for the camera frustum
+            if 'rgb_image' in selected:
+                # Use original resolution for camera frustum display
+                # MOVi-A is 128x128, which should display clearly
+                small_rgb = selected['rgb_image']
+                
+                # Add camera frustum with image
+                current_viz['camera_frustum'] = server.scene.add_camera_frustum(
+                    "/camera_frustum",
+                    fov=fov,
+                    aspect=1.0,  # Square aspect ratio for MOVi
+                    scale=2.0,   # Size of frustum visualization
+                    wxyz=wxyz,
+                    position=cam_pos_normalized,
+                    image=small_rgb,
+                )
+            else:
+                # Add camera frustum without image
+                current_viz['camera_frustum'] = server.scene.add_camera_frustum(
+                    "/camera_frustum",
+                    fov=fov,
+                    aspect=1.0,
+                    scale=2.0,
+                    wxyz=wxyz,
+                    position=cam_pos_normalized,
+                    color=(255, 255, 0),
+                )
+            
+            # Add camera label
+            if show_labels.value:
+                current_viz['camera_label'] = server.scene.add_label(
+                    "/camera_label",
+                    text=f"Camera Frame {selected['frame']}",
+                    position=cam_pos_normalized + np.array([0, 0.5, 0]),
+                )
+        
+        # Create instance toggles
+        with instance_folder:
+            for inst in selected['instances']:
+                inst_info = inst['info']
+                toggle = server.gui.add_checkbox(
+                    f"Instance {inst['id']}: {inst_info['shape']} ({inst_info['color']})",
+                    initial_value=True
+                )
+                current_viz['instance_toggles'].append(toggle)
+        
+        # Add instances
+        for i, inst in enumerate(selected['instances']):
+            inst_id = inst['id']
+            inst_info = inst['info']
+            
+            # Check if this instance should be shown
+            show_this = i < len(current_viz['instance_toggles']) and current_viz['instance_toggles'][i].value
+            
+            if not show_this:
+                continue
+            
+            # Get color
+            shape_name = inst_info['shape']
+            color = object_colors.get(shape_name, (128, 128, 128))
+            
+            # Add point cloud
+            if show_points.value:
+                pc = server.scene.add_point_cloud(
+                    f"/instance_{inst_id}/points",
+                    points=inst['points'],
+                    colors=np.array([color] * len(inst['points']), dtype=np.uint8),
+                    point_size=point_size.value,
+                )
+                current_viz['points'][inst_id] = pc
+            
+            # Add superquadric
+            if show_quadrics.value:
+                try:
+                    vertices, faces = generate_superquadric_mesh(inst['quadric'], num_samples=20)
+                    
+                    mesh = server.scene.add_mesh_simple(
+                        f"/instance_{inst_id}/mesh",
+                        vertices=vertices,
+                        faces=faces,
+                        color=color,
+                        opacity=mesh_opacity.value,
+                    )
+                    current_viz['meshes'][inst_id] = mesh
+                    
+                    if show_labels.value:
+                        sq = inst['quadric']
+                        label_text = f"{inst_info['shape']}\n"
+                        label_text += f"ε₁={sq.shape[0]:.2f}, ε₂={sq.shape[1]:.2f}\n"
+                        label_text += f"Inliers: {inst['inlier_ratio']:.1%}\n"
+                        label_text += f"Outliers: {(1 - inst['inlier_ratio']):.1%}"
+                        
+                        label = server.scene.add_label(
+                            f"/instance_{inst_id}/label",
+                            text=label_text,
+                            position=sq.translation,
+                        )
+                        current_viz['labels'][inst_id] = label
+                    
+                except Exception as e:
+                    print(f"Error visualizing instance {inst_id}: {e}")
+            
+            # Update info
+            info_text += f"\n**Instance {inst_id}:**\n"
+            info_text += f"- Shape: {inst_info['shape']}\n"
+            info_text += f"- Size: {inst_info['size']}\n"
+            info_text += f"- Color: {inst_info['color']}\n"
+            info_text += f"- Points: {len(inst['points'])}\n"
+            info_text += f"- ε₁={inst['quadric'].shape[0]:.3f}, ε₂={inst['quadric'].shape[1]:.3f}\n"
+            info_text += f"- Inliers: {inst['inlier_ratio']:.1%}\n"
+        
+        info_display.value = info_text
+    
+    # Set up callbacks
+    @current_sample.on_update
+    def _(_):
+        # Update frame slider max value based on selected sample
+        for s in samples_info:
+            if s['name'] == current_sample.value:
+                frame_slider.max = s['num_frames'] - 1
+                frame_slider.value = 0  # Reset to first frame
+                break
+        load_frame()  # Automatically load when sample changes
+    
+    @frame_slider.on_update
+    def _(_):
+        if not current_viz['is_playing']:  # Only load if not playing (playback will handle it)
+            load_frame()
+    
+    # Playback functions
+    import threading
+    playback_thread = None
+    
+    def playback_loop():
+        """Playback loop in separate thread"""
+        while current_viz['is_playing']:
+            # Move to next frame
+            current_frame = int(frame_slider.value)
+            next_frame = (current_frame + 1) % (frame_slider.max + 1)
+            frame_slider.value = next_frame
+            load_frame()
+            
+            # Sleep based on FPS
+            time.sleep(1.0 / fps_slider.value)
+    
+    @play_button.on_click
+    def _(_):
+        if not current_viz['is_playing']:
+            current_viz['is_playing'] = True
+            play_button.disabled = True
+            pause_button.disabled = False
+            # Start playback thread
+            playback_thread = threading.Thread(target=playback_loop)
+            playback_thread.start()
+            status_display.value = "**Status:** Playing..."
+    
+    @pause_button.on_click
+    def _(_):
+        if current_viz['is_playing']:
+            current_viz['is_playing'] = False
+            play_button.disabled = False
+            pause_button.disabled = True
+            status_display.value = "**Status:** Paused"
+    
+    @show_scene.on_update
+    def _(_):
+        update_scene()
+    
+    @use_rgb_colors.on_update
+    def _(_):
+        if show_scene.value:
+            update_scene()
+    
+    @show_points.on_update
+    def _(_):
+        for pc in current_viz['points'].values():
+            pc.visible = show_points.value
+    
+    @show_quadrics.on_update
+    def _(_):
+        for mesh in current_viz['meshes'].values():
+            mesh.visible = show_quadrics.value
+    
+    @show_labels.on_update
+    def _(_):
+        for label in current_viz['labels'].values():
+            label.visible = show_labels.value
+    
+    @point_size.on_update
+    def _(event):
+        if current_viz['scene_cloud'] is not None:
+            current_viz['scene_cloud'].point_size = event.target.value
+        for pc in current_viz['points'].values():
+            pc.point_size = event.target.value
+    
+    @mesh_opacity.on_update
+    def _(event):
+        for mesh in current_viz['meshes'].values():
+            mesh.opacity = event.target.value
+    
+    @match_camera_view.on_click
+    def _(event: viser.GuiEvent):
+        """Set the viewer camera to match the current frame's camera"""
+        if current_viz['current_result'] is None:
+            status_display.value = "**Status:** Please process a frame first"
+            return
+            
+        result = current_viz['current_result']
+        if 'camera_position' not in result:
+            status_display.value = "**Status:** No camera data available"
+            return
+        
+        # Get normalized camera position and orientation
+        cam_pos = result['camera_position']
+        cam_quat = result['camera_quaternion']  # [x, y, z, w] format
+        
+        # Apply scene normalization to camera position
+        scale = result.get('scene_scale', 1.0)
+        center = result.get('scene_center', np.zeros(3))
+        cam_pos_normalized = (cam_pos - center) * scale
+        
+        # Convert quaternion from xyzw to wxyz format for viser
+        wxyz = np.array([cam_quat[3], cam_quat[0], cam_quat[1], cam_quat[2]])
+        
+        # Set camera position and orientation
+        client = event.client
+        if client is not None:
+            client.camera.position = cam_pos_normalized
+            client.camera.wxyz = wxyz
+            
+            # Also update the up direction based on camera orientation
+            from scipy.spatial.transform import Rotation
+            rot = Rotation.from_quat(cam_quat)  # xyzw format
+            # In MOVi, +Y is up in camera space, but viser might need -Y
+            # Try negative Y to fix the upside-down issue
+            camera_up = rot.apply([0, -1, 0])
+            client.camera.up_direction = camera_up
+            
+            status_display.value = f"**Status:** Matched camera view for frame {result['frame']}"
+    
+    # Initial setup
+    pause_button.disabled = True  # Initially disabled
+    
+    # Load first frame of first sample
+    if len(samples_info) > 0:
+        load_frame()
+    else:
+        info_display.value = "**Sample Info:**\n\nNo samples available"
+    
+    # Set initial viewer camera to a reasonable position
+    # Look at the origin from a distance
+    server.scene.set_up_direction("+y")
+    
+    # Keep server running
+    try:
+        while True:
+            time.sleep(0.1)
+    except KeyboardInterrupt:
+        print("\nShutting down server...")
+        server.stop()
+
+
+if __name__ == "__main__":
+    main()

EMS-superquadric_fitting_inference/process_viser_hierarchical.py

@@ -0,0 +1,486 @@
+#!/usr/bin/env python3
+"""
+Hierarchical multi-superquadric fitting with viser visualization
+Based on the hierarchical_ems algorithm from multiquadric_test.py
+"""
+
+import numpy as np
+import sys
+import os
+import time
+import viser
+from sklearn.cluster import DBSCAN
+
+# Add the src directory to Python path
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src'))
+
+from EMS.EMS_recovery import EMS_recovery
+
+
+def hierarchical_ems(
+    point,
+    OutlierRatio=0.5,           # Reduced for better initial fit
+    MaxIterationEM=20,
+    ToleranceEM=1e-3,
+    RelativeToleranceEM=2e-1,
+    MaxOptiIterations=2,
+    Sigma=0.3,
+    MaxiSwitch=2,
+    AdaptiveUpperBound=True,
+    Rescale=False,
+    MaxLayer=3,                 # Reduced for faster processing
+    Eps=0.1,                    # Adjusted for normalized point clouds
+    MinPoints=50,               # Minimum points to form a cluster
+):
+    """
+    Hierarchical EMS for extracting multiple superquadrics from a point cloud
+    """
+    point_seg = {key: [] for key in list(range(0, MaxLayer+1))}
+    point_outlier = {key: [] for key in list(range(0, MaxLayer+1))}
+    point_seg[0] = [point]
+    list_quadrics = []
+    quadric_info = []  # Store additional info about each quadric
+    
+    for h in range(MaxLayer):
+        if len(point_seg[h]) == 0:
+            break
+            
+        for c in range(len(point_seg[h])):
+            current_points = point_seg[h][c]
+            if len(current_points) < MinPoints * 2:
+                continue
+                
+            print(f"  Layer {h}, Segment {c}: Processing {len(current_points)} points")
+            
+            try:
+                # Fit superquadric
+                x_raw, p_raw = EMS_recovery(
+                    current_points,
+                    OutlierRatio,
+                    MaxIterationEM,
+                    ToleranceEM,
+                    RelativeToleranceEM,
+                    MaxOptiIterations,
+                    Sigma,
+                    MaxiSwitch,
+                    AdaptiveUpperBound,
+                    Rescale,
+                )
+                
+                # Calculate fitting quality
+                inlier_mask = p_raw > 0.5
+                inlier_ratio = np.sum(inlier_mask) / len(p_raw)
+                
+                if inlier_ratio > 0.3:  # Accept if at least 30% inliers
+                    list_quadrics.append(x_raw)
+                    quadric_info.append({
+                        'layer': h,
+                        'segment': c,
+                        'inlier_ratio': inlier_ratio,
+                        'num_points': len(current_points),
+                        'inlier_points': current_points[inlier_mask]
+                    })
+                    print(f"    → Fitted superquadric with {inlier_ratio:.1%} inliers")
+                
+                # Separate outliers for next layer
+                outlier_mask = p_raw < 0.1
+                outlier = current_points[outlier_mask]
+                
+                # If many outliers and not last layer, try clustering
+                if len(outlier) > MinPoints * 2 and h < MaxLayer - 1:
+                    clustering = DBSCAN(eps=Eps, min_samples=MinPoints).fit(outlier)
+                    labels = list(set(clustering.labels_))
+                    labels = [item for item in labels if item >= 0]
+                    
+                    if len(labels) >= 1:
+                        print(f"    → Found {len(labels)} clusters in outliers")
+                        for i in range(len(labels)):
+                            cluster_points = outlier[clustering.labels_ == labels[i]]
+                            if len(cluster_points) > MinPoints:
+                                point_seg[h + 1].append(cluster_points)
+                
+            except Exception as e:
+                print(f"    → Error: {e}")
+                continue
+    
+    return list_quadrics, quadric_info
+
+
+def generate_superquadric_mesh(sq, num_samples=25):
+    """Generate mesh vertices and faces for superquadric surface"""
+    eta = np.linspace(-np.pi/2, np.pi/2, num_samples)
+    omega = np.linspace(-np.pi, np.pi, num_samples)
+    
+    vertices = []
+    faces = []
+    
+    # Generate vertices
+    for i, e in enumerate(eta):
+        for j, w in enumerate(omega):
+            # Superquadric parametric equations
+            cos_eta = np.sign(np.cos(e)) * np.abs(np.cos(e))**sq.shape[0]
+            sin_eta = np.sign(np.sin(e)) * np.abs(np.sin(e))**sq.shape[0]
+            cos_omega = np.sign(np.cos(w)) * np.abs(np.cos(w))**sq.shape[1]
+            sin_omega = np.sign(np.sin(w)) * np.abs(np.sin(w))**sq.shape[1]
+            
+            # Local coordinates
+            x_local = sq.scale[0] * cos_eta * cos_omega
+            y_local = sq.scale[1] * cos_eta * sin_omega
+            z_local = sq.scale[2] * sin_eta
+            
+            # Apply rotation and translation
+            point_local = np.array([x_local, y_local, z_local])
+            point_global = sq.RotM @ point_local + sq.translation
+            
+            vertices.append(point_global)
+    
+    vertices = np.array(vertices)
+    
+    # Generate faces (triangles)
+    for i in range(num_samples - 1):
+        for j in range(num_samples - 1):
+            # Current vertex indices
+            idx1 = i * num_samples + j
+            idx2 = i * num_samples + (j + 1) % num_samples
+            idx3 = (i + 1) * num_samples + j
+            idx4 = (i + 1) * num_samples + (j + 1) % num_samples
+            
+            # Two triangles per quad
+            faces.append([idx1, idx2, idx3])
+            faces.append([idx2, idx4, idx3])
+    
+    return vertices, np.array(faces)
+
+
+def main():
+    # Import utilities for reading PLY files
+    from EMS.utilities import read_ply
+    
+    all_samples = []
+    sample_idx = 0
+    
+    print("Loading and processing samples with hierarchical multi-quadric fitting...")
+    
+    # 1. Load repository example PLY files
+    example_data_dir = "/research/cbim/vast/sf895/code/EMS-superquadric_fitting/MATLAB/example_scripts/data"
+    
+    # Single superquadric examples
+    single_ply_files = [
+        "single_superquadric/noisy_pointCloud_example_1.ply",
+        "single_superquadric/noisy_pointCloud_example_2.ply",
+        "single_superquadric/partial_pointCloud_example_1.ply",
+    ]
+    
+    # Multi superquadric examples
+    multi_ply_files = [
+        "multi_superquadrics/cat.ply",
+        "multi_superquadrics/dog.ply",
+        "multi_superquadrics/turtle.ply",
+    ]
+    
+    # Process single superquadric files
+    for ply_file in single_ply_files:
+        file_path = os.path.join(example_data_dir, ply_file)
+        if os.path.exists(file_path):
+            print(f"\nProcessing {ply_file}...")
+            try:
+                # Load PLY data
+                point_cloud = read_ply(file_path)
+                
+                # Single quadric fitting
+                from EMS.EMS_recovery import EMS_recovery
+                sq, p = EMS_recovery(point_cloud, OutlierRatio=0.2, AdaptiveUpperBound=True)
+                
+                all_samples.append({
+                    'name': os.path.basename(ply_file),
+                    'idx': sample_idx,
+                    'points': point_cloud,
+                    'quadrics': [sq],
+                    'quadric_info': [{
+                        'layer': 0,
+                        'segment': 0,
+                        'inlier_ratio': np.sum(p > 0.5) / len(p),
+                        'num_points': len(point_cloud),
+                        'inlier_points': point_cloud[p > 0.5]
+                    }]
+                })
+                sample_idx += 1
+                
+                print(f"  Success! Shape: {sq.shape}, Scale: {sq.scale}")
+                
+            except Exception as e:
+                print(f"  Failed: {e}")
+    
+    # Process multi superquadric files
+    for ply_file in multi_ply_files:
+        file_path = os.path.join(example_data_dir, ply_file)
+        if os.path.exists(file_path):
+            print(f"\nProcessing {ply_file} (multi-quadric)...")
+            try:
+                # Load PLY data
+                point_cloud = read_ply(file_path)
+                
+                # Hierarchical multi-quadric fitting
+                # Adjust parameters for these specific examples
+                quadrics, quadric_info = hierarchical_ems(
+                    point_cloud,
+                    OutlierRatio=0.9,    # Higher for multi-object scenes
+                    Eps=1.7,             # Larger for non-normalized data
+                    MinPoints=60,        # Standard minimum
+                    Rescale=True         # Enable rescaling for raw PLY data
+                )
+                
+                all_samples.append({
+                    'name': os.path.basename(ply_file),
+                    'idx': sample_idx,
+                    'points': point_cloud,
+                    'quadrics': quadrics,
+                    'quadric_info': quadric_info
+                })
+                sample_idx += 1
+                
+                print(f"Summary: Found {len(quadrics)} superquadrics")
+                for j, (sq, info) in enumerate(zip(quadrics, quadric_info)):
+                    print(f"  SQ{j+1}: Shape={sq.shape}, Scale={sq.scale}, "
+                          f"Inliers={info['inlier_ratio']:.1%}")
+                
+            except Exception as e:
+                print(f"  Failed: {e}")
+    
+    # 2. Also load normalized point cloud samples if they exist
+    normalized_dir = "/research/cbim/vast/sf895/code/EMS-superquadric_fitting/20250811_231035_step10_stage0_waves1"
+    if os.path.exists(normalized_dir):
+        print("\n--- Processing normalized point cloud samples ---")
+        for i in range(2):  # Just load first 2 samples
+            sample_name = f"sample_{i}_normalized_points.npz"
+            sample_path = os.path.join(normalized_dir, sample_name)
+            
+            if os.path.exists(sample_path):
+                print(f"\nProcessing {sample_name}...")
+                try:
+                    # Load data
+                    data = np.load(sample_path)
+                    point_cloud = data['points'][0]  # First frame
+                    
+                    # Hierarchical multi-quadric fitting
+                    quadrics, quadric_info = hierarchical_ems(point_cloud)
+                    
+                    all_samples.append({
+                        'name': sample_name,
+                        'idx': sample_idx,
+                        'points': point_cloud,
+                        'quadrics': quadrics,
+                        'quadric_info': quadric_info
+                    })
+                    sample_idx += 1
+                    
+                    print(f"Summary: Found {len(quadrics)} superquadrics")
+                    
+                except Exception as e:
+                    print(f"  Failed: {e}")
+    
+    # Start viser server
+    server = viser.ViserServer(port=8080)
+    print(f"\n{'='*60}")
+    print(f"Viser server started at http://localhost:8080")
+    print("Open this URL in your browser to view the 3D visualization")
+    print("Press Ctrl+C to stop the server")
+    print('='*60)
+    
+    # Colors for different superquadrics
+    quadric_colors = [
+        (255, 0, 0),    # Red
+        (0, 255, 0),    # Green
+        (0, 0, 255),    # Blue
+        (255, 255, 0),  # Yellow
+        (255, 0, 255),  # Magenta
+        (0, 255, 255),  # Cyan
+    ]
+    
+    # Create GUI
+    with server.gui.add_folder("Controls"):
+        # Sample selector
+        sample_names = [s['name'] for s in all_samples if s['points'] is not None]
+        current_sample = server.gui.add_dropdown(
+            "Select Sample",
+            options=sample_names,
+            initial_value=sample_names[0] if sample_names else None
+        )
+        
+        # Visibility controls
+        show_points = server.gui.add_checkbox("Show Points", initial_value=True)
+        show_all_quadrics = server.gui.add_checkbox("Show All Quadrics", initial_value=True)
+        show_labels = server.gui.add_checkbox("Show Labels", initial_value=True)
+        
+        # Individual quadric toggles will be added dynamically
+        quadric_toggles_folder = server.gui.add_folder("Individual Quadrics")
+        
+        # Visual parameters
+        point_size = server.gui.add_slider(
+            "Point Size",
+            min=0.001,
+            max=0.02,
+            step=0.001,
+            initial_value=0.003
+        )
+        
+        mesh_opacity = server.gui.add_slider(
+            "Mesh Opacity",
+            min=0.0,
+            max=1.0,
+            step=0.1,
+            initial_value=0.5
+        )
+        
+        # Info display
+        info_display = server.gui.add_markdown("**Sample Info:**\n\nSelect a sample to view")
+    
+    # Store current visualization handles
+    current_viz = {
+        'points': None,
+        'meshes': [],
+        'labels': [],
+        'quadric_toggles': []
+    }
+    
+    def update_scene():
+        """Update the 3D scene based on current selection"""
+        # Clear existing visualization
+        if current_viz['points'] is not None:
+            current_viz['points'].remove()
+            current_viz['points'] = None
+        
+        for mesh in current_viz['meshes']:
+            mesh.remove()
+        current_viz['meshes'] = []
+        
+        for label in current_viz['labels']:
+            label.remove()
+        current_viz['labels'] = []
+        
+        # Clear quadric toggles
+        for toggle in current_viz['quadric_toggles']:
+            toggle.remove()
+        current_viz['quadric_toggles'] = []
+        
+        # Find selected sample
+        selected = None
+        for sample in all_samples:
+            if sample['name'] == current_sample.value:
+                selected = sample
+                break
+        
+        if selected is None or selected['points'] is None:
+            info_display.value = "**No valid sample selected**"
+            return
+        
+        # Update info
+        info_text = f"**{selected['name']}**\n\n"
+        info_text += f"Total points: {len(selected['points'])}\n"
+        info_text += f"Superquadrics found: {len(selected['quadrics'])}\n\n"
+        
+        if len(selected['quadrics']) > 0:
+            info_text += "**Superquadric Details:**\n"
+            for i, (sq, info) in enumerate(zip(selected['quadrics'], selected['quadric_info'])):
+                info_text += f"\n**SQ{i+1}** (Layer {info['layer']}):\n"
+                info_text += f"- Shape: ε₁={sq.shape[0]:.3f}, ε₂={sq.shape[1]:.3f}\n"
+                info_text += f"- Scale: ({sq.scale[0]:.2f}, {sq.scale[1]:.2f}, {sq.scale[2]:.2f})\n"
+                info_text += f"- Inliers: {info['inlier_ratio']:.1%} ({info['num_points']} points)\n"
+        
+        info_display.value = info_text
+        
+        # Add point cloud
+        if show_points.value:
+            current_viz['points'] = server.scene.add_point_cloud(
+                "/current/points",
+                points=selected['points'],
+                colors=np.array([(128, 128, 128)] * len(selected['points']), dtype=np.uint8),
+                point_size=point_size.value,
+            )
+        
+        # Add individual quadric toggles
+        with quadric_toggles_folder:
+            for i in range(len(selected['quadrics'])):
+                toggle = server.gui.add_checkbox(
+                    f"Quadric {i+1}",
+                    initial_value=True
+                )
+                current_viz['quadric_toggles'].append(toggle)
+        
+        # Add superquadrics
+        for i, (sq, info) in enumerate(zip(selected['quadrics'], selected['quadric_info'])):
+            color = quadric_colors[i % len(quadric_colors)]
+            
+            # Check if this quadric should be shown
+            show_this = show_all_quadrics.value
+            if i < len(current_viz['quadric_toggles']):
+                show_this = show_this and current_viz['quadric_toggles'][i].value
+            
+            if show_this:
+                try:
+                    vertices, faces = generate_superquadric_mesh(sq, num_samples=20)
+                    
+                    mesh = server.scene.add_mesh_simple(
+                        f"/current/mesh_{i}",
+                        vertices=vertices,
+                        faces=faces,
+                        color=color,
+                        opacity=mesh_opacity.value,
+                    )
+                    current_viz['meshes'].append(mesh)
+                    
+                    if show_labels.value:
+                        label = server.scene.add_label(
+                            f"/current/label_{i}",
+                            text=f"SQ{i+1}: ε₁={sq.shape[0]:.2f}, ε₂={sq.shape[1]:.2f}",
+                            position=sq.translation,
+                        )
+                        current_viz['labels'].append(label)
+                    
+                except Exception as e:
+                    print(f"Error visualizing quadric {i}: {e}")
+    
+    # Set up callbacks
+    @current_sample.on_update
+    def _(_):
+        update_scene()
+    
+    @show_points.on_update
+    def _(_):
+        if current_viz['points'] is not None:
+            current_viz['points'].visible = show_points.value
+    
+    @show_all_quadrics.on_update
+    def _(_):
+        for mesh in current_viz['meshes']:
+            mesh.visible = show_all_quadrics.value
+    
+    @show_labels.on_update
+    def _(_):
+        for label in current_viz['labels']:
+            label.visible = show_labels.value
+    
+    @point_size.on_update
+    def _(event):
+        if current_viz['points'] is not None:
+            current_viz['points'].point_size = event.target.value
+    
+    @mesh_opacity.on_update
+    def _(event):
+        for mesh in current_viz['meshes']:
+            mesh.opacity = event.target.value
+    
+    # Initial scene
+    update_scene()
+    
+    # Keep server running
+    try:
+        while True:
+            time.sleep(0.1)
+    except KeyboardInterrupt:
+        print("\nShutting down server...")
+        server.stop()
+
+
+if __name__ == "__main__":
+    main()

EMS-superquadric_fitting_inference/process_viser_single.py

@@ -0,0 +1,263 @@
+#!/usr/bin/env python3
+"""
+Simple viser visualization for superquadric fitting results
+"""
+
+import numpy as np
+import sys
+import os
+import time
+import viser
+
+# Add the src directory to Python path
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src'))
+
+from EMS.EMS_recovery import EMS_recovery
+from sklearn.cluster import DBSCAN
+
+
+def generate_superquadric_mesh(sq, num_samples=30):
+    """Generate mesh vertices and faces for superquadric surface"""
+    eta = np.linspace(-np.pi/2, np.pi/2, num_samples)
+    omega = np.linspace(-np.pi, np.pi, num_samples)
+    
+    vertices = []
+    faces = []
+    
+    # Generate vertices
+    for i, e in enumerate(eta):
+        for j, w in enumerate(omega):
+            # Superquadric parametric equations
+            cos_eta = np.sign(np.cos(e)) * np.abs(np.cos(e))**sq.shape[0]
+            sin_eta = np.sign(np.sin(e)) * np.abs(np.sin(e))**sq.shape[0]
+            cos_omega = np.sign(np.cos(w)) * np.abs(np.cos(w))**sq.shape[1]
+            sin_omega = np.sign(np.sin(w)) * np.abs(np.sin(w))**sq.shape[1]
+            
+            # Local coordinates
+            x_local = sq.scale[0] * cos_eta * cos_omega
+            y_local = sq.scale[1] * cos_eta * sin_omega
+            z_local = sq.scale[2] * sin_eta
+            
+            # Apply rotation and translation
+            point_local = np.array([x_local, y_local, z_local])
+            point_global = sq.RotM @ point_local + sq.translation
+            
+            vertices.append(point_global)
+    
+    vertices = np.array(vertices)
+    
+    # Generate faces (triangles)
+    for i in range(num_samples - 1):
+        for j in range(num_samples - 1):
+            # Current vertex indices
+            idx1 = i * num_samples + j
+            idx2 = i * num_samples + (j + 1) % num_samples
+            idx3 = (i + 1) * num_samples + j
+            idx4 = (i + 1) * num_samples + (j + 1) % num_samples
+            
+            # Two triangles per quad
+            faces.append([idx1, idx2, idx3])
+            faces.append([idx2, idx4, idx3])
+    
+    return vertices, np.array(faces)
+
+
+def main():
+    # Base directory with samples
+    base_dir = "/research/cbim/vast/sf895/code/EMS-superquadric_fitting/20250811_231035_step10_stage0_waves1"
+    
+    # Process samples
+    all_samples = []
+    
+    print("Loading and processing samples...")
+    for i in range(5):  # Process samples 0-4
+        sample_name = f"sample_{i}_normalized_points.npz"
+        sample_path = os.path.join(base_dir, sample_name)
+        
+        if os.path.exists(sample_path):
+            try:
+                # Load data
+                data = np.load(sample_path)
+                point_cloud = data['points'][0]  # First frame
+                
+                # Try to fit superquadric
+                print(f"\nProcessing {sample_name}...")
+                sq, p = EMS_recovery(point_cloud, OutlierRatio=0.2, AdaptiveUpperBound=True)
+                
+                all_samples.append({
+                    'name': sample_name,
+                    'idx': i,
+                    'points': point_cloud,
+                    'quadric': sq,
+                    'probs': p
+                })
+                print(f"  Success! Shape: {sq.shape}, Scale: {sq.scale}")
+                
+            except Exception as e:
+                print(f"  Failed: {e}")
+                all_samples.append({
+                    'name': sample_name,
+                    'idx': i,
+                    'points': point_cloud if 'point_cloud' in locals() else None,
+                    'quadric': None,
+                    'probs': None
+                })
+    
+    # Start viser server
+    server = viser.ViserServer(port=8080)
+    print(f"\n{'='*60}")
+    print(f"Viser server started at http://localhost:8080")
+    print("Open this URL in your browser to view the 3D visualization")
+    print("Press Ctrl+C to stop the server")
+    print('='*60)
+    
+    # Colors for different samples
+    colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0), (255, 0, 255)]
+    
+    # Create GUI
+    with server.gui.add_folder("Controls"):
+        # Sample selector
+        sample_names = [s['name'] for s in all_samples if s['points'] is not None]
+        current_sample = server.gui.add_dropdown(
+            "Select Sample",
+            options=sample_names,
+            initial_value=sample_names[0] if sample_names else None
+        )
+        
+        # Visibility controls
+        show_points = server.gui.add_checkbox("Show Points", initial_value=True)
+        show_quadric = server.gui.add_checkbox("Show Quadric", initial_value=True)
+        
+        # Visual parameters
+        point_size = server.gui.add_slider(
+            "Point Size",
+            min=0.001,
+            max=0.02,
+            step=0.001,
+            initial_value=0.005
+        )
+        
+        mesh_opacity = server.gui.add_slider(
+            "Mesh Opacity",
+            min=0.0,
+            max=1.0,
+            step=0.1,
+            initial_value=0.5
+        )
+        
+        # Info display
+        info_display = server.gui.add_markdown("**Sample Info:**\n\nSelect a sample to view")
+    
+    # Store current visualization handles
+    current_viz = {'points': None, 'mesh': None, 'label': None}
+    
+    def update_scene():
+        """Update the 3D scene based on current selection"""
+        # Clear existing visualization
+        if current_viz['points'] is not None:
+            current_viz['points'].remove()
+            current_viz['points'] = None
+        if current_viz['mesh'] is not None:
+            current_viz['mesh'].remove()
+            current_viz['mesh'] = None
+        if current_viz['label'] is not None:
+            current_viz['label'].remove()
+            current_viz['label'] = None
+        
+        # Find selected sample
+        selected = None
+        for sample in all_samples:
+            if sample['name'] == current_sample.value:
+                selected = sample
+                break
+        
+        if selected is None or selected['points'] is None:
+            info_display.value = "**No valid sample selected**"
+            return
+        
+        # Update info
+        info_text = f"**{selected['name']}**\n\n"
+        info_text += f"Points: {len(selected['points'])}\n\n"
+        
+        if selected['quadric'] is not None:
+            sq = selected['quadric']
+            info_text += f"**Superquadric Parameters:**\n"
+            info_text += f"- Shape: ε₁={sq.shape[0]:.3f}, ε₂={sq.shape[1]:.3f}\n"
+            info_text += f"- Scale: ({sq.scale[0]:.2f}, {sq.scale[1]:.2f}, {sq.scale[2]:.2f})\n"
+            info_text += f"- Translation: ({sq.translation[0]:.2f}, {sq.translation[1]:.2f}, {sq.translation[2]:.2f})\n"
+        else:
+            info_text += "**No superquadric fitted**"
+        
+        info_display.value = info_text
+        
+        color = colors[selected['idx'] % len(colors)]
+        
+        # Add point cloud
+        if show_points.value:
+            current_viz['points'] = server.scene.add_point_cloud(
+                "/current/points",
+                points=selected['points'],
+                colors=np.array([color] * len(selected['points']), dtype=np.uint8),
+                point_size=point_size.value,
+            )
+        
+        # Add superquadric
+        if show_quadric.value and selected['quadric'] is not None:
+            try:
+                vertices, faces = generate_superquadric_mesh(selected['quadric'], num_samples=25)
+                
+                current_viz['mesh'] = server.scene.add_mesh_simple(
+                    "/current/mesh",
+                    vertices=vertices,
+                    faces=faces,
+                    color=color,
+                    opacity=mesh_opacity.value,
+                )
+                
+                # Add label
+                current_viz['label'] = server.scene.add_label(
+                    "/current/label",
+                    text=f"ε₁={selected['quadric'].shape[0]:.2f}, ε₂={selected['quadric'].shape[1]:.2f}",
+                    position=selected['quadric'].translation,
+                )
+                
+            except Exception as e:
+                print(f"Error visualizing quadric: {e}")
+    
+    # Set up callbacks
+    @current_sample.on_update
+    def _(_):
+        update_scene()
+    
+    @show_points.on_update
+    def _(_):
+        update_scene()
+        
+    @show_quadric.on_update
+    def _(_):
+        update_scene()
+    
+    @point_size.on_update
+    def _(event):
+        if current_viz['points'] is not None:
+            current_viz['points'].point_size = event.target.value
+    
+    @mesh_opacity.on_update
+    def _(event):
+        if current_viz['mesh'] is not None:
+            current_viz['mesh'].opacity = event.target.value
+    
+    # Initial scene
+    update_scene()
+    
+    # Keep server running
+    try:
+        while True:
+            time.sleep(0.1)
+    except KeyboardInterrupt:
+        print("\nShutting down server...")
+        server.stop()
+
+
+if __name__ == "__main__":
+    main()

EMS-superquadric_fitting_inference/pyproject.toml

@@ -0,0 +1,3 @@
+[build-system]
+requires = ["setuptools>=42"]
+build-backend = "setuptools.build_meta"

EMS-superquadric_fitting_inference/setup.py

@@ -0,0 +1,32 @@
+import setuptools
+
+with open("README.md", "r", encoding="utf-8") as fh:
+    long_description = fh.read()
+
+setuptools.setup(
+    name='EMS',
+    version='0.0.1',
+    description='EMS: a package for probabilistic recovery of superquadrics from point clouds',
+    url='https://github.com/bmlklwx/EMS-probabilistic_superquadric_fitting.git',
+    author='Weixiao Liu, Yuwei Wu, Sipu Ruan, Gregory Chirikjian',
+    author_email='wliu72@jhu.edu',
+    long_description=long_description,
+    long_description_content_type="text/markdown",
+    
+    install_requires=[
+        'numpy',
+        'scipy',
+        'plyfile',
+        'mayavi',
+        'numba'
+    ],
+    classifiers=[
+            "Programming Language :: Python :: 3",
+            "License :: OSI Approved :: MIT License",
+            "Operating System :: OS Independent",
+    ],
+    
+    package_dir={"": "src"},
+    packages=setuptools.find_packages(where="src"),
+    python_requires='>=3.6'
+)

EMS-superquadric_fitting_inference/src/EMS/EMS_recovery.py

@@ -0,0 +1,378 @@
+import numpy as np
+from numba import njit
+from scipy.optimize import least_squares
+
+from EMS.superquadrics import rotations, superquadric
+
+def EMS_recovery(
+        point, OutlierRatio=0.1, MaxIterationEM=20,
+        ToleranceEM=1e-3, RelativeToleranceEM=1e-1,
+        MaxOptiIterations=3, Sigma=0, MaxiSwitch=2,
+        AdaptiveUpperBound=False, Rescale=True):
+    # The function conducting probabilistic superquadric recovery.
+    # Input: point - point cloud np array of N * 3
+    #
+
+    # ---------------------------------------INITIALIZATIONS--------------------------------------------
+    # translate the points to the center of mass
+    point = np.array(point, dtype=float)
+    t0 = np.mean(point, 0)
+    point = point - t0
+
+    # rescale
+    if Rescale is True:
+        max_length = np.max(point)
+        scale = max_length / 10
+        point = point / scale
+
+    # eigen analysis for rotation initialization
+    EigVec = EigenAnalysis(point)
+    R0 = rotations()
+    R0.RotM = np.array([-EigVec[:, 0], -EigVec[:, 2],
+                       np.cross(EigVec[:, 0], EigVec[:, 2])]).T
+    euler0 = R0.euler
+
+    # scale initialization
+    point_rot0 = point @ R0.RotM
+    s0 = np.median(np.abs(point_rot0), 0)
+
+    # initialize configuration
+    x0 = np.array([1.0, 1.0, s0[0], s0[1], s0[2],
+                  euler0[0], euler0[1], euler0[2], 0, 0, 0])
+
+    # set lower and upper bounds for the superquadrics
+    upper = 4 * np.max(np.abs(point))
+    lb = np.array([0, 0, 0.001, 0.001, 0.001, -2 * np.pi, -2 *
+                  np.pi, -2 * np.pi, -upper, -upper, -upper])
+    ub = np.array([2.0, 2.0, upper, upper, upper, 2 * np.pi,
+                  2 * np.pi, 2 * np.pi, upper, upper, upper])
+
+    # calculate bounding volume of ourlier space
+    V = BoundVolume(point_rot0)
+
+    # set prior outlier density
+    p0 = 1 / V
+
+    # initialize variance
+    if Sigma == 0:
+        sigma2 = V ** (1 / 3) / 10
+    else:
+        sigma2 = Sigma
+
+    # initialize EMS
+    x = x0
+    cost = 0.0
+    num_switch = int(0)
+    p = np.ones(point.shape[0])
+
+    # ---------------------------------------EMS ALGORITHM--------------------------------------------
+    for iterEM in range(MaxIterationEM):
+        # evaluating distance from points to superquadric
+        dist = Distance(point, x)
+
+        # inferring the postierior outlier probability (E-step)
+        if OutlierRatio != 0:
+            p = OutlierProb(dist, sigma2, OutlierRatio, p0)
+
+        # calculate adaptive upper bound
+        if AdaptiveUpperBound is True:
+            R_cur = Euler2RotM(x[5: 8])
+            point_cur = point @ R_cur - x[8: 11] @ R_cur
+            ub_a = 1.1 * np.max(np.abs(point_cur), 0)
+            ub[2: 5] = ub_a
+            ub[8: 11] = ub_a
+            lb[8: 11] = -ub_a
+
+        # Optimize the superquadric configuration (M-step)
+        optfunc = least_squares(CostFunc, x, bounds=(
+            lb, ub), max_nfev=MaxOptiIterations, args=(point, p, sigma2))
+        x_n = optfunc.x
+        cost_n = 2 * optfunc.cost
+
+        # update sigma
+        sigma2_n = cost_n / (3 * np.sum(p))       
+
+        # evaluate raletive decreasing of cost
+        relative_cost = (cost - cost_n) / cost_n
+
+        # check optimality for termination
+        if (cost_n < ToleranceEM and iterEM > 0) or \
+                (relative_cost < RelativeToleranceEM and num_switch >= MaxiSwitch and iterEM > 4):
+            x = x_n
+            break
+
+        # check for entering similarity switch
+        if relative_cost < RelativeToleranceEM and iterEM > 0:
+            # entering similarity switch (S-step)
+            # initialize swith success flag
+            switch_success = False
+
+            # search for similarity candidates
+            x_candidate = SimilarityCandidates(x)
+
+            # evaluating switch (S-step)
+            x, cost, sigma2, switch_success = Switch(
+                x_candidate, point, p, AdaptiveUpperBound, ub, lb, MaxOptiIterations, \
+                sigma2, sigma2_n, cost, cost_n, x_n, switch_success
+            )
+
+            num_switch = num_switch + 1
+
+        else:
+            # update parameter and prepare for the next EM iteration
+            cost = cost_n
+            sigma2 = sigma2_n
+            x = x_n
+    
+    if Rescale is True:
+        x[2 : 5] = x[2 : 5] * scale
+        x[8 : 11] = x[8 : 11] * scale
+    
+    x[8 : 11] = x[8 : 11] + t0
+
+    sq = superquadric(x[0 : 2], x[2 : 5], x[5 : 8], x[8 : 11])
+
+    return sq, p
+
+# ---------------------------------------UTILITIES-------------------------------------------
+@njit(cache=True)
+def SimilarityCandidates(x):
+    # axis mismatch similarity
+    axis_0 = Euler2RotM(x[5: 8])
+    axis_1 = axis_0[:, np.array([1, 2, 0])]
+    axis_2 = axis_0[:, np.array([2, 0, 1])]
+    eul_1 = RotM2Euler(axis_1)
+    eul_2 = RotM2Euler(axis_2)
+    x_axis = np.array(
+        [[x[1], x[0], x[3], x[4], x[2], eul_1[0], eul_1[1], eul_1[2], x[8], x[9], x[10]],
+         [x[1], x[0], x[4], x[2], x[3], eul_2[0], eul_2[1], eul_2[2], x[8], x[9], x[10]]]
+    )   
+
+    # duality similarities
+    scale_ratio = x[np.array([3, 4, 2])] / x[2 : 5]
+    scale_idx = np.argwhere(np.logical_and(scale_ratio > 0.6, scale_ratio < 1.4))
+    x_rot = np.zeros((scale_idx.shape[0], 11))
+    
+    for idx in range(scale_idx.shape[0]):
+        if scale_idx[idx, 0] == 0:
+            eul_rot = RotM2Euler(axis_0 @ Euler2RotM(np.array([np.pi / 4, 0.0, 0.0])))
+            if x[1] <= 1:
+                x_rot[idx, :] = np.array(
+                    [x[0], 2 - x[1], 
+                    ((1 - np.sqrt(2)) * x[1] + np.sqrt(2)) * min(x[2], x[3]),
+                    ((1 - np.sqrt(2)) * x[1] + np.sqrt(2)) * min(x[2], x[3]),
+                    x[4], eul_rot[0], eul_rot[1], eul_rot[2],
+                    x[8], x[9], x[10]]
+                )
+            else:
+                x_rot[idx, :] = np.array(
+                    [x[0], 2 - x[1], 
+                    ((np.sqrt(2)/2 - 1) * x[1] + 2 - np.sqrt(2) / 2) * min(x[2], x[3]),
+                    ((np.sqrt(2)/2 - 1) * x[1] + 2 - np.sqrt(2) / 2) * min(x[2], x[3]),
+                    x[4], eul_rot[0], eul_rot[1], eul_rot[2],
+                    x[8], x[9], x[10]]
+                )
+
+        elif scale_idx[idx, 0] == 1:
+            eul_rot = RotM2Euler(axis_1 @ Euler2RotM(np.array([np.pi / 4, 0.0, 0.0])))
+            if x[0] <= 1:
+                x_rot[idx, :] = np.array(
+                    [x[1], 2 - x[0], 
+                    ((1 - np.sqrt(2)) * x[0] + np.sqrt(2)) * min(x[3], x[4]),
+                    ((1 - np.sqrt(2)) * x[0] + np.sqrt(2)) * min(x[3], x[4]),
+                    x[2], eul_rot[0], eul_rot[1], eul_rot[2],
+                    x[8], x[9], x[10]]
+                )
+            else:
+                x_rot[idx, :] = np.array(
+                    [x[1], 2 - x[0], 
+                    ((np.sqrt(2)/2 - 1) * x[0] + 2 - np.sqrt(2)/2) * min(x[3], x[4]),
+                    ((np.sqrt(2)/2 - 1) * x[0] + 2 - np.sqrt(2)/2) * min(x[3], x[4]),
+                    x[2], eul_rot[0], eul_rot[1], eul_rot[2],
+                    x[8], x[9], x[10]]
+                )
+
+        elif scale_idx[idx, 0] == 2:
+            eul_rot = RotM2Euler(axis_2 @ Euler2RotM(np.array([np.pi / 4, 0.0, 0.0])))
+            if x[0] <= 1:
+                x_rot[idx, :] = np.array(
+                    [x[1], 2 - x[0], 
+                    ((1 - np.sqrt(2)) * x[0] + np.sqrt(2)) * min(x[4], x[2]),
+                    ((1 - np.sqrt(2)) * x[0] + np.sqrt(2)) * min(x[4], x[2]),
+                    x[2], eul_rot[0], eul_rot[1], eul_rot[2],
+                    x[8], x[9], x[10]]
+                )
+            else:
+                x_rot[idx, :] = np.array(
+                    [x[1], 2 - x[0], 
+                    ((np.sqrt(2)/2 - 1) * x[0] + 2 - np.sqrt(2)/2) * min(x[4], x[2]),
+                    ((np.sqrt(2)/2 - 1) * x[0] + 2 - np.sqrt(2)/2) * min(x[4], x[2]),
+                    x[2], eul_rot[0], eul_rot[1], eul_rot[2],
+                    x[8], x[9], x[10]]
+                )
+    
+    x_candidate = np.zeros((2 + x_rot.shape[0], 11))
+    x_candidate[0 : 2] = x_axis
+    if scale_idx.shape[0] > 0:
+        x_candidate[2 : 2 + scale_idx.shape[0]] = x_rot
+    
+    return x_candidate
+
+def Switch(
+    x_candidate, point, p, AdaptiveUpperBound, ub, lb, MaxOptiIterations, \
+        sigma2, sigma2_n, cost, cost_n, x_n, switch_success
+):
+    cost_candidate = SwitchCost(x_candidate, point, p)
+    idx_nan = np.argwhere(
+        np.logical_and(~np.isnan(cost_candidate), ~np.isinf(cost_candidate))
+    ).reshape(1, -1)[0]
+
+    cost_candidate = cost_candidate[idx_nan]
+    idx = np.argsort(cost_candidate)
+
+    for i in idx:
+        if AdaptiveUpperBound is True:
+            R_cur = Euler2RotM(x_candidate[i, 5: 8])
+            point_cur = point @ R_cur - x_candidate[i, 8: 11] @ R_cur
+            ub_a = 1.1 * np.max(np.abs(point_cur), 0)
+            ub[2: 5] = ub_a
+            ub[8: 11] = ub_a
+            lb[8: 11] = -ub_a
+
+        x_candidate[i] = np.minimum(x_candidate[i], ub)
+        x_candidate[i] = np.maximum(x_candidate[i], lb)
+
+        optfunc = least_squares(CostFunc, x_candidate[i], bounds=(
+                lb, ub), max_nfev=MaxOptiIterations, args=(point, p, sigma2))
+        x_switch = optfunc.x
+        cost_switch = 2 * optfunc.cost
+
+        if cost_switch < min(cost_n, cost):
+            x = x_switch
+            cost = cost_switch
+
+            sigma2 = cost_switch / (3 * sum(p))
+            switch_success = True
+            break
+    
+    if switch_success == False:
+        cost = cost_n
+        sigma2 = sigma2_n
+        x = x_n
+
+    return x, cost, sigma2, switch_success
+    
+@njit(cache=True)
+def SwitchCost(x_candidate, point, p):
+    val = np.zeros(x_candidate.shape[0])
+    for i in range(x_candidate.shape[0]):
+        val[i] = np.sum(p * (Distance(point, x_candidate[i]) ** 2))
+    return val
+
+@njit(cache=True)
+def EigenAnalysis(point):
+    CovM = point.T @ point / point.shape[0]
+    EVal, EVec = np.linalg.eig(CovM)
+    idx = np.flip(np.argsort(EVal))
+    return EVec[:, idx]
+
+@njit(cache=True)
+def BoundVolume(point):
+    V = (np.max(point[:, 0]) - np.min(point[:, 0])) * \
+        (np.max(point[:, 1]) - np.min(point[:, 1])) * \
+        (np.max(point[:, 2]) - np.min(point[:, 2]))
+    return V
+
+@njit(cache=True)
+def Distance(point, x):
+    # approximate the distance from a point to its nearest point on the superquadric surface
+    # extract transformation from superquadric parameters
+    R = Euler2RotM(x[5: 8])
+    t = x[8: 11]
+
+    # transform to the canonical frame
+    point_c = point @ R - t @ R
+
+    # calculating radial distance
+    # r_norm = np.linalg.norm(point_c, axis=1)
+    r_norm = np.sqrt(np.sum(point_c ** 2, 1))
+    
+    dist = r_norm * np.abs((
+        (((point_c[:, 0] / x[2]) ** 2) ** (1 / x[1]) +
+         ((point_c[:, 1] / x[3]) ** 2) ** (1 / x[1])) ** (x[1] / x[0]) +
+        ((point_c[:, 2] / x[4]) ** 2) ** (1 / x[0])) ** (-x[0] / 2) - 1
+    )
+    return dist
+
+@njit(cache=True)
+def CostFunc(x, point, p, sigma2):
+    if sigma2 > 1e-10:
+        value = p ** 0.5 * Distance(point, x)
+    else:
+        value = np.abs((p * Distance(point, x) ** 2 + 2 *
+                       sigma2 * np.log(SurfaceArea(x)))) ** 0.5
+    return value
+
+@njit(cache=True)
+def OutlierProb(dist, sigma2, w, p0):
+    c = (2 * np.pi * sigma2) ** (- 3 / 2)
+    const = (w * p0) / (c * (1 - w))
+    p = np.exp(-1 / (2 * sigma2) * dist ** 2)
+    p = p / (const + p)
+    return p
+
+@njit(cache=True)
+def SurfaceArea(x):
+    a00 = 8 * (x[2] * x[3] + x[3] * x[4] + x[2] * x[4])
+    a02 = 8 * (x[2] ** 2 + x[3] ** 2) ** 0.5 * x[4] + 4 * x[2] * x[3]
+    a20 = 4 * (x[2] * (x[3] ** 2 + x[4] ** 2) ** 0.5 +
+               x[3] * (x[2] ** 2 + x[4] ** 2) ** 0.5)
+    a = (x[2] ** 2 + x[3] ** 2) ** 0.5
+    b = (x[3] ** 2 + x[4] ** 2) ** 0.5
+    c = (x[2] ** 2 + x[4] ** 2) ** 0.5
+    s = (a + b + c) / 2
+    a22 = 8 * (s * (s - a) * (s - b) * (s - c)) ** (1/2)
+    area = np.array([[1 - x[0] / 2, x[0] / 2]]) @ np.array([[a00, a02],
+                                                            [a20, a22]]) @ np.array([[1 - x[1] / 2], [x[1] / 2]])
+    return area[0, 0]
+
+@njit(cache=True)
+def Euler2RotM(euler):
+    # from euler angles to rotation matrix (ZYX_intrinsic)
+
+    RotZ = np.array(
+        [[np.cos(euler[0]), -np.sin(euler[0]), 0.0],
+         [np.sin(euler[0]), np.cos(euler[0]), 0.0],
+         [0.0, 0.0, 1.0]]
+    )
+
+    RotY = np.array(
+        [[np.cos(euler[1]), 0.0, np.sin(euler[1])],
+         [0.0, 1.0, 0.0],
+         [-np.sin(euler[1]), 0.0, np.cos(euler[1])]]
+    )
+
+    RotX = np.array(
+        [[1.0, 0.0, 0.0],
+         [0.0, np.cos(euler[2]), -np.sin(euler[2])],
+         [0.0, np.sin(euler[2]), np.cos(euler[2])]]
+    )
+
+    return RotZ @ RotY @ RotX
+
+@njit(cache=True)
+def RotM2Euler(R):
+
+    s = np.sqrt(R[0, 0] * R[0, 0] + R[1, 0] * R[1, 0])
+    singular = s < 1e-6
+
+    if not singular:
+        x = np.arctan2(R[2, 1], R[2, 2])
+        y = np.arctan2(-R[2, 0], s)
+        z = np.arctan2(R[1, 0], R[0, 0])
+    else:
+        x = np.arctan2(-R[1, 2], R[1, 1])
+        y = np.arctan2(-R[2, 0], s)
+        z = 0
+
+    return np.array([z, y, x])

EMS-superquadric_fitting_inference/src/EMS/__pycache__/EMS_recovery.CostFunc-307.py311.1.nbc

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:575177a225271df742d51f983c9267f0ff29006fac837a569792fecf97e09396
+size 268721

EMS-superquadric_fitting_inference/src/EMS/__pycache__/EMS_recovery.CostFunc-307.py312.1.nbc

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cd12b1b67c67b3ee8c3af1db109a8a96a6071a0bdbbbf051dabbe3fc95e89025
+size 195624

EMS-superquadric_fitting_inference/src/EMS/__pycache__/EMS_recovery.Distance-286.py311.1.nbc

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a38d3fe9a0f6a6853e9ab94bb66c3319aa0f2bdcf9b1995165f5ef8fc5d3a676
+size 218890

EMS-superquadric_fitting_inference/src/EMS/__pycache__/EMS_recovery.Distance-286.py312.1.nbc

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:49a3e75f660f91481c1410a7388b106a817a348d2ca3093c9635cf93a4dc7b59
+size 175972

EMS-superquadric_fitting_inference/src/EMS/__pycache__/EMS_recovery.EigenAnalysis-272.py311.1.nbc

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:350b381765ffc2a6d02a6380dec17deb4fbf73636116b9208bd800ffef369dfb
+size 136198

EMS-superquadric_fitting_inference/src/EMS/__pycache__/EMS_recovery.EigenAnalysis-272.py312.1.nbc

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:748f9d90b283eaf074de14248d68a842979537c268b1a7277597b620ee161f71
+size 147907

EMS-superquadric_fitting_inference/src/EMS/__pycache__/EMS_recovery.SimilarityCandidates-138.py311.1.nbc

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d2451c48b493ecdcf4fe2b870a7dc7f96e4160c6011eed929c6e8fff5a480353
+size 488688

EMS-superquadric_fitting_inference/src/EMS/__pycache__/EMS_recovery.SimilarityCandidates-138.py312.1.nbc

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:05a72a147c53aa27bd81fc144108c22d9fb3aa30221ee4c3f9c069116d0a5d0e
+size 454977