add bbox support

Dev.ipynb

@@ -26,6 +26,46 @@
     "slicer.util.selectModule(\"SammBase\")"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "id": "642430aa-9f67-4af0-940f-540047fe92a3",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "vtkPoints (0x62a4f80)\n",
+      "  Debug: Off\n",
+      "  Modified Time: 1776089\n",
+      "  Reference Count: 1\n",
+      "  Registered Events: (none)\n",
+      "  Data: 0x102c7160\n",
+      "  Data Array Name: Points\n",
+      "  Number Of Points: 4\n",
+      "  Bounds: \n",
+      "    Xmin,Xmax: (-32.0076, 95.4545)\n",
+      "    Ymin,Ymax: (-34.8485, 63.2576)\n",
+      "    Zmin,Zmax: (0, 0)\n",
+      "\n",
+      "\n",
+      "0.0\n"
+     ]
+    }
+   ],
+   "source": [
+    "import vtk\n",
+    "\n",
+    "logic = SammBaseLogic()\n",
+    "logic._parameterNode.SetNodeReferenceID(\"sammPrompt2DBox\", \"vtkMRMLMarkupsPlaneNode1\")\n",
+    "plane = logic._parameterNode.GetNodeReference(\"sammPrompt2DBox\")\n",
+    "points = vtk.vtkPoints() \n",
+    "plane.GetPlaneCornerPoints(points)\n",
+    "print(points)\n",
+    "print(points.GetPoint(0)[2])\n"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 3,
@@ -643,17 +683,49 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 6,
    "id": "92be4fc6-f838-47d0-ae4e-47c2598eb353",
    "metadata": {},
-   "outputs": [],
-   "source": []
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "array([1, 2, 3, 4])"
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "import numpy as np\n",
+    "np.array([1,2,3,4]).reshape([1,4])[0]"
+   ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 20,
    "id": "4e83cab3-759b-4bf4-a2ff-2f0134331451",
    "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "False\n"
+     ]
+    }
+   ],
+   "source": [
+    "print([]==None)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "512d3635-580d-443a-b011-f51e4bc4f952",
+   "metadata": {},
    "outputs": [],
    "source": []
   }

README.md

