Spaces:

TobiasPitters
/

bids-neuroimaging

Sleeping

File size: 15,231 Bytes

aad6f22

from fastapi import FastAPI
from fastapi.responses import HTMLResponse
from datasets import load_dataset
import base64
import asyncio
from contextlib import asynccontextmanager

ds_iter = None
first_sample = None
nifti_base_url = "data:application/octet-stream;base64,{}"
initial_metadata = {}
num_iterations = 0

async def load_dataset_async():
    global ds_iter
    dataset = load_dataset("TobiasPitters/ds004884-mini", streaming=True)
    ds_iter = iter(dataset["train"])

@asynccontextmanager
async def lifespan(app: FastAPI):
    global first_sample, initial_metadata

    asyncio.create_task(load_dataset_async())

    yield

app = FastAPI(title="BIDS Neuroimaging Viewer", lifespan=lifespan)

@app.get("/", response_class=HTMLResponse)
async def root():
    """Main page with NiiVue viewer"""
    html_content = f"""<!DOCTYPE html>
<html>
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>BIDS Neuroimaging Viewer</title>
    <style>
        * {{
            margin: 0;
            padding: 0;
            box-sizing: border-box;
        }}
        body {{
            font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
            background: #1a1a1a;
            color: #fff;
        }}
        .header {{
            background: #2d2d2d;
            padding: 20px;
            border-bottom: 2px solid #3498db;
        }}
        .header h1 {{
            margin: 0;
            color: #3498db;
        }}
        .container {{
            max-width: 1400px;
            margin: 0 auto;
            padding: 20px;
        }}
        .main-content {{
            display: flex;
            gap: 20px;
        }}
        .viewer-section {{
            flex: 1;
            min-width: 0;
        }}
        .metadata-section {{
            width: 300px;
            flex-shrink: 0;
        }}
        .metadata-card {{
            background: #2d2d2d;
            border-radius: 8px;
            padding: 20px;
            position: sticky;
            top: 20px;
        }}
        .metadata-card h3 {{
            margin: 0 0 15px 0;
            color: #3498db;
            font-size: 18px;
        }}
        .metadata-item {{
            margin-bottom: 12px;
            padding-bottom: 12px;
            border-bottom: 1px solid #404040;
        }}
        .metadata-item:last-child {{
            border-bottom: none;
            margin-bottom: 0;
            padding-bottom: 0;
        }}
        .metadata-label {{
            font-size: 11px;
            text-transform: uppercase;
            color: #888;
            letter-spacing: 0.5px;
            margin-bottom: 4px;
        }}
        .metadata-value {{
            font-size: 14px;
            color: #fff;
            word-break: break-word;
        }}
        .viewer-container {{
            width: 100%;
            height: 80vh;
            min-height: 600px;
            position: relative;
            background: #000;
            border-radius: 8px;
            overflow: hidden;
        }}
        #niivue-canvas {{
            width: 100%;
            height: 100%;
            display: block;
        }}
        #status {{
            position: absolute;
            top: 10px;
            left: 10px;
            background: rgba(0, 0, 0, 0.8);
            color: #fff;
            font: 14px monospace;
            padding: 10px 15px;
            border-radius: 5px;
            z-index: 100;
        }}
        .info {{
            margin-top: 20px;
            padding: 15px;
            background: #2d2d2d;
            border-radius: 5px;
        }}
        .tabs {{
            display: flex;
            gap: 10px;
            margin-bottom: 20px;
            border-bottom: 2px solid #3498db;
        }}
        .tab {{
            padding: 12px 24px;
            background: #2d2d2d;
            border: none;
            color: #fff;
            cursor: pointer;
            font-size: 16px;
            border-radius: 8px 8px 0 0;
            transition: all 0.3s;
        }}
        .tab:hover {{
            background: #3d3d3d;
        }}
        .tab.active {{
            background: #3498db;
            color: #fff;
        }}
        .controls {{
            display: flex;
            gap: 10px;
            margin-top: 20px;
        }}
        .btn {{
            padding: 12px 32px;
            background: #3498db;
            border: none;
            color: #fff;
            cursor: pointer;
            font-size: 16px;
            border-radius: 5px;
            transition: all 0.3s;
            font-weight: 600;
        }}
        .btn:hover {{
            background: #2980b9;
            transform: translateY(-2px);
            box-shadow: 0 4px 8px rgba(52, 152, 219, 0.3);
        }}
        .btn:active {{
            transform: translateY(0);
        }}
    </style>
</head>
<body>
    <div class="header">
        <h1>🧠 BIDS Neuroimaging Viewer</h1>
        <p>Aphasia Recovery Cohort (ARC) Dataset - Mini Sample</p>
    </div>

    <div class="container">
        <div class="main-content">
            <div class="viewer-section">
                <div class="tabs">
                    <button class="tab active" id="tab-multiplanar">Multiplanar + 3D</button>
                    <button class="tab" id="tab-render3d">3D Render Only</button>
                </div>

                <div class="viewer-container">
                    <canvas id="niivue-canvas"></canvas>
                    <div id="status">Loading NiiVue...</div>
                </div>

                <div class="controls">
                    <button class="btn" id="next-btn">Next</button>
                </div>

                <div class="info">
                    <h3 id="info-title">Controls:</h3>
                    <ul id="info-list">
                        <li><strong>Slice views (3 panels):</strong> Click to move crosshair, drag to pan, scroll to zoom</li>
                        <li><strong>3D view (bottom-right):</strong> Left-click drag to rotate, scroll to zoom</li>
                    </ul>
                </div>
            </div>

            <div class="metadata-section">
                <div class="metadata-card">
                    <h3>Sample Info</h3>
                    <div id="metadata-content">
                        <div class="metadata-item">
                            <div class="metadata-label">Loading...</div>
                        </div>
                    </div>
                </div>
            </div>
        </div>
    </div>

    <script type="module">
    let nv;

    function updateMetadata(metadata) {{
        const metadataContent = document.getElementById('metadata-content');
        const fields = ['subject', 'session', 'datatype', 'suffix', 'task', 'run'];

        metadataContent.innerHTML = fields.map(field => `
            <div class="metadata-item">
                <div class="metadata-label">${{field}}</div>
                <div class="metadata-value">${{metadata[field] || 'N/A'}}</div>
            </div>
        `).join('');
    }}

    function switchView(mode, clickedElement) {{
        if (!nv) return; // Not initialized yet

        // Update tab styles
        if (clickedElement) {{
            document.querySelectorAll('.tab').forEach(tab => tab.classList.remove('active'));
            clickedElement.classList.add('active');
        }}

        // Switch view mode
        if (mode === 'multiplanar') {{
            nv.setSliceType(nv.sliceTypeMultiplanar);
            if (nv.setMultiplanarLayout) {{
                nv.setMultiplanarLayout(2);
            }}
            nv.opts.show3Dcrosshair = true;

            document.getElementById('info-list').innerHTML = `
                <li><strong>Slice views (3 panels):</strong> Click to move crosshair, drag to pan, scroll to zoom</li>
                <li><strong>3D view (bottom-right):</strong> Left-click drag to rotate, right-click drag to pan, scroll to zoom</li>
                <li><strong>Double click:</strong> Reset view</li>
            `;
        }} else if (mode === 'render3d') {{
            nv.setSliceType(nv.sliceTypeRender);

            // Enable clip plane in 3D render
            nv.setClipPlane([0, 270, 0]);  // Show clipping plane
            nv.opts.clipPlaneHotKey = 'c';  // Press 'c' to toggle clip plane

            document.getElementById('info-list').innerHTML = `
                <li><strong>Left click + drag:</strong> Rotate the 3D volume</li>
                <li><strong>Right click + drag:</strong> Pan/move</li>
                <li><strong>Mouse wheel:</strong> Adjust clip plane depth</li>
            `;
        }}

        nv.drawScene();
    }}

    // Setup tab event listeners
    document.getElementById('tab-multiplanar').addEventListener('click', function() {{
        switchView('multiplanar', this);
    }});

    document.getElementById('tab-render3d').addEventListener('click', function() {{
        switchView('render3d', this);
    }});

    // Next button handler
    document.getElementById('next-btn').addEventListener('click', async function() {{
        if (!nv) {{
            alert('Viewer not initialized yet');
            return;
        }}

        const statusEl = document.getElementById('status');
        try {{
            statusEl.style.display = 'block';
            statusEl.textContent = 'Loading next sample...';
            statusEl.style.background = 'rgba(0, 0, 0, 0.8)';

            const response = await fetch('/next');
            const data = await response.json();

            if (data.status === 'error') {{
                statusEl.textContent = 'Error: ' + data.message;
                statusEl.style.background = 'rgba(160, 0, 0, 0.8)';
                setTimeout(() => {{ statusEl.style.display = 'none'; }}, 3000);
                return;
            }}

            console.log('Loading next sample...');
            await nv.loadVolumes([{{
                url: data.data_url,
                name: 'volume.nii.gz'
            }}]);

            // Update metadata
            updateMetadata(data.metadata);

            nv.drawScene();
            statusEl.style.display = 'none';
            console.log('✓ Next sample loaded!');
        }} catch (err) {{
            console.error('Error loading next sample:', err);
            statusEl.textContent = 'Error: ' + err.message;
            statusEl.style.background = 'rgba(160, 0, 0, 0.8)';
            setTimeout(() => {{ statusEl.style.display = 'none'; }}, 3000);
        }}
    }});

    (async () => {{
        const statusEl = document.getElementById('status');

        try {{
            statusEl.textContent = 'Fetching NiiVue library...';

            const niivueModule = await import('https://unpkg.com/@niivue/niivue@0.57.0/dist/index.js');
            const Niivue = niivueModule.Niivue;

            statusEl.textContent = 'Initializing viewer...';
            nv = new Niivue({{
                logging: true,
                show3Dcrosshair: true,
                textHeight: 0.04
            }});

            await nv.attachTo('niivue-canvas');

            statusEl.textContent = 'Loading initial sample...';

            // Fetch initial sample and metadata
            const initResponse = await fetch('/initial');
            const initData = await initResponse.json();

            const volumes = [{{
                url: initData.data_url,
                name: 'volume.nii.gz'
            }}];

            await nv.loadVolumes(volumes);

            // Update metadata display
            updateMetadata(initData.metadata);

            // Start with multiplanar + 3D
            nv.setSliceType(nv.sliceTypeMultiplanar);
            if (nv.setMultiplanarLayout) {{
                nv.setMultiplanarLayout(2);
            }}
            nv.setRenderAzimuthElevation(120, 10);
            nv.opts.show3Dcrosshair = true;
            nv.opts.crosshairWidth = 2;

            statusEl.textContent = 'Rendering...';
            setTimeout(() => {{
                nv.updateGLVolume();
                nv.drawScene();
                statusEl.style.display = 'none';
                console.log('✓✓✓ Viewer ready!');
            }}, 300);

        }} catch (err) {{
            console.error('ERROR:', err);
            console.error('Error message:', err.message);
            console.error('Error stack:', err.stack);
            statusEl.textContent = 'Error: ' + err.message;
            statusEl.style.background = 'rgba(160, 0, 0, 0.8)';
        }}
    }})();
    </script>
</body>
</html>
"""
    return html_content

