import gradio as gr
# import nibabel as nib
# import pydicom
import os
import dicom2nifti

def convert(input):
    path = input.name
    print("Zip file path:", path)
    assert path.endswith('.zip'), 'File must be a zip file'
    os.system(f'unzip -q "{path}" -d "{path[:-4]}"')
    path = path[:-4]
    files = os.listdir(path)
    if len(files) == 1 and os.path.isdir(os.path.join(path, files[0])):
        path = os.path.join(path, files[0])
        print("Dicom folder path:", path)
    dicom_files = os.listdir(path)
    assert [f.endswith('.dcm') for f in dicom_files], 'Dicom folder must contain only .dcm files'
    
    output_file = os.path.join(path, os.path.basename(path) + '.nii.gz')
    dicom2nifti.dicom_series_to_nifti(path, output_file, reorient_nifti=True)
    print("Nifti file path:", output_file)
    return output_file
    

app = gr.Interface(
    fn=convert,
    inputs=gr.File(type="file", label="Dicom folder (.zip)", file_types=["zip"]),
    outputs="file",
    title="Convert a Dicom folder to a Nifti (.nii.gz) file",
)
app.queue(concurrency_count=1)
app.launch()