@@ -4,8 +4,6 @@
 
 [Laboratory of Biomechanical and Image Guided Surgical Systems](https://bigss.lcsr.jhu.edu/), [Johns Hopkins University](https://www.jhu.edu/)
 
-[Yihao Liu](https://yihao.one/), [Jeremy Zhang](https://jeremyzz830.github.io/), Zhangcong She
-
 ## Known issues
 #13 Resolved with a not-so-elegant way. Tested work on MRHead, DZ-MR, MRBrainTumor2 and CTChest data from Slicer. Please report a bug if your data is not working. Note for the first few seconds when you start "Mask Sync", the server is not so stable, wait a few seconds then slide it up and down, the mask then will be updated. Note now you can only work on the RED view. Will update later to support all 3 views.
 

samm-python-terminal/sam_server.py

@@ -68,10 +68,10 @@ def looping(self):
                 self.cleanup()
             except zmq.error.Again:
                 continue
-            except Exception as e:
-                logging.error(traceback.format_exc())
-                time.sleep(self.interv)
-                continue
+            # except Exception as e:
+            #     logging.error(traceback.format_exc())
+            #     time.sleep(self.interv)
+            #     continue
 
             time.sleep(self.interv)
 

samm-python-terminal/utl_sam_msg.py

@@ -109,6 +109,8 @@ def getDecodedData(msgbyte):
 
 '''
 n : int
+view : char
+bbox : int 32, 4
 positivePrompts, int 32, n * 2
 negativePrompts, int 32
 '''
@@ -118,12 +120,14 @@ def getEncodedData(self):
 
         n = self.msg["n"]
         view = SammViewMapper[self.msg["view"]]
+        bbox2D = self.msg["bbox2D"]
         positivePoints = self.msg["positivePrompts"]
         negativePoints = self.msg["negativePrompts"]
 
         msg = b''
         msg += np.array([n], dtype='int32').tobytes()
         msg += np.array([view], dtype='int32').tobytes()
+        msg += bbox2D.astype("int32").tobytes()
 
         if positivePoints is not None and positivePoints.shape[0] > 0:
             msg += np.array([positivePoints.shape[0]], dtype='int32').tobytes()
@@ -150,6 +154,9 @@ def getDecodedData(msgbyte):
         msg["view"] = SammViewMapper["DICT"][np.frombuffer(msgbyte[pt:pt+4], dtype="int32").reshape([1])[0]]
         pt += 4
 
+        msg["bbox2D"] = np.frombuffer(msgbyte[pt:pt+4*4], dtype="int32").reshape([1,4])[0]
+        pt += 4*4
+
         positivePromptNum = np.frombuffer(msgbyte[pt:pt+4], dtype="int32").reshape([1])[0]
         pt += 4
 
@@ -164,7 +171,7 @@ def getDecodedData(msgbyte):
         pt += 4
 
         if negativePromptNum > 0:
-            negativePrompt = np.frombuffer(msgbyte[pt:pt+4*2*negativePromptNum], dtype="int32").reshape([positivePromptNum, 2])
+            negativePrompt = np.frombuffer(msgbyte[pt:pt+4*2*negativePromptNum], dtype="int32").reshape([negativePromptNum, 2])
             pt += 4 * 2 * negativePromptNum
             msg["negativePrompts"] = negativePrompt
         else:

samm-python-terminal/utl_sam_server.py

@@ -158,10 +158,21 @@ def CalculateEmbeddings(msg):
 
     print("[SAMM INFO] Embeddings Cached.")
 
+def helperPredict(dataNode, msg, points, labels, bbox2d):
+    dataNode.samPredictor[msg["view"]].features = dataNode.features[msg["view"]][msg["n"]].to("cuda")
+    seg, _, _ = dataNode.samPredictor[msg["view"]].predict(
+        point_coords = points,
+        point_labels = labels,
+        box = bbox2d,
+        multimask_output = False,)
+    seg = seg[0]
+    return seg
+
 def sammProcessingCallBack_INFERENCE(msg):
     dataNode = SammParameterNode()
     positivePoints = msg["positivePrompts"]
     negativePoints = msg["negativePrompts"]
+    bbox2d = msg["bbox2D"]
 
     points = []
     labels = []
@@ -176,20 +187,26 @@ def sammProcessingCallBack_INFERENCE(msg):
             labels.append(0)
 
     seg = None
-    if len(points) > 0:
-        if msg["view"] == "R":
-            tempsize = [dataNode.imageSize[1], dataNode.imageSize[2]]
-        if msg["view"] == "G":
-            tempsize = [dataNode.imageSize[0], dataNode.imageSize[2]]
-        if msg["view"] == "Y":
-            tempsize = [dataNode.imageSize[0], dataNode.imageSize[1]]
+    if len(points) > 0 and (bbox2d[0]!=-404):
+
+        points = np.array(points)
+        point_labels = np.array(labels)
+        bbox2d = np.array(bbox2d)
+        seg = helperPredict(dataNode, msg, points, point_labels, bbox2d)
+
+    elif len(points) == 0 and (bbox2d[0]!=-404):
+
+        points = None
+        point_labels = None
+        bbox2d = np.array(bbox2d)
+        seg = helperPredict(dataNode, msg, points, point_labels, bbox2d)
+
+    elif len(points) > 0 and (bbox2d[0]==-404):
 
-        dataNode.samPredictor[msg["view"]].features = dataNode.features[msg["view"]][msg["n"]].to("cuda")
-        seg, _, _ = dataNode.samPredictor[msg["view"]].predict(
-            point_coords = np.array(points),
-            point_labels = np.array(labels),
-            multimask_output = False,)
-        seg = seg[0]
+        points = np.array(points)
+        point_labels = np.array(labels)
+        bbox2d = None
+        seg = helperPredict(dataNode, msg, points, point_labels, bbox2d)
 
     else:
         if msg["view"] == "R":