diff --git a/face_parsing/face_parsing_loader.py b/face_parsing/face_parsing_loader.py index 7b8118c..fbb64c1 100644 --- a/face_parsing/face_parsing_loader.py +++ b/face_parsing/face_parsing_loader.py @@ -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 \ No newline at end of file + + return face_parser diff --git a/face_parsing/face_parsing_model.py b/face_parsing/face_parsing_model.py index 718ea14..198f949 100644 --- a/face_parsing/face_parsing_model.py +++ b/face_parsing/face_parsing_model.py @@ -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__() diff --git a/face_parsing/resnet.py b/face_parsing/resnet.py index bd30bfd..ba9002d 100644 --- a/face_parsing/resnet.py +++ b/face_parsing/resnet.py @@ -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__() diff --git a/partswap_inference.py b/partswap_inference.py index 15a5c2c..b062e1a 100644 --- a/partswap_inference.py +++ b/partswap_inference.py @@ -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, diff --git a/run.py b/run.py index 3775eb5..4fdde2d 100644 --- a/run.py +++ b/run.py @@ -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, diff --git a/seg_viz.py b/seg_viz.py index d78663c..953d998 100644 --- a/seg_viz.py +++ b/seg_viz.py @@ -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 \ No newline at end of file + return result_tensor