fixed noise in mask issue

app/main.py

@@ -208,7 +208,9 @@ async def label_image(image: str, mask_labels: MaskLabels) -> Response:
     for i in range(len(mask_labels.masks)):
         mask_i = mask_labels.masks[i]
         mask_i = mask_i.replace("data:image/png;base64,", "")
-        save_masks.append(Image.open(io.BytesIO(base64.b64decode(mask_i))).convert("L"))
+        mask_i = Image.open(io.BytesIO(base64.b64decode(mask_i))).convert("L")
+        mask_i = mask_i.point(lambda p: 0 if <= 1 else p)
+        save_masks.append(mask_i)
     bboxes = (
         batched_mask_to_box(
             torch.as_tensor(np.array([np.array(m) for m in save_masks]))

app/per_sam/model.py

@@ -8,6 +8,7 @@ import cv2
 from torchvision.ops.boxes import batched_nms
 from app.mobile_sam import SamPredictor
 from app.mobile_sam.utils import batched_mask_to_box
+from app.sam.postprocessing import clean_mask_torch
 
 
 def point_selection(mask_sim, topk: int = 1):
@@ -139,7 +140,10 @@ def fast_inference(
         # Weighted sum three-scale masks
         logits_high = logits_high * weights.unsqueeze(-1)
         logit_high = logits_high.sum(0)
-        mask = (logit_high > 0).detach().cpu().numpy()
+        # mask = (logit_high > 0).detach().cpu().numpy()
+
+        mask = (logit_high > 0)
+        mask = clean_mask_torch(mask).bool()[0, 0, :, :].detach().cpu().numpy()
 
         logits = logits * weights_np[..., None]
         logit = logits.sum(0)

app/sam/predictor.py

@@ -10,6 +10,7 @@ import torch
 from typing import Optional, Tuple
 from .sam import Sam
 from .transforms import ResizeLongestSide
+from .postprocess import clean_mask_torch
 
 
 class SamPredictor:
@@ -237,6 +238,7 @@ class SamPredictor:
 
         if not return_logits:
             masks = masks > self.model.mask_threshold
+            masks = clean_mask_torch(masks.int()).bool()
 
         return masks, iou_predictions, low_res_masks
 

instance-labeler/app/canvas.tsx

@@ -28,7 +28,7 @@ const maskFilter = (imageData: ImageData, color: RGB) => {
     const g = imageData.data[i + 1];
     const b = imageData.data[i + 2];
 
-    if (r === 0 && g === 0 && b === 0) {
+    if (r <= 1 && g <= 1 && b <= 1) {
       imageData.data[i + 3] = 0;
     } else {
       imageData.data[i] = color.r;
@@ -261,8 +261,8 @@ export default function Canvas({ imageUrl, imageName }: { imageUrl: string, imag
 
   const predLabels = async () => {
     const length = useLatest ? groupRef.current.length : groupRef.current.length - 1;
-    if (groupRef.current.length === 0) {
-      alert('Please pin an instance');
+    if (groupRef.current.length === 0 || classList.length === 0) {
+      alert('Please pin an instance first');
       return
     }
     const mask = groupRef.current[length - 1].toDataURL({ x: 0, y: 0, width: image?.width, height: image?.height });

requirements.txt

@@ -6,3 +6,4 @@ Pillow
 fastapi
 uvicorn
 timm
+kornia