Update app.py

app.py

@@ -1,16 +1,14 @@
-
-import matplotlib.pyplot as plt
-import numpy as np
+import os
+import subprocess
+import sys
+import requests
+import zipfile
+import gradio as gr
 import torch
+import numpy as np
 from torchvision.transforms import ToTensor
 from PIL import Image
-import io
 import cv2
-import gradio as gr
-import os
-import requests
-import zipfile
-import subprocess
 
 # Ensure the necessary model files are available
 def download_file(url, destination):
@@ -18,12 +16,30 @@ def download_file(url, destination):
     with open(destination, 'wb') as f:
         f.write(response.content)
 
-# Install the necessary packages
-subprocess.run(["pip", "install", "git+https://github.com/facebookresearch/segment-anything.git"])
-subprocess.run(["git", "clone", "https://github.com/yformer/EfficientSAM.git"])
+# Download SAM model
+if not os.path.exists("weights/sam_vit_h_4b8939.pth"):
+    os.makedirs("weights", exist_ok=True)
+    download_file("https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth", "weights/sam_vit_h_4b8939.pth")
+
+# Add EfficientSAM to Python path
+sys.path.append(os.path.abspath("EfficientSAM-main"))
 
-# Change directory to EfficientSAM
-os.chdir("EfficientSAM")
+from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
+from efficient_sam.build_efficient_sam_vits import build_efficient_sam_vits
+
+# Constants
+DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+MODEL_TYPE = "vit_h"
+CHECKPOINT_PATH = "weights/sam_vit_h_4b8939.pth"
+
+# Load SAM model
+sam = sam_model_registry[MODEL_TYPE](checkpoint=CHECKPOINT_PATH).to(device=DEVICE)
+mask_generator_sam = SamAutomaticMaskGenerator(sam)
+
+# Load EfficientSAM model
+with zipfile.ZipFile("weights/efficient_sam_vits.pt.zip", 'r') as zip_ref:
+    zip_ref.extractall("weights")
+efficient_sam_vits_model = build_efficient_sam_vits()
 
 from segment_anything.utils.amg import (
     batched_mask_to_box,
@@ -47,6 +63,7 @@ def process_small_region(rles):
         unchanged = unchanged and not changed
         new_masks.append(torch.as_tensor(mask).unsqueeze(0))
         scores.append(float(unchanged))
+
     masks = torch.cat(new_masks, dim=0)
     boxes = batched_mask_to_box(masks)
     keep_by_nms = batched_nms(
@@ -137,18 +154,6 @@ def process_image(image):
 
     return [image, sam_annotated_image, efficient_sam_annotated_image]
 
-# Download EfficientSAM model
-if not os.path.exists("weights/efficient_sam_vits.pt.zip"):
-    download_file("https://example.com/path/to/efficient_sam_vits.pt.zip", "weights/efficient_sam_vits.pt.zip")
-
-# Extract EfficientSAM model
-with zipfile.ZipFile("weights/efficient_sam_vits.pt.zip", 'r') as zip_ref:
-    zip_ref.extractall("weights")
-
-from efficient_sam.build_efficient_sam import build_efficient_sam_vits
-efficient_sam_vits_model = build_efficient_sam_vits()
-efficient_sam_vits_model.eval()
-
 # Gradio interface
 interface = gr.Interface(
     fn=process_image,
@@ -161,3 +166,4 @@ interface = gr.Interface(
 interface.launch(debug=True)
 
 
+