narrowfy
This commit is contained in:
Ubuntu
+7
-3
@@ -15,7 +15,7 @@ from scaled_paste import main_scaled_paste
|
||||
from scaled_paste import main_scaled_paste_2
|
||||
from simple_bg_swap import (simple_bg_swap, get_threshold_for_bg_swap, RGB_2_LAB, LAB_2_RGB, get_mean_and_standard_deviation, renormalize_array)
|
||||
from distribution_reshape import (simple_rescale_histogram, get_histogram_limits)
|
||||
from utility_nodes import TRI3D_clean_mask, TRI3D_extract_pose_part, TRI3D_position_pose_part, TRI3D_fill_mask, TRI3D_is_only_trouser
|
||||
from utility_nodes import TRI3D_clean_mask, TRI3D_extract_pose_part, TRI3D_position_pose_part, TRI3D_fill_mask, TRI3D_is_only_trouser, TRI3D_extract_facer_mask,
|
||||
from utility_nodes import TRI3D_extract_facer_mask
|
||||
|
||||
from .AEMatter import (load_AEMatter_Model, run_AEMatter_inference)
|
||||
@@ -3694,7 +3694,7 @@ class TRI3D_BGREMOVE_MEGA():
|
||||
|
||||
|
||||
from photoroom import TRI3D_photoroom_bgremove_api
|
||||
from smart_box import TRI3D_SmartBox, TRI3D_Skip_HeadMask, TRI3D_Skip_HeadMask_AddNeck, TRI3D_Image_extend
|
||||
from smart_box import TRI3D_SmartBox, TRI3D_Skip_HeadMask, TRI3D_Skip_HeadMask_AddNeck, TRI3D_Image_extend, TRI3D_Smart_Depth, TRI3D_NarrowfyImage
|
||||
from nsfw import TRI3DNSFWFilter
|
||||
|
||||
# A dictionary that contains all nodes you want to export with their names
|
||||
@@ -3760,11 +3760,13 @@ NODE_CLASS_MAPPINGS = {
|
||||
"tri3d_Skip_HeadMask": TRI3D_Skip_HeadMask,
|
||||
"tri3d_Skip_HeadMask_AddNeck": TRI3D_Skip_HeadMask_AddNeck,
|
||||
"tri3d_Image_extend": TRI3D_Image_extend,
|
||||
"tri3d_Smart_Depth": TRI3D_Smart_Depth,
|
||||
"tri3d_NSFWFilter": TRI3DNSFWFilter,
|
||||
"tri3d_NarrowfyImage": TRI3D_NarrowfyImage,
|
||||
}
|
||||
|
||||
|
||||
VERSION = "4.8.6"
|
||||
VERSION = "4.8.7"
|
||||
# A dictionary that contains the friendly/humanly readable titles for the nodes
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"tri3d-photoroom-bgremove-api": "Photoroom BG Remove" + " v" + VERSION,
|
||||
@@ -3830,4 +3832,6 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"tri3d_Skip_HeadMask_AddNeck": "Skip Head Mask and add neck" + " v" + VERSION,
|
||||
"tri3d_NSFWFilter": "TRI3D NSFW Filter" + " v" + VERSION,
|
||||
"tri3d_Image_extend": "Image extend" + " v" + VERSION,
|
||||
"tri3d_Smart_Depth": "Smart Depth" + " v" + VERSION,
|
||||
"tri3d_NarrowfyImage": "Narrowfy Image" + " v" + VERSION,
|
||||
}
|
||||
|
||||
+321
@@ -408,3 +408,324 @@ class TRI3D_Image_extend:
|
||||
torch_mask = torch_mask.unsqueeze(0)
|
||||
|
||||
return (torch_image, torch_mask)
|
||||
|
||||
|
||||
class TRI3D_Smart_Depth:
|
||||
|
||||
|
||||
def from_torch_image(self, image):
|
||||
image = image.cpu().numpy() * 255.0
|
||||
image = np.clip(image, 0, 255).astype(np.uint8)
|
||||
return image
|
||||
|
||||
def to_torch_image(self, image):
|
||||
image = image.astype(dtype=np.float32)
|
||||
image /= 255.0
|
||||
image = torch.from_numpy(image)
|
||||
return image
|
||||
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE", ),
|
||||
"keypoints_json": ("STRING", {"multiline": True}),
|
||||
},
|
||||
}
|
||||
|
||||
FUNCTION = "run"
|
||||
RETURN_TYPES = ("IMAGE", )
|
||||
CATEGORY = "TRI3D"
|
||||
|
||||
def extract_torso_keypoints(self, keypoints):
|
||||
# Indices for torso-related keypoints
|
||||
torso_indices = [8, 9, 10, 11, 12, 13]
|
||||
return [keypoints[i] for i in torso_indices]
|
||||
|
||||
|
||||
def run(self, image, keypoints_json):
|
||||
|
||||
kp_data = json.loads(open(keypoints_json, 'r').read())
|
||||
original_height, original_width = kp_data['height'], kp_data['width']
|
||||
torso_keypoints = self.extract_torso_keypoints(kp_data['keypoints'])
|
||||
|
||||
# Convert Torch image to OpenCV format
|
||||
cv_image = self.from_torch_image(image)
|
||||
|
||||
# Remove the batch dimension if present
|
||||
if len(cv_image.shape) == 4:
|
||||
cv_image = cv_image[0]
|
||||
|
||||
# Adjust keypoints to match the image dimensions
|
||||
adjusted_keypoints = self.adjust_keypoints(torso_keypoints, cv_image.shape, original_height, original_width)
|
||||
|
||||
|
||||
# Fill the area below the hip line
|
||||
filled_image = self.fill_below_hip(cv_image, adjusted_keypoints)
|
||||
|
||||
|
||||
# Convert back to Torch format
|
||||
torch_image = self.to_torch_image(filled_image)
|
||||
|
||||
# Add the batch dimension back
|
||||
torch_image = torch_image.unsqueeze(0)
|
||||
|
||||
|
||||
return (torch_image,)
|
||||
|
||||
def adjust_keypoints(self, keypoints, image_shape, original_height, original_width):
|
||||
image_height, image_width = image_shape[:2]
|
||||
scale_x = image_width / original_width
|
||||
scale_y = image_height / original_height
|
||||
|
||||
adjusted_keypoints = [
|
||||
(int(x * scale_x), int(y * scale_y)) for x, y in keypoints
|
||||
]
|
||||
return adjusted_keypoints
|
||||
|
||||
def fill_below_hip(self, image, keypoints):
|
||||
# Correct the indices for hip keypoints
|
||||
# Assuming indices 8 and 11 are for left and right hips
|
||||
# print(keypoints,"hip keypoints")
|
||||
try:
|
||||
valid_y_coords = [kp[1] for kp in keypoints if kp[1] >= 0]
|
||||
hip_y = min(valid_y_coords) if valid_y_coords else 0
|
||||
except:
|
||||
hip_y = 0
|
||||
|
||||
if hip_y == 0:
|
||||
return image
|
||||
|
||||
# Find the bounding box of the mask below the hip line
|
||||
mask = image[:, :, 0] # Assuming single-channel mask
|
||||
below_hip = mask[hip_y:, :]
|
||||
contours, _ = cv2.findContours(below_hip, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||||
cnt = 0
|
||||
|
||||
for contour in contours:
|
||||
x, y, w, h = cv2.boundingRect(contour)
|
||||
# print(cnt, x,y,w,h, cv2.contourArea(contour), "cnt,x,y,w,h,area")
|
||||
cnt+=1
|
||||
# cv2.rectangle(image, (x, y + hip_y), (x + w, y + h + hip_y), (255, 255, 255), -1)
|
||||
contours = [contour for contour in contours if cv2.contourArea(contour) > 0]
|
||||
|
||||
if len(contours) == 0:
|
||||
return image
|
||||
# Combine all contours into one
|
||||
all_contours = np.vstack(contours)
|
||||
|
||||
# Calculate a single bounding rectangle for all contours
|
||||
x, y, w, h = cv2.boundingRect(all_contours)
|
||||
# print(x,y,w,h, "x,y,w,h")
|
||||
cv2.rectangle(image, (x, y + hip_y), (x + w, y + h + hip_y), (0, 0, 0), -1)
|
||||
|
||||
return image
|
||||
|
||||
|
||||
class TRI3D_NarrowfyImage:
|
||||
def from_torch_image(self, image):
|
||||
image = image.cpu().numpy() * 255.0
|
||||
image = np.clip(image, 0, 255).astype(np.uint8)
|
||||
return image
|
||||
|
||||
def to_torch_image(self, image):
|
||||
image = image.astype(dtype=np.float32)
|
||||
image /= 255.0
|
||||
image = torch.from_numpy(image)
|
||||
return image
|
||||
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"face_mask": ("IMAGE", ),
|
||||
"image": ("IMAGE", ),
|
||||
"ratio": ("FLOAT", {"default": 1.5, "min": 1.2, "max": 2, "step": 0.01}),
|
||||
},
|
||||
}
|
||||
|
||||
FUNCTION = "run"
|
||||
RETURN_TYPES = ("IMAGE", "IMAGE", )
|
||||
RETURN_NAMES = ("image", "mask_image", )
|
||||
CATEGORY = "TRI3D"
|
||||
|
||||
def run(self, face_mask, image, ratio):
|
||||
cv_face_mask = self.from_torch_image(face_mask)
|
||||
cv_image = self.from_torch_image(image)
|
||||
|
||||
# Remove the batch dimension if present
|
||||
if len(cv_image.shape) == 4:
|
||||
cv_image = cv_image[0]
|
||||
if len(cv_face_mask.shape) == 4:
|
||||
cv_face_mask = cv_face_mask[0]
|
||||
mask = cv_face_mask[:, :, 0] # Assuming single-channel mask
|
||||
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||||
lowest_y = 0
|
||||
highest_y = cv_image.shape[0]
|
||||
for contour in contours:
|
||||
for point in contour:
|
||||
x, y = point[0]
|
||||
if y > lowest_y:
|
||||
lowest_y = y
|
||||
if y < highest_y:
|
||||
highest_y = y
|
||||
|
||||
y_below_face = cv_image.shape[0] - lowest_y
|
||||
y_face = lowest_y-highest_y
|
||||
|
||||
# Only extend if the space below face is less than 1.5 times face height
|
||||
target_below_face = int(y_face * ratio)
|
||||
# print("y_face", y_face)
|
||||
# print("lowest_y", lowest_y)
|
||||
# print("highest_y", highest_y)
|
||||
# print("target_below_face", target_below_face)
|
||||
# print("y_below_face", y_below_face)
|
||||
|
||||
original_height = cv_image.shape[0]
|
||||
original_width = cv_image.shape[1]
|
||||
|
||||
if y_below_face < target_below_face:
|
||||
y_extend = target_below_face - y_below_face
|
||||
|
||||
# Calculate how much to extend horizontally to maintain aspect ratio
|
||||
new_height = original_height + y_extend
|
||||
new_width = int(original_width * (new_height / original_height))
|
||||
x_extend = new_width - original_width
|
||||
x_extend_left = x_extend // 2
|
||||
x_extend_right = x_extend - x_extend_left
|
||||
|
||||
# Extend the image in all necessary directions
|
||||
cv_image = cv2.copyMakeBorder(
|
||||
cv_image,
|
||||
0, y_extend, # top, bottom
|
||||
x_extend_left, x_extend_right, # left, right
|
||||
cv2.BORDER_CONSTANT,
|
||||
value=[0, 0, 0]
|
||||
)
|
||||
|
||||
# Create extension mask
|
||||
extension_mask = np.zeros_like(cv_image)
|
||||
# Make extended portions white
|
||||
extension_mask[original_height:, :] = 255 # bottom extension
|
||||
extension_mask[:, :x_extend_left] = 255 # left extension
|
||||
extension_mask[:, -x_extend_right:] = 255 # right extension
|
||||
|
||||
else:
|
||||
extension_mask = np.zeros_like(cv_image)
|
||||
|
||||
# Convert both images back to torch format
|
||||
torch_image = self.to_torch_image(cv_image)
|
||||
torch_mask = self.to_torch_image(extension_mask)
|
||||
|
||||
# Add batch dimension to both
|
||||
torch_image = torch_image.unsqueeze(0)
|
||||
torch_mask = torch_mask.unsqueeze(0)
|
||||
|
||||
return (torch_image, torch_mask)
|
||||
|
||||
|
||||
class TRI3D_CropAndExtend:
|
||||
def from_torch_image(self, image):
|
||||
image = image.cpu().numpy() * 255.0
|
||||
image = np.clip(image, 0, 255).astype(np.uint8)
|
||||
return image
|
||||
|
||||
def to_torch_image(self, image):
|
||||
image = image.astype(dtype=np.float32)
|
||||
image /= 255.0
|
||||
image = torch.from_numpy(image)
|
||||
return image
|
||||
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"garment_image": ("IMAGE",),
|
||||
"garment_mask": ("IMAGE",),
|
||||
"human_image": ("IMAGE",),
|
||||
"human_mask": ("IMAGE",),
|
||||
"margin": ("INT", {"default": 10, "min": 0, "max": 50}),
|
||||
},
|
||||
}
|
||||
|
||||
FUNCTION = "run"
|
||||
RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "IMAGE",)
|
||||
RETURN_NAMES = ("cropped_garment", "cropped_garment_mask", "cropped_human", "cropped_human_mask",)
|
||||
CATEGORY = "TRI3D"
|
||||
|
||||
def process_image_and_mask(self, image, mask, margin):
|
||||
# Convert to CV format and remove batch dimension
|
||||
cv_image = self.from_torch_image(image)[0]
|
||||
cv_mask = self.from_torch_image(mask)[0]
|
||||
|
||||
# Find bounding box from mask
|
||||
mask_channel = cv_mask[:, :, 0]
|
||||
contours, _ = cv2.findContours(mask_channel, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||||
|
||||
if not contours:
|
||||
return image, mask
|
||||
|
||||
# Get bounding box with margin
|
||||
x, y, w, h = cv2.boundingRect(contours[0])
|
||||
x = max(0, x - margin)
|
||||
y = max(0, y - margin)
|
||||
w = min(cv_image.shape[1] - x, w + 2 * margin)
|
||||
h = min(cv_image.shape[0] - y, h + 2 * margin)
|
||||
|
||||
# Crop image and mask
|
||||
cropped_image = cv_image[y:y+h, x:x+w]
|
||||
cropped_mask = cv_mask[y:y+h, x:x+w]
|
||||
|
||||
# Calculate required height for aspect ratio 1/3
|
||||
min_height = w * 3
|
||||
if h < min_height:
|
||||
height_extend = min_height - h
|
||||
|
||||
# Extend image with black pixels
|
||||
extended_image = cv2.copyMakeBorder(
|
||||
cropped_image,
|
||||
0, int(height_extend), # top, bottom
|
||||
0, 0, # left, right
|
||||
cv2.BORDER_CONSTANT,
|
||||
value=[0, 0, 0]
|
||||
)
|
||||
|
||||
# Extend mask with white pixels for garment mask
|
||||
extended_mask = cv2.copyMakeBorder(
|
||||
cropped_mask,
|
||||
0, int(height_extend), # top, bottom
|
||||
0, 0, # left, right
|
||||
cv2.BORDER_CONSTANT,
|
||||
value=[255, 255, 255]
|
||||
)
|
||||
|
||||
cropped_image = extended_image
|
||||
cropped_mask = extended_mask
|
||||
|
||||
# Convert back to torch format and add batch dimension
|
||||
torch_image = self.to_torch_image(cropped_image).unsqueeze(0)
|
||||
torch_mask = self.to_torch_image(cropped_mask).unsqueeze(0)
|
||||
|
||||
return torch_image, torch_mask
|
||||
|
||||
def run(self, garment_image, garment_mask, human_image, human_mask, margin):
|
||||
# Process garment
|
||||
cropped_garment, cropped_garment_mask = self.process_image_and_mask(
|
||||
garment_image, garment_mask, margin
|
||||
)
|
||||
|
||||
# Process human
|
||||
cropped_human, cropped_human_mask = self.process_image_and_mask(
|
||||
human_image, human_mask, margin
|
||||
)
|
||||
|
||||
return (cropped_garment, cropped_garment_mask, cropped_human, cropped_human_mask)
|
||||
|
||||
Reference in New Issue
Block a user