working on sam encoder

Makefile

@@ -1,2 +1,6 @@
 dev:
 	uvicorn app.server:app --reload
+
+download_sam_model:
+	wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth
+	mv sam_vit_h_4b8939.pth sam_model/

app/cloth_segmentation/model.py

@@ -1,12 +1,18 @@
+import sys
+sys.path.append("..")
 from base64 import b64encode
 from dataclasses import dataclass
 from io import BytesIO
 from PIL import Image
+from segment_anything import sam_model_registry, SamPredictor
 
 from transformers import pipeline
 
 pipe = pipeline("image-segmentation", model="mattmdjaga/segformer_b2_clothes")
 
+sam_checkpoint = "app/sam_models/sam_vit_h_4b8939.pth"
+model_type = "vit_h"
+device = "cuda"
 
 @dataclass
 class Layer:
@@ -31,3 +37,12 @@ def segment(image: Image) -> [Layer]:
         result.append(Layer(t['label'], image_to_base64(t['mask'])))
 
     return result
+
+def sam_anything(image: Image) -> [Layer]:
+    sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
+    sam.to(device=device)
+
+    predictor = SamPredictor(sam)
+    pred = predictor.set_image(image)
+    print(f"Predicted {len(pred.pred_masks)} instances")
+    return pred

app/server.py

@@ -1,5 +1,5 @@
 from PIL import Image
-from app.cloth_segmentation.model import Layer, segment
+from app.cloth_segmentation.model import Layer, segment, sam_anything
 from typing import List
 
 from contextlib import asynccontextmanager
@@ -32,3 +32,8 @@ def index():
 def mask(upload: UploadFile) -> List[Layer]:
     image = Image.open(upload.file)
     return segment(image)
+
+@app.post("/encode")
+def encode(upload: UploadFile) -> List[Layer]:
+    image = Image.open(upload.file)
+    return sam_anything(image)

requirements.txt

@@ -8,3 +8,5 @@ numpy==1.24.3
 torch==2.0.*
 torchvision
 transformers==4.30.1
+
+segment-anything @ git+https://github.com/facebookresearch/segment-anything.git@6fdee8f2727f4506cfbbe553e23b895e27956588