import streamlit as st
import pydicom
import matplotlib.pyplot as plt
import zipfile
import io
import os

# Streamlit app title
st.title("DICOM Image Viewer")

# Upload a ZIP file containing DICOMs
uploaded_zip = st.file_uploader("Upload a ZIP file containing DICOMs", type=["zip"])

# Initialize variables for DICOM data and image list
dicom_images = []
num_slices = 0

if uploaded_zip is not None:
    # Unzip the uploaded ZIP file and extract DICOMs
    with zipfile.ZipFile(uploaded_zip, "r") as zip_ref:
        zip_ref.extractall("temp_zip_folder")

    # Get a list of DICOM files
    dicom_files = [f for f in os.listdir("temp_zip_folder") if f.lower().endswith(".dcm")]
    num_slices = len(dicom_files)

    # Read and display DICOM images
    for dicom_file in dicom_files:
        dicom_data = pydicom.dcmread(f"temp_zip_folder/{dicom_file}")
        dicom_images.append(dicom_data.pixel_array)

    # User selects the slice using a slider
    selected_slice = st.slider("Select a DICOM slice", 0, num_slices - 1)

    # Display the selected DICOM image
    plt.imshow(dicom_images[selected_slice], cmap=plt.cm.bone)
    plt.axis("off")
    plt.title(f"DICOM Image - Slice {selected_slice + 1}/{num_slices}")
    plt.tight_layout()

    # Show the image in the Streamlit app
    st.pyplot(plt)

# Remove the temporary folder and its contents
if num_slices > 0:
    for dicom_file in dicom_files:
        os.remove(f"temp_zip_folder/{dicom_file}")
    os.rmdir("temp_zip_folder")

st.write("Upload a ZIP file containing DICOMs to view and scroll through the slices.")