From 435d1433c55b0145ca098066203e3d4d3c6fd9a9 Mon Sep 17 00:00:00 2001
From: Tobias Pitters <tobias.pitters@gmail.com>
Date: Wed, 19 Nov 2025 22:23:24 +0100
Subject: [PATCH 1/3] WIP: nifti vis working, now improve

---
 src/datasets/features/nifti.py | 71 +++++++++++++++++++++++++++++++---
 1 file changed, 65 insertions(+), 6 deletions(-)

diff --git a/src/datasets/features/nifti.py b/src/datasets/features/nifti.py
index f63b8cf6aa1..c5511a00404 100644
--- a/src/datasets/features/nifti.py
+++ b/src/datasets/features/nifti.py
@@ -1,6 +1,7 @@
+import base64
 import os
+import uuid
 from dataclasses import dataclass, field
-from io import BytesIO
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, ClassVar, Dict, Optional, Union
 
@@ -18,6 +19,62 @@
 
     from .features import FeatureType
 
+if config.NIBABEL_AVAILABLE:
+    import nibabel as nib
+
+    class Nifti1ImageWrapper(nib.nifti1.Nifti1Image):
+        """
+        A wrapper around nibabel's Nifti1Image to customize its representation.
+        """
+
+        def __init__(self, nifti_image: nib.nifti1.Nifti1Image):
+            super().__init__(
+                dataobj=nifti_image.get_fdata(),
+                affine=nifti_image.affine,
+                header=nifti_image.header,
+                extra=nifti_image.extra,
+                file_map=nifti_image.file_map,
+                dtype=nifti_image.get_data_dtype(),
+            )
+            self.nifti_image = nifti_image
+
+        def _repr_html_(self):
+            bytes_ = self.nifti_image.to_bytes()
+            b64 = base64.b64encode(bytes_).decode("utf-8")
+
+            self.nifti_data_url = f"data:application/octet-stream;base64,{b64}"
+            viewer_id = f"papaya-{uuid.uuid4().hex[:8]}"
+
+            html = f"""
+    <div id="{viewer_id}" style="width: 100%; height: 400px;"></div>
+    <link rel="stylesheet" type="text/css" href="https://cdn.jsdelivr.net/npm/papaya-viewer@1.0.1455/release/current/standard/papaya.css" />
+    <script src="https://cdn.jsdelivr.net/npm/papaya-viewer@1.0.1455/release/current/standard/papaya.js"></script>
+    <script type="text/javascript">
+    (function() {{
+        // Wait for Papaya to load
+        function initPapaya() {{
+            if (typeof papaya === 'undefined' || typeof papaya.Container === 'undefined') {{
+                setTimeout(initPapaya, 100);
+                return;
+            }}
+
+            // Papaya loaded - manually initialize
+            var params = {{}};
+            params["images"] = ["{self.nifti_data_url}"];
+            params["kioskMode"] = false;
+            params["showControls"] = true;
+
+            // Manual initialization
+            papaya.Container.startPapaya();
+            papaya.Container.addViewer("{viewer_id}", params);
+        }}
+
+        initPapaya();
+    }})();
+    </script>
+    """
+            return html
+
 
 @dataclass
 class Nifti:
@@ -106,7 +163,7 @@ def encode_example(self, value: Union[str, bytes, bytearray, dict, "nib.Nifti1Im
                 f"A nifti sample should be a string, bytes, Path, nibabel image, or dict, but got {type(value)}."
             )
 
-    def decode_example(self, value: dict, token_per_repo_id=None) -> "nib.nifti1.Nifti1Image":
+    def decode_example(self, value: dict, token_per_repo_id=None) -> "Nifti1ImageWrapper":
         """Decode example NIfTI file into nibabel image object.
 
         Args:
@@ -165,11 +222,13 @@ def decode_example(self, value: dict, token_per_repo_id=None) -> "nib.nifti1.Nif
             ):  # gzip magic number, see https://stackoverflow.com/a/76055284/9534390 or "Magic number" on https://en.wikipedia.org/wiki/Gzip
                 bytes_ = gzip.decompress(bytes_)
 
-            bio = BytesIO(bytes_)
-            fh = nib.FileHolder(fileobj=bio)
-            nifti = nib.Nifti1Image.from_file_map({"header": fh, "image": fh})
+            nifti = nib.Nifti1Image.from_bytes(bytes_)
+            # bio = BytesIO(bytes_)
+            # fh = nib.FileHolder(fileobj=bio)
+            # nifti = nib.Nifti1Image.from_file_map({"header": fh, "image": fh})
 
-        return nifti
+        # return nifti
+        return Nifti1ImageWrapper(nifti)
 
     def embed_storage(self, storage: pa.StructArray, token_per_repo_id=None) -> pa.StructArray:
         """Embed NifTI files into the Arrow array.

From b15c72410e48a49ecadca92b63758c56ba5e2a64 Mon Sep 17 00:00:00 2001
From: Tobias Pitters <tobias.pitters@gmail.com>
Date: Wed, 19 Nov 2025 22:32:01 +0100
Subject: [PATCH 2/3] seems to work fine, tests not there yet

---
 src/datasets/features/nifti.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/datasets/features/nifti.py b/src/datasets/features/nifti.py
index c5511a00404..a35122e3f95 100644
--- a/src/datasets/features/nifti.py
+++ b/src/datasets/features/nifti.py
@@ -46,7 +46,7 @@ def _repr_html_(self):
             viewer_id = f"papaya-{uuid.uuid4().hex[:8]}"
 
             html = f"""
-    <div id="{viewer_id}" style="width: 100%; height: 400px;"></div>
+    <div id="{viewer_id}" style="width: 100%; height: 800px;"></div>
     <link rel="stylesheet" type="text/css" href="https://cdn.jsdelivr.net/npm/papaya-viewer@1.0.1455/release/current/standard/papaya.css" />
     <script src="https://cdn.jsdelivr.net/npm/papaya-viewer@1.0.1455/release/current/standard/papaya.js"></script>
     <script type="text/javascript">

From 8b9a7681953786a813575fb11dc3004b44a837e9 Mon Sep 17 00:00:00 2001
From: Tobias Pitters <tobias.pitters@gmail.com>
Date: Wed, 19 Nov 2025 22:36:06 +0100
Subject: [PATCH 3/3] remove uncommented lines

---
 src/datasets/features/nifti.py | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/src/datasets/features/nifti.py b/src/datasets/features/nifti.py
index a35122e3f95..2cfbefdfaad 100644
--- a/src/datasets/features/nifti.py
+++ b/src/datasets/features/nifti.py
@@ -223,11 +223,7 @@ def decode_example(self, value: dict, token_per_repo_id=None) -> "Nifti1ImageWra
                 bytes_ = gzip.decompress(bytes_)
 
             nifti = nib.Nifti1Image.from_bytes(bytes_)
-            # bio = BytesIO(bytes_)
-            # fh = nib.FileHolder(fileobj=bio)
-            # nifti = nib.Nifti1Image.from_file_map({"header": fh, "image": fh})
 
-        # return nifti
         return Nifti1ImageWrapper(nifti)
 
     def embed_storage(self, storage: pa.StructArray, token_per_repo_id=None) -> pa.StructArray: