Update app.py

app.py

@@ -1,65 +1,29 @@
 import streamlit as st
 import cv2
 import numpy as np
-from PIL import Image
-import torch
-
-
-!git clone --recursive https://github.com/frank-xwang/UnSAM.git
-!python -m pip install 'git+https://github.com/MaureenZOU/detectron2-xyz.git'
-%cd UnSAM
-!python -m pip install -r requirements.txt
-
-# uncomment the following lines if you want to run with GPU
-%cd whole_image_segmentation/mask2former/modeling/pixel_decoder/ops
-!sh make.sh
-# clone and install Mask2Former
-!git clone https://github.com/facebookresearch/Mask2Former.git
-%cd Mask2Former
-!pip install -U opencv-python
-!pip install git+https://github.com/cocodataset/panopticapi.git
-!pip install -r requirements.txt
-%cd mask2former/modeling/pixel_decoder/ops
-!python setup.py build install
-%cd ../../../../
-%cd /kaggle/working/Mask2Former
-#%cd /kaggle/working/UnSAM/whole_image_segmentation/mask2former/modeling/pixel_decoder/ops/Mask2Former
-
-#%cd /kaggle/working/UnSAM/whole_image_segmentation/mask2former
-import detectron2
-from detectron2.utils.logger import setup_logger
-setup_logger()
-setup_logger(name="mask2former")
-
-# import some common libraries
-import numpy as np
-import cv2
 import torch
-
-
-
-
-from detectron2.engine import DefaultPredictor
+from PIL import Image
 from detectron2.config import get_cfg
 from detectron2.projects.deeplab import add_deeplab_config
-from detectron2.utils.colormap import random_color
+from detectron2.engine import DefaultPredictor
 from mask2former import add_maskformer2_config
-from tqdm import tqdm
+from detectron2.utils.colormap import random_color
+import os
 
-def setup_predictor(config_file, weights_path, device='cpu'):
+@st.cache_resource
+def setup_config(weights_path):
     cfg = get_cfg()
     cfg.set_new_allowed(True)
     add_deeplab_config(cfg)
     add_maskformer2_config(cfg)
-    cfg.merge_from_file(config_file)
+    cfg.merge_from_file("configs/maskformer2_R50_bs16_50ep.yaml")
     cfg.MODEL.WEIGHTS = weights_path
-    cfg.MODEL.DEVICE = device
-    predictor = DefaultPredictor(cfg)
-    return predictor
+    cfg.MODEL.DEVICE = "cpu"  # Use CPU for inference
+    cfg.freeze()
+    return cfg
 
 def area(mask):
-    if mask.size == 0: 
-        return 0
+    if mask.size == 0: return 0
     return np.count_nonzero(mask) / mask.size
 
 def vis_mask(input, mask, mask_color):
@@ -71,7 +35,7 @@ def vis_mask(input, mask, mask_color):
 def show_image(I, pool):
     already_painted = np.zeros(np.array(I).shape[:2])
     input = I.copy()
-    for mask in tqdm(pool):
+    for mask in pool:
         already_painted += mask.astype(np.uint8)
         overlap = (already_painted == 2)
         if np.sum(overlap) != 0:
@@ -79,40 +43,56 @@ def show_image(I, pool):
             already_painted -= overlap
         input = vis_mask(input, mask, random_color(rgb=True))
     return input
+import gdown
 
-# Load UnSAM and UnSAM+ predictors
-unsam_predictor = setup_predictor(
-    "/kaggle/working/UnSAM/whole_image_segmentation/configs/maskformer2_R50_bs16_50ep.yaml", 
-    "/kaggle/working/Mask2Former/unsam_sa1b_4perc_ckpt_200k.pth"
-)
-unsam_plus_predictor = setup_predictor(
-    "/kaggle/working/UnSAM/whole_image_segmentation/configs/maskformer2_R50_bs16_50ep.yaml", 
-    "/kaggle/working/Mask2Former/unsam_plus_sa1b_1perc_ckpt_50k.pth"
-)
+gdown.download("https://drive.google.com/uc?id=1sCZM5j2pQr34-scSEkgG7VmUaHJc8n4d", "unsam_plus_sa1b_1perc_ckpt_50k.pth", quiet=False)
+gdown.download("https://drive.google.com/uc?id=1qUdZ2ELU_5SNTsmx3Q0wSA87u4SebiO4", "unsam_sa1b_4perc_ckpt_200k.pth", quiet=False)
 
-st.title("Image Segmentation with UnSAM and UnSAM+")
+@st.cache_data
+def process_image(image, model_type):
+    if model_type == "UNSAM+":
+        weights_path = "unsam_plus_sa1b_1perc_ckpt_50k.pth"
+    else:  # UNSAM
+        weights_path = "unsam_sa1b_4perc_ckpt_200k.pth"
+    
+    cfg = setup_config(weights_path)
+    predictor = DefaultPredictor(cfg)
+    
+    inputs = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+    outputs = predictor(inputs)['instances']
+    
+    masks = []
+    for score, mask in zip(outputs.scores, outputs.pred_masks):
+        if score < 0.5: continue
+        masks.append(mask.cpu().numpy())
+    
+    sorted_masks = sorted(masks, key=lambda m: area(m), reverse=True)
+    result_image = show_image(np.array(image), sorted_masks)
+    
+    return result_image
+
+st.title("UNSAM and UNSAM+ Image Segmentation")
 
-# Upload image
-uploaded_file = st.file_uploader("Choose an image...", type="png")
+uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
 
 if uploaded_file is not None:
-    # Read the image
-    image = np.array(Image.open(uploaded_file))
+    image = Image.open(uploaded_file)
     
-    # Display the original image
-    st.image(image, caption='Original Image', use_column_width=True)
+    col1, col2, col3 = st.columns(3)
     
-    # Run predictions for UnSAM+
-    unsam_plus_outputs = unsam_plus_predictor(image)['instances']
-    unsam_plus_masks = [mask.cpu().numpy() for mask in unsam_plus_outputs.pred_masks]
-    sorted_unsam_plus_masks = sorted(unsam_plus_masks, key=lambda m: area(m), reverse=True)
-    unsam_plus_image = show_image(image, sorted_unsam_plus_masks)
+    with col1:
+        st.header("Original Image")
+        st.image(image, use_column_width=True)
     
-    # Run predictions for UnSAM
-    unsam_outputs = unsam_predictor(image)['instances']
-    unsam_masks = [mask.cpu().numpy() for mask in unsam_outputs.pred_masks]
-    sorted_unsam_masks = sorted(unsam_masks, key=lambda m: area(m), reverse=True)
-    unsam_image = show_image(image, sorted_unsam_masks)
+    with col2:
+        st.header("UNSAM+ Output")
+        unsam_plus_output = process_image(image, "UNSAM+")
+        st.image(unsam_plus_output, use_column_width=True)
     
-    # Display the images side by side
-    st.image([image, unsam_plus_image, unsam_image], caption=['Original Image', 'UnSAM+ Output', 'UnSAM Output'], use_column_width=True)
+    with col3:
+        st.header("UNSAM Output")
+        unsam_output = process_image(image, "UNSAM")
+        st.image(unsam_output, use_column_width=True)
+
+else:
+    st.write("Please upload an image to see the segmentation results.")