Update streamlit_app.py

streamlit_app.py

@@ -6,7 +6,7 @@ then displays results and provides download links.
 """
 
 # imports
-import streamlit as st, logging, shutil
+import streamlit as st, logging, shutil, torch
 from PIL import Image
 from pathlib import Path
 from utils.constants import *
@@ -47,6 +47,12 @@ if uploaded:
         ImageSplitter(source_dir=PNG_IMAGES_DIR, output_dir=SPLIT_IMAGES_DIR, sub_image_width=IMG_WIDTH, sub_image_height=IMG_HEIGHT).split_all()
     # generate - cellpose masks (detect step using a pre-trained model)
     with st.spinner("Running Cellpose segmentation..."):
+        # from huggingface_hub import snapshot_download, hf_hub_download
+        # model_dir = snapshot_download(repo_id="unikill066/drg_cellpose_sam_model", cache_dir="/mnt/models_cache")
+        # # weights_path = hf_hub_download(repo_id="unikill066/drg_cellpose_sam_model",filename="model.pt")
+        # # model = torch.load(weights_path, map_location="cpu")
+        # print(model.eval())
+        # MODEL = "/mnt/models_cache/unikill066/drg_cellpose_sam_model"
         cellpose_sam_detect_images_eval(model_path=MODEL, image_input_dir=SPLIT_IMAGES_DIR, image_output_dir=CELLPOSE_MASKS_DIR)
     # generate - stitched masks (.npy files)
     with st.spinner("Stitching masks..."):