finish

app/cloth_segmentation/model.py

@@ -9,7 +9,7 @@ from PIL import Image
 from transformers import pipeline
 
 from app.simple_segmentation.network import U2NET
-from app.utils import image_to_base64, check_or_download_model
+from app.utils import image_to_base64
 
 pipe = pipeline("image-segmentation", model="mattmdjaga/segformer_b2_clothes")
 

app/segment_anything/model.py

@@ -5,8 +5,10 @@ from segment_anything import sam_model_registry, SamPredictor
 from base64 import b64encode
 
 from huggingface_hub import hf_hub_download
+from app.utils import check_or_download_model
 
-chkpt_path = hf_hub_download("ybelkada/segment-anything", "checkpoints/sam_vit_h_4b8939.pth")
+MODEL_URL = "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth"
+chkpt_path = ".checkpoint/sam_vit_h_4b8939.pth"
 model_type = "default"
 
 def initialize_model(model_type, chkpt_path):
@@ -18,6 +20,9 @@ def initialize_model(model_type, chkpt_path):
         print("GPU not available. Using CPU.")
     return sam
 
+def load_segment_model(checkpoint_path, device='cpu'):
+    check_or_download_model(MODEL_URL, checkpoint_path)
+
 def preprocess_image(image: Image) -> np.ndarray:
     # Implement any necessary preprocessing steps
     # Ensure the image is in the expected format

app/server.py

@@ -10,6 +10,7 @@ from pydantic.dataclasses import dataclass
 from app.cloth_segmentation.model import Layer, segment
 from app.segment_anything.model import predict
 from app.simple_segmentation.model import Mode, binary_segment, load_seg_model, Result
+from app.segment_anything.model import load_segment_model
 from app.utils import image_to_base64
 
 # === Context ===
@@ -30,6 +31,7 @@ class MaskResponse:
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     ml_models["cloth_segmentation"] = load_seg_model(".checkpoint/model.pth")
+    ml_models["segment_anything"] = load_segment_model('.checkpoint/sam_vit_h_4b8939.pth')
     yield
     ml_models.clear()
 

app/simple_segmentation/model.py

@@ -9,7 +9,7 @@ from PIL import Image
 from torchvision import transforms
 
 from app.simple_segmentation.network import U2NET
-from app.utils import check_or_download_model
+from app.utils import check_or_download_model, load_checkpoint
 
 MODEL_URL = "https://huggingface.co/spaces/wildoctopus/cloth-segmentation/resolve/main/model/cloth_segm.pth"
 
@@ -23,18 +23,6 @@ class Result:
     lower_body: Optional[Image.Image] = None
     full_body: Optional[Image.Image] = None
 
-def load_checkpoint(model, checkpoint_path):
-    model_state_dict = torch.load(
-        checkpoint_path, map_location=torch.device("cpu"))
-    new_state_dict = OrderedDict()
-    for k, v in model_state_dict.items():
-        name = k[7:]  # remove `module.`
-        new_state_dict[name] = v
-
-    model.load_state_dict(new_state_dict)
-    print("----checkpoints loaded from path: {}----".format(checkpoint_path))
-    return model
-
 def load_seg_model(checkpoint_path, device='cpu'):
     net = U2NET(in_ch=3, out_ch=4)
     check_or_download_model(MODEL_URL, checkpoint_path)

app/utils.py

@@ -1,10 +1,11 @@
-import os
+import os, torch
 from io import BytesIO
 from base64 import b64encode
 from urllib.request import urlretrieve
 from urllib.parse import urlparse
 from PIL import Image
 from typing import Optional
+from collections import OrderedDict
 
 def image_to_base64(image: Image.Image | None):
     if image == None:
@@ -19,16 +20,52 @@ def is_valid_url(url):
         return all([result.scheme, result.netloc])
     except ValueError:
         return False
+    
+def download_with_progress(model_url, file_path):
+    try:
+        response, _ = urlretrieve(model_url, file_path, reporthook=download_progress)
+    except Exception as e:
+        print(f"Error downloading the model: {e}")
+        return False
+    else:
+        return True
+
+def download_progress(block_num, block_size, total_size):
+    progress = min(1.0, block_num * block_size / total_size)
+    bar_length = 50
+    block = int(round(bar_length * progress))
+    progress_percent = progress * 100
+    progress_bar = f"[{'=' * block}{' ' * (bar_length - block)}] {progress_percent:.2f}%\r"
+    print(progress_bar, end='', flush=True)
 
 def check_or_download_model(model_url, file_path):
     if not is_valid_url(model_url):
         print("Invalid model URL.")
         return
 
-    if not os.path.exists(file_path):
-        print("No model found, downloading model")
-        os.makedirs(os.path.dirname(file_path), exist_ok=True)
-        urlretrieve(model_url, file_path)
-        print("Model downloaded successfully.")
-    else:
+    if os.path.exists(file_path):
         print("Model already exists at:", file_path)
+    else:
+        print("No model found, downloading model.")
+        os.makedirs(os.path.dirname(file_path), exist_ok=True)
+        download_with_progress(model_url, file_path)
+        print("\nModel downloaded successfully.")
+
+def load_checkpoint(model, checkpoint_path):
+    # Load model checkpoint
+    model_state_dict = torch.load(
+        checkpoint_path, map_location=torch.device("cuda" if torch.cuda.is_available() else "cpu"))
+
+    # Create a new state dictionary without the 'module.' prefix
+    new_state_dict = OrderedDict()
+    for k, v in model_state_dict.items():
+        name = k[7:]  # remove `module.`
+        new_state_dict[name] = v
+
+    # Load the new state dictionary into the model
+    model.load_state_dict(new_state_dict)
+
+    # Print a confirmation message
+    print("---- Checkpoint loaded from path: {} ----".format(checkpoint_path))
+
+    return model