diff --git a/src/methods_segmentation/custom_segmentation/script.py b/src/methods_segmentation/custom_segmentation/script.py
index c108117f..d29cb9ab 100644
--- a/src/methods_segmentation/custom_segmentation/script.py
+++ b/src/methods_segmentation/custom_segmentation/script.py
@@ -2,6 +2,7 @@
 import anndata as ad
 import os
 import shutil
+import pandas as pd
 
 ## VIASH START
 par = {
@@ -20,14 +21,18 @@
 assert par["labels_key"] in sdata.labels, f"Key '{par['labels_key']}' not found in input data."
 
 print(f"Copy segmentation from '{par['labels_key']}'", flush=True)
+metadata = sdata.tables["metadata"]
+# Select only the columns that exist — Xenium provides cell_id and region,
+# Vizgen uses different column names (or an empty obs) so we take what's available.
+obs_cols = [c for c in ["cell_id", "region"] if c in metadata.obs.columns]
 sdata_segmentation_only = sd.SpatialData(
   labels={
     "segmentation": sdata[par["labels_key"]]
   },
   tables={
     "table": ad.AnnData(
-      obs=sdata.tables["metadata"].obs[["cell_id", "region"]],
-      var=sdata.tables["metadata"].var[[]]
+      obs=metadata.obs[obs_cols],
+      var=metadata.var[[]]
     )
   }
 )