import folder_paths import os import numpy as np import torch import re import json from segment_anything import sam_model_registry, SamPredictor from einops import rearrange, repeat from PIL import Image import mediapipe as mp from math import sqrt BaseOptions = mp.tasks.BaseOptions FaceLandmarker = mp.tasks.vision.FaceLandmarker FaceLandmarkerOptions = mp.tasks.vision.FaceLandmarkerOptions PoseLandmarker = mp.tasks.vision.PoseLandmarker PoseLandmarkerOptions = mp.tasks.vision.PoseLandmarkerOptions VisionRunningMode = mp.tasks.vision.RunningMode global_predictor = None face_landmarker = None pose_landmarker = None # For auto-segmentation layerMapping = { "L_eye": { "useMiddle": False, "positiveOffsetX": 0, "positiveOffsetY": 0, "negativeOffsetX": 0, "negativeOffsetY": 0, "positiveScale": 0, "negativeScale": 0.5, "indices": mp.solutions.face_mesh.FACEMESH_LEFT_EYE, }, "R_eye": { "useMiddle": False, "positiveOffsetX": 0, "positiveOffsetY": 0, "negativeOffsetX": 0, "negativeOffsetY": 0, "positiveScale": 0, "negativeScale": 0.5, "indices": mp.solutions.face_mesh.FACEMESH_RIGHT_EYE, }, "L_iris": { "useMiddle": False, "positiveOffsetX": 0, "positiveOffsetY": 0, "negativeOffsetX": 0, "negativeOffsetY": 0, "positiveScale": -0.2, "negativeScale": 0.5, "indices": mp.solutions.face_mesh.FACEMESH_LEFT_IRIS, }, "R_iris": { "useMiddle": False, "positiveOffsetX": 0, "positiveOffsetY": 0, "negativeOffsetX": 0, "negativeOffsetY": 0, "positiveScale": -0.2, "negativeScale": 0.5, "indices": mp.solutions.face_mesh.FACEMESH_RIGHT_IRIS, }, "face": { "useMiddle": False, "positiveOffsetX": 0, "positiveOffsetY": 40, "negativeOffsetX": 0, "negativeOffsetY": 60, "positiveScale": 0.2, "negativeScale": 0.6, "indices": mp.solutions.face_mesh.FACEMESH_FACE_OVAL, }, "mouth": { "useMiddle": False, "positiveOffsetX": 0, "positiveOffsetY": 0, "negativeOffsetX": 0, "negativeOffsetY": 0, "positiveScale": -0.3, "negativeScale": 0.3, # https://stackoverflow.com/questions/66649492/how-to-get-specific-landmark-of-face-like-lips-or-eyes-using-tensorflow-js-face "indices": [[x, x] for x in [61, 37, 270, 91, 314]], }, "mouth_in": { "useMiddle": False, "positiveOffsetX": 0, "positiveOffsetY": 0, "negativeOffsetX": 0, "negativeOffsetY": 0, "positiveScale": -0.5, "negativeScale": 0.5, # https://stackoverflow.com/questions/66649492/how-to-get-specific-landmark-of-face-like-lips-or-eyes-using-tensorflow-js-face "indices": [[x, x] for x in [310, 88]], }, } class SAMMultiLayer: def __init__(self): self.predictor = None self.output_dir = folder_paths.get_output_directory() @classmethod def INPUT_TYPES(s): return { "required": { "image": ("IMAGE",), "ckpt": (folder_paths.get_filename_list("sams"),), "embedding_id": ( "STRING", {"multiline": False, "default": "embedding"}, ), "image_prompts_json": ("STRING", {"multiline": False, "default": "[]"}), }, } CATEGORY = "image" RETURN_TYPES = ("SAM_PROMPT",) FUNCTION = "load_image" def load_models(self, ckpt, model_type): global global_predictor, face_landmarker, pose_landmarker ckpt = folder_paths.get_full_path("sams", ckpt) sam = sam_model_registry[model_type](checkpoint=ckpt) # .to("cuda") global_predictor = SamPredictor(sam) face_landmarker_model_path = os.path.join( os.path.dirname(__file__), "../mediapipe_models/face_landmarker.task" ) face_landmarker_options = FaceLandmarkerOptions( base_options=BaseOptions(model_asset_path=face_landmarker_model_path), running_mode=VisionRunningMode.IMAGE, ) face_landmarker = FaceLandmarker.create_from_options(face_landmarker_options) pose_landmarker_model_path = os.path.join( os.path.dirname(__file__), "../mediapipe_models/pose_landmarker_full.task" ) pose_landmarker_options = PoseLandmarkerOptions( base_options=BaseOptions(model_asset_path=pose_landmarker_model_path), running_mode=VisionRunningMode.IMAGE, ) pose_landmarker = PoseLandmarker.create_from_options(pose_landmarker_options) return global_predictor, face_landmarker, pose_landmarker def auto_segment(self, image, face_landmarks, pose_landmarks): H, W, C = image.shape imagePromptsMulti = {} boxesMulti = {} for key, value in layerMapping.items(): positivePoints = [] middlePoints = [] negativePoints = [] for index in value["indices"]: start, end = index startPoint = face_landmarks[start] startX = startPoint.x * W startY = startPoint.y * H if len(middlePoints) == 0: middlePoints.append({"x": startX, "y": startY, "label": 1}) else: middlePoints[0]["x"] += startX middlePoints[0]["y"] += startY positivePoints.append({"x": startX, "y": startY, "label": 1}) len_indices = len(value["indices"]) middlePoints[0]["x"] /= len_indices middlePoints[0]["y"] /= len_indices if value["useMiddle"]: imagePromptsMulti[key] = middlePoints else: for i, index in enumerate(value["indices"]): start, end = index startPoint = face_landmarks[start] startX = startPoint.x * W startY = startPoint.y * H middlePoint = middlePoints[0] directionVector = { "x": middlePoint["x"] - startX, "y": middlePoint["y"] - startY, } directionVectorLength = sqrt( directionVector["x"] * directionVector["x"] + directionVector["y"] * directionVector["y"] ) if value["negativeScale"] != 0: negativePointDistance = ( value["negativeScale"] * directionVectorLength ) negativePoint = { "x": startX - (negativePointDistance * directionVector["x"]) / directionVectorLength - value["negativeOffsetX"], "y": startY - (negativePointDistance * directionVector["y"]) / directionVectorLength - value["negativeOffsetY"], "label": 0, } negativePoints.append(negativePoint) positivePointDistance = ( value["positiveScale"] * directionVectorLength ) positivePoints[i] = { "x": positivePoints[i]["x"] - (positivePointDistance * directionVector["x"]) / directionVectorLength - value["positiveOffsetX"], "y": positivePoints[i]["y"] - (positivePointDistance * directionVector["y"]) / directionVectorLength - value["positiveOffsetY"], "label": 1, } imagePromptsMulti[key] = positivePoints + negativePoints points = negativePoints if len(negativePoints) > 0 else positivePoints box = np.array( [ min(x["x"] for x in points), min(x["y"] for x in points), max(x["x"] for x in points), max(x["y"] for x in points), ] ) boxesMulti[key] = box if pose_landmarks is not None: positiveBreathX = ( (pose_landmarks[11].x + pose_landmarks[12].x) / 2 ) * W positiveBreathY = ( (pose_landmarks[11].y + pose_landmarks[12].y) / 2 ) * H negativeBreathX1 = pose_landmarks[0].x * W negativeBreathY1 = pose_landmarks[0].y * H negativeBreathX2 = pose_landmarks[9].x * W negativeBreathY2 = pose_landmarks[9].y * H negativeBreathX3 = pose_landmarks[10].x * W negativeBreathY3 = pose_landmarks[10].y * H imagePromptsMulti["breath"] = [ {"x": positiveBreathX, "y": positiveBreathY, "label": 1}, {"x": negativeBreathX1, "y": negativeBreathY1, "label": 0}, {"x": negativeBreathX2, "y": negativeBreathY2, "label": 0}, {"x": negativeBreathX3, "y": negativeBreathY3, "label": 0}, ] return imagePromptsMulti, boxesMulti def detect_face(self, np_image): global face_landmarker, pose_landmarker mp_image = mp.Image( image_format=mp.ImageFormat.SRGB, data=(np_image * 255).astype(np.uint8) ) face_landmarks = face_landmarker.detect(mp_image).face_landmarks face_landmarks = face_landmarks[0] if len(face_landmarks) > 0 else None pose_landmarks = pose_landmarker.detect(mp_image).pose_landmarks pose_landmarks = pose_landmarks[0] if len(pose_landmarks) > 0 else None imagePromptsMulti, boxesMulti = self.auto_segment( np_image, face_landmarks, pose_landmarks ) return imagePromptsMulti, boxesMulti def load_image(self, image, ckpt, embedding_id, image_prompts_json): image_prompts = json.loads(image_prompts_json.replace("'", '"')) order_file = f"{self.output_dir}/segments_{embedding_id}/order.json" if os.path.exists(order_file): # Frontend uploads segments images to backend => backend reads all segments images and passes them to next nodes with open(order_file) as f: order = json.load(f) result = [image_prompts] for segment in order: image = Image.open( f"{self.output_dir}/segments_{embedding_id}/{segment}.png" ) image = np.array(image).astype(np.float32) / 255.0 image = torch.from_numpy(image)[None,] result.append(image) return result else: # Frontend uploads clicks coordinates to backend => backend runs SAM and passes the segments to next nodes model_type = re.findall(r"vit_[lbh]", ckpt)[0] global global_predictor if global_predictor is None: global_predictor, _, _ = self.load_models(ckpt, model_type) if image.shape[3] == 4: image = image[:, :, :, :3] emb_filename = f"{self.output_dir}/{embedding_id}_{model_type}.npy" if not os.path.exists(emb_filename): image_np = (image[0].numpy() * 255).astype(np.uint8) global_predictor.set_image(image_np) emb = global_predictor.get_image_embedding().cpu().numpy() np.save(emb_filename, emb) with open( f"{self.output_dir}/{embedding_id}_{model_type}.json", "w" ) as f: data = { "input_size": global_predictor.input_size, "original_size": global_predictor.original_size, } json.dump(data, f) else: emb = np.load(emb_filename) with open(f"{self.output_dir}/{embedding_id}_{model_type}.json") as f: data = json.load(f) global_predictor.input_size = data["input_size"] global_predictor.features = torch.from_numpy(emb) global_predictor.is_image_set = True global_predictor.original_size = data["original_size"] imagePromptsMulti, boxesMulti = self.detect_face(image[0].numpy()) image_prompts = json.loads(image_prompts_json.replace("'", '"')) result = [image_prompts] if isinstance(image_prompts, list): pass elif all(isinstance(item, list) for item in image_prompts.values()): for key, item in image_prompts.items(): if len(item) == 0: h, w, c = image[0].shape result.append(torch.zeros(1, h, w, c)) continue points = ( imagePromptsMulti[key] if key in imagePromptsMulti else item ) point_coords = np.array([[p["x"], p["y"]] for p in points]) point_labels = np.array([p["label"] for p in points]) masks, _, _ = global_predictor.predict( point_coords=point_coords, point_labels=point_labels, box=boxesMulti[key] if key in boxesMulti else None, ) masks = torch.from_numpy(masks) masks = rearrange(masks[0], "h w -> 1 h w") out_image = repeat(masks, "1 h w -> 1 h w c", c=3) * image result.append(out_image) return result NODE_CLASS_MAPPINGS = {"SAM MultiLayer": SAMMultiLayer} NODE_DISPLAY_NAME_MAPPINGS = {"SAM MultiLayer": "SAM MultiLayer"}