fix the resizing issue during concurrent visiting

app.py

@@ -109,6 +109,8 @@ def reset(session_state):
     session_state['ori_image'] = None
     session_state['image_with_prompt'] = None
     session_state['feature'] = None
+    session_state['input_size'] = None
+    session_state['original_size'] = None
     return None, None, None, session_state
 
 
@@ -119,6 +121,8 @@ def reset_all(session_state):
     session_state['ori_image'] = None
     session_state['image_with_prompt'] = None
     session_state['feature'] = None
+    session_state['input_size'] = None
+    session_state['original_size'] = None
     return None, None, None, None, None, None, session_state
 
 
@@ -149,7 +153,7 @@ def on_image_upload(
     session_state['image_with_prompt'] = copy.deepcopy(image)
     print("Image changed")
     nd_image = np.array(image)
-    session_state['feature'] = predictor.set_image(nd_image)
+    session_state['feature'], session_state['input_size'], session_state['original_size'] = predictor.set_image(nd_image)
 
     return image, None, None, session_state
 
@@ -190,12 +194,16 @@ def segment_with_points(
         fill=point_color,
     )
     image = session_state['image_with_prompt']
+    print(f"image: {image.size}")
+    nd_image = np.array(session_state['ori_image'])
 
     if ENABLE_ONNX:
         coord_np = np.array(session_state['coord_list'])[None]
         label_np = np.array(session_state['label_list'])[None]
         masks, scores, _ = predictor.predict(
             features=session_state['feature'],
+            input_size=session_state['input_size'],
+            original_size=session_state['original_size'],
             point_coords=coord_np,
             point_labels=label_np,
         )
@@ -206,6 +214,8 @@ def segment_with_points(
         label_np = np.array(session_state['label_list'])
         masks, scores, logits = predictor.predict(
             features=session_state['feature'],
+            input_size=session_state['input_size'],
+            original_size=session_state['original_size'],
             point_coords=coord_np,
             point_labels=label_np,
             num_multimask_outputs=4,
@@ -233,7 +243,7 @@ def segment_with_points(
     binary_mask = np.where(annotations[0] > 0.5, 255, 0).astype(np.uint8)
     mask = Image.fromarray(binary_mask)
     binary_mask = np.expand_dims(binary_mask, axis=2)
-    crop = Image.fromarray(np.concatenate((session_state['ori_image'], binary_mask), axis=2), "RGBA")
+    crop = Image.fromarray(np.concatenate((nd_image, binary_mask), axis=2), "RGBA")
     return seg, mask, crop, session_state
 
 
@@ -282,6 +292,8 @@ def segment_with_box(
             point_labels = np.array([2, 3])[None]
             masks, _, _ = predictor.predict(
                 features=session_state['feature'],
+                input_size=session_state['input_size'],
+                original_size=session_state['original_size'],
                 point_coords=point_coords,
                 point_labels=point_labels,
             )
@@ -289,6 +301,8 @@ def segment_with_box(
         else:
             masks, scores, _ = predictor.predict(
                 features=session_state['feature'],
+                input_size=session_state['input_size'],
+                original_size=session_state['original_size'],
                 box=box_np,
                 num_multimask_outputs=1,
             )

segment_anything/onnx/predictor_onnx.py

@@ -53,34 +53,30 @@ class SamPredictorONNX:
         input_image = self.transform.apply_image(image)
         input_image = input_image.transpose(2, 0, 1)[None, :, :, :]
         self.reset_image()
-        self.original_size = image.shape[:2]
-        self.input_size = tuple(input_image.shape[-2:])
+        original_size = image.shape[:2]
+        input_size = tuple(input_image.shape[-2:])
         input_image = self.preprocess(input_image).astype(np.float32)
         outputs = self.encoder.run(None, {'image': input_image})
-        self.features = outputs[0]
-        self.is_image_set = True
+        features = outputs[0]
 
-        return self.features
+        return features, input_size, original_size
 
     def predict(
             self,
-            features: np.ndarray = None,
+            features: np.ndarray,
+            input_size: Tuple[int, int],
+            original_size: Tuple[int, int],
             point_coords: Optional[np.ndarray] = None,
             point_labels: Optional[np.ndarray] = None,
     ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
-        if features is None and not self.is_image_set:
-            raise RuntimeError("An image must be set with .set_image(...) before mask prediction.")
-        if features is None:
-            features = self.features
-
-        point_coords = self.transform.apply_coords(point_coords, self.original_size)
+        point_coords = self.transform.apply_coords(point_coords, original_size)
         outputs = self.decoder.run(None, {
             'image_embeddings': features,
             'point_coords': point_coords.astype(np.float32),
             'point_labels': point_labels.astype(np.float32)
         })
         scores, low_res_masks = outputs[0], outputs[1]
-        masks = self.postprocess_masks(low_res_masks)
+        masks = self.postprocess_masks(low_res_masks, input_size, original_size)
         masks = masks > self.mask_threshold
 
         return masks, scores, low_res_masks
@@ -102,10 +98,10 @@ class SamPredictorONNX:
         x = np.pad(x, ((0, 0), (0, 0), (0, padh), (0, padw)), mode='constant', constant_values=0)
         return x
 
-    def postprocess_masks(self, mask: np.ndarray):
+    def postprocess_masks(self, mask: np.ndarray, input_size: Tuple[int, int], original_size: Tuple[int, int]) -> np.ndarray:
         mask = mask.squeeze(0).transpose(1, 2, 0)
         mask = cv2.resize(mask, (self.img_size, self.img_size), interpolation=cv2.INTER_LINEAR)
-        mask = mask[:self.input_size[0], :self.input_size[1], :]
-        mask = cv2.resize(mask, (self.original_size[1], self.original_size[0]), interpolation=cv2.INTER_LINEAR)
+        mask = mask[:input_size[0], :input_size[1], :]
+        mask = cv2.resize(mask, (original_size[1], original_size[0]), interpolation=cv2.INTER_LINEAR)
         mask = mask.transpose(2, 0, 1)[None, :, :, :]
         return mask

segment_anything/predictor.py

@@ -59,13 +59,15 @@ class SamPredictor:
         input_image_torch = torch.as_tensor(input_image, device=self.device)
         input_image_torch = input_image_torch.permute(2, 0, 1).contiguous()[None, :, :, :]
 
-        return self.set_torch_image(input_image_torch, image.shape[:2])
+        input_size = tuple(input_image_torch.shape[-2:])
+        original_size = image.shape[:2]
+
+        return self.set_torch_image(input_image_torch), input_size, original_size
 
     @torch.no_grad()
     def set_torch_image(
         self,
         transformed_image: torch.Tensor,
-        original_image_size: Tuple[int, ...],
     ) -> torch.Tensor:
         """
         Calculates the image embeddings for the provided image, allowing
@@ -75,8 +77,6 @@ class SamPredictor:
         Arguments:
           transformed_image (torch.Tensor): The input image, with shape
             1x3xHxW, which has been transformed with ResizeLongestSide.
-          original_image_size (tuple(int, int)): The size of the image
-            before transformation, in (H, W) format.
         """
         assert (
             len(transformed_image.shape) == 4
@@ -85,24 +85,23 @@ class SamPredictor:
         ), f"set_torch_image input must be BCHW with long side {self.model.image_encoder.img_size}."
         self.reset_image()
 
-        self.original_size = original_image_size
-        self.input_size = tuple(transformed_image.shape[-2:])
         input_image = self.model.preprocess(transformed_image)
-        self.features = self.model.image_encoder(input_image)
-        self.is_image_set = True
+        features = self.model.image_encoder(input_image)
 
-        return self.features
+        return features
 
     def predict(
         self,
-        features: torch.Tensor = None,
+        features: torch.Tensor,
+        input_size: Tuple[int, int],
+        original_size: Tuple[int, int],
         point_coords: Optional[np.ndarray] = None,
         point_labels: Optional[np.ndarray] = None,
         box: Optional[np.ndarray] = None,
         mask_input: Optional[np.ndarray] = None,
         num_multimask_outputs: int = 3,
         return_logits: bool = False,
-        use_stability_score: bool = False
+        use_stability_score: bool = False,
     ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
         """
         Predict masks for the given input prompts, using the currently set image.
@@ -134,24 +133,18 @@ class SamPredictor:
             of masks and H=W=256. These low resolution logits can be passed to
             a subsequent iteration as mask input.
         """
-        if features is None and not self.is_image_set:
-            raise RuntimeError("An image must be set with .set_image(...) before mask prediction.")
-
-        if features is None:
-            features = self.features
-
         # Transform input prompts
         coords_torch, labels_torch, box_torch, mask_input_torch = None, None, None, None
         if point_coords is not None:
             assert (
                 point_labels is not None
             ), "point_labels must be supplied if point_coords is supplied."
-            point_coords = self.transform.apply_coords(point_coords, self.original_size)
+            point_coords = self.transform.apply_coords(point_coords, original_size)
             coords_torch = torch.as_tensor(point_coords, dtype=torch.float, device=self.device)
             labels_torch = torch.as_tensor(point_labels, dtype=torch.int, device=self.device)
             coords_torch, labels_torch = coords_torch[None, :, :], labels_torch[None, :]
         if box is not None:
-            box = self.transform.apply_boxes(box, self.original_size)
+            box = self.transform.apply_boxes(box, original_size)
             box_torch = torch.as_tensor(box, dtype=torch.float, device=self.device)
             box_torch = box_torch[None, :]
         if mask_input is not None:
@@ -160,6 +153,8 @@ class SamPredictor:
 
         masks, iou_predictions, low_res_masks = self.predict_torch(
             features,
+            input_size,
+            original_size,
             coords_torch,
             labels_torch,
             box_torch,
@@ -178,6 +173,8 @@ class SamPredictor:
     def predict_torch(
         self,
         features: torch.Tensor,
+        input_size: Tuple[int, int],
+        original_size: Tuple[int, int],
         point_coords: Optional[torch.Tensor],
         point_labels: Optional[torch.Tensor],
         boxes: Optional[torch.Tensor] = None,
@@ -249,7 +246,7 @@ class SamPredictor:
             )
 
         # Upscale the masks to the original image resolution
-        masks = self.model.postprocess_masks(low_res_masks, self.input_size, self.original_size)
+        masks = self.model.postprocess_masks(low_res_masks, input_size, original_size)
 
         if not return_logits:
             masks = masks > self.model.mask_threshold