diff --git a/nodes.py b/nodes.py index 6e2b9f0..57b6cc2 100644 --- a/nodes.py +++ b/nodes.py @@ -144,7 +144,6 @@ class Sam2RealtimeSegmentation: }, "optional": { "coordinates_positive": ("STRING", ), - # "point_labels": ("STRING",), "coordinates_negative": ("STRING", ), # "bboxes": ("BBOX", ), # "individual_objects": ("BOOLEAN", {"default": False}), @@ -152,7 +151,7 @@ class Sam2RealtimeSegmentation: }, } - RETURN_NAMES = ("PROCESSED_IMAGES","MASK",) + RETURN_NAMES = ("PROCESSED_IMAGES", "MASK",) RETURN_TYPES = ("IMAGE", "MASK",) FUNCTION = "segment_images" CATEGORY = "SAM2-Realtime" @@ -161,6 +160,32 @@ class Sam2RealtimeSegmentation: self.predictor = None self.if_init = False + def _process_coordinate_input(self, coordinates, label): + """Helper function to process coordinate inputs safely""" + if not coordinates: + return [], [] + try: + coord_list = ast.literal_eval(coordinates) + points = [tuple(map(int, point)) for point in coord_list] + labels = [label] * len(points) + return points, labels + except (ValueError, SyntaxError) as e: + print(f"Error processing coordinates: {e}") + return [], [] + + def _process_mask_logits(self, out_mask_logits, frame_shape, device): + """Helper function to process mask logits""" + if out_mask_logits.shape[0] > 0: + mask = (out_mask_logits[0, 0] > 0.5).byte() + mask = torch.nn.functional.interpolate( + mask.unsqueeze(0).unsqueeze(0).float(), + size=frame_shape[:2], + mode='nearest' + ).squeeze().byte().to(device) + else: + mask = torch.ones(frame_shape[:2], device=device, dtype=torch.uint8) + return mask + def segment_images( self, images, @@ -175,8 +200,6 @@ class Sam2RealtimeSegmentation: # mask=None, ): model = sam2_model["model"] - #device = sam2_model["device"] - device = torch.device("cuda") model.to(device) @@ -184,7 +207,6 @@ class Sam2RealtimeSegmentation: mask_list = [] # The `model` variable is now ready and equivalent to `predictor` returned by sam2.build_sam.build_sam2_camera_predictor - if reset_tracking: self.if_init = False self.predictor = None @@ -192,48 +214,25 @@ class Sam2RealtimeSegmentation: if self.predictor is None: self.predictor = model - def process_frame(frame, frame_idx): - def process_coordinate_input(coordinates, label): - if not coordinates: - return [], [] - coord_list = ast.literal_eval(coordinates) - points = [tuple(map(int, point)) for point in coord_list] - labels = [label] * len(points) - return points, labels + # Process coordinates once, outside the frame loop + pos_points, pos_labels = self._process_coordinate_input(coordinates_positive, 1) + neg_points, neg_labels = self._process_coordinate_input(coordinates_negative, 0) + all_points = pos_points + neg_points + all_labels = pos_labels + neg_labels - def process_mask_logits(out_mask_logits, frame_shape): - if out_mask_logits.shape[0] > 0: - out_mask_logits = out_mask_logits.to(device) - mask = (out_mask_logits[0, 0] > 0.5).byte() - mask = torch.nn.functional.interpolate( - mask.unsqueeze(0).unsqueeze(0).float(), - size=(frame_shape[0], frame_shape[1]), - mode='nearest' - ).squeeze(0).squeeze(0).byte().to(device) - else: - mask = torch.ones((frame_shape[0], frame_shape[1]), device=device, dtype=torch.uint8) - return mask + if all_points: + points_tensor = torch.tensor([all_points], device=device) + labels_tensor = torch.tensor([all_labels], device=device) - with torch.inference_mode(), torch.autocast("cuda", dtype=torch.float16): + with torch.inference_mode(), torch.autocast("cuda", dtype=torch.float16): + for frame_idx, frame in enumerate(images): frame = frame.to(device).float() if not self.if_init: self.predictor.load_first_frame(frame) self.if_init = True - # Process both coordinate lists - pos_points, pos_labels = process_coordinate_input(coordinates_positive, 1) - neg_points, neg_labels = process_coordinate_input(coordinates_negative, 0) - - # Combine points and labels - all_points = pos_points + neg_points - all_labels = pos_labels + neg_labels - if all_points: - points_tensor = torch.tensor([all_points], device=device) - labels_tensor = torch.tensor([all_labels], device=device) - - #NOTE: single object tracking only for now. TODO switch coordinate input to dict to track multiple objects _, _, out_mask_logits = self.predictor.add_new_prompt( frame_idx=0, obj_id=1, @@ -244,23 +243,21 @@ class Sam2RealtimeSegmentation: out_mask_logits = torch.zeros((0,), device=device) else: out_obj_ids, out_mask_logits = self.predictor.track(frame) + # Process mask logits - mask = process_mask_logits(out_mask_logits, frame.shape) + mask = self._process_mask_logits(out_mask_logits, frame.shape, device) # Create colored overlay for processed frames - mask_colored = torch.stack([mask] * 3, dim=2).to(device) - overlayed_frame = torch.add(frame * 0.7, mask_colored * 0.3).to(device) + mask_colored = torch.stack([mask] * 3, dim=2) + overlayed_frame = torch.add(frame * 0.7, mask_colored * 0.3) + processed_frames.append(overlayed_frame) - - # Store the raw binary mask for mask output mask_list.append(mask) - for frame_idx, img in enumerate(images): - process_frame(img, frame_idx) - # Stack masks and frames stacked_masks = torch.stack(mask_list, dim=0) stacked_frames = torch.stack(processed_frames, dim=0) + return (stacked_frames, stacked_masks) NODE_CLASS_MAPPINGS = {