import streamlit as st
import huggingface_hub as hf_hub
import monai
import os
import zipfile
import torch

hf_hub.login(token=st.secrets["HF_TOKEN"])

with st.spinner("Downloading Dataset"):
    data_path = hf_hub.hf_hub_download(repo_id="osbm/prostate158", filename="data.zip", repo_type="dataset")

st.write(data_path)
with st.spinner("Unzipping..."):
    with zipfile.ZipFile(data_path, 'r') as zip_ref:
        zip_ref.extractall(".")


# st.write(os.listdir(os.getcwd()))
# st.write(os.getcwd())


model = monai.networks.nets.UNet(
    in_channels=1,
    out_channels=3,
    spatial_dims=3,
    channels=[16, 32, 64, 128, 256, 512],
    strides=[2, 2, 2, 2, 2],
    num_res_units=4,
    act="PRELU",
    norm="BATCH",
    dropout=0.15,
)

# load this model using anatomy.pt 
model.load_state_dict(torch.load('anatomy.pt', map_location=torch.device('cpu')))

print(model)