nodes/FuouM-ComfyUI-FirstOrderMM
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

reformat

Browse Source
This commit is contained in:
FuouM
2024-07-01 14:45:00 +07:00
parent a6143c8858
commit ee622c22dc
6 changed files with 62 additions and 98 deletions
Download Patch File Download Diff File Expand all files Collapse all files
face_parsing/face_parsing_loader.py
+14 -7
View File
@@ -3,23 +3,30 @@ import numpy as np
from .face_parsing_model import BiSeNet
def load_face_parser_model(base_dir: str, face_parser_checkpoint_name: str, cpu=False):
print("Loading face_parser model")
face_parser = BiSeNet(n_classes=19)
face_parser_path = f"{base_dir}/face_parsing/{face_parser_checkpoint_name}"
mean = torch.Tensor(np.array([0.485, 0.456, 0.406], dtype=np.float32)).view(1, 3, 1, 1)
std = torch.Tensor(np.array([0.229, 0.224, 0.225], dtype=np.float32)).view(1, 3, 1, 1)
mean = torch.Tensor(np.array([0.485, 0.456, 0.406], dtype=np.float32)).view(
1, 3, 1, 1
)
std = torch.Tensor(np.array([0.229, 0.224, 0.225], dtype=np.float32)).view(
1, 3, 1, 1
)
if not cpu:
face_parser.cuda()
face_parser.load_state_dict(torch.load(face_parser_path))
face_parser.mean = mean.cuda()
face_parser.std = std.cuda()
else:
face_parser.load_state_dict(torch.load(face_parser_path, map_location=torch.device('cpu')))
face_parser.load_state_dict(
torch.load(face_parser_path, map_location=torch.device("cpu"))
)
face_parser.mean = mean
face_parser.std = std
print("face_parser model loaded!")
return face_parser
return face_parser
face_parsing/face_parsing_model.py
+1
View File
@@ -4,6 +4,7 @@ import torch.nn.functional as F
from .resnet import Resnet18
class ConvBNReLU(nn.Module):
def __init__(self, in_chan, out_chan, ks=3, stride=1, padding=1, *args, **kwargs):
super(ConvBNReLU, self).__init__()
face_parsing/resnet.py
-3
View File
@@ -1,11 +1,9 @@
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.model_zoo as modelzoo
import torch
from pathlib import Path
def conv3x3(in_planes, out_planes, stride=1):
"""3x3 convolution with padding"""
return nn.Conv2d(
@@ -50,7 +48,6 @@ def create_layer_basic(in_chan, out_chan, bnum, stride=1):
return nn.Sequential(*layers)
class Resnet18(nn.Module):
def __init__(self):
super(Resnet18, self).__init__()
partswap_inference.py
+1
View File
@@ -17,6 +17,7 @@ def face_parse_seg(
seg = F.interpolate(seg, size=out_size)
return seg
def partswap_inference(
swap_indices: list[int],
source_image,
run.py
+13 -9
View File
@@ -13,14 +13,15 @@ from .constants import (
config_folder,
default_model_name,
default_partswap_model_name,
face_parser_checkpoint_name,
partswap_face_parser_enforce_models,
partswap_fomm_model_names,
partswap_model_config_dict,
partswap_model_length_dict,
partswap_model_names,
supported_model_names,
face_parser_checkpoint_name,
partswap_face_parser_enforce_models
)
from .face_parsing.face_parsing_loader import load_face_parser_model
from .fomm_inference import inference, inference_best_frame
from .fomm_loader import load_checkpoint
from .partswap_inference import partswap_inference
@@ -34,7 +35,6 @@ from .utils import (
out_video,
reshape_image,
)
from .face_parsing.face_parsing_loader import load_face_parser_model
base_dir = Path(__file__).resolve().parent
@@ -225,7 +225,7 @@ class FOMM_Partswap:
"use_face_parser": (
"BOOLEAN",
{"default": False},
),
),
"chosen_seg_indices": ("STRING", {}),
"viz_alpha": (
"FLOAT",
@@ -289,18 +289,22 @@ class FOMM_Partswap:
reconstruction_module, segmentation_module = load_partswap_checkpoint(
config_path, checkpoint_path, blend_scale, use_fomm=use_fomm
)
face_parser_model = None
if model_name in partswap_face_parser_enforce_models:
assert use_face_parser, "Special supervised model, requires face_parser"
if use_face_parser:
face_parser_model = load_face_parser_model(base_dir, face_parser_checkpoint_name)
face_parser_model = load_face_parser_model(
base_dir, face_parser_checkpoint_name
)
face_parser_model.cuda()
face_parser_model.eval()
print(f"Chosen {chosen_seg_indices}")
swap_indices = deserialize_integers(chosen_seg_indices)
swap_indices = [x for x in swap_indices if x < partswap_model_length_dict[model_name]]
swap_indices = [
x for x in swap_indices if x < partswap_model_length_dict[model_name]
]
print(f"Swapping {swap_indices}")
params = {
@@ -324,7 +328,7 @@ class FOMM_Partswap:
output_images = out_video(predictions)
del face_parser_model
return (
seg_src_viz,
seg_tgt_viz,
seg_viz.py
+33 -79
View File
@@ -6,129 +6,83 @@ import numpy as np
import torch
import torch.nn.functional as F
def preprocess_image_and_mask(image, seg_source):
image = image.squeeze(0).cpu() # Remove batch dimension
mask = seg_source["segmentation"]
if mask.shape[2:] != image.shape[1:]:
mask = F.interpolate(mask, size=image.shape[1:], mode='bilinear')
mask = F.interpolate(mask, size=image.shape[1:], mode="bilinear")
return image, mask
def apply_hard_edges(mask):
return (torch.max(mask, dim=1, keepdim=True)[0] == mask).float()
def create_color_mask(mask, colormap='gist_rainbow'):
def create_color_mask(mask, colormap="gist_rainbow"):
colormap = plt.get_cmap(colormap)
num_segments = mask.shape[1]
mask = mask.squeeze(0).permute(1, 2, 0).cpu().numpy()
color_mask = np.zeros((256, 256, 3), dtype=np.float32)
patches = []
for i in range(num_segments):
color = np.array(colormap((i - 1) / (num_segments - 1)))[:3] if i != 0 else np.array((0, 0, 0))
color = (
np.array(colormap((i - 1) / (num_segments - 1)))[:3]
if i != 0
else np.array((0, 0, 0))
)
patches.append(mpatches.Patch(color=color, label=str(i)))
color_mask += mask[..., i:(i+1)] * color.reshape(1, 1, 3)
color_mask += mask[..., i : (i + 1)] * color.reshape(1, 1, 3)
return color_mask, patches
def overlay_mask(image, color_mask, alpha=0.6):
image_np = image.permute(1, 2, 0).numpy()
if image_np.max() > 1:
image_np = image_np / 255.0
overlaid_image = image_np * (1 - alpha) + color_mask * alpha
return np.clip(overlaid_image, 0, 1)
def create_figure(overlaid_image, patches):
fig = Figure(figsize=(10, 10))
canvas = FigureCanvasAgg(fig)
ax = fig.add_subplot(111)
ax.imshow(overlaid_image)
ax.axis('off')
ax.legend(handles=patches, loc='upper right', bbox_to_anchor=(1, 1), fontsize='small')
ax.axis("off")
ax.legend(
handles=patches, loc="upper right", bbox_to_anchor=(1, 1), fontsize="small"
)
fig.tight_layout(pad=0)
return fig, canvas
def render_figure_to_tensor(canvas):
canvas.draw()
result_image = np.frombuffer(canvas.tostring_rgb(), dtype='uint8')
result_image = np.frombuffer(canvas.tostring_rgb(), dtype="uint8")
result_image = result_image.reshape(canvas.get_width_height()[::-1] + (3,))
result_tensor = torch.from_numpy(result_image).unsqueeze(0).float() / 255.0
return result_tensor
def visualize_frame(image, seg_source, hard_edges=True, colormap='gist_rainbow', alpha=0.6):
def visualize_frame(
image, seg_source, hard_edges=True, colormap="gist_rainbow", alpha=0.6
):
image, mask = preprocess_image_and_mask(image, seg_source)
if hard_edges:
mask = apply_hard_edges(mask)
color_mask, patches = create_color_mask(mask, colormap)
overlaid_image = overlay_mask(image, color_mask, alpha)
fig, canvas = create_figure(overlaid_image, patches)
result_tensor = render_figure_to_tensor(canvas)
return result_tensor
# def visualize_frame(image, seg_source, hard_edges=True, colormap='gist_rainbow', alpha=0.6):
# # Ensure image is the right shape and on CPU
# image = image.squeeze(0).cpu() # Remove batch dimension
# # Extract segmentation from seg_source dictionary
# mask = seg_source["segmentation"]
# # Interpolate mask to match image size if necessary
# if mask.shape[2:] != image.shape[1:]:
# mask = F.interpolate(mask, size=image.shape[1:], mode='bilinear')
# if hard_edges:
# mask = (torch.max(mask, dim=1, keepdim=True)[0] == mask).float()
# colormap = plt.get_cmap(colormap)
# num_segments = mask.shape[1]
# mask = mask.squeeze(0).permute(1, 2, 0).cpu().numpy()
# color_mask = np.zeros((256, 256, 3), dtype=np.float32)
# patches = []
# for i in range(num_segments):
# if i != 0:
# color = np.array(colormap((i - 1) / (num_segments - 1)))[:3]
# else:
# color = np.array((0, 0, 0))
# patches.append(mpatches.Patch(color=color, label=str(i)))
# color_mask += mask[..., i:(i+1)] * color.reshape(1, 1, 3)
# # Convert image to numpy and normalize to [0, 1] range if necessary
# image_np = image.permute(1, 2, 0).numpy()
# if image_np.max() > 1:
# image_np = image_np / 255.0
# # Overlay the color mask on the original image
# overlaid_image = image_np * (1 - alpha) + color_mask * alpha
# # Ensure all values are in [0, 1] range
# overlaid_image = np.clip(overlaid_image, 0, 1)
# # Create a figure and axis
# fig = Figure(figsize=(10, 10))
# canvas = FigureCanvasAgg(fig)
# ax = fig.add_subplot(111)
# # Display the image
# ax.imshow(overlaid_image)
# ax.axis('off')
# # Add legend
# ax.legend(handles=patches, loc='upper right', bbox_to_anchor=(1, 1), fontsize='small')
# # Render the figure to numpy array
# fig.tight_layout(pad=0)
# canvas.draw()
# result_image = np.frombuffer(canvas.tostring_rgb(), dtype='uint8')
# result_image = result_image.reshape(canvas.get_width_height()[::-1] + (3,))
# # Convert to torch tensor and add batch dimension to get [1, 256, 256, 3]
# result_tensor = torch.from_numpy(result_image).unsqueeze(0).float() / 255.0
# return result_tensor
return result_tensor