change node input to json

nodes.py

@@ -31,7 +31,7 @@ class DownloadAndLoadSAM2RealtimeModel:
     def INPUT_TYPES(s):
         return {"required": {
             "model": ([ 
-                    'sam2_hiera_tiny.pt', 'sam2_hiera_small.pt',
+                    'sam2_hiera_tiny.pt',
                     ],),
             "segmentor": (
                     ['realtime'],
@@ -70,8 +70,7 @@ class DownloadAndLoadSAM2RealtimeModel:
 
         if not os.path.exists(model_path):
             print(f"Downloading SAM2 model to: {model_path}")
-            base_url = "https://dl.fbaipublicfiles.com/segment_anything_2/072824/"
-            url = f"{base_url}{model}"
+            url = "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_tiny.pt"
             response = requests.get(url, stream=True)
             response.raise_for_status()
 
@@ -84,9 +83,8 @@ class DownloadAndLoadSAM2RealtimeModel:
 
         config_dir = os.path.join(script_directory, "sam2_configs") 
 
-        model_cfg = model.replace(".pt", ".yaml")
-
         # Code ripped out of sam2.build_sam.build_sam2_camera_predictor to appease Hydra
+        model_cfg = "sam2_hiera_t.yaml" #TODO: remove hardcoded config and path
         with initialize_config_dir(config_dir=config_dir, version_base=None):
             cfg = compose(config_name=model_cfg)
 
@@ -142,12 +140,12 @@ class Sam2RealtimeSegmentation:
             "required": {
                 "images": ("IMAGE",),
                 "sam2_model": ("SAM2MODEL",),
+                "reset_tracking": ("BOOLEAN", {"default": False}),
                 # "keep_model_loaded": ("BOOLEAN", {"default": True}),
             },
             "optional": {
                 "coordinates_positive": ("STRING", ),
                 "coordinates_negative": ("STRING", ),
-                "reset_tracking": ("BOOLEAN", {"default": False}),
                 # "bboxes": ("BBOX", ),
                 # "individual_objects": ("BOOLEAN", {"default": False}),
                 # "mask": ("MASK", ),
@@ -194,9 +192,9 @@ class Sam2RealtimeSegmentation:
         images,
         sam2_model,
         # keep_model_loaded,
+        reset_tracking,
         coordinates_positive=None,
         coordinates_negative=None,
-        reset_tracking=False,
         #point_labels=None,
         # bboxes=None,
         # individual_objects=False,
@@ -252,7 +250,6 @@ class Sam2RealtimeSegmentation:
 
                 # Create colored overlay for processed frames
                 mask_colored = torch.stack([mask] * 3, dim=2)
-
                 overlayed_frame = torch.add(frame * 0.7, mask_colored * 0.3)
                 
                 processed_frames.append(overlayed_frame)
@@ -264,12 +261,44 @@ class Sam2RealtimeSegmentation:
         
         return (stacked_frames, stacked_masks)
 
+class BoundingBoxToCenter:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "bbox_data": ("JSON",),
+            }
+        }
+
+    RETURN_TYPES = ("STRING",)
+    RETURN_NAMES = ("center_coordinates",)
+    FUNCTION = "convert_bbox_to_center"
+    CATEGORY = "SAM2-Realtime"
+
+    def convert_bbox_to_center(self, bbox_data):
+        try:
+            bbox_list = ast.literal_eval(bbox_data)
+            tlx, tly, brx, bry = bbox_list[0][0]
+        
+            center_x = int((tlx + brx) / 2)
+            center_y = int((tly + bry) / 2)
+        
+            center_coords = f"[[{center_x}, {center_y}]]"
+            
+            return (center_coords,)
+            
+        except (ValueError, SyntaxError, IndexError) as e:
+            print(f"Error processing bounding box data: {e}")
+            return ("[[0, 0]]",)
+
 NODE_CLASS_MAPPINGS = {
     "DownloadAndLoadSAM2RealtimeModel": DownloadAndLoadSAM2RealtimeModel,
-    "Sam2RealtimeSegmentation": Sam2RealtimeSegmentation
+    "Sam2RealtimeSegmentation": Sam2RealtimeSegmentation,
+    "BoundingBoxToCenter": BoundingBoxToCenter
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
     "DownloadAndLoadSAM2RealtimeModel": "(Down)Load sam2_realtime Model",
-    "Sam2RealtimeSegmentation": "Sam2RealtimeSegmentation"
+    "Sam2RealtimeSegmentation": "Sam2RealtimeSegmentation",
+    "BoundingBoxToCenter": "BoundingBox To Center"
 }

requirements.txt

@@ -4,4 +4,4 @@ tqdm>=4.66.1
 hydra-core>=1.3.2
 iopath>=0.1.10
 pillow>=9.4.0
-git+https://github.com/pschroedl/ComfyUI-SAM2-Realtime.git@main#egg=sam2_realtime
+-e .

sam2_configs/sam2_hiera_small.yaml

@@ -1,116 +0,0 @@
-# @package _global_
-
-# Model
-model:
-  _target_: sam2_realtime.modeling.sam2_base.SAM2Base
-  image_encoder:
-    _target_: sam2_realtime.modeling.backbones.image_encoder.ImageEncoder
-    scalp: 1
-    trunk:
-      _target_: sam2_realtime.modeling.backbones.hieradet.Hiera
-      embed_dim: 96
-      num_heads: 1
-      stages: [1, 2, 11, 2]
-      global_att_blocks: [7, 10, 13]
-      window_pos_embed_bkg_spatial_size: [7, 7]
-    neck:
-      _target_: sam2_realtime.modeling.backbones.image_encoder.FpnNeck
-      position_encoding:
-        _target_: sam2_realtime.modeling.position_encoding.PositionEmbeddingSine
-        num_pos_feats: 256
-        normalize: true
-        scale: null
-        temperature: 10000
-      d_model: 256
-      backbone_channel_list: [768, 384, 192, 96]
-      fpn_top_down_levels: [2, 3]  # output level 0 and 1 directly use the backbone features
-      fpn_interp_model: nearest
-
-  memory_attention:
-    _target_: sam2_realtime.modeling.memory_attention.MemoryAttention
-    d_model: 256
-    pos_enc_at_input: true
-    layer:
-      _target_: sam2_realtime.modeling.memory_attention.MemoryAttentionLayer
-      activation: relu
-      dim_feedforward: 2048
-      dropout: 0.1
-      pos_enc_at_attn: false
-      self_attention:
-        _target_: sam2_realtime.modeling.sam.transformer.RoPEAttention
-        rope_theta: 10000.0
-        feat_sizes: [32, 32]
-        embedding_dim: 256
-        num_heads: 1
-        downsample_rate: 1
-        dropout: 0.1
-      d_model: 256
-      pos_enc_at_cross_attn_keys: true
-      pos_enc_at_cross_attn_queries: false
-      cross_attention:
-        _target_: sam2_realtime.modeling.sam.transformer.RoPEAttention
-        rope_theta: 10000.0
-        feat_sizes: [32, 32]
-        rope_k_repeat: True
-        embedding_dim: 256
-        num_heads: 1
-        downsample_rate: 1
-        dropout: 0.1
-        kv_in_dim: 64
-    num_layers: 4
-
-  memory_encoder:
-      _target_: sam2_realtime.modeling.memory_encoder.MemoryEncoder
-      out_dim: 64
-      position_encoding:
-        _target_: sam2_realtime.modeling.position_encoding.PositionEmbeddingSine
-        num_pos_feats: 64
-        normalize: true
-        scale: null
-        temperature: 10000
-      mask_downsampler:
-        _target_: sam2_realtime.modeling.memory_encoder.MaskDownSampler
-        kernel_size: 3
-        stride: 2
-        padding: 1
-      fuser:
-        _target_: sam2_realtime.modeling.memory_encoder.Fuser
-        layer:
-          _target_: sam2_realtime.modeling.memory_encoder.CXBlock
-          dim: 256
-          kernel_size: 7
-          padding: 3
-          layer_scale_init_value: 1e-6
-          use_dwconv: True  # depth-wise convs
-        num_layers: 2
-
-  num_maskmem: 7
-  image_size: 512
-  # apply scaled sigmoid on mask logits for memory encoder, and directly feed input mask as output mask
-  sigmoid_scale_for_mem_enc: 20.0
-  sigmoid_bias_for_mem_enc: -10.0
-  use_mask_input_as_output_without_sam: true
-  # Memory
-  directly_add_no_mem_embed: true
-  # use high-resolution feature map in the SAM mask decoder
-  use_high_res_features_in_sam: true
-  # output 3 masks on the first click on initial conditioning frames
-  multimask_output_in_sam: true
-  # SAM heads
-  iou_prediction_use_sigmoid: True
-  # cross-attend to object pointers from other frames (based on SAM output tokens) in the encoder
-  use_obj_ptrs_in_encoder: true
-  add_tpos_enc_to_obj_ptrs: false
-  only_obj_ptrs_in_the_past_for_eval: true
-  # object occlusion prediction
-  pred_obj_scores: true
-  pred_obj_scores_mlp: true
-  fixed_no_obj_ptr: true
-  # multimask tracking settings
-  multimask_output_for_tracking: true
-  use_multimask_token_for_obj_ptr: true
-  multimask_min_pt_num: 0
-  multimask_max_pt_num: 1
-  use_mlp_for_obj_ptr_proj: true
-  # Compilation flag
-  compile_image_encoder: False

sam2_configs/sam2_hiera_t.yaml

@@ -85,7 +85,7 @@ model:
         num_layers: 2
 
   num_maskmem: 7
-  image_size: 512
+  image_size: 1024
   # apply scaled sigmoid on mask logits for memory encoder, and directly feed input mask as output mask
   # SAM decoder
   sigmoid_scale_for_mem_enc: 20.0

sam2_realtime/sam2_tensor_predictor.py

@@ -14,6 +14,7 @@ from tqdm import tqdm
 from sam2_realtime.modeling.sam2_base import NO_OBJ_SCORE, SAM2Base
 from sam2_realtime.utils.misc import concat_points, fill_holes_in_mask_scores, load_video_frames
 
+
 class SAM2TensorPredictor(SAM2Base):
     """The predictor class to handle user interactions and manage inference states."""
 
@@ -54,9 +55,7 @@ class SAM2TensorPredictor(SAM2Base):
             img = img.float()
         else:
             raise ValueError("Input must be a numpy array or a PyTorch tensor")
-        #save original height/width
-        orig_h, orig_w = img.shape[1:]
-        
+
         # Resize to the target size (supports tensor resizing)
         img = torch.nn.functional.interpolate(
             img.unsqueeze(0), size=(image_size, image_size), mode="bilinear", align_corners=False
@@ -69,29 +68,28 @@ class SAM2TensorPredictor(SAM2Base):
         img /= img_std
 
         height, width = img.shape[1:]  # CHW format
-        return img, width, height, orig_w, orig_h
+        return img, width, height
 
     @torch.inference_mode()
     def load_first_frame(self, img):
         if isinstance(img, torch.Tensor):
             img = img.to(self.device)  # Ensure the tensor is on the correct device
-            
+
         self.condition_state = self._init_state(
             offload_video_to_cpu=False, offload_state_to_cpu=False
         )
-        img, width, height, orig_w, orig_h = self.prepare_data(img, image_size=self.image_size)
-        self._orig_hw = (orig_w, orig_h)
+        img, width, height = self.prepare_data(img, image_size=self.image_size)
         self.condition_state["images"] = [img]
         self.condition_state["num_frames"] = len(self.condition_state["images"])
         self.condition_state["video_height"] = height
         self.condition_state["video_width"] = width
         self._get_image_feature(frame_idx=0, batch_size=1)
-    
+
     def add_conditioning_frame(self, img):
         if isinstance(img, torch.Tensor):
             img = img.to(self.device)  # Ensure the tensor is on the correct device
 
-        img, width, height, _, _ = self.prepare_data(img, image_size=self.image_size)
+        img, width, height = self.prepare_data(img, image_size=self.image_size)
         self.condition_state["images"].append(img)
         self.condition_state["num_frames"] = len(self.condition_state["images"])
         self._get_image_feature(
@@ -237,15 +235,14 @@ class SAM2TensorPredictor(SAM2Base):
                     points = torch.cat([box_coords, points], dim=1)
                     labels = torch.cat([box_labels, labels], dim=1)
         if normalize_coords:
-            #video_H = self.condition_state["video_height"]
-            #video_W = self.condition_state["video_width"]
-            orig_w, orig_h = self._orig_hw
-            
-            points = points / torch.tensor([orig_w, orig_h]).to(points.device)
+            video_H = self.condition_state["video_height"]
+            video_W = self.condition_state["video_width"]
+            points = points / torch.tensor([video_W, video_H]).to(points.device)
         # scale the (normalized) coordinates by the model's internal image size
         points = points * self.image_size
         points = points.to(self.condition_state["device"])
         labels = labels.to(self.condition_state["device"])
+
         if not clear_old_points:
             point_inputs = point_inputs_per_frame.get(frame_idx, None)
         else:
@@ -345,16 +342,14 @@ class SAM2TensorPredictor(SAM2Base):
         if labels.dim() == 1:
             labels = labels.unsqueeze(0)  # add batch dimension
         if normalize_coords:
-            #video_H = self.condition_state["video_height"]
-            #video_W = self.condition_state["video_width"]
-            orig_w, orig_h = self._orig_hw
-            
-            points = points / torch.tensor([orig_w, orig_h]).to(points.device)
-            
+            video_H = self.condition_state["video_height"]
+            video_W = self.condition_state["video_width"]
+            points = points / torch.tensor([video_W, video_H]).to(points.device)
         # scale the (normalized) coordinates by the model's internal image size
         points = points * self.image_size
         points = points.to(self.condition_state["device"])
         labels = labels.to(self.condition_state["device"])
+
         if not clear_old_points:
             point_inputs = point_inputs_per_frame.get(frame_idx, None)
         else:
@@ -774,7 +769,7 @@ class SAM2TensorPredictor(SAM2Base):
         if isinstance(img, torch.Tensor):
             img = img.to(self.device)  # Ensure the tensor is on the correct device
 
-        img, _, _ , _, _ = self.prepare_data(img, image_size=self.image_size)
+        img, _, _ = self.prepare_data(img, image_size=self.image_size)
 
         output_dict = self.condition_state["output_dict"]
         obj_ids = self.condition_state["obj_ids"]

setup.py

@@ -4,10 +4,6 @@
 # This source code is licensed under the license found in the
 # LICENSE file in the root directory of this source tree.
 
-import os
-# Set the CUDA architecture list
-os.environ["TORCH_CUDA_ARCH_LIST"] = "8.0 8.6+PTX 8.7 9.0 9.0a"
-
 from setuptools import find_packages, setup
 from torch.utils.cpp_extension import BuildExtension, CUDAExtension
 
@@ -57,7 +53,7 @@ setup(
     license=LICENSE,
     packages=find_packages(),
     install_requires=REQUIRED_PACKAGES,
-    python_requires=">=3.10.15",
+    python_requires=">=3.11.10",
     ext_modules=get_extensions(),
     cmdclass={"build_ext": BuildExtension.with_options(no_python_abi_suffix=True)},
 )