files2

.gitignore

@@ -0,0 +1,9 @@
+/__pycache__
+/__pycache__/*
+*/__pycache__/*
+*/__pycache__*
+
+
+/.ipynb_checkpoints
+/.ipynb_checkpoints/*
+/notebook/.ipynb_checkpoints

configs/config.py

@@ -0,0 +1,14 @@
+import torch
+import os
+
+DEVICE = torch.device("cpu")
+
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+BACKEND_DIR = os.path.abspath(os.path.join(BASE_DIR, '..'))
+
+MODEL_PATH = os.path.join(BACKEND_DIR, 'models', 'checkpoint')
+
+FAST_SAM_IMGSZ = 1024
+FAST_SAM_CONF = 0.5
+FAST_SAM_IOU = 0.6
+FAST_SAM_RETINA_MASKS = True

models/checkpoint/FastSAM-s.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c9f78716a81c7aff0d608ccc73e1b82ab3aaad86005049f6a92106a0be6d0844
+size 23851578

models/checkpoint/FastSAM-x.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:752cadc2828edb1cd4bc4f9eb587100631af06ea2108f4c9ed56df4755701e76
+size 144972346

models/checkpoint/FastSAM.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c0be4e7ddbe4c15333d15a859c676d053c486d0a746a3be6a7a9790d52a9b6d7
+size 144943063

models/model.py

@@ -0,0 +1,37 @@
+from typing import Any, List
+from fastsam import FastSAM
+from PIL import Image
+from fastsam import FastSAM, FastSAMPrompt
+import numpy as np
+from configs.config import DEVICE, FAST_SAM_CONF, FAST_SAM_IMGSZ, FAST_SAM_IOU, FAST_SAM_RETINA_MASKS, MODEL_PATH
+from PIL import Image
+import os
+
+
+def loadModel(model_name: str = "FastSAM-x.pt") -> Any:
+    path: str = MODEL_PATH + "/" + model_name
+    if not os.path.exists(path):
+        downloadModel(path)
+    return FastSAM(f"{MODEL_PATH}/{model_name}")
+
+
+def getMask(image_path: Image, fast_sam: FastSAM, point: List[List[int]], point_label: List[int]) -> Any:
+    result: Any = fast_sam(
+        source=image_path,
+        device=DEVICE,
+        retina_masks=FAST_SAM_RETINA_MASKS,
+        imgsz=FAST_SAM_IMGSZ,
+        conf=FAST_SAM_CONF,
+        iou=FAST_SAM_IOU,
+    )
+    prompt_process = FastSAMPrompt(image_path, result, device=DEVICE)
+    return prompt_process.point_prompt(points=point, pointlabel=point_label)
+
+
+def downloadModel(model_name):
+    import requests
+    url = "https://firebasestorage.googleapis.com/v0/b/lexicons-5.appspot.com/o/FastSam-Models%2FFastSAM-x.pt?alt=media&token=64b65560-17d6-47b0-8a2b-8e2ee096da64"
+    r = requests.get(url)
+    with open(model_name, 'wb') as f:
+        f.write(r.content)
+    return model_name

models/preprocess.py

@@ -0,0 +1,17 @@
+from typing import Any, Dict, List
+from fastsam import FastSAM, FastSAMPrompt
+import numpy as np
+from configs.config import DEVICE, FAST_SAM_CONF, FAST_SAM_IMGSZ, FAST_SAM_IOU, FAST_SAM_RETINA_MASKS
+
+def preprocess(points_data: List[Dict]) -> Any:
+    
+    input_points = []
+    input_labels = []
+    
+    for point in points_data:
+        input_points.append([int(point['x_']), int(point['y_'])])
+        input_labels.append(int(point['flag_']))
+    
+    return input_points, input_labels
+
+    

tools/tools.py

@@ -0,0 +1,10 @@
+
+import base64
+import io
+from typing import Any
+
+def convertToBuffer(image : Any) -> bytes:
+    buffered : io.BytesIO = io.BytesIO()
+    image.save(buffered, format="PNG")
+    img_base64 : bytes = base64.b64encode(buffered.getvalue()).decode('utf-8')
+    return img_base64

visualize/visualize.py

@@ -0,0 +1,25 @@
+from PIL import Image
+import cv2
+import numpy as np
+import torch
+from typing import Any, Dict, List
+
+def removeBgFromSegmentImage(og_image : Any, og_mask : Any, color: tuple = (0, 0, 255, 255),  opacity: float = 0.2) -> Image:
+    og_image = np.array(og_image.convert('RGB'))
+    mask = og_mask.astype(np.uint8) * 255  # Convert to 0 or 255
+
+    rgba_image = np.zeros((og_image.shape[0], og_image.shape[1], 4), dtype=np.uint8)
+
+    color_with_opacity = (color[0], color[1], color[2], int(color[3] * opacity))
+    rgba_image[mask > 0] = color_with_opacity
+
+    return Image.fromarray(rgba_image)
+
+def removeOnlyBg(og_image : Any, og_mask : Any) -> Image:
+    img = np.array(og_image.convert('RGB'))
+    mask = cv2.resize(og_mask.astype(np.uint8), (img.shape[1], img.shape[0]), interpolation=cv2.INTER_NEAREST)
+    rgba_image = np.zeros((img.shape[0], img.shape[1], 4), dtype=np.uint8)
+    rgba_image[..., :3] = img
+    rgba_image[..., 3] = mask * 255
+
+    return Image.fromarray(rgba_image)