chore: remove unused nodes
This commit is contained in:
Radionic
+1
-1
@@ -72,7 +72,7 @@ def download_sam_model():
|
||||
|
||||
download_sam_model()
|
||||
|
||||
paths = ["blender", "sam", "common"]
|
||||
paths = ["blender", "sam"]
|
||||
files = []
|
||||
|
||||
for path in paths:
|
||||
|
||||
@@ -1,35 +0,0 @@
|
||||
class CombineMesh:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"obj_1": ("BPY_OBJS",),
|
||||
"obj_2": ("BPY_OBJS",),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("BPY_OBJS",)
|
||||
RETURN_NAMES = ("bpy_objs",)
|
||||
|
||||
FUNCTION = "process"
|
||||
|
||||
CATEGORY = "mesh"
|
||||
|
||||
def process(self, obj_1, obj_2):
|
||||
import global_bpy
|
||||
bpy = global_bpy.get_bpy()
|
||||
|
||||
override = bpy.context.copy()
|
||||
override["active_object"] = obj_1[0]
|
||||
override["selected_editable_objects"] = list([*obj_1, *obj_2])
|
||||
bpy.ops.object.join(override)
|
||||
joined_object = override["active_object"]
|
||||
return ([joined_object], )
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"CombineMesh": CombineMesh
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"CombineMesh": "Combine Mesh"
|
||||
}
|
||||
@@ -1,63 +0,0 @@
|
||||
class CreateVertexBounds():
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"bpy_objs_target": ("BPY_OBJS",),
|
||||
"name": ("STRING", {
|
||||
"multiline": False,
|
||||
"default": "Group"
|
||||
}),
|
||||
"extrude": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01, "display": "number"}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("BPY_OBJS",)
|
||||
RETURN_NAMES = ("bpy_objs",)
|
||||
|
||||
FUNCTION = "process"
|
||||
|
||||
CATEGORY = "mesh"
|
||||
|
||||
def process(self, bpy_objs_target, name, extrude):
|
||||
import global_bpy
|
||||
bpy = global_bpy.get_bpy()
|
||||
|
||||
if len(bpy_objs_target) == 0:
|
||||
# throw error
|
||||
return
|
||||
|
||||
target_object = bpy_objs_target[0]
|
||||
|
||||
# deselect all objects
|
||||
bpy.ops.object.select_all(action='DESELECT')
|
||||
# select only the target object
|
||||
bpy.context.view_layer.objects.active = target_object
|
||||
# enter enter edit mode and select all faces of the object to fill
|
||||
bpy.ops.object.mode_set(mode='EDIT')
|
||||
bpy.ops.mesh.select_all(action='SELECT')
|
||||
|
||||
bpy.ops.mesh.extrude_region_move(MESH_OT_extrude_region={
|
||||
"use_normal_flip": False, "use_dissolve_ortho_edges": False, "mirror": False})
|
||||
|
||||
bpy.ops.transform.resize(
|
||||
value=(extrude, extrude, 1), constraint_axis=(False, False, False))
|
||||
|
||||
# remove the extruded vertex from the vertex group "name"
|
||||
|
||||
bpy.ops.object.vertex_group_remove_from()
|
||||
|
||||
# bpy.ops.mesh.delete(type='FACE')
|
||||
bpy.ops.object.mode_set(mode='OBJECT')
|
||||
|
||||
return ([target_object],)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"CreateVertexBounds": CreateVertexBounds
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"CreateVertexBounds": "Vertex Bounds"
|
||||
}
|
||||
@@ -1,56 +0,0 @@
|
||||
class CreateVertexGroup():
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"bpy_objs_target": ("BPY_OBJS",),
|
||||
"name": ("STRING", {
|
||||
"multiline": False,
|
||||
"default": "Group"
|
||||
}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("BPY_OBJS",)
|
||||
RETURN_NAMES = ("bpy_objs",)
|
||||
|
||||
FUNCTION = "process"
|
||||
|
||||
CATEGORY = "mesh"
|
||||
|
||||
def process(self, bpy_objs_target, name):
|
||||
import global_bpy
|
||||
bpy = global_bpy.get_bpy()
|
||||
|
||||
if len(bpy_objs_target) == 0:
|
||||
# throw error
|
||||
return
|
||||
|
||||
target_object = bpy_objs_target[0]
|
||||
|
||||
# deselect all objects
|
||||
bpy.ops.object.select_all(action='DESELECT')
|
||||
# select only the target object
|
||||
bpy.context.view_layer.objects.active = target_object
|
||||
# enter enter edit mode and select all faces of the object to fill
|
||||
bpy.ops.object.mode_set(mode='EDIT')
|
||||
bpy.ops.mesh.select_all(action='SELECT')
|
||||
|
||||
# create vertex group and assign all selected faces to it (the name of the vertex group is the same as the name of the object)
|
||||
bpy.context.object.vertex_groups.new(name=name)
|
||||
bpy.ops.object.vertex_group_assign()
|
||||
|
||||
# bpy.ops.mesh.delete(type='FACE')
|
||||
bpy.ops.object.mode_set(mode='OBJECT')
|
||||
|
||||
return ([target_object],)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"CreateVertexGroup": CreateVertexGroup
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"CreateVertexGroup": "Vertex Group (Old)"
|
||||
}
|
||||
@@ -1,61 +0,0 @@
|
||||
class DissolveFaces():
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"bpy_objs_target": ("BPY_OBJS",),
|
||||
"target_vertex_group": ("STRING", {
|
||||
"multiline": False,
|
||||
"default": ""
|
||||
}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("BPY_OBJS",)
|
||||
RETURN_NAMES = ("bpy_objs",)
|
||||
|
||||
FUNCTION = "process"
|
||||
|
||||
CATEGORY = "mesh"
|
||||
|
||||
def process(self, bpy_objs_target, target_vertex_group):
|
||||
import global_bpy
|
||||
bpy = global_bpy.get_bpy()
|
||||
|
||||
target_object = bpy_objs_target[0]
|
||||
|
||||
# deselect all objects
|
||||
bpy.ops.object.select_all(action='DESELECT')
|
||||
# select only the target object
|
||||
bpy.context.view_layer.objects.active = target_object
|
||||
# enter enter edit mode and select all faces of the object to fill
|
||||
bpy.ops.object.mode_set(mode='EDIT')
|
||||
|
||||
bpy.ops.mesh.select_all(action='DESELECT')
|
||||
|
||||
# if we have vertex group, select that instead
|
||||
if target_vertex_group and target_vertex_group.strip() != "":
|
||||
group = target_object.vertex_groups.get(target_vertex_group)
|
||||
if group:
|
||||
bpy.ops.object.vertex_group_set_active(group=group.name)
|
||||
bpy.ops.object.vertex_group_select()
|
||||
else:
|
||||
bpy.ops.mesh.select_all(action='SELECT')
|
||||
|
||||
# dissolve the faces
|
||||
# bpy.ops.mesh.dissolve_faces()
|
||||
bpy.ops.mesh.delete(type='ONLY_FACE')
|
||||
# bpy.ops.mesh.delete(type='FACE')
|
||||
bpy.ops.object.mode_set(mode='OBJECT')
|
||||
|
||||
return ([target_object],)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"DissolveFaces": DissolveFaces
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"DissolveFaces": "Dissolve Faces (Old)"
|
||||
}
|
||||
@@ -1,46 +0,0 @@
|
||||
class FillFace:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"bpy_objs_target": ("BPY_OBJS",),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("BPY_OBJS",)
|
||||
RETURN_NAMES = ("bpy_objs",)
|
||||
|
||||
FUNCTION = "process"
|
||||
|
||||
CATEGORY = "mesh"
|
||||
|
||||
def process(self, bpy_objs_target):
|
||||
import global_bpy
|
||||
bpy = global_bpy.get_bpy()
|
||||
|
||||
target_object = bpy_objs_target[0]
|
||||
|
||||
# deselect all objects
|
||||
bpy.ops.object.select_all(action='DESELECT')
|
||||
# select only the target object
|
||||
bpy.context.view_layer.objects.active = target_object
|
||||
# enter enter edit mode and select all faces of the object to fill
|
||||
bpy.ops.object.mode_set(mode='EDIT')
|
||||
bpy.ops.mesh.select_all(action='SELECT')
|
||||
|
||||
# fill the faces
|
||||
bpy.ops.mesh.fill()
|
||||
|
||||
bpy.ops.object.mode_set(mode='OBJECT')
|
||||
|
||||
return ([target_object],)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"FillFace": FillFace
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"FillFace": "Fill Face (Old)"
|
||||
}
|
||||
@@ -1,30 +0,0 @@
|
||||
class GroupObject():
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"bpy_objs_target": ("BPY_OBJS",),
|
||||
"bpy_objs_target_2": ("BPY_OBJS",),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("BPY_OBJS",)
|
||||
RETURN_NAMES = ("bpy_objs",)
|
||||
|
||||
FUNCTION = "process"
|
||||
|
||||
CATEGORY = "mesh"
|
||||
|
||||
def process(self, bpy_objs_target, bpy_objs_target_2):
|
||||
combined_list = bpy_objs_target + bpy_objs_target_2
|
||||
return (combined_list,)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"GroupObject": GroupObject
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"GroupObject": "Group Object"
|
||||
}
|
||||
@@ -1,108 +0,0 @@
|
||||
class GroupVertexInside():
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"bpy": ("BPY",),
|
||||
"bpy_objs_target": ("BPY_OBJS",),
|
||||
"name": ("STRING", {
|
||||
"multiline": False,
|
||||
"default": "Group"
|
||||
}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("BPY_OBJS",)
|
||||
RETURN_NAMES = ("bpy_objs",)
|
||||
|
||||
FUNCTION = "process"
|
||||
|
||||
CATEGORY = "mesh"
|
||||
|
||||
def process(self, bpy_objs_target, name):
|
||||
import global_bpy
|
||||
bpy = global_bpy.get_bpy()
|
||||
|
||||
target_object = bpy_objs_target[0]
|
||||
|
||||
bpy.ops.object.mode_set(mode='OBJECT')
|
||||
|
||||
# deselect all objects
|
||||
bpy.ops.object.select_all(action='DESELECT')
|
||||
# select only the target object
|
||||
bpy.context.view_layer.objects.active = target_object
|
||||
# enter enter edit mode and select all faces of the object to fill
|
||||
bpy.ops.object.mode_set(mode='EDIT')
|
||||
bpy.ops.mesh.select_all(action='DESELECT')
|
||||
|
||||
# select the vetex group's vertices
|
||||
bpy.ops.object.vertex_group_select()
|
||||
|
||||
# get vertex group
|
||||
vertex_group = target_object.vertex_groups[name]
|
||||
|
||||
verts = target_object.data.vertices
|
||||
|
||||
# all_vertices = [(v.co.x, v.co.y, i) for i, v in target_object.data.vertices]
|
||||
|
||||
# get the verts in the group
|
||||
verts_in_group = [v for v in verts if vertex_group.index in [vg.group for vg in v.groups]]
|
||||
|
||||
# get all selected vertices's xy
|
||||
# get the vertices of the vertex group
|
||||
vertex_group_vertices = [(v.co.x, v.co.y) for v in verts_in_group]
|
||||
|
||||
# get all vertices
|
||||
# all_vertices = [(v.co.x, v.co.y) for v in target_object.data.vertices]
|
||||
all_vertices_without_vertex_group = [(v.co.x, v.co.y, i) for i, v in enumerate(target_object.data.vertices) if v not in verts_in_group]
|
||||
|
||||
from scipy.spatial import ConvexHull, Delaunay
|
||||
|
||||
# create a convex hull of the vertex group's vertices
|
||||
hull = ConvexHull(vertex_group_vertices)
|
||||
|
||||
# create a Delaunay triangulation of the convex hull
|
||||
tri = Delaunay(hull.points[hull.vertices])
|
||||
|
||||
# get the indices of the vertices that are inside the convex hull
|
||||
inside_vertices_indices = [i for (x,y,i) in all_vertices_without_vertex_group if tri.find_simplex((x,y))>=0]
|
||||
|
||||
# print(inside_vertices_indices)
|
||||
# print(len(vertex_group_vertices))
|
||||
|
||||
bpy.ops.mesh.select_all(action='DESELECT')
|
||||
|
||||
# select the vertices that are inside the convex hull
|
||||
# for i in inside_vertices_indices:
|
||||
# target_object.data.vertices[i].select = True
|
||||
|
||||
# # select vertex from 1 to 10
|
||||
# for i in range(1, 100):
|
||||
# target_object.data.vertices[i].select = True
|
||||
|
||||
# create a new vertex group
|
||||
# bpy.context.object.vertex_groups.new(name="fuck")
|
||||
# bpy.ops.object.vertex_group_assign()
|
||||
bpy.ops.object.mode_set(mode='OBJECT')
|
||||
# create a new vertex group
|
||||
# new_group = bpy.context.object.vertex_groups.new(name="new_group")
|
||||
vertex_group.add(inside_vertices_indices, weight=1.0, type='REPLACE')
|
||||
|
||||
# assign the selected vertices to the new vertex group
|
||||
# new_group.add([v.index for v in target_object.data.vertices if v.select], weight=1.0, type='REPLACE')
|
||||
# new_group.add([20, 100], weight=1.0, type='REPLACE')
|
||||
# assign the selected vertices to the active vertex group
|
||||
bpy.ops.object.mode_set(mode='OBJECT')
|
||||
|
||||
|
||||
return ([target_object],)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"GroupVertexInside": GroupVertexInside
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"GroupVertexInside": "Group Vertex Inside"
|
||||
}
|
||||
@@ -1,121 +0,0 @@
|
||||
class KnifeProjection:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"bpy_objs_knife": ("BPY_OBJS",),
|
||||
"bpy_objs_target": ("BPY_OBJS",),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("BPY_OBJS",)
|
||||
RETURN_NAMES = ("bpy_objs",)
|
||||
|
||||
FUNCTION = "process"
|
||||
|
||||
CATEGORY = "mesh"
|
||||
|
||||
def process(self, bpy_objs_knife, bpy_objs_target):
|
||||
import mathutils
|
||||
import math
|
||||
import global_bpy
|
||||
bpy = global_bpy.get_bpy()
|
||||
|
||||
knife_object = bpy_objs_knife[0]
|
||||
target_object = bpy_objs_target[0]
|
||||
|
||||
bpy.ops.object.select_all(action='DESELECT')
|
||||
bpy.context.view_layer.objects.active = target_object
|
||||
# knife_object.select_set(True)
|
||||
# target_object.select_set(True)
|
||||
|
||||
# bpy.ops.object.mode_set(mode='EDIT')
|
||||
|
||||
area = next(
|
||||
area for area in bpy.context.screen.areas if area.type == 'VIEW_3D')
|
||||
space = next(space for space in area.spaces if space.type == 'VIEW_3D')
|
||||
region = next(
|
||||
region for region in area.regions if region.type == 'WINDOW')
|
||||
|
||||
# space.region_3d.view_perspective = 'ORTHO'
|
||||
# space.region_3d.view_rotation = mathutils.Euler(
|
||||
# (0, 0, 0)).to_quaternion()
|
||||
|
||||
# set blender to orthographic view without using numpad
|
||||
|
||||
# bpy.ops.view3d.view_axis(type='TOP')
|
||||
|
||||
override = bpy.context.copy()
|
||||
|
||||
# override["active_object"] = target_object
|
||||
# override["edit_object"] = target_object
|
||||
# override["selected_objects"] = [knife_object, target_object]
|
||||
# override["selected_editable_objects"] = [knife_object, target_object]
|
||||
|
||||
override["region"] = region
|
||||
override["area"] = area
|
||||
override["space"] = space
|
||||
|
||||
override["active_object"] = target_object
|
||||
override["edit_object"] = target_object
|
||||
override["selected_objects"] = [knife_object, target_object]
|
||||
override["selected_editable_objects"] = [knife_object, target_object]
|
||||
|
||||
knife_object.select_set(False)
|
||||
target_object.select_set(True)
|
||||
# move the kinfe object up a bit
|
||||
bpy.ops.transform.translate(value=(0, 0, -0.1))
|
||||
bpy.ops.object.mode_set(mode='EDIT')
|
||||
knife_object.select_set(False)
|
||||
|
||||
# Get the viewport rotation
|
||||
viewport_rotation = space.region_3d.view_rotation
|
||||
|
||||
with bpy.context.temp_override(**override):
|
||||
# # Create a rotation matrix representing a top-down view
|
||||
# rotation_matrix = mathutils.Matrix.Rotation(math.radians(90.0), 4, 'X')
|
||||
|
||||
# # Create a rotation matrix to make the object face the user
|
||||
# rotation_matrix_face_user = mathutils.Matrix.Rotation(math.radians(180.0), 4, 'Z')
|
||||
|
||||
# # Create a rotation matrix for the viewport rotation
|
||||
# rotation_matrix_viewport = viewport_rotation.to_matrix().to_4x4()
|
||||
|
||||
# # Remember the original locations
|
||||
# original_location_knife = knife_object.location.copy()
|
||||
# original_location_target = target_object.location.copy()
|
||||
|
||||
# # Move the objects to the origin
|
||||
# knife_object.location = mathutils.Vector((0, 0, 0))
|
||||
# target_object.location = mathutils.Vector((0, 0, 0))
|
||||
|
||||
# # Rotate the objects
|
||||
# knife_object.matrix_world = rotation_matrix @ rotation_matrix_face_user @ rotation_matrix_viewport @ knife_object.matrix_world
|
||||
# target_object.matrix_world = rotation_matrix @ rotation_matrix_face_user @ rotation_matrix_viewport @ target_object.matrix_world
|
||||
|
||||
bpy.ops.view3d.view_axis(type='TOP', align_active=False)
|
||||
bpy.ops.view3d.view_persportho()
|
||||
bpy.ops.wm.redraw_timer(type='DRAW_WIN_SWAP', iterations=1)
|
||||
bpy.ops.mesh.knife_project()
|
||||
|
||||
# # Rotate the objects back
|
||||
# knife_object.matrix_world = rotation_matrix.inverted() @ knife_object.matrix_world
|
||||
# target_object.matrix_world = rotation_matrix.inverted() @ target_object.matrix_world
|
||||
|
||||
# # Move the objects back to their original locations
|
||||
# knife_object.location = original_location_knife
|
||||
# target_object.location = original_location_target
|
||||
|
||||
bpy.ops.object.mode_set(mode='OBJECT')
|
||||
|
||||
return ([target_object],)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"KnifeProjection": KnifeProjection
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"KnifeProjection": "Knife Projection"
|
||||
}
|
||||
@@ -1,104 +0,0 @@
|
||||
class MeshFromTexture:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE",),
|
||||
# For disabling cache
|
||||
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "BPY_OBJS")
|
||||
RETURN_NAMES = ("image", "bpy_objs")
|
||||
|
||||
FUNCTION = "process"
|
||||
|
||||
CATEGORY = "mesh"
|
||||
|
||||
def process(self, image, seed):
|
||||
import torch
|
||||
import cv2
|
||||
import numpy as np
|
||||
import global_bpy
|
||||
bpy = global_bpy.get_bpy()
|
||||
|
||||
image = np.copy(image[0].numpy())
|
||||
|
||||
gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
|
||||
gray = (gray * 255).astype(np.uint8)
|
||||
|
||||
# Find contours
|
||||
contours, _ = cv2.findContours(
|
||||
gray, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||||
|
||||
# Get the largest contour
|
||||
areas = [cv2.contourArea(contour) for contour in contours]
|
||||
|
||||
max_area_index = areas.index(max(areas))
|
||||
largest_contour = contours[max_area_index]
|
||||
contours = [largest_contour]
|
||||
|
||||
def normalize_vertices(vertices, max_value):
|
||||
return vertices / float(max_value) * 2 - 1
|
||||
|
||||
# Get the image width and height
|
||||
height, width = image.shape[:2]
|
||||
|
||||
# Normalize the vertices
|
||||
normalized_contours = []
|
||||
for contour in contours:
|
||||
normalized_contour = []
|
||||
for vertex in contour:
|
||||
normalized_vertex = [normalize_vertices(
|
||||
vertex[0][0], width), normalize_vertices(vertex[0][1], height) * -1]
|
||||
normalized_contour.append(normalized_vertex)
|
||||
normalized_contours.append(
|
||||
np.array(normalized_contour, dtype=np.float32))
|
||||
|
||||
meshes = []
|
||||
# print(len(normalized_contours))
|
||||
for i, contour in enumerate(normalized_contours):
|
||||
# Create a new mesh for each contour
|
||||
mesh = bpy.data.meshes.new(name=f"NewMesh{i}")
|
||||
# Create a new object for each mesh
|
||||
obj = bpy.data.objects.new(f"NewObject{i}", mesh)
|
||||
# Link the object to the current collection
|
||||
bpy.context.collection.objects.link(obj)
|
||||
|
||||
# Add a z coordinate to each vertex
|
||||
ordered_vertices = [(*vertex, 0) for vertex in contour]
|
||||
# Create a face from the vertices
|
||||
face = list(range(len(ordered_vertices)))
|
||||
|
||||
# Create the mesh from the vertices and face
|
||||
mesh.from_pydata(ordered_vertices, [], [face])
|
||||
|
||||
# Create a default shape key for the mesh
|
||||
sk_basis = obj.shape_key_add(name='Basis')
|
||||
|
||||
meshes.append(obj) # Add the object to the list of meshes
|
||||
|
||||
# Draw contours on the original image
|
||||
if not image.flags['C_CONTIGUOUS']:
|
||||
image = np.ascontiguousarray(image)
|
||||
cv2.drawContours(image, contours, -1, (0, 255, 0), 3)
|
||||
|
||||
# Convert image back to tensor
|
||||
image = [torch.from_numpy(image)]
|
||||
return (image, meshes)
|
||||
|
||||
|
||||
# A dictionary that contains all nodes you want to export with their names
|
||||
# NOTE: names should be globally unique
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"MeshFromTexture": MeshFromTexture
|
||||
}
|
||||
|
||||
# A dictionary that contains the friendly/humanly readable titles for the nodes
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"MeshFromTexture": "Mesh from texture"
|
||||
}
|
||||
@@ -1,48 +0,0 @@
|
||||
import subprocess
|
||||
import folder_paths
|
||||
import os
|
||||
from mesh_utils import open_in_blender
|
||||
|
||||
|
||||
class OpenInBlender:
|
||||
def __init__(self):
|
||||
self.my_blender_process = None
|
||||
self.output_dir = folder_paths.get_output_directory()
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"bpy_objs": ("BPY_OBJS",),
|
||||
"blender_path": ("STRING", {
|
||||
"multiline": False,
|
||||
"default": "blender"
|
||||
}),
|
||||
"shading": (["Material", "Solid", "Rendered", "Wireframe"],),
|
||||
"camera_location": ("VECTOR3D",),
|
||||
"camera_rotation": ("VECTOR3D",),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ()
|
||||
FUNCTION = "process"
|
||||
|
||||
OUTPUT_NODE = True
|
||||
|
||||
CATEGORY = "mesh"
|
||||
|
||||
def process(self, bpy_objs, blender_path, shading, camera_location, camera_rotation):
|
||||
output_file = self.output_dir + '/tmp.blend'
|
||||
p = open_in_blender(self.my_blender_process, blender_path=blender_path, output_file=output_file, camera_location=camera_location,
|
||||
camera_rotation=camera_rotation, shading=shading)
|
||||
self.my_blender_process = p
|
||||
return {"ui": {}}
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"OpenInBlender": OpenInBlender
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"OpenInBlender": "Open in Blender"
|
||||
}
|
||||
@@ -1,74 +0,0 @@
|
||||
class Plane:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"division": ("INT", {
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 256,
|
||||
"step": 1,
|
||||
"display": "number"
|
||||
}),
|
||||
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), # For disabling cache
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("BPY_OBJS",)
|
||||
RETURN_NAMES = ("bpy_objs",)
|
||||
|
||||
FUNCTION = "process"
|
||||
|
||||
CATEGORY = "mesh"
|
||||
|
||||
def process(self, division, seed):
|
||||
import global_bpy
|
||||
bpy = global_bpy.get_bpy()
|
||||
|
||||
# Define coordinates for the plane to object mode
|
||||
if bpy.context.active_object is not None:
|
||||
bpy.ops.object.mode_set(mode='OBJECT')
|
||||
coords = [(-1, -1, 0), (1, -1, 0), (-1, 1, 0), (1, 1, 0)]
|
||||
|
||||
# Create a mesh object
|
||||
mesh = bpy.data.meshes.new(name="PlaneMesh")
|
||||
mesh.from_pydata(coords,[],[(0,1,3,2)])
|
||||
|
||||
# Create a new object with the mesh
|
||||
object_a = bpy.data.objects.new(name="PlaneObject", object_data=mesh)
|
||||
bpy.ops.object.select_all(action='DESELECT')
|
||||
bpy.context.collection.objects.link(object_a)
|
||||
bpy.context.view_layer.objects.active = object_a
|
||||
object_a.select_set(True)
|
||||
|
||||
# Invert the Y axis of the UV to fix the vertical inversion
|
||||
bpy.ops.transform.resize(value=(-1, -1, 1))
|
||||
|
||||
# Create a new UV map
|
||||
uv_map = mesh.uv_layers.new(name="newUV")
|
||||
mesh.uv_layers.active = uv_map
|
||||
|
||||
# Unwrap the mesh
|
||||
bpy.ops.object.mode_set(mode='EDIT')
|
||||
bpy.ops.uv.unwrap()
|
||||
bpy.ops.object.mode_set(mode='OBJECT')
|
||||
|
||||
# Subdivide the plane
|
||||
if division > 0:
|
||||
original_mode = bpy.context.object.mode
|
||||
bpy.ops.object.mode_set(mode='EDIT')
|
||||
bpy.ops.mesh.subdivide(number_cuts=division)
|
||||
bpy.ops.object.mode_set(mode=original_mode)
|
||||
|
||||
sk_basis = object_a.shape_key_add(name='Basis')
|
||||
|
||||
return ([object_a],)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"Plane": Plane
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"Plane": "Plane (Old)"
|
||||
}
|
||||
@@ -1,55 +0,0 @@
|
||||
class ScaleVertex():
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"bpy_objs_target": ("BPY_OBJS",),
|
||||
"extrude": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01, "display": "number"}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("BPY_OBJS",)
|
||||
RETURN_NAMES = ("bpy_objs",)
|
||||
|
||||
FUNCTION = "process"
|
||||
|
||||
CATEGORY = "mesh"
|
||||
|
||||
def process(self, bpy_objs_target, extrude):
|
||||
import global_bpy
|
||||
bpy = global_bpy.get_bpy()
|
||||
|
||||
if len(bpy_objs_target) == 0:
|
||||
# throw error
|
||||
return
|
||||
|
||||
target_object = bpy_objs_target[0]
|
||||
|
||||
# deselect all objects
|
||||
bpy.ops.object.select_all(action='DESELECT')
|
||||
# select only the target object
|
||||
bpy.context.view_layer.objects.active = target_object
|
||||
# enter enter edit mode and select all faces of the object to fill
|
||||
bpy.ops.object.mode_set(mode='EDIT')
|
||||
bpy.ops.mesh.select_all(action='SELECT')
|
||||
|
||||
bpy.ops.transform.resize(
|
||||
value=(extrude, extrude, 1), constraint_axis=(False, False, False))
|
||||
|
||||
# remove the extruded vertex from the vertex group "name"
|
||||
# bpy.ops.object.vertex_group_remove_from()
|
||||
|
||||
# bpy.ops.mesh.delete(type='FACE')
|
||||
bpy.ops.object.mode_set(mode='OBJECT')
|
||||
|
||||
return ([target_object],)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"ScaleVertex": ScaleVertex
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"ScaleVertex": "Scale Vertex"
|
||||
}
|
||||
@@ -1,92 +0,0 @@
|
||||
class Texture:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"bpy_objs": ("BPY_OBJS",),
|
||||
"texture": ("IMAGE",),
|
||||
"texture_name": ("STRING", {
|
||||
"multiline": False,
|
||||
"default": "my_image"
|
||||
}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("BPY_OBJS",)
|
||||
RETURN_NAMES = ("bpy_objs",)
|
||||
|
||||
FUNCTION = "process"
|
||||
|
||||
# OUTPUT_NODE = False
|
||||
|
||||
CATEGORY = "mesh"
|
||||
|
||||
def process(self, bpy_objs, texture, texture_name):
|
||||
import numpy as np
|
||||
import global_bpy
|
||||
bpy = global_bpy.get_bpy()
|
||||
|
||||
# Convert image to numpy
|
||||
texture = texture[0].numpy()
|
||||
|
||||
# Flip the image vertically
|
||||
texture = np.flipud(texture)
|
||||
|
||||
# Create an image with the required dimensions
|
||||
img = bpy.data.images.new(texture_name, width=texture.shape[1], height=texture.shape[0])
|
||||
|
||||
# If there is no alpha channel, append one full of 1's
|
||||
if texture.shape[2] == 3:
|
||||
alpha_channel = np.ones((*texture.shape[:2], 1))
|
||||
texture = np.concatenate((texture, alpha_channel), axis=2)
|
||||
|
||||
# Flatten image data and rearrange color channels for blender
|
||||
img.pixels = texture.ravel()
|
||||
|
||||
# Pack image to store it within .blend file
|
||||
img.pack()
|
||||
|
||||
# Save image to a file
|
||||
# img.filepath_raw = 'test.png'
|
||||
# img.file_format = 'PNG'
|
||||
# img.save()
|
||||
|
||||
# Get the active object
|
||||
obj = bpy_objs[0]
|
||||
|
||||
# Create a material
|
||||
mat = bpy.data.materials.new("MaterialName")
|
||||
mat.use_nodes = True
|
||||
nodes = mat.node_tree.nodes
|
||||
for node in nodes:
|
||||
nodes.remove(node)
|
||||
|
||||
# Add a new texture node
|
||||
texture_node = nodes.new(type='ShaderNodeTexImage')
|
||||
texture_node.image = img
|
||||
|
||||
# Add a new BSDF node
|
||||
bsdf_node = nodes.new(type='ShaderNodeBsdfPrincipled')
|
||||
|
||||
# Add a new output node
|
||||
output_node = nodes.new(type='ShaderNodeOutputMaterial')
|
||||
|
||||
# Link nodes together
|
||||
links = mat.node_tree.links
|
||||
links.new(bsdf_node.inputs['Base Color'], texture_node.outputs['Color'])
|
||||
links.new(output_node.inputs['Surface'], bsdf_node.outputs['BSDF'])
|
||||
|
||||
# Assign the material to the active object
|
||||
if obj.data.materials:
|
||||
obj.data.materials[0] = mat
|
||||
else:
|
||||
obj.data.materials.append(mat)
|
||||
return (bpy_objs,)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"Texture": Texture
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"Texture": "Texture"
|
||||
}
|
||||
@@ -1,45 +0,0 @@
|
||||
class TransfromObject():
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"bpy_objs_target": ("BPY_OBJS",),
|
||||
"x": ("FLOAT", {"default": 0.0, "min": -10.0, "max": 10.0, "step": 0.01, "display": "number"}),
|
||||
"y": ("FLOAT", {"default": 0.0, "min": -10.0, "max": 10.0, "step": 0.01, "display": "number"}),
|
||||
"z": ("FLOAT", {"default": 0.0, "min": -10.0, "max": 10.0, "step": 0.01, "display": "number"}),
|
||||
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("BPY_OBJS",)
|
||||
RETURN_NAMES = ("bpy_objs",)
|
||||
|
||||
FUNCTION = "process"
|
||||
|
||||
CATEGORY = "mesh"
|
||||
|
||||
def process(self, bpy_objs_target, x, y, z):
|
||||
import global_bpy
|
||||
bpy = global_bpy.get_bpy()
|
||||
|
||||
target_object = bpy_objs_target[0]
|
||||
|
||||
bpy.ops.object.select_all(action='DESELECT')
|
||||
bpy.context.view_layer.objects.active = target_object
|
||||
|
||||
# move the object in level transformation
|
||||
target_object.location.x += x
|
||||
target_object.location.y += y
|
||||
target_object.location.z += z
|
||||
|
||||
return ([target_object],)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"TransfromObject": TransfromObject
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"TransfromObject": "Transfrom Object"
|
||||
}
|
||||
@@ -1,52 +0,0 @@
|
||||
import torch
|
||||
from einops import rearrange
|
||||
from nodes import PreviewImage
|
||||
|
||||
class PointVisualizer(PreviewImage):
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE",),
|
||||
"point": ("VECTOR3D",),
|
||||
"point_size": ("INT", {
|
||||
"default": 1,
|
||||
"min": 10,
|
||||
"max": 20,
|
||||
"step": 1,
|
||||
"display": "number"
|
||||
}),
|
||||
},
|
||||
}
|
||||
|
||||
FUNCTION = "process"
|
||||
|
||||
CATEGORY = "image"
|
||||
|
||||
def process(self, images, point, point_size):
|
||||
x, y, z = point
|
||||
h, w, c = images[0].shape
|
||||
point_image = torch.zeros(images[0].shape)
|
||||
|
||||
center_x, center_y, _ = point
|
||||
top_left_x = center_x - point_size // 2
|
||||
top_left_y = center_y - point_size // 2
|
||||
|
||||
# Make sure the square is within the image boundaries
|
||||
top_left_x = max(0, min(w - point_size, top_left_x))
|
||||
top_left_y = max(0, min(h - point_size, top_left_y))
|
||||
|
||||
point_image[top_left_y:top_left_y+point_size, top_left_x:top_left_x+point_size] = 1
|
||||
point_image = rearrange(point_image, 'h w c -> 1 h w c')
|
||||
|
||||
# Blending
|
||||
images = images * 0.2 + point_image * 0.8
|
||||
return self.save_images(images)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"PointVisualizer": PointVisualizer
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"PointVisualizer": "Point Visualizer"
|
||||
}
|
||||
@@ -1,45 +0,0 @@
|
||||
class VECTOR3D:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"x": ("INT", {
|
||||
"default": 0,
|
||||
"min": -1024,
|
||||
"max": 1024,
|
||||
"step": 1,
|
||||
"display": "number"
|
||||
}),
|
||||
"y": ("INT", {
|
||||
"default": 0,
|
||||
"min": -1024,
|
||||
"max": 1024,
|
||||
"step": 1,
|
||||
"display": "number"
|
||||
}),
|
||||
"z": ("INT", {
|
||||
"default": 0,
|
||||
"min": -1024,
|
||||
"max": 1024,
|
||||
"step": 1,
|
||||
"display": "number"
|
||||
}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("VECTOR3D",)
|
||||
|
||||
FUNCTION = "run"
|
||||
|
||||
CATEGORY = "utils"
|
||||
|
||||
def run(self, x, y, z):
|
||||
return ([x, y, z],)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"VECTOR3D": VECTOR3D
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"VECTOR3D": "Vector 3D"
|
||||
}
|
||||
@@ -1,67 +0,0 @@
|
||||
import torch
|
||||
from einops import rearrange, repeat
|
||||
from torchvision import transforms
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
import cv2
|
||||
import base64
|
||||
import requests
|
||||
import json
|
||||
import io
|
||||
|
||||
class AvatarSegmentation:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "IMAGE", "IMAGE")
|
||||
RETURN_NAMES = ("back_hair", "front_hair","eyes", "eyelashes", "mouth")
|
||||
|
||||
FUNCTION = "segment"
|
||||
|
||||
CATEGORY = "image"
|
||||
|
||||
def segment(self, image):
|
||||
# Convert tensor to PIL image
|
||||
image = image[0]
|
||||
image = rearrange(image, 'h w c -> c h w')
|
||||
image = transforms.ToPILImage()(image)
|
||||
|
||||
# Save image to in-memory file
|
||||
image_buff = io.BytesIO()
|
||||
image.save(image_buff, format="PNG")
|
||||
image_string = base64.b64encode(image_buff.getvalue()).decode("utf-8")
|
||||
|
||||
# Send request to API
|
||||
url = 'https://q41iq6s6t8.execute-api.ap-southeast-1.amazonaws.com/seg-cpu'
|
||||
headers = {
|
||||
'content-type': 'application/json'
|
||||
}
|
||||
data = {
|
||||
'img_str': image_string
|
||||
}
|
||||
response = requests.post(url, headers=headers, data=json.dumps(data))
|
||||
|
||||
# Convert response to tensors
|
||||
outs = []
|
||||
masks = response.json()
|
||||
for segment_name, img_str in masks.items():
|
||||
mask = base64.b64decode(img_str.encode('utf-8'))
|
||||
mask = Image.open(io.BytesIO(mask))
|
||||
mask = np.array(mask)
|
||||
mask = torch.from_numpy(mask) # shape: H, W, 3
|
||||
mask = rearrange(mask, 'h w c -> 1 h w c')
|
||||
outs.append(mask)
|
||||
return outs
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"AvatarSegmentation": AvatarSegmentation
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"AvatarSegmentation": "Avatar Segmentation"
|
||||
}
|
||||
@@ -1,52 +0,0 @@
|
||||
import folder_paths
|
||||
import torch
|
||||
from segment_anything import SamAutomaticMaskGenerator, sam_model_registry
|
||||
from einops import rearrange, repeat
|
||||
|
||||
|
||||
class SAM_Load_Embedding:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"filename": ("STRING", {
|
||||
"multiline": False,
|
||||
"default": "embeddings"
|
||||
}),
|
||||
"embedding_id": ("STRING", {
|
||||
"multiline": False,
|
||||
"default": "embedding"
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("EMBEDDINGS",)
|
||||
RETURN_NAMES = ("EMBEDDINGS",)
|
||||
|
||||
OUTPUT_NODE = True
|
||||
|
||||
FUNCTION = "process"
|
||||
|
||||
CATEGORY = "image"
|
||||
|
||||
def process(self, filename, embedding_id):
|
||||
import json
|
||||
import numpy as np
|
||||
|
||||
data = {}
|
||||
with open(filename, 'r') as f:
|
||||
data = json.load(f)
|
||||
|
||||
# Convert list to numpy ndarray
|
||||
data['image_embedding'] = np.array(data['image_embedding'])
|
||||
|
||||
return (data, )
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"SAM_Load_Embedding": SAM_Load_Embedding
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"SAM_Load_Embedding": "SAM_Load_Embedding "
|
||||
}
|
||||
@@ -1,42 +0,0 @@
|
||||
import folder_paths
|
||||
import torch
|
||||
from segment_anything import SamAutomaticMaskGenerator, sam_model_registry
|
||||
from einops import rearrange, repeat
|
||||
import requests
|
||||
from PIL import Image
|
||||
import io
|
||||
import numpy as np
|
||||
from segment_anything import SamAutomaticMaskGenerator, sam_model_registry, SamPredictor
|
||||
from einops import rearrange, repeat
|
||||
|
||||
class SAM_Remote_Emb:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model_type": (["vit_h", "vit_l", "vit_b"],),
|
||||
"ckpt": (folder_paths.get_filename_list("sams"),),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("SAM", "SAMPREDICTOR")
|
||||
RETURN_NAMES = ("sam", "predictor")
|
||||
|
||||
FUNCTION = "segment"
|
||||
|
||||
CATEGORY = "image"
|
||||
|
||||
def segment(self, model_type, ckpt):
|
||||
ckpt = folder_paths.get_full_path("sams", ckpt)
|
||||
sam = sam_model_registry[model_type](checkpoint=ckpt)
|
||||
predictor = SamPredictor(sam)
|
||||
|
||||
return (sam, predictor)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"SAM_Loader": SAM_Remote_Emb
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"SAM_Loader": "SAM Loader"
|
||||
}
|
||||
@@ -1,110 +0,0 @@
|
||||
import torch
|
||||
from einops import rearrange, repeat
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
class SAM_Predict:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE",),
|
||||
"embeddings": ("EMBEDDINGS",),
|
||||
"predictor": ("SAMPREDICTOR",),
|
||||
},
|
||||
"optional": {
|
||||
"prompt": ("SAM_PROMPT", ),
|
||||
"mask": ("IMAGE",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE","IMAGE", "IMAGE")
|
||||
RETURN_NAMES = ("image","out_image", "mask")
|
||||
|
||||
FUNCTION = "segment"
|
||||
|
||||
CATEGORY = "image"
|
||||
|
||||
def segment(self, image, embeddings, predictor, prompt, mask=None):
|
||||
image_embedding_list = embeddings['image_embedding']
|
||||
shape = tuple(embeddings['shape'])
|
||||
input_size = tuple(embeddings['input_size'])
|
||||
|
||||
# Convert the list back to a numpy array with the original shape
|
||||
image_embedding_np = np.array(image_embedding_list, dtype=np.single).reshape(shape)
|
||||
|
||||
# Convert the numpy array to a PyTorch tensor
|
||||
image_embedding_tensor = torch.from_numpy(image_embedding_np)
|
||||
|
||||
# Set the image embeddings for the model
|
||||
# predictor.set_torch_image(image_embedding_tensor, image.shape[:2])
|
||||
# print(image[0].shape[:2])
|
||||
predictor.input_size = input_size
|
||||
predictor.features = image_embedding_tensor
|
||||
predictor.is_image_set = True
|
||||
predictor.original_size = image[0].shape[:2]
|
||||
|
||||
# prompt = [{"x":364,"y":153,"label":1},{"x":296,"y":189,"label":1},{"x":277,"y":246,"label":1}]
|
||||
|
||||
# if point_1 != None:
|
||||
# x, y, z = point_1
|
||||
|
||||
# point_coords = np.array([[x, y]])
|
||||
# point_labels = np.array([1])
|
||||
|
||||
if prompt == None or len(prompt) == 0:
|
||||
return (image, image, None)
|
||||
|
||||
point_coords = np.array([[p['x'], p['y']] for p in prompt])
|
||||
point_labels = np.array([p['label'] for p in prompt])
|
||||
|
||||
masks, iou_predictions, low_res_masks = predictor.predict(
|
||||
point_coords=point_coords,
|
||||
point_labels=point_labels,
|
||||
)
|
||||
|
||||
if mask != None:
|
||||
# scale the mask to 256x256
|
||||
cv2_mask = cv2.resize(np.array(mask[0]), (256, 256))
|
||||
cv2_mask = cv2_mask[np.newaxis, :, :]
|
||||
cv2_mask = (cv2_mask * 255).astype(int)
|
||||
|
||||
true_locations = np.array(np.where(cv2_mask[0] == 255))
|
||||
if true_locations.shape[1] > 0:
|
||||
# Randomly select a point in the mask
|
||||
rand_index = np.random.randint(true_locations.shape[1])
|
||||
y, x = true_locations[:, rand_index]
|
||||
point_coords = np.array([[x, y]])
|
||||
point_labels = np.array([1])
|
||||
|
||||
masks, iou_predictions, low_res_masks = predictor.predict(
|
||||
point_coords=point_coords,
|
||||
point_labels=point_labels,
|
||||
mask_input=cv2_mask
|
||||
)
|
||||
else:
|
||||
# No detected mask
|
||||
h, w, c = image[0].shape
|
||||
masks = np.zeros((1, h, w))
|
||||
|
||||
masks = torch.from_numpy(masks)
|
||||
|
||||
masks = rearrange(masks[0], 'h w -> 1 h w')
|
||||
# masks = rearrange(masks, 'c h w -> 1 c h w')
|
||||
out_image = repeat(masks, '1 h w -> 1 h w c', c=3) * image
|
||||
|
||||
print(masks.shape, torch.max(masks), torch.min(masks))
|
||||
print(image.shape, torch.max(image), torch.min(image))
|
||||
|
||||
# print(emb)
|
||||
|
||||
# print(masks, out_image)
|
||||
return (image, out_image, masks)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"SAM_Predict": SAM_Predict
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"SAM_Predict": "SAM Predict"
|
||||
}
|
||||
@@ -1,88 +0,0 @@
|
||||
import requests
|
||||
from PIL import Image
|
||||
import io
|
||||
import numpy as np
|
||||
import folder_paths
|
||||
|
||||
class SAM_Embedding:
|
||||
def __init__(self):
|
||||
self.output_dir = folder_paths.get_output_directory()
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE",),
|
||||
"embedding_id": ("STRING", {
|
||||
"multiline": False,
|
||||
"default": "embedding"
|
||||
}),
|
||||
},
|
||||
"optional": {
|
||||
"predictor": ("SAMPREDICTOR",)
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("EMBEDDINGS", )
|
||||
RETURN_NAMES = ("embeddings", )
|
||||
|
||||
FUNCTION = "segment"
|
||||
|
||||
CATEGORY = "image"
|
||||
|
||||
def segment(self, image, embedding_id, predictor=None):
|
||||
# Convert PyTorch tensor to numpy array
|
||||
image_np = (image[0].numpy() * 255).astype(np.uint8)
|
||||
|
||||
if predictor != None:
|
||||
predictor.set_image(image_np)
|
||||
emb = predictor.get_image_embedding().cpu().numpy()
|
||||
output = {
|
||||
"image_embedding": emb,
|
||||
"shape": emb.shape,
|
||||
"input_size": predictor.input_size
|
||||
}
|
||||
else:
|
||||
# Convert numpy array to PIL Image
|
||||
img = Image.fromarray(image_np)
|
||||
|
||||
# Create an in-memory bytes buffer
|
||||
img_byte_arr = io.BytesIO()
|
||||
|
||||
# Save the PIL Image to the bytes buffer in PNG format
|
||||
img.save(img_byte_arr, format='PNG')
|
||||
|
||||
# Get the bytes value of the buffer
|
||||
img_byte_arr = img_byte_arr.getvalue()
|
||||
|
||||
# Create a dictionary with the image bytes
|
||||
files = {'image': ('image.png', img_byte_arr)}
|
||||
|
||||
# Send the POST request
|
||||
response = requests.post('https://avatechgg--segment-anything-entrypoint.modal.run', files=files)
|
||||
|
||||
# Check if the request was successful
|
||||
if response.status_code == 200:
|
||||
# Parse the JSON response
|
||||
output = response.json()
|
||||
else:
|
||||
print(f"Request failed with status code {response.status_code}")
|
||||
output = None
|
||||
|
||||
# sam = sam_model_registry[model_type](checkpoint=ckpt)
|
||||
# predictor = SamPredictor(sam)
|
||||
# masks = predictor.set_torch_image
|
||||
# masks = predictor.predict
|
||||
|
||||
# print(output)
|
||||
|
||||
np.save(f"{self.output_dir}/{embedding_id}.npy", output["image_embedding"])
|
||||
return (output, )
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"SAM_Embedding": SAM_Embedding
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"SAM_Embedding": "SAM Embedding"
|
||||
}
|
||||
@@ -1,68 +0,0 @@
|
||||
import folder_paths
|
||||
import os
|
||||
import numpy as np
|
||||
import re
|
||||
from segment_anything import sam_model_registry, SamPredictor
|
||||
|
||||
|
||||
class SAM_Prompt_Image:
|
||||
def __init__(self):
|
||||
self.output_dir = folder_paths.get_output_directory()
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
input_dir = folder_paths.get_input_directory()
|
||||
files = [
|
||||
f
|
||||
for f in os.listdir(input_dir)
|
||||
if os.path.isfile(os.path.join(input_dir, f))
|
||||
]
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE",),
|
||||
"ckpt": (folder_paths.get_filename_list("sams"),),
|
||||
"embedding_id": (
|
||||
"STRING",
|
||||
{"multiline": False, "default": "embedding"},
|
||||
),
|
||||
# "image": (sorted(files), ),
|
||||
"image_prompts_json": ("STRING", {"multiline": False, "default": "[]"}),
|
||||
},
|
||||
}
|
||||
|
||||
CATEGORY = "image"
|
||||
|
||||
RETURN_TYPES = ("SAM_PROMPT",)
|
||||
FUNCTION = "load_image"
|
||||
|
||||
def load_image(self, image, ckpt, embedding_id, image_prompts_json):
|
||||
import json
|
||||
|
||||
emb_filename = f"{self.output_dir}/{embedding_id}.npy"
|
||||
if not os.path.exists(emb_filename):
|
||||
ckpt = folder_paths.get_full_path("sams", ckpt)
|
||||
model_type = re.findall(r'vit_[lbh]', ckpt)[0]
|
||||
sam = sam_model_registry[model_type](checkpoint=ckpt)
|
||||
predictor = SamPredictor(sam)
|
||||
|
||||
image_np = (image[0].numpy() * 255).astype(np.uint8)
|
||||
predictor.set_image(image_np)
|
||||
emb = predictor.get_image_embedding().cpu().numpy()
|
||||
np.save(emb_filename, emb)
|
||||
|
||||
image_prompts = json.loads(image_prompts_json)
|
||||
|
||||
result = (image_prompts,)
|
||||
|
||||
if isinstance(image_prompts, list):
|
||||
pass
|
||||
elif all(isinstance(item, list) for item in image_prompts.values()):
|
||||
for item in image_prompts.values():
|
||||
result += (item,)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {"SAM_Prompt_Image": SAM_Prompt_Image}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {"SAM_Prompt_Image": "SAM_Prompt_Image "}
|
||||
@@ -1,64 +0,0 @@
|
||||
import folder_paths
|
||||
import torch
|
||||
import os
|
||||
from segment_anything import SamAutomaticMaskGenerator, sam_model_registry
|
||||
from einops import rearrange, repeat
|
||||
|
||||
class SAM_Save_Embedding:
|
||||
def __init__(self):
|
||||
self.output_dir = folder_paths.get_output_directory()
|
||||
self.type = "output"
|
||||
self.prefix_append = ""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"embeddings": ("EMBEDDINGS",),
|
||||
"filename": ("STRING", {
|
||||
"multiline": False,
|
||||
"default": "embeddings"
|
||||
}),
|
||||
"write_mode": (["Overwrite", "Increment"],),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ()
|
||||
RETURN_NAMES = ()
|
||||
|
||||
OUTPUT_NODE = True
|
||||
|
||||
FUNCTION = "process"
|
||||
|
||||
CATEGORY = "image"
|
||||
|
||||
def process(self, embeddings, filename, write_mode):
|
||||
import json
|
||||
import numpy as np
|
||||
|
||||
filepath = self.output_dir + "/" + filename + '.json'
|
||||
|
||||
if write_mode == "Increment":
|
||||
count = 0
|
||||
# while file exists, increment count
|
||||
while os.path.exists(self.output_dir + "/" + filename + '_' + str(count) + '.json'):
|
||||
count += 1
|
||||
|
||||
filepath = self.output_dir + "/" + filename + '_' + str(count) + '.json'
|
||||
|
||||
# print(embeddings)
|
||||
if isinstance(embeddings['image_embedding'], np.ndarray):
|
||||
embeddings['image_embedding'] = embeddings['image_embedding'].tolist()
|
||||
with open(filepath, 'w') as f:
|
||||
json.dump(embeddings, f)
|
||||
|
||||
return { "ui" : { "file": { filepath.replace(f"{self.output_dir}/", "") } } }
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"SAM_Save_Embedding": SAM_Save_Embedding
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"SAM_Save_Embedding": "SAM_Save_Embedding "
|
||||
}
|
||||
Reference in New Issue
Block a user