@app.get("/initial")
async def get_initial():
    print("Loading dataset...")
    # dataset = load_dataset("TobiasPitters/ds004884-mini", streaming=True)
    global ds_iter
    dataset = load_dataset("parquet", data_dir="./data/" ,streaming=True)
    ds_iter = iter(dataset["train"])

    first_sample = next(ds_iter)

    if isinstance(first_sample['nifti'], dict):
        nifti_bytes = first_sample['nifti']['bytes']
    else:
        nifti_bytes = first_sample['nifti'].to_bytes()
    nifti_b64 = base64.b64encode(nifti_bytes).decode("utf-8")
    data_url = nifti_base_url.format(nifti_b64)

    initial_metadata = {
        'subject': first_sample.get('subject', 'N/A'),
        'session': first_sample.get('session', 'N/A'),
        'datatype': first_sample.get('datatype', 'N/A'),
        'suffix': first_sample.get('suffix', 'N/A'),
        'task': first_sample.get('task', 'N/A'),
        'run': first_sample.get('run', 'N/A'),
        'path': first_sample.get('path', 'N/A')
    }

    return {
        "status": "success",
        "data_url": data_url,
        "metadata": initial_metadata
    }

async def load_dataset_iterator():
    global num_iterations
    dataset = load_dataset("TobiasPitters/ds004884-mini", streaming=True)
    ds_iter_new = iter(dataset["train"])
    for _ in range(num_iterations):
        next(ds_iter_new)

    global ds_iter
    ds_iter = ds_iter_new

@app.get("/next")
async def next_sample():
    """Load next sample from dataset"""
    global ds_iter
    global num_iterations
    try:
        ex = next(ds_iter)
        if isinstance(ex['nifti'], dict):
            nifti_bytes = ex['nifti']['bytes']
        else:
            nifti_bytes = ex['nifti'].to_bytes()
        nifti_b64 = base64.b64encode(nifti_bytes).decode("utf-8")
        new_data_url = nifti_base_url.format(nifti_b64)

        metadata = {
            'subject': ex.get('subject', 'N/A'),
            'session': ex.get('session', 'N/A'),
            'datatype': ex.get('datatype', 'N/A'),
            'suffix': ex.get('suffix', 'N/A'),
            'task': ex.get('task', 'N/A'),
            'run': ex.get('run', 'N/A'),
            'path': ex.get('path', 'N/A')
        }

        num_iterations += 1
        return {
            "status": "success",
            "data_url": new_data_url,
            "metadata": metadata
        }
    except StopIteration:
        return {"status": "error", "message": "No more samples in dataset"}

@app.get("/health")
async def health():
    """Health check endpoint"""
    return {"status": "healthy", "dataset_loaded": True}