create_samples.py

"""
Create synthetic CT phantom samples for HuggingFace Space demo
"""
import numpy as np
import nibabel as nib
import os

def create_synthetic_samples():
    examples_dir = 'examples'
    os.makedirs(examples_dir, exist_ok=True)
    
    np.random.seed(42)
    
    # Sample 1: Chest phantom
    print('Creating sample_ct_chest.nii.gz...')
    shape = (192, 192, 128)
    data = np.ones(shape, dtype=np.float32) * -1000  # Air background
    
    cx, cy, cz = shape[0]//2, shape[1]//2, shape[2]//2
    
    # Create coordinate grids
    X, Y, Z = np.meshgrid(
        np.arange(shape[0]) - cx,
        np.arange(shape[1]) - cy,
        np.arange(shape[2]) - cz,
        indexing='ij'
    )
    
    # Body contour (ellipsoid)
    body = (X**2 / 70**2 + Y**2 / 50**2 + Z**2 / 55**2) < 1
    data[body] = 40 + np.random.randn(*shape)[body] * 30
    
    # Spine
    spine = (X**2 + (Y + 35)**2) < 100
    data[spine] = 400 + np.random.randn(*shape)[spine] * 50
    
    # Liver
    liver = ((X - 30)**2 / 25**2 + (Y + 10)**2 / 30**2 + (Z + 20)**2 / 35**2) < 1
    data[liver] = 60 + np.random.randn(*shape)[liver] * 15
    
    # Lungs
    lung_l = ((X + 25)**2 / 20**2 + (Y - 20)**2 / 25**2 + Z**2 / 40**2) < 1
    lung_r = ((X - 25)**2 / 20**2 + (Y - 20)**2 / 25**2 + Z**2 / 40**2) < 1
    data[lung_l | lung_r] = -700 + np.random.randn(*shape)[lung_l | lung_r] * 100
    
    # Heart
    heart = (X**2 / 20**2 + (Y - 10)**2 / 18**2 + (Z + 10)**2 / 25**2) < 1
    data[heart] = 50 + np.random.randn(*shape)[heart] * 20
    
    # Kidneys
    kidney_l = ((X + 35)**2 / 8**2 + (Y + 15)**2 / 12**2 + (Z + 10)**2 / 18**2) < 1
    kidney_r = ((X - 35)**2 / 8**2 + (Y + 15)**2 / 12**2 + (Z + 10)**2 / 18**2) < 1
    data[kidney_l | kidney_r] = 30 + np.random.randn(*shape)[kidney_l | kidney_r] * 10
    
    # Ribs
    for i in range(-5, 6):
        rib_z = i * 10
        rib_mask = ((X - 45)**2 + (Y - 15)**2 < 20) & (np.abs(Z - rib_z) < 3)
        data[rib_mask] = 350 + np.random.randn(*shape)[rib_mask] * 40
        rib_mask = ((X + 45)**2 + (Y - 15)**2 < 20) & (np.abs(Z - rib_z) < 3)
        data[rib_mask] = 350 + np.random.randn(*shape)[rib_mask] * 40
    
    affine = np.diag([1.5, 1.5, 1.5, 1.0])
    img = nib.Nifti1Image(data.astype(np.int16), affine)
    nib.save(img, os.path.join(examples_dir, 'sample_ct_chest.nii.gz'))
    print('  ✓ Created sample_ct_chest.nii.gz')
    
    # Sample 2: Abdomen phantom
    print('Creating sample_ct_abdomen.nii.gz...')
    shape2 = (160, 160, 96)
    data2 = np.ones(shape2, dtype=np.float32) * -1000
    
    cx2, cy2, cz2 = shape2[0]//2, shape2[1]//2, shape2[2]//2
    
    X2, Y2, Z2 = np.meshgrid(
        np.arange(shape2[0]) - cx2,
        np.arange(shape2[1]) - cy2,
        np.arange(shape2[2]) - cz2,
        indexing='ij'
    )
    
    # Body
    body2 = (X2**2 / 60**2 + Y2**2 / 45**2) < 1
    data2[body2] = 35 + np.random.randn(*shape2)[body2] * 25
    
    # Spine
    spine2 = (X2**2 + (Y2 + 30)**2) < 80
    data2[spine2] = 380 + np.random.randn(*shape2)[spine2] * 45
    
    # Liver
    liver2 = ((X2 - 25)**2 / 30**2 + (Y2 + 5)**2 / 28**2 + (Z2 + 15)**2 / 30**2) < 1
    data2[liver2] = 55 + np.random.randn(*shape2)[liver2] * 12
    
    # Spleen
    spleen = ((X2 + 35)**2 / 15**2 + (Y2 + 10)**2 / 18**2 + (Z2 + 5)**2 / 20**2) < 1
    data2[spleen] = 45 + np.random.randn(*shape2)[spleen] * 10
    
    # Kidneys
    kidney_l2 = ((X2 + 28)**2 / 10**2 + (Y2 + 12)**2 / 14**2 + Z2**2 / 18**2) < 1
    kidney_r2 = ((X2 - 28)**2 / 10**2 + (Y2 + 12)**2 / 14**2 + Z2**2 / 18**2) < 1
    data2[kidney_l2 | kidney_r2] = 32 + np.random.randn(*shape2)[kidney_l2 | kidney_r2] * 8
    
    # Pancreas
    pancreas = (X2**2 / 25**2 + (Y2 + 8)**2 / 6**2 + (Z2 - 5)**2 / 10**2) < 1
    data2[pancreas] = 40 + np.random.randn(*shape2)[pancreas] * 10
    
    # Aorta
    aorta = (X2**2 + (Y2 + 20)**2) < 36
    data2[aorta] = 120 + np.random.randn(*shape2)[aorta] * 25
    
    # Stomach
    stomach = ((X2 + 10)**2 / 20**2 + (Y2 - 5)**2 / 15**2 + (Z2 + 25)**2 / 18**2) < 1
    data2[stomach] = -50 + np.random.randn(*shape2)[stomach] * 30
    
    affine2 = np.diag([2.0, 2.0, 2.0, 1.0])
    img2 = nib.Nifti1Image(data2.astype(np.int16), affine2)
    nib.save(img2, os.path.join(examples_dir, 'sample_ct_abdomen.nii.gz'))
    print('  ✓ Created sample_ct_abdomen.nii.gz')
    
    print('\nDone! Synthetic CT samples created.')
    return ['sample_ct_chest.nii.gz', 'sample_ct_abdomen.nii.gz']

if __name__ == '__main__':
    create_synthetic_samples()