diff --git a/.gitignore b/.gitignore index 61735cc..890e731 100644 --- a/.gitignore +++ b/.gitignore @@ -2,6 +2,7 @@ pretrained character src dst +.idea # Byte-compiled / optimized / DLL files __pycache__/ diff --git a/tha2/__init__.py b/tha2/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha2/app/__init__.py b/tha2/app/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha2/app/ifacialmocap_puppeteer.py b/tha2/app/ifacialmocap_puppeteer.py new file mode 100644 index 0000000..6e81938 --- /dev/null +++ b/tha2/app/ifacialmocap_puppeteer.py @@ -0,0 +1,461 @@ +import errno +import json +import os +import socket +import sys +import threading +import time + +sys.path.append(os.getcwd()) + +import numpy +import torch +import wx + +from tha2.poser.poser import Poser +from tha2.mocap.ifacialmocap_constants import * +from tha2.mocap.ifacialmocap_pose_converter import IFacialMocapPoseConverter +from tha2.util import extract_PIL_image_from_filelike, resize_PIL_image, extract_pytorch_image_from_PIL_image, convert_output_image_from_torch_to_numpy + + +class CaptureData: + def __init__(self): + self.lock = threading.Lock() + self.data = self.create_default_data() + + def write_data(self, data): + self.lock.acquire() + self.data = data + self.lock.release() + + def read_data(self): + self.lock.acquire() + output = self.data + self.lock.release() + return output + + @staticmethod + def create_default_data(): + data = {} + + for blendshape_name in BLENDSHAPE_NAMES: + data[blendshape_name] = 0.0 + + data[HEAD_BONE_X] = 0.0 + data[HEAD_BONE_Y] = 0.0 + data[HEAD_BONE_Z] = 0.0 + data[HEAD_BONE_QUAT] = [0.0, 0.0, 0.0, 1.0] + + data[LEFT_EYE_BONE_X] = 0.0 + data[LEFT_EYE_BONE_Y] = 0.0 + data[LEFT_EYE_BONE_Z] = 0.0 + data[LEFT_EYE_BONE_QUAT] = [0.0, 0.0, 0.0, 1.0] + + data[RIGHT_EYE_BONE_X] = 0.0 + data[RIGHT_EYE_BONE_Y] = 0.0 + data[RIGHT_EYE_BONE_Z] = 0.0 + data[RIGHT_EYE_BONE_QUAT] = [0.0, 0.0, 0.0, 1.0] + + return data + + +class ClientThread(threading.Thread): + def __init__(self, capture_data: CaptureData): + super().__init__() + self.capture_data = capture_data + self.should_terminate = False + self.address = "0.0.0.0" + self.port = 50002 + + def run(self): + self.socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + self.socket.setblocking(False) + self.socket.bind((self.address, self.port)) + while True: + if self.should_terminate: + break + try: + socket_bytes = self.socket.recv(8192) + except socket.error as e: + err = e.args[0] + if err == errno.EAGAIN or err == errno.EWOULDBLOCK: + continue + else: + raise e + socket_string = socket_bytes.decode("utf-8") + blender_data = json.loads(socket_string) + data = self.convert_from_blender_data(blender_data) + self.capture_data.write_data(data) + self.socket.close() + + @staticmethod + def convert_from_blender_data(blender_data): + data = {} + + shape_key_values = blender_data["shapeKeyValues"] + for blendshape_name in BLENDSHAPE_NAMES: + data[blendshape_name] = shape_key_values[blendshape_name] + + head_bone = blender_data["boneValues"]["armatures"]["bones"]["headBone"] + data[HEAD_BONE_X] = head_bone["e_rx"] + data[HEAD_BONE_Y] = head_bone["e_ry"] + data[HEAD_BONE_Z] = head_bone["e_rz"] + data[HEAD_BONE_QUAT] = [ + head_bone["q_rx"], + head_bone["q_ry"], + head_bone["q_rz"], + head_bone["q_w"], + ] + + right_eye_bone = blender_data["boneValues"]["armatures"]["bones"]["rightEyeBone"] + data[RIGHT_EYE_BONE_X] = right_eye_bone["e_rx"] + data[RIGHT_EYE_BONE_Y] = right_eye_bone["e_ry"] + data[RIGHT_EYE_BONE_Z] = right_eye_bone["e_rz"] + data[RIGHT_EYE_BONE_QUAT] = [ + right_eye_bone["q_rx"], + right_eye_bone["q_ry"], + right_eye_bone["q_rz"], + right_eye_bone["q_w"], + ] + + left_eye_bone = blender_data["boneValues"]["armatures"]["bones"]["leftEyeBone"] + data[LEFT_EYE_BONE_X] = left_eye_bone["e_rx"] + data[LEFT_EYE_BONE_Y] = left_eye_bone["e_ry"] + data[LEFT_EYE_BONE_Z] = left_eye_bone["e_rz"] + data[LEFT_EYE_BONE_QUAT] = [ + left_eye_bone["q_rx"], + left_eye_bone["q_ry"], + left_eye_bone["q_rz"], + left_eye_bone["q_w"], + ] + + return data + + +class MainFrame(wx.Frame): + def __init__(self, poser: Poser, pose_converter: IFacialMocapPoseConverter, device: torch.device): + super().__init__(None, wx.ID_ANY, "iFacialMocap Puppeteer") + self.pose_converter = pose_converter + self.poser = poser + self.device = device + self.capture_data = CaptureData() + self.client_thread = ClientThread(self.capture_data) + self.create_ui() + self.Bind(wx.EVT_CLOSE, self.on_close) + + self.capture_timer = wx.Timer(self, wx.ID_ANY) + self.Bind(wx.EVT_TIMER, self.update_capture_panel, id=self.capture_timer.GetId()) + + self.animation_timer = wx.Timer(self, wx.ID_ANY) + self.Bind(wx.EVT_TIMER, self.update_result_image_panel, id=self.animation_timer.GetId()) + + self.wx_source_image = None + self.torch_source_image = None + self.last_pose = None + + self.client_thread.start() + self.capture_timer.Start(33) + self.animation_timer.Start(33) + + self.source_image_string = "Nothing yet!" + + def on_close(self, event: wx.Event): + self.client_thread.should_terminate = True + self.client_thread.join() + self.capture_timer.Stop() + self.animation_timer.Stop() + event.Skip() + + def create_animation_panel(self, parent): + self.animation_panel = wx.Panel(parent, style=wx.RAISED_BORDER) + self.animation_panel_sizer = wx.BoxSizer(wx.HORIZONTAL) + self.animation_panel.SetSizer(self.animation_panel_sizer) + self.animation_panel.SetAutoLayout(1) + + if True: + self.input_panel = wx.Panel(self.animation_panel, size=(256, 368), style=wx.SIMPLE_BORDER) + self.input_panel_sizer = wx.BoxSizer(wx.VERTICAL) + self.input_panel.SetSizer(self.input_panel_sizer) + self.input_panel.SetAutoLayout(1) + self.animation_panel_sizer.Add(self.input_panel, 0, wx.FIXED_MINSIZE) + + self.source_image_panel = wx.Panel(self.input_panel, size=(256, 256), style=wx.SIMPLE_BORDER) + self.source_image_panel.Bind(wx.EVT_PAINT, self.paint_source_image_panel) + self.input_panel_sizer.Add(self.source_image_panel, 0, wx.FIXED_MINSIZE) + + self.load_image_button = wx.Button(self.input_panel, wx.ID_ANY, "Load Image") + self.input_panel_sizer.Add(self.load_image_button, 1, wx.EXPAND) + self.load_image_button.Bind(wx.EVT_BUTTON, self.load_image) + + self.input_panel_sizer.Fit(self.input_panel) + + if True: + self.pose_converter.init_pose_converter_panel(self.animation_panel) + + if True: + self.animation_left_panel = wx.Panel(self.animation_panel, style=wx.SIMPLE_BORDER) + self.animation_left_panel_sizer = wx.BoxSizer(wx.VERTICAL) + self.animation_left_panel.SetSizer(self.animation_left_panel_sizer) + self.animation_left_panel.SetAutoLayout(1) + self.animation_panel_sizer.Add(self.animation_left_panel, 0, wx.EXPAND) + + self.result_image_panel = wx.Panel(self.animation_left_panel, size=(256, 256), style=wx.SIMPLE_BORDER) + self.result_image_panel.Bind(wx.EVT_PAINT, self.paint_result_image_panel) + self.animation_left_panel_sizer.Add(self.result_image_panel, 0, wx.FIXED_MINSIZE) + + self.output_index_choice = wx.Choice(self.animation_left_panel, + choices=[str(i) for i in range(self.poser.get_output_length())]) + self.output_index_choice.SetSelection(0) + self.animation_left_panel_sizer.Add(self.output_index_choice, 0, wx.EXPAND) + + separator = wx.StaticLine(self.animation_left_panel, -1, size=(256, 5)) + self.animation_left_panel_sizer.Add(separator, 0, wx.EXPAND) + + background_text = wx.StaticText(self.animation_left_panel, label="--- Background ---", + style=wx.ALIGN_CENTER) + self.animation_left_panel_sizer.Add(background_text, 0, wx.EXPAND) + + self.output_background_choice = wx.Choice( + self.animation_left_panel, + choices=[ + "TRANSPARENT", + "GREEN", + "BLUE", + "BLACK", + "WHITE" + ]) + self.output_background_choice.SetSelection(0) + self.animation_left_panel_sizer.Add(self.output_background_choice, 0, wx.EXPAND) + + separator = wx.StaticLine(self.animation_left_panel, -1, size=(256, 5)) + self.animation_left_panel_sizer.Add(separator, 0, wx.EXPAND) + + self.fps_text = wx.StaticText(self.animation_left_panel, label="") + self.animation_left_panel_sizer.Add(self.fps_text, wx.SizerFlags().Border()) + + self.animation_left_panel_sizer.Fit(self.animation_left_panel) + + self.animation_panel_sizer.Fit(self.animation_panel) + + def create_ui(self): + self.main_sizer = wx.BoxSizer(wx.VERTICAL) + self.SetSizer(self.main_sizer) + self.SetAutoLayout(1) + + self.capture_pose_lock = threading.Lock() + + self.create_animation_panel(self) + self.main_sizer.Add(self.animation_panel, wx.SizerFlags(0).Expand().Border(wx.ALL, 5)) + + self.create_capture_panel(self) + self.main_sizer.Add(self.capture_panel, wx.SizerFlags(0).Expand().Border(wx.ALL, 5)) + + self.main_sizer.Fit(self) + + def create_capture_panel(self, parent): + self.capture_panel = wx.Panel(parent, style=wx.RAISED_BORDER) + self.capture_panel_sizer = wx.FlexGridSizer(cols=5) + for i in range(5): + self.capture_panel_sizer.AddGrowableCol(i) + self.capture_panel.SetSizer(self.capture_panel_sizer) + self.capture_panel.SetAutoLayout(1) + + self.blendshape_labels = {} + self.blendshape_value_labels = {} + self.blendshape_gauges = {} + blendshape_column_0 = self.create_blendshapes_column(self.capture_panel, COLUMN_0_BLENDSHAPES) + self.capture_panel_sizer.Add(blendshape_column_0, wx.SizerFlags(0).Expand().Border(wx.ALL, 3)) + blendshape_column_1 = self.create_blendshapes_column(self.capture_panel, COLUMN_1_BLENDSHAPES) + self.capture_panel_sizer.Add(blendshape_column_1, wx.SizerFlags(0).Expand().Border(wx.ALL, 3)) + blendshape_column_2 = self.create_blendshapes_column(self.capture_panel, COLUMN_2_BLENDSHAPES) + self.capture_panel_sizer.Add(blendshape_column_2, wx.SizerFlags(0).Expand().Border(wx.ALL, 3)) + blendshape_column_3 = self.create_blendshapes_column(self.capture_panel, COLUMN_3_BLENDSHAPES) + self.capture_panel_sizer.Add(blendshape_column_3, wx.SizerFlags(0).Expand().Border(wx.ALL, 3)) + blendshape_column_4 = self.create_blendshapes_column(self.capture_panel, COLUMN_4_BLENDSHAPES) + self.capture_panel_sizer.Add(blendshape_column_4, wx.SizerFlags(0).Expand().Border(wx.ALL, 3)) + + self.rotation_labels = {} + self.rotation_value_labels = {} + rotation_column_0 = self.create_rotation_column(self.capture_panel, RIGHT_EYE_BONE_ROTATIONS) + self.capture_panel_sizer.Add(rotation_column_0, wx.SizerFlags(0).Expand().Border(wx.ALL, 3)) + rotation_column_1 = self.create_rotation_column(self.capture_panel, LEFT_EYE_BONE_ROTATIONS) + self.capture_panel_sizer.Add(rotation_column_1, wx.SizerFlags(0).Expand().Border(wx.ALL, 3)) + rotation_column_2 = self.create_rotation_column(self.capture_panel, HEAD_BONE_ROTATIONS) + self.capture_panel_sizer.Add(rotation_column_2, wx.SizerFlags(0).Expand().Border(wx.ALL, 3)) + + def create_blendshapes_column(self, parent, blendshape_names): + column_panel = wx.Panel(parent, style=wx.SIMPLE_BORDER) + column_panel_sizer = wx.FlexGridSizer(cols=3) + column_panel_sizer.AddGrowableCol(1) + column_panel.SetSizer(column_panel_sizer) + column_panel.SetAutoLayout(1) + + for blendshape_name in blendshape_names: + self.blendshape_labels[blendshape_name] = wx.StaticText( + column_panel, label=blendshape_name, style=wx.ALIGN_RIGHT) + column_panel_sizer.Add(self.blendshape_labels[blendshape_name], + wx.SizerFlags(1).Expand().Border(wx.ALL, 3)) + + self.blendshape_gauges[blendshape_name] = wx.Gauge( + column_panel, style=wx.GA_HORIZONTAL, size=(100, -1)) + column_panel_sizer.Add(self.blendshape_gauges[blendshape_name], wx.SizerFlags(1).Expand().Border(wx.ALL, 3)) + + self.blendshape_value_labels[blendshape_name] = wx.TextCtrl( + column_panel, style=wx.TE_RIGHT, size=(40, -1)) + self.blendshape_value_labels[blendshape_name].SetValue("0.00") + self.blendshape_value_labels[blendshape_name].Disable() + column_panel_sizer.Add(self.blendshape_value_labels[blendshape_name], + wx.SizerFlags(0).Border(wx.ALL, 3)) + + column_panel.GetSizer().Fit(column_panel) + return column_panel + + def create_rotation_column(self, parent, rotation_names): + column_panel = wx.Panel(parent, style=wx.SIMPLE_BORDER) + column_panel_sizer = wx.FlexGridSizer(cols=2) + column_panel_sizer.AddGrowableCol(1) + column_panel.SetSizer(column_panel_sizer) + column_panel.SetAutoLayout(1) + + for rotation_name in rotation_names: + self.rotation_labels[rotation_name] = wx.StaticText( + column_panel, label=rotation_name, style=wx.ALIGN_RIGHT) + column_panel_sizer.Add(self.rotation_labels[rotation_name], + wx.SizerFlags(1).Expand().Border(wx.ALL, 3)) + + self.rotation_value_labels[rotation_name] = wx.TextCtrl( + column_panel, style=wx.TE_RIGHT) + self.rotation_value_labels[rotation_name].SetValue("0.00") + self.rotation_value_labels[rotation_name].Disable() + column_panel_sizer.Add(self.rotation_value_labels[rotation_name], + wx.SizerFlags(1).Expand().Border(wx.ALL, 3)) + + column_panel.GetSizer().Fit(column_panel) + return column_panel + + def paint_capture_panel(self, event: wx.Event): + self.update_capture_panel(event) + + def update_capture_panel(self, event: wx.Event): + data = self.capture_data.read_data() + for blendshape_name in BLENDSHAPE_NAMES: + value = data[blendshape_name] + self.blendshape_gauges[blendshape_name].SetValue(MainFrame.convert_to_100(value)) + self.blendshape_value_labels[blendshape_name].SetValue("%0.2f" % value) + for rotation_name in ROTATION_NAMES: + value = data[rotation_name] + self.rotation_value_labels[rotation_name].SetValue("%0.2f" % value) + + @staticmethod + def convert_to_100(x): + return int(max(0.0, min(1.0, x)) * 100) + + def paint_source_image_panel(self, event: wx.Event): + if self.wx_source_image is None: + self.draw_source_image_string(self.source_image_panel, use_paint_dc=True) + else: + dc = wx.PaintDC(self.source_image_panel) + dc.Clear() + dc.DrawBitmap(self.wx_source_image, 0, 0, True) + + def draw_source_image_string(self, widget, use_paint_dc: bool = True): + if use_paint_dc: + dc = wx.PaintDC(widget) + else: + dc = wx.ClientDC(widget) + dc.Clear() + font = wx.Font(wx.FontInfo(14).Family(wx.FONTFAMILY_SWISS)) + dc.SetFont(font) + w, h = dc.GetTextExtent(self.source_image_string) + dc.DrawText(self.source_image_string, 128 - w // 2, 128 - h // 2) + + def paint_result_image_panel(self, event: wx.Event): + self.last_pose = None + + def update_result_image_panel(self, event: wx.Event): + tic = time.perf_counter() + + ifacialmocap_pose = self.capture_data.read_data() + current_pose = self.pose_converter.convert(ifacialmocap_pose) + if self.last_pose is not None \ + and self.last_pose == current_pose \ + and self.last_output_index == self.output_index_choice.GetSelection(): + return + self.last_pose = current_pose + self.last_output_index = self.output_index_choice.GetSelection() + + if self.torch_source_image is None: + self.draw_source_image_string(self.result_image_panel, use_paint_dc=False) + return + + pose = torch.tensor(current_pose, device=self.device) + output_index = self.output_index_choice.GetSelection() + output_image = self.poser.pose(self.torch_source_image, pose, output_index)[0].detach().cpu() + numpy_image = convert_output_image_from_torch_to_numpy(output_image) + + background_choice = self.output_background_choice.GetSelection() + if background_choice != 0: + background = numpy.zeros((numpy_image.shape[0], numpy_image.shape[1], numpy_image.shape[2])) + background[:, :, 3] = 1.0 + if background_choice == 1: + background[:, :, 1] = 1.0 + elif background_choice == 2: + background[:, :, 2] = 1.0 + elif background_choice == 4: + background[:, :, 0:3] = 1.0 + numpy_image = self.blend_with_background(numpy_image, background) + + numpy_image = numpy.uint8(numpy.rint(numpy_image * 255.0)) + wx_image = wx.ImageFromBuffer( + numpy_image.shape[0], + numpy_image.shape[1], + numpy_image[:, :, 0:3].tobytes(), + numpy_image[:, :, 3].tobytes()) + wx_bitmap = wx_image.ConvertToBitmap() + + dc = wx.ClientDC(self.result_image_panel) + dc.Clear() + dc.DrawBitmap(wx_bitmap, (256 - numpy_image.shape[0]) // 2, (256 - numpy_image.shape[1]) // 2, True) + + toc = time.perf_counter() + elapsed_time = toc - tic + fps = min(1.0 / elapsed_time, 1000.0 / 33.0) + self.fps_text.SetLabelText("FPS = %0.2f" % fps) + + def blend_with_background(self, numpy_image, background): + alpha = numpy_image[:, :, 3:4] + color = numpy_image[:, :, 0:3] + new_color = color * alpha + (1.0 - alpha) * background[:, :, 0:3] + return numpy.concatenate([new_color, background[:, :, 3:4]], axis=2) + + def load_image(self, event: wx.Event): + dir_name = "data/illust" + file_dialog = wx.FileDialog(self, "Choose an image", dir_name, "", "*.png", wx.FD_OPEN) + if file_dialog.ShowModal() == wx.ID_OK: + image_file_name = os.path.join(file_dialog.GetDirectory(), file_dialog.GetFilename()) + pil_image = resize_PIL_image(extract_PIL_image_from_filelike(image_file_name)) + w, h = pil_image.size + if pil_image.mode != 'RGBA': + self.source_image_string = "Image must have alpha channel!" + self.wx_source_image = None + self.torch_source_image = None + else: + self.wx_source_image = wx.Bitmap.FromBufferRGBA(w, h, pil_image.convert("RGBA").tobytes()) + self.torch_source_image = extract_pytorch_image_from_PIL_image(pil_image).to(self.device) + + self.Refresh() + file_dialog.Destroy() + + +if __name__ == "__main__": + import tha2.poser.modes.mode_20 + import tha2.poser.modes.mode_20_wx + + cuda = torch.device('cuda') + poser = tha2.poser.modes.mode_20.create_poser(cuda) + pose_converter = tha2.poser.modes.mode_20_wx.create_ifacialmocap_pose_converter() + + app = wx.App() + main_frame = MainFrame(poser, pose_converter, cuda) + main_frame.Show(True) + app.MainLoop() diff --git a/tha2/app/manual_poser.py b/tha2/app/manual_poser.py new file mode 100644 index 0000000..6064471 --- /dev/null +++ b/tha2/app/manual_poser.py @@ -0,0 +1,379 @@ +import logging +import os +import sys +from typing import List + +sys.path.append(os.getcwd()) + +import numpy +import torch +import wx +import PIL.Image + +from tha2.poser.poser import Poser, PoseParameterCategory, PoseParameterGroup +from tha2.util import extract_PIL_image_from_filelike, resize_PIL_image, extract_pytorch_image_from_PIL_image, convert_output_image_from_torch_to_numpy + + +class MorphCategoryControlPanel(wx.Panel): + def __init__(self, + parent, + title: str, + pose_param_category: PoseParameterCategory, + param_groups: List[PoseParameterGroup]): + super().__init__(parent, style=wx.SIMPLE_BORDER) + self.pose_param_category = pose_param_category + self.sizer = wx.BoxSizer(wx.VERTICAL) + self.SetSizer(self.sizer) + self.SetAutoLayout(1) + + title_text = wx.StaticText(self, label=title, style=wx.ALIGN_CENTER) + self.sizer.Add(title_text, 0, wx.EXPAND) + + self.param_groups = [group for group in param_groups if group.get_category() == pose_param_category] + self.choice = wx.Choice(self, choices=[group.get_group_name() for group in self.param_groups]) + if len(self.param_groups) > 0: + self.choice.SetSelection(0) + self.choice.Bind(wx.EVT_CHOICE, self.on_choice_updated) + self.sizer.Add(self.choice, 0, wx.EXPAND) + + self.left_slider = wx.Slider(self, minValue=-1000, maxValue=1000, value=-1000, style=wx.HORIZONTAL) + self.sizer.Add(self.left_slider, 0, wx.EXPAND) + + self.right_slider = wx.Slider(self, minValue=-1000, maxValue=1000, value=-1000, style=wx.HORIZONTAL) + self.sizer.Add(self.right_slider, 0, wx.EXPAND) + + self.checkbox = wx.CheckBox(self, label="Show") + self.checkbox.SetValue(True) + self.sizer.Add(self.checkbox, 0, wx.SHAPED | wx.ALIGN_CENTER) + + self.update_ui() + + self.sizer.Fit(self) + + def update_ui(self): + param_group = self.param_groups[self.choice.GetSelection()] + if param_group.is_discrete(): + self.left_slider.Enable(False) + self.right_slider.Enable(False) + self.checkbox.Enable(True) + elif param_group.get_arity() == 1: + self.left_slider.Enable(True) + self.right_slider.Enable(False) + self.checkbox.Enable(False) + else: + self.left_slider.Enable(True) + self.right_slider.Enable(True) + self.checkbox.Enable(False) + + def on_choice_updated(self, event: wx.Event): + param_group = self.param_groups[self.choice.GetSelection()] + if param_group.is_discrete(): + self.checkbox.SetValue(True) + self.update_ui() + + def set_param_value(self, pose: List[float]): + if len(self.param_groups) == 0: + return + selected_morph_index = self.choice.GetSelection() + param_group = self.param_groups[selected_morph_index] + param_index = param_group.get_parameter_index() + if param_group.is_discrete(): + if self.checkbox.GetValue(): + for i in range(param_group.get_arity()): + pose[param_index + i] = 1.0 + else: + param_range = param_group.get_range() + alpha = (self.left_slider.GetValue() + 1000) / 2000.0 + pose[param_index] = param_range[0] + (param_range[1] - param_range[0]) * alpha + if param_group.get_arity() == 2: + alpha = (self.right_slider.GetValue() + 1000) / 2000.0 + pose[param_index + 1] = param_range[0] + (param_range[1] - param_range[0]) * alpha + + +class RotationControlPanel(wx.Panel): + def __init__(self, parent, + pose_param_category: PoseParameterCategory, + param_groups: List[PoseParameterGroup]): + super().__init__(parent, style=wx.SIMPLE_BORDER) + self.sizer = wx.BoxSizer(wx.VERTICAL) + self.SetSizer(self.sizer) + self.SetAutoLayout(1) + + self.param_groups = [group for group in param_groups if group.get_category() == pose_param_category] + for param_group in self.param_groups: + assert not param_group.is_discrete() + assert param_group.get_arity() == 1 + + self.sliders = [] + for param_group in self.param_groups: + static_text = wx.StaticText(self, label="--- %s ---" % param_group.get_group_name(), style=wx.ALIGN_CENTER) + self.sizer.Add(static_text, 0, wx.EXPAND) + slider = wx.Slider(self, minValue=-1000, maxValue=1000, value=0, style=wx.HORIZONTAL) + self.sizer.Add(slider, 0, wx.EXPAND) + self.sliders.append(slider) + + self.sizer.Fit(self) + + def set_param_value(self, pose: List[float]): + if len(self.param_groups) == 0: + return + for param_group_index in range(len(self.param_groups)): + param_group = self.param_groups[param_group_index] + slider = self.sliders[param_group_index] + param_range = param_group.get_range() + param_index = param_group.get_parameter_index() + alpha = (slider.GetValue() + 1000) / 2000.0 + pose[param_index] = param_range[0] + (param_range[1] - param_range[0]) * alpha + + +class MainFrame(wx.Frame): + def __init__(self, poser: Poser, device: torch.device): + super().__init__(None, wx.ID_ANY, "Poser") + self.poser = poser + self.device = device + + self.wx_source_image = None + self.torch_source_image = None + self.source_image_string = "Nothing yet!" + + self.main_sizer = wx.BoxSizer(wx.HORIZONTAL) + self.SetSizer(self.main_sizer) + self.SetAutoLayout(1) + self.init_left_panel() + self.init_control_panel() + self.init_right_panel() + self.main_sizer.Fit(self) + + self.timer = wx.Timer(self, wx.ID_ANY) + self.Bind(wx.EVT_TIMER, self.update_result_image_panel, self.timer) + + save_image_id = wx.NewIdRef() + self.Bind(wx.EVT_MENU, self.on_save_image, id=save_image_id) + accelerator_table = wx.AcceleratorTable([ + (wx.ACCEL_CTRL, ord('S'), save_image_id) + ]) + self.SetAcceleratorTable(accelerator_table) + + self.last_pose = None + self.last_output_index = self.output_index_choice.GetSelection() + self.last_output_numpy_image = None + + def init_left_panel(self): + self.control_panel = wx.Panel(self, style=wx.SIMPLE_BORDER, size=(256, -1)) + self.left_panel = wx.Panel(self, style=wx.SIMPLE_BORDER) + left_panel_sizer = wx.BoxSizer(wx.VERTICAL) + self.left_panel.SetSizer(left_panel_sizer) + self.left_panel.SetAutoLayout(1) + + self.source_image_panel = wx.Panel(self.left_panel, size=(256, 256), style=wx.SIMPLE_BORDER) + self.source_image_panel.Bind(wx.EVT_PAINT, self.paint_source_image_panel) + left_panel_sizer.Add(self.source_image_panel, 0, wx.FIXED_MINSIZE) + + self.load_image_button = wx.Button(self.left_panel, wx.ID_ANY, "\nLoad Image\n\n") + left_panel_sizer.Add(self.load_image_button, 1, wx.EXPAND) + self.load_image_button.Bind(wx.EVT_BUTTON, self.load_image) + + left_panel_sizer.Fit(self.left_panel) + self.main_sizer.Add(self.left_panel, 0, wx.FIXED_MINSIZE) + + def init_control_panel(self): + self.control_panel_sizer = wx.BoxSizer(wx.VERTICAL) + self.control_panel.SetSizer(self.control_panel_sizer) + + morph_categories = [ + PoseParameterCategory.EYEBROW, + PoseParameterCategory.EYE, + PoseParameterCategory.MOUTH, + PoseParameterCategory.IRIS_MORPH + ] + morph_category_titles = { + PoseParameterCategory.EYEBROW: "--- Eyebrow ---", + PoseParameterCategory.EYE: "--- Eye ---", + PoseParameterCategory.MOUTH: "--- Mouth ---", + PoseParameterCategory.IRIS_MORPH: "--- Iris morphs ---", + } + self.morph_control_panels = {} + for category in morph_categories: + param_groups = self.poser.get_pose_parameter_groups() + filtered_param_groups = [group for group in param_groups if group.get_category() == category] + if len(filtered_param_groups) == 0: + continue + control_panel = MorphCategoryControlPanel( + self.control_panel, + morph_category_titles[category], + category, + self.poser.get_pose_parameter_groups()) + self.morph_control_panels[category] = control_panel + self.control_panel_sizer.Add(control_panel, 0, wx.EXPAND) + + self.rotation_control_panels = {} + rotation_categories = [ + PoseParameterCategory.IRIS_ROTATION, + PoseParameterCategory.FACE_ROTATION + ] + for category in rotation_categories: + param_groups = self.poser.get_pose_parameter_groups() + filtered_param_groups = [group for group in param_groups if group.get_category() == category] + if len(filtered_param_groups) == 0: + continue + control_panel = RotationControlPanel( + self.control_panel, + category, + self.poser.get_pose_parameter_groups()) + self.rotation_control_panels[category] = control_panel + self.control_panel_sizer.Add(control_panel, 0, wx.EXPAND) + + self.control_panel_sizer.Fit(self.control_panel) + self.main_sizer.Add(self.control_panel, 1, wx.EXPAND) + + def init_right_panel(self): + self.right_panel = wx.Panel(self, style=wx.SIMPLE_BORDER) + right_panel_sizer = wx.BoxSizer(wx.VERTICAL) + self.right_panel.SetSizer(right_panel_sizer) + self.right_panel.SetAutoLayout(1) + + self.result_image_panel = wx.Panel(self.right_panel, size=(256, 256), style=wx.SIMPLE_BORDER) + self.result_image_panel.Bind(wx.EVT_PAINT, self.paint_result_image_panel) + self.output_index_choice = wx.Choice( + self.right_panel, + choices=[str(i) for i in range(self.poser.get_output_length())]) + self.output_index_choice.SetSelection(0) + right_panel_sizer.Add(self.result_image_panel, 0, wx.FIXED_MINSIZE) + right_panel_sizer.Add(self.output_index_choice, 0, wx.EXPAND) + + self.save_image_button = wx.Button(self.right_panel, wx.ID_ANY, "\nSave Image\n\n") + right_panel_sizer.Add(self.save_image_button, 1, wx.EXPAND) + self.save_image_button.Bind(wx.EVT_BUTTON, self.on_save_image) + + right_panel_sizer.Fit(self.right_panel) + self.main_sizer.Add(self.right_panel, 0, wx.FIXED_MINSIZE) + + def create_param_category_choice(self, param_category: PoseParameterCategory): + params = [] + for param_group in self.poser.get_pose_parameter_groups(): + if param_group.get_category() == param_category: + params.append(param_group.get_group_name()) + choice = wx.Choice(self.control_panel, choices=params) + if len(params) > 0: + choice.SetSelection(0) + return choice + + def load_image(self, event: wx.Event): + dir_name = "data/illust" + file_dialog = wx.FileDialog(self, "Choose an image", dir_name, "", "*.png", wx.FD_OPEN) + if file_dialog.ShowModal() == wx.ID_OK: + image_file_name = os.path.join(file_dialog.GetDirectory(), file_dialog.GetFilename()) + pil_image = resize_PIL_image(extract_PIL_image_from_filelike(image_file_name)) + w, h = pil_image.size + if pil_image.mode != 'RGBA': + self.source_image_string = "Image must have alpha channel!" + self.wx_source_image = None + self.torch_source_image = None + else: + self.wx_source_image = wx.Bitmap.FromBufferRGBA(w, h, pil_image.convert("RGBA").tobytes()) + self.torch_source_image = extract_pytorch_image_from_PIL_image(pil_image).to(self.device) + self.Refresh() + file_dialog.Destroy() + + def paint_source_image_panel(self, event: wx.Event): + if self.wx_source_image is None: + self.draw_source_image_string(self.source_image_panel, use_paint_dc=True) + else: + dc = wx.PaintDC(self.source_image_panel) + dc.Clear() + dc.DrawBitmap(self.wx_source_image, 0, 0, True) + + def paint_result_image_panel(self, event: wx.Event): + self.last_pose = None + + def draw_source_image_string(self, widget, use_paint_dc: bool = True): + if use_paint_dc: + dc = wx.PaintDC(widget) + else: + dc = wx.ClientDC(widget) + dc.Clear() + font = wx.Font(wx.FontInfo(14).Family(wx.FONTFAMILY_SWISS)) + dc.SetFont(font) + w, h = dc.GetTextExtent(self.source_image_string) + dc.DrawText(self.source_image_string, 128 - w // 2, 128 - h // 2) + + def get_current_pose(self): + current_pose = [0.0 for i in range(self.poser.get_num_parameters())] + for morph_control_panel in self.morph_control_panels.values(): + morph_control_panel.set_param_value(current_pose) + for rotation_control_panel in self.rotation_control_panels.values(): + rotation_control_panel.set_param_value(current_pose) + return current_pose + + def update_result_image_panel(self, event: wx.Event): + current_pose = self.get_current_pose() + if self.last_pose is not None \ + and self.last_pose == current_pose \ + and self.last_output_index == self.output_index_choice.GetSelection(): + return + self.last_pose = current_pose + self.last_output_index = self.output_index_choice.GetSelection() + + if self.torch_source_image is None: + self.draw_source_image_string(self.result_image_panel, use_paint_dc=False) + return + + pose = torch.tensor(current_pose, device=self.device) + output_index = self.output_index_choice.GetSelection() + output_image = self.poser.pose(self.torch_source_image, pose, output_index)[0].detach().cpu() + numpy_image = numpy.uint8(numpy.rint(convert_output_image_from_torch_to_numpy(output_image) * 255.0)) + self.last_output_numpy_image = numpy_image + wx_image = wx.ImageFromBuffer( + numpy_image.shape[0], + numpy_image.shape[1], + numpy_image[:, :, 0:3].tobytes(), + numpy_image[:, :, 3].tobytes()) + wx_bitmap = wx_image.ConvertToBitmap() + + dc = wx.ClientDC(self.result_image_panel) + dc.Clear() + dc.DrawBitmap(wx_bitmap, (256 - numpy_image.shape[0]) // 2, (256 - numpy_image.shape[1]) // 2, True) + + def on_save_image(self, event: wx.Event): + if self.last_output_numpy_image is None: + logging.info("There is no output image to save!!!") + return + + dir_name = "data/illust" + file_dialog = wx.FileDialog(self, "Save image", dir_name, "", "*.png", wx.FD_SAVE) + if file_dialog.ShowModal() == wx.ID_OK: + image_file_name = os.path.join(file_dialog.GetDirectory(), file_dialog.GetFilename()) + try: + if os.path.exists(image_file_name): + message_dialog = wx.MessageDialog(self, f"Override {image_file_name}", "Manual Poser", + wx.YES_NO | wx.ICON_QUESTION) + result = message_dialog.ShowModal() + if result == wx.ID_YES: + self.save_last_numpy_image(image_file_name) + message_dialog.Destroy() + else: + self.save_last_numpy_image(image_file_name) + except: + message_dialog = wx.MessageDialog(self, f"Could not save {image_file_name}", "Manual Poser", wx.OK) + message_dialog.ShowModal() + message_dialog.Destroy() + file_dialog.Destroy() + + def save_last_numpy_image(self, image_file_name): + numpy_image = self.last_output_numpy_image + pil_image = PIL.Image.fromarray(numpy_image, mode='RGBA') + os.makedirs(os.path.dirname(image_file_name), exist_ok=True) + pil_image.save(image_file_name) + + +if __name__ == "__main__": + cuda = torch.device('cuda') + import tha2.poser.modes.mode_20 + + poser = tha2.poser.modes.mode_20.create_poser(cuda) + + app = wx.App() + main_frame = MainFrame(poser, cuda) + main_frame.Show(True) + main_frame.timer.Start(30) + app.MainLoop() diff --git a/tha2/compute/__init__.py b/tha2/compute/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha2/compute/cached_computation_func.py b/tha2/compute/cached_computation_func.py new file mode 100644 index 0000000..8bcf874 --- /dev/null +++ b/tha2/compute/cached_computation_func.py @@ -0,0 +1,10 @@ +from typing import Callable, Dict, List + +from torch import Tensor + +from tha2.nn.batch_module.batch_input_module import BatchInputModule + +TensorCachedComputationFunc = Callable[ + [Dict[str, BatchInputModule], List[Tensor], Dict[str, List[Tensor]]], Tensor] +TensorListCachedComputationFunc = Callable[ + [Dict[str, BatchInputModule], List[Tensor], Dict[str, List[Tensor]]], List[Tensor]] diff --git a/tha2/compute/cached_computation_protocol.py b/tha2/compute/cached_computation_protocol.py new file mode 100644 index 0000000..db8360f --- /dev/null +++ b/tha2/compute/cached_computation_protocol.py @@ -0,0 +1,43 @@ +from abc import ABC, abstractmethod +from typing import Dict, List + +from torch import Tensor + +from tha2.nn.batch_module.batch_input_module import BatchInputModule +from tha2.compute.cached_computation_func import TensorCachedComputationFunc, TensorListCachedComputationFunc + + +class CachedComputationProtocol(ABC): + def get_output(self, + key: str, + modules: Dict[str, BatchInputModule], + batch: List[Tensor], + outputs: Dict[str, List[Tensor]]): + if key in outputs: + return outputs[key] + else: + output = self.compute_output(key, modules, batch, outputs) + outputs[key] = output + return outputs[key] + + @abstractmethod + def compute_output(self, + key: str, + modules: Dict[str, BatchInputModule], + batch: List[Tensor], + outputs: Dict[str, List[Tensor]]) -> List[Tensor]: + pass + + def get_output_tensor_func(self, key: str, index: int) -> TensorCachedComputationFunc: + def func(modules: Dict[str, BatchInputModule], + batch: List[Tensor], + outputs: Dict[str, List[Tensor]]): + return self.get_output(key, modules, batch, outputs)[index] + return func + + def get_output_tensor_list_func(self, key: str) -> TensorListCachedComputationFunc: + def func(modules: Dict[str, BatchInputModule], + batch: List[Tensor], + outputs: Dict[str, List[Tensor]]): + return self.get_output(key, modules, batch, outputs) + return func \ No newline at end of file diff --git a/tha2/mocap/__init__.py b/tha2/mocap/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha2/mocap/ifacialmocap_constants.py b/tha2/mocap/ifacialmocap_constants.py new file mode 100644 index 0000000..3d81fda --- /dev/null +++ b/tha2/mocap/ifacialmocap_constants.py @@ -0,0 +1,239 @@ +EYE_LOOK_IN_LEFT = "eyeLookInLeft" +EYE_LOOK_OUT_LEFT = "eyeLookOutLeft" +EYE_LOOK_DOWN_LEFT = "eyeLookDownLeft" +EYE_LOOK_UP_LEFT = "eyeLookUpLeft" +EYE_BLINK_LEFT = "eyeBlinkLeft" +EYE_SQUINT_LEFT = "eyeSquintLeft" +EYE_WIDE_LEFT = "eyeWideLeft" +EYE_LOOK_IN_RIGHT = "eyeLookInRight" +EYE_LOOK_OUT_RIGHT = "eyeLookOutRight" +EYE_LOOK_DOWN_RIGHT = "eyeLookDownRight" +EYE_LOOK_UP_RIGHT = "eyeLookUpRight" +EYE_BLINK_RIGHT = "eyeBlinkRight" +EYE_SQUINT_RIGHT = "eyeSquintRight" +EYE_WIDE_RIGHT = "eyeWideRight" +BROW_DOWN_LEFT = "browDownLeft" +BROW_OUTER_UP_LEFT = "browOuterUpLeft" +BROW_DOWN_RIGHT = "browDownRight" +BROW_OUTER_UP_RIGHT = "browOuterUpRight" +BROW_INNER_UP = "browInnerUp" +NOSE_SNEER_LEFT = "noseSneerLeft" +NOSE_SNEER_RIGHT = "noseSneerRight" +CHEEK_SQUINT_LEFT = "cheekSquintLeft" +CHEEK_SQUINT_RIGHT = "cheekSquintRight" +CHEEK_PUFF = "cheekPuff" +MOUTH_LEFT = "mouthLeft" +MOUTH_DIMPLE_LEFT = "mouthDimpleLeft" +MOUTH_FROWN_LEFT = "mouthFrownLeft" +MOUTH_LOWER_DOWN_LEFT = "mouthLowerDownLeft" +MOUTH_PRESS_LEFT = "mouthPressLeft" +MOUTH_SMILE_LEFT = "mouthSmileLeft" +MOUTH_STRETCH_LEFT = "mouthStretchLeft" +MOUTH_UPPER_UP_LEFT = "mouthUpperUpLeft" +MOUTH_RIGHT = "mouthRight" +MOUTH_DIMPLE_RIGHT = "mouthDimpleRight" +MOUTH_FROWN_RIGHT = "mouthFrownRight" +MOUTH_LOWER_DOWN_RIGHT = "mouthLowerDownRight" +MOUTH_PRESS_RIGHT = "mouthPressRight" +MOUTH_SMILE_RIGHT = "mouthSmileRight" +MOUTH_STRETCH_RIGHT = "mouthStretchRight" +MOUTH_UPPER_UP_RIGHT = "mouthUpperUpRight" +MOUTH_CLOSE = "mouthClose" +MOUTH_FUNNEL = "mouthFunnel" +MOUTH_PUCKER = "mouthPucker" +MOUTH_ROLL_LOWER = "mouthRollLower" +MOUTH_ROLL_UPPER = "mouthRollUpper" +MOUTH_SHRUG_LOWER = "mouthShrugLower" +MOUTH_SHRUG_UPPER = "mouthShrugUpper" +JAW_LEFT = "jawLeft" +JAW_RIGHT = "jawRight" +JAW_FORWARD = "jawForward" +JAW_OPEN = "jawOpen" +TONGUE_OUT = "tongueOut" + +BLENDSHAPE_NAMES = [ + EYE_LOOK_IN_LEFT, # 0 + EYE_LOOK_OUT_LEFT, # 1 + EYE_LOOK_DOWN_LEFT, # 2 + EYE_LOOK_UP_LEFT, # 3 + EYE_BLINK_LEFT, # 4 + EYE_SQUINT_LEFT, # 5 + EYE_WIDE_LEFT, # 6 + EYE_LOOK_IN_RIGHT, # 7 + EYE_LOOK_OUT_RIGHT, # 8 + EYE_LOOK_DOWN_RIGHT, # 9 + EYE_LOOK_UP_RIGHT, # 10 + EYE_BLINK_RIGHT, # 11 + EYE_SQUINT_RIGHT, # 12 + EYE_WIDE_RIGHT, # 13 + BROW_DOWN_LEFT, # 14 + BROW_OUTER_UP_LEFT, # 15 + BROW_DOWN_RIGHT, # 16 + BROW_OUTER_UP_RIGHT, # 17 + BROW_INNER_UP, # 18 + NOSE_SNEER_LEFT, # 19 + NOSE_SNEER_RIGHT, # 20 + CHEEK_SQUINT_LEFT, # 21 + CHEEK_SQUINT_RIGHT, # 22 + CHEEK_PUFF, # 23 + MOUTH_LEFT, # 24 + MOUTH_DIMPLE_LEFT, # 25 + MOUTH_FROWN_LEFT, # 26 + MOUTH_LOWER_DOWN_LEFT, # 27 + MOUTH_PRESS_LEFT, # 28 + MOUTH_SMILE_LEFT, # 29 + MOUTH_STRETCH_LEFT, # 30 + MOUTH_UPPER_UP_LEFT, # 31 + MOUTH_RIGHT, # 32 + MOUTH_DIMPLE_RIGHT, # 33 + MOUTH_FROWN_RIGHT, # 34 + MOUTH_LOWER_DOWN_RIGHT, # 35 + MOUTH_PRESS_RIGHT, # 36 + MOUTH_SMILE_RIGHT, # 37 + MOUTH_STRETCH_RIGHT, # 38 + MOUTH_UPPER_UP_RIGHT, # 39 + MOUTH_CLOSE, # 40 + MOUTH_FUNNEL, # 41 + MOUTH_PUCKER, # 42 + MOUTH_ROLL_LOWER, # 43 + MOUTH_ROLL_UPPER, # 44 + MOUTH_SHRUG_LOWER, # 45 + MOUTH_SHRUG_UPPER, # 46 + JAW_LEFT, # 47 + JAW_RIGHT, # 48 + JAW_FORWARD, # 49 + JAW_OPEN, # 50 + TONGUE_OUT, # 51 +] + +EYE_LEFT_BLENDSHAPES = [ + EYE_LOOK_IN_LEFT, # 0 + EYE_LOOK_OUT_LEFT, # 1 + EYE_LOOK_DOWN_LEFT, # 2 + EYE_LOOK_UP_LEFT, # 3 + EYE_BLINK_LEFT, # 4 + EYE_SQUINT_LEFT, # 5 + EYE_WIDE_LEFT, # 6 +] + +EYE_RIGHT_BLENDSHAPES = [ + EYE_LOOK_IN_RIGHT, # 7 + EYE_LOOK_OUT_RIGHT, # 8 + EYE_LOOK_DOWN_RIGHT, # 9 + EYE_LOOK_UP_RIGHT, # 10 + EYE_BLINK_RIGHT, # 11 + EYE_SQUINT_RIGHT, # 12 + EYE_WIDE_RIGHT, # 13 +] + +BROW_LEFT_BLENDSHAPES = [ + BROW_DOWN_LEFT, # 14 + BROW_OUTER_UP_LEFT, # 15 + +] + +BROW_RIGHT_BLENDSHAPES = [ + BROW_DOWN_RIGHT, # 16 + BROW_OUTER_UP_RIGHT, # 17 + +] + +BROW_BOTH_BLENDSHAPES = [ + BROW_INNER_UP, # 18 +] + +NOSE_BLENDSHAPES = [ + NOSE_SNEER_LEFT, # 19 + NOSE_SNEER_RIGHT, # 20 +] + +CHECK_BLENDSHAPES = [ + CHEEK_SQUINT_LEFT, # 21 + CHEEK_SQUINT_RIGHT, # 22 + CHEEK_PUFF, # 23 +] + +MOUTH_LEFT_BLENDSHAPES = [ + MOUTH_LEFT, # 24 + MOUTH_DIMPLE_LEFT, # 25 + MOUTH_FROWN_LEFT, # 26 + MOUTH_LOWER_DOWN_LEFT, # 27 + MOUTH_PRESS_LEFT, # 28 + MOUTH_SMILE_LEFT, # 29 + MOUTH_STRETCH_LEFT, # 30 + MOUTH_UPPER_UP_LEFT, # 31 +] + +MOUTH_RIGHT_BLENDSHAPES = [ + MOUTH_RIGHT, # 32 + MOUTH_DIMPLE_RIGHT, # 33 + MOUTH_FROWN_RIGHT, # 34 + MOUTH_LOWER_DOWN_RIGHT, # 35 + MOUTH_PRESS_RIGHT, # 36 + MOUTH_SMILE_RIGHT, # 37 + MOUTH_STRETCH_RIGHT, # 38 + MOUTH_UPPER_UP_RIGHT, # 39 +] + +MOUTH_BOTH_BLENDSHAPES = [ + MOUTH_CLOSE, # 40 + MOUTH_FUNNEL, # 41 + MOUTH_PUCKER, # 42 + MOUTH_ROLL_LOWER, # 43 + MOUTH_ROLL_UPPER, # 44 + MOUTH_SHRUG_LOWER, # 45 + MOUTH_SHRUG_UPPER, # 46 +] + +JAW_BLENDSHAPES = [ + JAW_LEFT, # 47 + JAW_RIGHT, # 48 + JAW_FORWARD, # 49 + JAW_OPEN, # 50 +] + +TONGUE_BLENDSHAPES = [ + TONGUE_OUT, # 51 +] + +COLUMN_0_BLENDSHAPES = EYE_RIGHT_BLENDSHAPES + BROW_RIGHT_BLENDSHAPES + [NOSE_SNEER_RIGHT, CHEEK_SQUINT_RIGHT] +COLUMN_1_BLENDSHAPES = EYE_LEFT_BLENDSHAPES + BROW_LEFT_BLENDSHAPES + [NOSE_SNEER_LEFT, CHEEK_SQUINT_LEFT] +COLUMN_2_BLENDSHAPES = MOUTH_RIGHT_BLENDSHAPES + [JAW_RIGHT] +COLUMN_3_BLENDSHAPES = MOUTH_LEFT_BLENDSHAPES + [JAW_LEFT] +COLUMN_4_BLENDSHAPES = [BROW_INNER_UP, CHEEK_PUFF] + MOUTH_BOTH_BLENDSHAPES + [JAW_FORWARD, JAW_OPEN, TONGUE_OUT] + +BLENDSHAPE_COLUMNS = [ + COLUMN_0_BLENDSHAPES, + COLUMN_1_BLENDSHAPES, + COLUMN_2_BLENDSHAPES, + COLUMN_3_BLENDSHAPES, + COLUMN_4_BLENDSHAPES, +] + +RIGHT_EYE_BONE_X = "rightEyeBoneX" +RIGHT_EYE_BONE_Y = "rightEyeBoneX" +RIGHT_EYE_BONE_Z = "rightEyeBoneX" +RIGHT_EYE_BONE_ROTATIONS = [RIGHT_EYE_BONE_X, RIGHT_EYE_BONE_Y, RIGHT_EYE_BONE_Z] + +LEFT_EYE_BONE_X = "leftEyeBoneX" +LEFT_EYE_BONE_Y = "leftEyeBoneY" +LEFT_EYE_BONE_Z = "leftEyeBoneZ" +LEFT_EYE_BONE_ROTATIONS = [LEFT_EYE_BONE_X, LEFT_EYE_BONE_Y, LEFT_EYE_BONE_Z] + +HEAD_BONE_X = "headBoneX" +HEAD_BONE_Y = "headBoneY" +HEAD_BONE_Z = "headBoneZ" +HEAD_BONE_ROTATIONS = [HEAD_BONE_X, HEAD_BONE_Y, HEAD_BONE_Z] + +ROTATION_NAMES = RIGHT_EYE_BONE_ROTATIONS + LEFT_EYE_BONE_ROTATIONS + HEAD_BONE_ROTATIONS + +RIGHT_EYE_BONE_QUAT = "rightEyeBoneQuat" +LEFT_EYE_BONE_QUAT = "leftEyeBoneQuat" +HEAD_BONE_QUAT = "headBoneQuat" +QUATERNION_NAMES = [ + RIGHT_EYE_BONE_QUAT, + LEFT_EYE_BONE_QUAT, + HEAD_BONE_QUAT +] + +IFACIALMOCAP_DATETIME_FORMAT = "%Y/%m/%d-%H:%M:%S.%f" diff --git a/tha2/mocap/ifacialmocap_pose_converter.py b/tha2/mocap/ifacialmocap_pose_converter.py new file mode 100644 index 0000000..390460b --- /dev/null +++ b/tha2/mocap/ifacialmocap_pose_converter.py @@ -0,0 +1,12 @@ +from abc import ABC, abstractmethod +from typing import Dict, List + + +class IFacialMocapPoseConverter(ABC): + @abstractmethod + def convert(self, ifacialmocap_pose: Dict[str, float]) -> List[float]: + pass + + @abstractmethod + def init_pose_converter_panel(self, parent): + pass \ No newline at end of file diff --git a/tha2/nn/__init__.py b/tha2/nn/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha2/nn/backbone/__init__.py b/tha2/nn/backbone/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha2/nn/backbone/poser_args.py b/tha2/nn/backbone/poser_args.py new file mode 100644 index 0000000..1c4fc7d --- /dev/null +++ b/tha2/nn/backbone/poser_args.py @@ -0,0 +1,68 @@ +from typing import Optional + +from torch.nn import Sigmoid, Sequential, Tanh + +from tha2.nn.base.conv import create_conv3, create_conv3_from_block_args +from tha2.nn.base.nonlinearity_factory import ReLUFactory +from tha2.nn.base.normalization import InstanceNorm2dFactory +from tha2.nn.base.util import BlockArgs + + +class PoserArgs00: + def __init__(self, + image_size: int, + input_image_channels: int, + output_image_channels: int, + start_channels: int, + num_pose_params: int, + block_args: Optional[BlockArgs] = None): + self.num_pose_params = num_pose_params + self.start_channels = start_channels + self.output_image_channels = output_image_channels + self.input_image_channels = input_image_channels + self.image_size = image_size + if block_args is None: + self.block_args = BlockArgs( + normalization_layer_factory=InstanceNorm2dFactory(), + nonlinearity_factory=ReLUFactory(inplace=True)) + else: + self.block_args = block_args + + def create_alpha_block(self): + from torch.nn import Sequential + return Sequential( + create_conv3( + in_channels=self.start_channels, + out_channels=1, + bias=True, + initialization_method=self.block_args.initialization_method, + use_spectral_norm=False), + Sigmoid()) + + def create_all_channel_alpha_block(self): + from torch.nn import Sequential + return Sequential( + create_conv3( + in_channels=self.start_channels, + out_channels=self.output_image_channels, + bias=True, + initialization_method=self.block_args.initialization_method, + use_spectral_norm=False), + Sigmoid()) + + def create_color_change_block(self): + return Sequential( + create_conv3_from_block_args( + in_channels=self.start_channels, + out_channels=self.output_image_channels, + bias=True, + block_args=self.block_args), + Tanh()) + + def create_grid_change_block(self): + return create_conv3( + in_channels=self.start_channels, + out_channels=2, + bias=False, + initialization_method='zero', + use_spectral_norm=False) \ No newline at end of file diff --git a/tha2/nn/backbone/poser_encoder_decoder_00.py b/tha2/nn/backbone/poser_encoder_decoder_00.py new file mode 100644 index 0000000..2de75eb --- /dev/null +++ b/tha2/nn/backbone/poser_encoder_decoder_00.py @@ -0,0 +1,122 @@ +import math +from typing import Optional, List + +import torch +from torch import Tensor +from torch.nn import ModuleList, Module + +from tha2.nn.backbone.poser_args import PoserArgs00 +from tha2.nn.base.conv import create_conv3_block_from_block_args, create_downsample_block_from_block_args, \ + create_upsample_block_from_block_args +from tha2.nn.base.nonlinearity_factory import ReLUFactory +from tha2.nn.base.normalization import InstanceNorm2dFactory +from tha2.nn.base.resnet_block import ResnetBlock +from tha2.nn.base.util import BlockArgs + + +class PoserEncoderDecoder00Args(PoserArgs00): + def __init__(self, + image_size: int, + input_image_channels: int, + output_image_channels: int, + num_pose_params: int , + start_channels: int, + bottleneck_image_size, + num_bottleneck_blocks, + max_channels: int, + block_args: Optional[BlockArgs] = None): + super().__init__( + image_size, input_image_channels, output_image_channels, start_channels, num_pose_params, block_args) + self.max_channels = max_channels + self.num_bottleneck_blocks = num_bottleneck_blocks + self.bottleneck_image_size = bottleneck_image_size + assert bottleneck_image_size > 1 + + if block_args is None: + self.block_args = BlockArgs( + normalization_layer_factory=InstanceNorm2dFactory(), + nonlinearity_factory=ReLUFactory(inplace=True)) + else: + self.block_args = block_args + + +class PoserEncoderDecoder00(Module): + def __init__(self, args: PoserEncoderDecoder00Args): + super().__init__() + self.args = args + + self.num_levels = int(math.log2(args.image_size // args.bottleneck_image_size)) + 1 + + self.downsample_blocks = ModuleList() + self.downsample_blocks.append( + create_conv3_block_from_block_args( + args.input_image_channels, + args.start_channels, + args.block_args)) + + current_image_size = args.image_size + current_num_channels = args.start_channels + while current_image_size > args.bottleneck_image_size: + next_image_size = current_image_size // 2 + next_num_channels = self.get_num_output_channels_from_image_size(next_image_size) + self.downsample_blocks.append(create_downsample_block_from_block_args( + in_channels=current_num_channels, + out_channels=next_num_channels, + is_output_1x1=False, + block_args=args.block_args)) + current_image_size = next_image_size + current_num_channels = next_num_channels + assert len(self.downsample_blocks) == self.num_levels + + self.bottleneck_blocks = ModuleList() + self.bottleneck_blocks.append(create_conv3_block_from_block_args( + in_channels=current_num_channels + args.num_pose_params, + out_channels=current_num_channels, + block_args=args.block_args)) + for i in range(1, args.num_bottleneck_blocks): + self.bottleneck_blocks.append( + ResnetBlock.create( + num_channels=current_num_channels, + is1x1=False, + block_args=args.block_args)) + + self.upsample_blocks = ModuleList() + while current_image_size < args.image_size: + next_image_size = current_image_size * 2 + next_num_channels = self.get_num_output_channels_from_image_size(next_image_size) + self.upsample_blocks.append(create_upsample_block_from_block_args( + in_channels=current_num_channels, + out_channels=next_num_channels, + block_args=args.block_args)) + current_image_size = next_image_size + current_num_channels = next_num_channels + + def get_num_output_channels_from_level(self, level: int): + return self.get_num_output_channels_from_image_size(self.args.image_size // (2 ** level)) + + def get_num_output_channels_from_image_size(self, image_size: int): + return min(self.args.start_channels * (self.args.image_size // image_size), self.args.max_channels) + + def forward(self, image: Tensor, pose: Optional[Tensor] = None) -> List[Tensor]: + if self.args.num_pose_params != 0: + assert pose is not None + else: + assert pose is None + outputs = [] + feature = image + outputs.append(feature) + for block in self.downsample_blocks: + feature = block(feature) + outputs.append(feature) + if pose is not None: + n, c = pose.shape + pose = pose.view(n, c, 1, 1).repeat(1, 1, self.args.bottleneck_image_size, self.args.bottleneck_image_size) + feature = torch.cat([feature, pose], dim=1) + for block in self.bottleneck_blocks: + feature = block(feature) + outputs.append(feature) + for block in self.upsample_blocks: + feature = block(feature) + outputs.append(feature) + outputs.reverse() + return outputs diff --git a/tha2/nn/backcomp/__init__.py b/tha2/nn/backcomp/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha2/nn/backcomp/nn/__init__.py b/tha2/nn/backcomp/nn/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha2/nn/backcomp/nn/conv.py b/tha2/nn/backcomp/nn/conv.py new file mode 100644 index 0000000..e13176f --- /dev/null +++ b/tha2/nn/backcomp/nn/conv.py @@ -0,0 +1,36 @@ +from torch.nn import Conv2d, Module, Sequential, InstanceNorm2d, ReLU, ConvTranspose2d + +from tha2.nn.backcomp.nn.init_function import create_init_function + + +def Conv7(in_channels: int, out_channels: int, initialization_method='he') -> Module: + init = create_init_function(initialization_method) + return init(Conv2d(in_channels, out_channels, kernel_size=7, stride=1, padding=3, bias=False)) + + +def Conv3(in_channels: int, out_channels: int, initialization_method='he') -> Module: + init = create_init_function(initialization_method) + return init(Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False)) + + +def Conv7Block(in_channels: int, out_channels: int, initialization_method='he') -> Module: + return Sequential( + Conv7(in_channels, out_channels, initialization_method), + InstanceNorm2d(out_channels, affine=True), + ReLU(inplace=True)) + + +def DownsampleBlock(in_channels: int, initialization_method='he') -> Module: + init = create_init_function(initialization_method) + return Sequential( + init(Conv2d(in_channels, in_channels * 2, kernel_size=4, stride=2, padding=1, bias=False)), + InstanceNorm2d(in_channels * 2, affine=True), + ReLU(inplace=True)) + + +def UpsampleBlock(in_channels: int, out_channels: int, initialization_method='he') -> Module: + init = create_init_function(initialization_method) + return Sequential( + init(ConvTranspose2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1, bias=False)), + InstanceNorm2d(out_channels, affine=True), + ReLU(inplace=True)) diff --git a/tha2/nn/backcomp/nn/encoder_decoder_module.py b/tha2/nn/backcomp/nn/encoder_decoder_module.py new file mode 100644 index 0000000..42ebea7 --- /dev/null +++ b/tha2/nn/backcomp/nn/encoder_decoder_module.py @@ -0,0 +1,45 @@ +from torch.nn import Module, ModuleList + +from tha2.nn.backcomp.nn.conv import Conv7Block, DownsampleBlock, UpsampleBlock +from tha2.nn.backcomp.nn.resnet_block import ResNetBlock + +# image_size: int = 256, +# image_channels: int = 4, +# pose_size: int = 3, +# intermediate_channels: int = 64, +# bottleneck_image_size: int = 32, +# bottleneck_block_count: int = 6, +# initialization_method: str = 'he', +# align_corners: bool = True): + +class EncoderDecoderModule(Module): + def __init__(self, + image_size: int, + image_channels: int, + output_channels: int, + bottleneck_image_size: int, + bottleneck_block_count: int, + initialization_method: str = 'he'): + super().__init__() + + self.module_list = ModuleList() + self.module_list.append(Conv7Block(image_channels, output_channels)) + current_size = image_size + current_channels = output_channels + while current_size > bottleneck_image_size: + self.module_list.append(DownsampleBlock(current_channels, initialization_method)) + current_size //= 2 + current_channels *= 2 + + for i in range(bottleneck_block_count): + self.module_list.append(ResNetBlock(current_channels, initialization_method)) + + while current_size < image_size: + self.module_list.append(UpsampleBlock(current_channels, current_channels // 2, initialization_method)) + current_size *= 2 + current_channels //= 2 + + def forward(self, x): + for module in self.module_list: + x = module(x) + return x diff --git a/tha2/nn/backcomp/nn/init_function.py b/tha2/nn/backcomp/nn/init_function.py new file mode 100644 index 0000000..cac751b --- /dev/null +++ b/tha2/nn/backcomp/nn/init_function.py @@ -0,0 +1,18 @@ +from torch.nn import Module +from torch.nn.init import kaiming_normal_, xavier_normal_ + + +def create_init_function(method: str = 'none'): + def init(module: Module): + if method == 'none': + return module + elif method == 'he': + kaiming_normal_(module.weight) + return module + elif method == 'xavier': + xavier_normal_(module.weight) + return module + else: + raise ("Invalid initialization method %s" % method) + + return init \ No newline at end of file diff --git a/tha2/nn/backcomp/nn/resnet_block.py b/tha2/nn/backcomp/nn/resnet_block.py new file mode 100644 index 0000000..346c370 --- /dev/null +++ b/tha2/nn/backcomp/nn/resnet_block.py @@ -0,0 +1,16 @@ +from torch import relu +from torch.nn import Module, InstanceNorm2d + +from tha2.nn.backcomp.nn.conv import Conv3 + + +class ResNetBlock(Module): + def __init__(self, num_channels: int, initialization_method: str = 'he'): + super().__init__() + self.conv1 = Conv3(num_channels, num_channels, initialization_method) + self.norm1 = InstanceNorm2d(num_features=num_channels, affine=True) + self.conv2 = Conv3(num_channels, num_channels, initialization_method) + self.norm2 = InstanceNorm2d(num_features=num_channels, affine=True) + + def forward(self, x): + return x + self.norm2(self.conv2(relu(self.norm1(self.conv1(x))))) diff --git a/tha2/nn/backcomp/nn/u_net_module.py b/tha2/nn/backcomp/nn/u_net_module.py new file mode 100644 index 0000000..11059e5 --- /dev/null +++ b/tha2/nn/backcomp/nn/u_net_module.py @@ -0,0 +1,63 @@ +import torch +from torch.nn import Module, ModuleList + +from tha2.nn.backcomp.nn.conv import Conv7Block, DownsampleBlock, UpsampleBlock +from tha2.nn.backcomp.nn.resnet_block import ResNetBlock + + +class UNetModule(Module): + def __init__(self, + image_size: int, + image_channels: int, + output_channels: int, + bottleneck_image_size: int, + bottleneck_block_count: int, + initialization_method: str = 'he'): + super().__init__() + self.downward_modules = ModuleList() + self.downward_module_channel_count = {} + + self.downward_modules.append(Conv7Block(image_channels, output_channels, initialization_method)) + self.downward_module_channel_count[image_size] = output_channels + + # Downsampling + current_channels = output_channels + current_image_size = image_size + while current_image_size > bottleneck_image_size: + self.downward_modules.append(DownsampleBlock(current_channels, initialization_method)) + current_channels = current_channels * 2 + current_image_size = current_image_size // 2 + self.downward_module_channel_count[current_image_size] = current_channels + + # Bottleneck + self.bottleneck_modules = ModuleList() + for i in range(bottleneck_block_count): + self.bottleneck_modules.append(ResNetBlock(current_channels, initialization_method)) + + # Upsampling + self.upsampling_modules = ModuleList() + while current_image_size < image_size: + if current_image_size == bottleneck_image_size: + input_channels = current_channels + else: + input_channels = current_channels + self.downward_module_channel_count[current_image_size] + self.upsampling_modules.insert(0, + UpsampleBlock(input_channels, current_channels // 2, initialization_method)) + current_channels = current_channels // 2 + current_image_size = current_image_size * 2 + + self.upsampling_modules.insert( + 0, Conv7Block(current_channels + output_channels, output_channels, initialization_method)) + + def forward(self, x): + downward_outputs = [] + for module in self.downward_modules: + x = module(x) + downward_outputs.append(x) + for module in self.bottleneck_modules: + x = module(x) + x = self.upsampling_modules[-1](x) + for i in range(len(self.upsampling_modules) - 2, -1, -1): + y = torch.cat([x, downward_outputs[i]], dim=1) + x = self.upsampling_modules[i](y) + return x diff --git a/tha2/nn/backcomp/tha/__init__.py b/tha2/nn/backcomp/tha/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha2/nn/backcomp/tha/combiner.py b/tha2/nn/backcomp/tha/combiner.py new file mode 100644 index 0000000..5f26a10 --- /dev/null +++ b/tha2/nn/backcomp/tha/combiner.py @@ -0,0 +1,81 @@ +import torch +from torch import Tensor +from torch.nn import Sequential, Sigmoid, Tanh + +from tha2.nn.backcomp.nn.conv import Conv7 +from tha2.nn.backcomp.nn.u_net_module import UNetModule +from tha2.nn.batch_module.batch_input_module import BatchInputModule, BatchInputModuleFactory + + +class Combiner(BatchInputModule): + def __init__(self, + image_size: int = 256, + image_channels: int = 4, + pose_size: int = 3, + intermediate_channels: int = 64, + bottleneck_image_size: int = 32, + bottleneck_block_count: int = 6, + initialization_method: str = 'he'): + super().__init__() + self.main_body = UNetModule( + image_size=image_size, + image_channels=2 * image_channels + pose_size, + output_channels=intermediate_channels, + bottleneck_image_size=bottleneck_image_size, + bottleneck_block_count=bottleneck_block_count, + initialization_method=initialization_method) + self.combine_alpha_mask = Sequential( + Conv7(intermediate_channels, image_channels, initialization_method), + Sigmoid()) + self.retouch_alpha_mask = Sequential( + Conv7(intermediate_channels, image_channels, initialization_method), + Sigmoid()) + self.retouch_color_change = Sequential( + Conv7(intermediate_channels, image_channels, initialization_method), + Tanh()) + + def forward(self, first_image: Tensor, second_image: Tensor, pose: Tensor): + pose = pose.unsqueeze(2).unsqueeze(3) + pose = pose.expand(pose.size(0), pose.size(1), first_image.size(2), first_image.size(3)) + + x = torch.cat([first_image, second_image, pose], dim=1) + y = self.main_body(x) + combine_alpha_mask = self.combine_alpha_mask(y) + combined_image = combine_alpha_mask * first_image + (1 - combine_alpha_mask) * second_image + retouch_alpha_mask = self.retouch_alpha_mask(y) + retouch_color_change = self.retouch_color_change(y) + final_image = retouch_alpha_mask * combined_image + (1 - retouch_alpha_mask) * retouch_color_change + # return [final_image, combined_image, combine_alpha_mask, retouch_alpha_mask, retouch_color_change] + return final_image + + def forward_from_batch(self, batch): + return self.forward(batch[0], batch[1], batch[2]) + + +class CombinerFactory(BatchInputModuleFactory): + def __init__(self, + image_size: int = 256, + image_channels: int = 4, + pose_size: int = 3, + intermediate_channels: int = 64, + bottleneck_image_size: int = 32, + bottleneck_block_count: int = 6, + initialization_method: str = 'he'): + super().__init__() + self.image_size = image_size + self.image_channels = image_channels + self.pose_size = pose_size + self.intermediate_channels = intermediate_channels + self.bottleneck_image_size = bottleneck_image_size + self.bottleneck_block_count = bottleneck_block_count + self.initialization_method = initialization_method + + def create(self) -> BatchInputModule: + return Combiner( + self.image_size, + self.image_channels, + self.pose_size, + self.intermediate_channels, + self.bottleneck_image_size, + self.bottleneck_block_count, + self.initialization_method) \ No newline at end of file diff --git a/tha2/nn/backcomp/tha/face_morpher.py b/tha2/nn/backcomp/tha/face_morpher.py new file mode 100644 index 0000000..9d1058a --- /dev/null +++ b/tha2/nn/backcomp/tha/face_morpher.py @@ -0,0 +1,74 @@ +import torch +from torch import Tensor +from torch.nn import Sequential, Tanh, Sigmoid + +from tha2.nn.backcomp.nn.conv import Conv7 +from tha2.nn.backcomp.nn.encoder_decoder_module import EncoderDecoderModule +from tha2.nn.batch_module.batch_input_module import BatchInputModuleFactory, BatchInputModule + + +class FaceMorpher(BatchInputModule): + def __init__(self, + image_size: int = 256, + image_channels: int = 4, + pose_size: int = 3, + intermediate_channels: int = 64, + bottleneck_image_size: int = 32, + bottleneck_block_count: int = 6, + initialization_method: str = 'he'): + super().__init__() + self.main_body = EncoderDecoderModule( + image_size=image_size, + image_channels=image_channels + pose_size, + output_channels=intermediate_channels, + bottleneck_image_size=bottleneck_image_size, + bottleneck_block_count=bottleneck_block_count, + initialization_method=initialization_method) + self.color_change = Sequential( + Conv7(intermediate_channels, image_channels, initialization_method), + Tanh()) + self.alpha_mask = Sequential( + Conv7(intermediate_channels, image_channels, initialization_method), + Sigmoid()) + + def forward(self, image: Tensor, pose: Tensor): + pose = pose.unsqueeze(2).unsqueeze(3) + pose = pose.expand(pose.size(0), pose.size(1), image.size(2), image.size(3)) + x = torch.cat([image, pose], dim=1) + y = self.main_body(x) + color = self.color_change(y) + alpha = self.alpha_mask(y) + output_image = alpha * image + (1 - alpha) * color + return [output_image, alpha, color] + + def forward_from_batch(self, batch): + return self.forward(batch[0], batch[1]) + + +class FaceMorpherSpec(BatchInputModuleFactory): + def __init__(self, + image_size: int = 256, + image_channels: int = 4, + pose_size: int = 3, + intermediate_channels: int = 64, + bottleneck_image_size: int = 32, + bottleneck_block_count: int = 6, + initialization_method: str = 'he'): + super().__init__() + self.image_size = image_size + self.image_channels = image_channels + self.pose_size = pose_size + self.intermediate_channels = intermediate_channels + self.bottleneck_image_size = bottleneck_image_size + self.bottleneck_block_count = bottleneck_block_count + self.initialization_method = initialization_method + + def create(self) -> BatchInputModule: + return FaceMorpher( + self.image_size, + self.image_channels, + self.pose_size, + self.intermediate_channels, + self.bottleneck_image_size, + self.bottleneck_block_count, + self.initialization_method) diff --git a/tha2/nn/backcomp/tha/two_algo_face_rotator.py b/tha2/nn/backcomp/tha/two_algo_face_rotator.py new file mode 100644 index 0000000..6949e7e --- /dev/null +++ b/tha2/nn/backcomp/tha/two_algo_face_rotator.py @@ -0,0 +1,101 @@ +import torch +from torch import Tensor +from torch.nn import Sequential, Tanh, Sigmoid +from torch.nn.functional import affine_grid, grid_sample + +from tha2.nn.backcomp.nn.conv import Conv7 +from tha2.nn.backcomp.nn.encoder_decoder_module import EncoderDecoderModule +from tha2.nn.batch_module.batch_input_module import BatchInputModule, BatchInputModuleFactory + + +class TwoAlgoFaceRotator(BatchInputModule): + COLOR_CHANGED_IMAGE_INDEX = 0 + RESAMPLED_IMAGE_INDEX = 1 + COLOR_CHANGE_INDEX = 2 + ALPHA_INDEX = 3 + GRID_CHANGE_INDEX = 4 + OUTPUT_LENGTH = 5 + + def __init__(self, + image_size: int = 256, + image_channels: int = 4, + pose_size: int = 3, + intermediate_channels: int = 64, + bottleneck_image_size: int = 32, + bottleneck_block_count: int = 6, + initialization_method: str = 'he', + align_corners: bool = True): + super().__init__() + self.align_corners = align_corners + self.main_body = EncoderDecoderModule( + image_size=image_size, + image_channels=image_channels + pose_size, + output_channels=intermediate_channels, + bottleneck_image_size=bottleneck_image_size, + bottleneck_block_count=bottleneck_block_count, + initialization_method=initialization_method) + self.pumarola_color_change = Sequential( + Conv7(intermediate_channels, image_channels, initialization_method), + Tanh()) + self.pumarola_alpha_mask = Sequential( + Conv7(intermediate_channels, image_channels, initialization_method), + Sigmoid()) + self.zhou_grid_change = Conv7(intermediate_channels, 2, initialization_method) + + def forward(self, image: Tensor, pose: Tensor): + n, c, h, w = image.shape + + pose = pose.unsqueeze(2).unsqueeze(3) + pose = pose.expand(pose.shape[0], pose.shape[1], image.shape[2], image.shape[3]) + x = torch.cat([image, pose], dim=1) + y = self.main_body(x) + + color_change = self.pumarola_color_change(y) + alpha_mask = self.pumarola_alpha_mask(y) + color_changed = alpha_mask * image + (1 - alpha_mask) * color_change + + original_grid_change = self.zhou_grid_change(y) + grid_change = torch.transpose(original_grid_change.view(n, 2, h * w), 1, 2).view(n, h, w, 2) + device = self.zhou_grid_change.weight.device + identity = torch.Tensor([[1, 0, 0], [0, 1, 0]]).to(device).unsqueeze(0).repeat(n, 1, 1) + base_grid = affine_grid(identity, [n, c, h, w], align_corners=self.align_corners) + grid = base_grid + grid_change + resampled = grid_sample(image, grid.half(), mode='bilinear', padding_mode='border', align_corners=self.align_corners) + + # return [color_changed, resampled, color_change, alpha_mask, original_grid_change] + return [color_changed, resampled] + + def forward_from_batch(self, batch): + return self.forward(batch[0], batch[1]) + + +class TwoAlgoFaceRotatorFactory(BatchInputModuleFactory): + def __init__(self, + image_size: int = 256, + image_channels: int = 4, + pose_size: int = 3, + intermediate_channels: int = 64, + bottleneck_image_size: int = 32, + bottleneck_block_count: int = 6, + initialization_method: str = 'he', + align_corners: bool = True): + super().__init__() + self.image_size = image_size + self.image_channels = image_channels + self.pose_size = pose_size + self.intermediate_channels = intermediate_channels + self.bottleneck_image_size = bottleneck_image_size + self.bottleneck_block_count = bottleneck_block_count + self.initialization_method = initialization_method + self.align_corners = align_corners + + def create(self) -> BatchInputModule: + return TwoAlgoFaceRotator( + self.image_size, + self.image_channels, + self.pose_size, + self.intermediate_channels, + self.bottleneck_image_size, + self.bottleneck_block_count, + self.initialization_method, + self.align_corners) \ No newline at end of file diff --git a/tha2/nn/base/__init__.py b/tha2/nn/base/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha2/nn/base/conv.py b/tha2/nn/base/conv.py new file mode 100644 index 0000000..0ab90d3 --- /dev/null +++ b/tha2/nn/base/conv.py @@ -0,0 +1,171 @@ +from typing import Optional + +from torch.nn import Conv2d, Module, Sequential, ConvTranspose2d + +from tha2.nn.base.module_factory import ModuleFactory +from tha2.nn.base.nonlinearity_factory import resolve_nonlinearity_factory +from tha2.nn.base.normalization import NormalizationLayerFactory +from tha2.nn.base.util import wrap_conv_or_linear_module, BlockArgs + + +def create_conv7(in_channels: int, out_channels: int, + bias: bool = False, + initialization_method='he', + use_spectral_norm: bool = False) -> Module: + return wrap_conv_or_linear_module( + Conv2d(in_channels, out_channels, kernel_size=7, stride=1, padding=3, bias=bias), + initialization_method, + use_spectral_norm) + +def create_conv7_from_block_args(in_channels: int, out_channels: int, + bias: bool = False, + block_args: Optional[BlockArgs] = None)-> Module: + if block_args is None: + block_args = BlockArgs() + return create_conv7(in_channels, out_channels, bias, + block_args.initialization_method, + block_args.use_spectral_norm) + + +def create_conv3(in_channels: int, out_channels: int, + bias: bool = False, + initialization_method='he', + use_spectral_norm: bool = False) -> Module: + return wrap_conv_or_linear_module( + Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=bias), + initialization_method, + use_spectral_norm) + + +def create_conv3_from_block_args(in_channels: int, out_channels: int, + bias: bool = False, + block_args: Optional[BlockArgs] = None): + if block_args is None: + block_args = BlockArgs() + return create_conv3(in_channels, out_channels, bias, + block_args.initialization_method, + block_args.use_spectral_norm) + + +def create_conv1(in_channels: int, out_channels: int, + initialization_method='he', + bias: bool = False, + use_spectral_norm: bool = False) -> Module: + return wrap_conv_or_linear_module( + Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=bias), + initialization_method, + use_spectral_norm) + + +def create_conv7_block(in_channels: int, out_channels: int, + initialization_method='he', + nonlinearity_factory: Optional[ModuleFactory] = None, + normalization_layer_factory: Optional[NormalizationLayerFactory] = None, + use_spectral_norm: bool = False) -> Module: + nonlinearity_factory = resolve_nonlinearity_factory(nonlinearity_factory) + return Sequential( + create_conv7(in_channels, out_channels, + bias=False, initialization_method=initialization_method, use_spectral_norm=use_spectral_norm), + NormalizationLayerFactory.resolve_2d(normalization_layer_factory).create(out_channels, affine=True), + resolve_nonlinearity_factory(nonlinearity_factory).create()) + + +def create_conv7_block_from_block_args( + in_channels: int, out_channels: int, + block_args: Optional[BlockArgs] = None) -> Module: + if block_args is None: + block_args = BlockArgs() + return create_conv7_block(in_channels, out_channels, + block_args.initialization_method, + block_args.nonlinearity_factory, + block_args.normalization_layer_factory, + block_args.use_spectral_norm) + + +def create_conv3_block(in_channels: int, out_channels: int, + initialization_method='he', + nonlinearity_factory: Optional[ModuleFactory] = None, + normalization_layer_factory: Optional[NormalizationLayerFactory] = None, + use_spectral_norm: bool = False) -> Module: + nonlinearity_factory = resolve_nonlinearity_factory(nonlinearity_factory) + return Sequential( + create_conv3(in_channels, out_channels, + bias=False, initialization_method=initialization_method, use_spectral_norm=use_spectral_norm), + NormalizationLayerFactory.resolve_2d(normalization_layer_factory).create(out_channels, affine=True), + resolve_nonlinearity_factory(nonlinearity_factory).create()) + + +def create_conv3_block_from_block_args( + in_channels: int, out_channels: int, block_args: Optional[BlockArgs] = None): + if block_args is None: + block_args = BlockArgs() + return create_conv3_block(in_channels, out_channels, + block_args.initialization_method, + block_args.nonlinearity_factory, + block_args.normalization_layer_factory, + block_args.use_spectral_norm) + + +def create_downsample_block(in_channels: int, out_channels: int, + is_output_1x1: bool = False, + initialization_method='he', + nonlinearity_factory: Optional[ModuleFactory] = None, + normalization_layer_factory: Optional[NormalizationLayerFactory] = None, + use_spectral_norm: bool = False) -> Module: + if is_output_1x1: + return Sequential( + wrap_conv_or_linear_module( + Conv2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1, bias=False), + initialization_method, + use_spectral_norm), + resolve_nonlinearity_factory(nonlinearity_factory).create()) + else: + return Sequential( + wrap_conv_or_linear_module( + Conv2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1, bias=False), + initialization_method, + use_spectral_norm), + NormalizationLayerFactory.resolve_2d(normalization_layer_factory).create(out_channels, affine=True), + resolve_nonlinearity_factory(nonlinearity_factory).create()) + + +def create_downsample_block_from_block_args(in_channels: int, out_channels: int, + is_output_1x1: bool = False, + block_args: Optional[BlockArgs] = None): + if block_args is None: + block_args = BlockArgs() + return create_downsample_block( + in_channels, out_channels, + is_output_1x1, + block_args.initialization_method, + block_args.nonlinearity_factory, + block_args.normalization_layer_factory, + block_args.use_spectral_norm) + + +def create_upsample_block(in_channels: int, + out_channels: int, + initialization_method='he', + nonlinearity_factory: Optional[ModuleFactory] = None, + normalization_layer_factory: Optional[NormalizationLayerFactory] = None, + use_spectral_norm: bool = False) -> Module: + nonlinearity_factory = resolve_nonlinearity_factory(nonlinearity_factory) + return Sequential( + wrap_conv_or_linear_module( + ConvTranspose2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1, bias=False), + initialization_method, + use_spectral_norm), + NormalizationLayerFactory.resolve_2d(normalization_layer_factory).create(out_channels, affine=True), + resolve_nonlinearity_factory(nonlinearity_factory).create()) + + +def create_upsample_block_from_block_args(in_channels: int, + out_channels: int, + block_args: Optional[BlockArgs] = None) -> Module: + if block_args is None: + block_args = BlockArgs() + return create_upsample_block(in_channels, out_channels, + block_args.initialization_method, + block_args.nonlinearity_factory, + block_args.normalization_layer_factory, + block_args.use_spectral_norm) diff --git a/tha2/nn/base/init_function.py b/tha2/nn/base/init_function.py new file mode 100644 index 0000000..62c99f3 --- /dev/null +++ b/tha2/nn/base/init_function.py @@ -0,0 +1,30 @@ +import torch +from torch import zero_ +from torch.nn import Module +from torch.nn.init import kaiming_normal_, xavier_normal_, normal_ + + +def create_init_function(method: str = 'none'): + def init(module: Module): + if method == 'none': + return module + elif method == 'he': + kaiming_normal_(module.weight) + return module + elif method == 'xavier': + xavier_normal_(module.weight) + return module + elif method == 'dcgan': + normal_(module.weight, 0.0, 0.02) + return module + elif method == 'dcgan_001': + normal_(module.weight, 0.0, 0.01) + return module + elif method == "zero": + with torch.no_grad(): + zero_(module.weight) + return module + else: + raise ("Invalid initialization method %s" % method) + + return init diff --git a/tha2/nn/base/module_factory.py b/tha2/nn/base/module_factory.py new file mode 100644 index 0000000..f6231c3 --- /dev/null +++ b/tha2/nn/base/module_factory.py @@ -0,0 +1,9 @@ +from abc import ABC, abstractmethod + +from torch.nn import Module + + +class ModuleFactory(ABC): + @abstractmethod + def create(self) -> Module: + pass \ No newline at end of file diff --git a/tha2/nn/base/nonlinearity_factory.py b/tha2/nn/base/nonlinearity_factory.py new file mode 100644 index 0000000..e367045 --- /dev/null +++ b/tha2/nn/base/nonlinearity_factory.py @@ -0,0 +1,38 @@ +from typing import Optional + +from torch.nn import Module, ReLU, LeakyReLU, ELU + +from tha2.nn.base.module_factory import ModuleFactory + + +class ReLUFactory(ModuleFactory): + def __init__(self, inplace: bool = False): + self.inplace = inplace + + def create(self) -> Module: + return ReLU(self.inplace) + + +class LeakyReLUFactory(ModuleFactory): + def __init__(self, inplace: bool = False, negative_slope: float = 1e-2): + self.negative_slope = negative_slope + self.inplace = inplace + + def create(self) -> Module: + return LeakyReLU(inplace=self.inplace, negative_slope=self.negative_slope) + + +class ELUFactory(ModuleFactory): + def __init__(self, inplace: bool = False, alpha: float = 1.0): + self.alpha = alpha + self.inplace = inplace + + def create(self) -> Module: + return ELU(inplace=self.inplace, alpha=self.alpha) + + +def resolve_nonlinearity_factory(nonlinearity_fatory: Optional[ModuleFactory]) -> ModuleFactory: + if nonlinearity_fatory is None: + return ReLUFactory(inplace=True) + else: + return nonlinearity_fatory diff --git a/tha2/nn/base/normalization.py b/tha2/nn/base/normalization.py new file mode 100644 index 0000000..9f465fc --- /dev/null +++ b/tha2/nn/base/normalization.py @@ -0,0 +1,102 @@ +from abc import ABC, abstractmethod +from typing import Optional + +import torch +from torch.nn import Module, BatchNorm2d, InstanceNorm2d, Parameter +from torch.nn.init import normal_, constant_ + +from tha2.nn.base.pass_through import PassThrough + + +class PixelNormalization(Module): + def __init__(self, epsilon=1e-8): + super().__init__() + self.epsilon = epsilon + + def forward(self, x): + return x / torch.sqrt((x ** 2).mean(dim=1, keepdim=True) + self.epsilon) + + +class NormalizationLayerFactory(ABC): + def __init__(self): + super().__init__() + + @abstractmethod + def create(self, num_features: int, affine: bool = True) -> Module: + pass + + @staticmethod + def resolve_2d(factory: Optional['NormalizationLayerFactory']) -> 'NormalizationLayerFactory': + if factory is None: + return InstanceNorm2dFactory() + else: + return factory + + +class Bias2d(Module): + def __init__(self, num_features: int): + super().__init__() + self.num_features = num_features + self.bias = Parameter(torch.zeros(1, num_features, 1, 1)) + + def forward(self, x): + return x + self.bias + + +class NoNorm2dFactory(NormalizationLayerFactory): + def __init__(self): + super().__init__() + + def create(self, num_features: int, affine: bool = True) -> Module: + if affine: + return Bias2d(num_features) + else: + return PassThrough() + + +class BatchNorm2dFactory(NormalizationLayerFactory): + def __init__(self, + weight_mean: Optional[float] = None, + weight_std: Optional[float] = None, + bias: Optional[float] = None): + super().__init__() + self.bias = bias + self.weight_std = weight_std + self.weight_mean = weight_mean + + def get_weight_mean(self): + if self.weight_mean is None: + return 1.0 + else: + return self.weight_mean + + def get_weight_std(self): + if self.weight_std is None: + return 0.02 + else: + return self.weight_std + + def create(self, num_features: int, affine: bool = True) -> Module: + module = BatchNorm2d(num_features=num_features, affine=affine) + if affine: + if self.weight_mean is not None or self.weight_std is not None: + normal_(module.weight, self.get_weight_mean(), self.get_weight_std()) + if self.bias is not None: + constant_(module.bias, self.bias) + return module + + +class InstanceNorm2dFactory(NormalizationLayerFactory): + def __init__(self): + super().__init__() + + def create(self, num_features: int, affine: bool = True) -> Module: + return InstanceNorm2d(num_features=num_features, affine=affine) + + +class PixelNormFactory(NormalizationLayerFactory): + def __init__(self): + super().__init__() + + def create(self, num_features: int, affine: bool = True) -> Module: + return PixelNormalization() \ No newline at end of file diff --git a/tha2/nn/base/pass_through.py b/tha2/nn/base/pass_through.py new file mode 100644 index 0000000..c64d678 --- /dev/null +++ b/tha2/nn/base/pass_through.py @@ -0,0 +1,9 @@ +from torch.nn import Module + + +class PassThrough(Module): + def __init__(self): + super().__init__() + + def forward(self, x): + return x \ No newline at end of file diff --git a/tha2/nn/base/resnet_block.py b/tha2/nn/base/resnet_block.py new file mode 100644 index 0000000..86d4647 --- /dev/null +++ b/tha2/nn/base/resnet_block.py @@ -0,0 +1,67 @@ +from typing import Optional + +import torch +from torch.nn import Module, Sequential, Parameter + +from tha2.nn.base.conv import create_conv1, create_conv3 +from tha2.nn.base.module_factory import ModuleFactory +from tha2.nn.base.nonlinearity_factory import resolve_nonlinearity_factory +from tha2.nn.base.normalization import NormalizationLayerFactory +from tha2.nn.base.util import BlockArgs + + +class ResnetBlock(Module): + @staticmethod + def create(num_channels: int, + is1x1: bool = False, + use_scale_parameters: bool = False, + block_args: Optional[BlockArgs] = None): + if block_args is None: + block_args = BlockArgs() + return ResnetBlock(num_channels, + is1x1, + block_args.initialization_method, + block_args.nonlinearity_factory, + block_args.normalization_layer_factory, + block_args.use_spectral_norm, + use_scale_parameters) + + def __init__(self, + num_channels: int, + is1x1: bool = False, + initialization_method: str = 'he', + nonlinearity_factory: ModuleFactory = None, + normalization_layer_factory: Optional[NormalizationLayerFactory] = None, + use_spectral_norm: bool = False, + use_scale_parameter: bool = False): + super().__init__() + self.use_scale_parameter = use_scale_parameter + if self.use_scale_parameter: + self.scale = Parameter(torch.zeros(1)) + nonlinearity_factory = resolve_nonlinearity_factory(nonlinearity_factory) + if is1x1: + self.resnet_path = Sequential( + create_conv1(num_channels, num_channels, initialization_method, + bias=True, + use_spectral_norm=use_spectral_norm), + nonlinearity_factory.create(), + create_conv1(num_channels, num_channels, initialization_method, + bias=True, + use_spectral_norm=use_spectral_norm)) + else: + self.resnet_path = Sequential( + create_conv3(num_channels, num_channels, + bias=False, initialization_method=initialization_method, + use_spectral_norm=use_spectral_norm), + NormalizationLayerFactory.resolve_2d(normalization_layer_factory).create(num_channels, affine=True), + nonlinearity_factory.create(), + create_conv3(num_channels, num_channels, + bias=False, initialization_method=initialization_method, + use_spectral_norm=use_spectral_norm), + NormalizationLayerFactory.resolve_2d(normalization_layer_factory).create(num_channels, affine=True)) + + def forward(self, x): + if self.use_scale_parameter: + return x + self.scale * self.resnet_path(x) + else: + return x + self.resnet_path(x) diff --git a/tha2/nn/base/spectral_norm.py b/tha2/nn/base/spectral_norm.py new file mode 100644 index 0000000..4f08dd4 --- /dev/null +++ b/tha2/nn/base/spectral_norm.py @@ -0,0 +1,9 @@ +from torch.nn import Module +from torch.nn.utils import spectral_norm + + +def apply_spectral_norm(module: Module, use_spectrial_norm: bool = False) -> Module: + if use_spectrial_norm: + return spectral_norm(module) + else: + return module diff --git a/tha2/nn/base/util.py b/tha2/nn/base/util.py new file mode 100644 index 0000000..96b877a --- /dev/null +++ b/tha2/nn/base/util.py @@ -0,0 +1,35 @@ +from typing import Optional + +from torch.nn import Module + +from tha2.nn.base.init_function import create_init_function +from tha2.nn.base.module_factory import ModuleFactory +from tha2.nn.base.nonlinearity_factory import resolve_nonlinearity_factory +from tha2.nn.base.normalization import NormalizationLayerFactory +from tha2.nn.base.spectral_norm import apply_spectral_norm + + +def wrap_conv_or_linear_module(module: Module, initialization_method: str, use_spectral_norm: bool): + init = create_init_function(initialization_method) + return apply_spectral_norm(init(module), use_spectral_norm) + + +class ImageArgs: + def __init__(self, size: int = 64, num_channels: int = 3): + self.num_channels = num_channels + self.size = size + + +class BlockArgs: + def __init__(self, + initialization_method: str = 'he', + use_spectral_norm: bool = False, + normalization_layer_factory: Optional[NormalizationLayerFactory] = None, + nonlinearity_factory: Optional[ModuleFactory] = None): + self.nonlinearity_factory = resolve_nonlinearity_factory(nonlinearity_factory) + self.normalization_layer_factory = normalization_layer_factory + self.use_spectral_norm = use_spectral_norm + self.initialization_method = initialization_method + + def wrap_module(self, module: Module) -> Module: + return wrap_conv_or_linear_module(module, self.initialization_method, self.use_spectral_norm) diff --git a/tha2/nn/base/view_change.py b/tha2/nn/base/view_change.py new file mode 100644 index 0000000..b98470f --- /dev/null +++ b/tha2/nn/base/view_change.py @@ -0,0 +1,44 @@ +from typing import List + +from torch import Tensor +from torch.nn import Module + + +class ViewChange(Module): + def __init__(self, new_size: List[int]): + super().__init__() + self.new_size = new_size + + def forward(self, x: Tensor): + n = x.shape[0] + return x.view([n] + self.new_size) + + +class ViewImageAsVector(Module): + def __init__(self): + super().__init__() + + def forward(self, x: Tensor): + assert x.dim() == 4 + n, c, w, h = x.shape + return x.view(n, c * w * h) + + +class ViewVectorAsMultiChannelImage(Module): + def __init__(self): + super().__init__() + + def forward(self, x: Tensor): + assert x.dim() == 2 + n, c = x.shape + return x.view(n, c, 1, 1) + + +class ViewVectorAsOneChannelImage(Module): + def __init__(self): + super().__init__() + + def forward(self, x: Tensor): + assert x.dim() == 2 + n, c = x.shape + return x.view(n, 1, c, 1) diff --git a/tha2/nn/batch_module/__init__.py b/tha2/nn/batch_module/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha2/nn/batch_module/batch_input_model_factory.py b/tha2/nn/batch_module/batch_input_model_factory.py new file mode 100644 index 0000000..54ccab2 --- /dev/null +++ b/tha2/nn/batch_module/batch_input_model_factory.py @@ -0,0 +1,20 @@ +from typing import Dict, Set + +from tha2.nn.batch_module.batch_input_module import BatchInputModule, BatchInputModuleFactory + + +class BatchInputModelFactory: + def __init__(self, module_factories: Dict[str, BatchInputModuleFactory]): + self.module_factories = module_factories + + def get_module_names(self) -> Set[str]: + return set(self.module_factories.keys()) + + def create(self) -> Dict[str, BatchInputModule]: + output = {} + for name in self.module_factories: + output[name] = self.module_factories[name].create() + return output + + def get_module_factory(self, module_name) -> BatchInputModuleFactory: + return self.module_factories[module_name] \ No newline at end of file diff --git a/tha2/nn/batch_module/batch_input_module.py b/tha2/nn/batch_module/batch_input_module.py new file mode 100644 index 0000000..88ded97 --- /dev/null +++ b/tha2/nn/batch_module/batch_input_module.py @@ -0,0 +1,25 @@ +from abc import ABC, abstractmethod +from typing import List + +from torch import Tensor +from torch.nn import Module + +from tha2.nn.base.module_factory import ModuleFactory + + +class BatchInputModule(Module, ABC): + def __init__(self): + super().__init__() + + @abstractmethod + def forward_from_batch(self, batch: List[Tensor]): + pass + + +class BatchInputModuleFactory(ModuleFactory): + def __init__(self): + super().__init__() + + @abstractmethod + def create(self) -> BatchInputModule: + pass diff --git a/tha2/nn/eyebrow/__init__.py b/tha2/nn/eyebrow/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha2/nn/eyebrow/eyebrow_decomposer_00.py b/tha2/nn/eyebrow/eyebrow_decomposer_00.py new file mode 100644 index 0000000..9379460 --- /dev/null +++ b/tha2/nn/eyebrow/eyebrow_decomposer_00.py @@ -0,0 +1,104 @@ +from typing import List, Optional + +import torch +from torch import Tensor + +from tha2.nn.backbone.poser_encoder_decoder_00 import PoserEncoderDecoder00Args, PoserEncoderDecoder00 +from tha2.nn.util import apply_color_change +from tha2.nn.batch_module.batch_input_module import BatchInputModule, BatchInputModuleFactory +from tha2.nn.base.nonlinearity_factory import ReLUFactory +from tha2.nn.base.normalization import InstanceNorm2dFactory +from tha2.nn.base.util import BlockArgs + + +class EyebrowDecomposer00Args(PoserEncoderDecoder00Args): + def __init__(self, + image_size: int = 128, + image_channels: int = 4, + start_channels: int = 64, + bottleneck_image_size=16, + num_bottleneck_blocks=6, + max_channels: int = 512, + block_args: Optional[BlockArgs] = None): + super().__init__( + image_size, + image_channels, + image_channels, + 0, + start_channels, + bottleneck_image_size, + num_bottleneck_blocks, + max_channels, + block_args) + + +class EyebrowDecomposer00(BatchInputModule): + def __init__(self, args: EyebrowDecomposer00Args): + super().__init__() + self.args = args + self.body = PoserEncoderDecoder00(args) + self.background_layer_alpha = self.args.create_alpha_block() + self.background_layer_color_change = self.args.create_color_change_block() + self.eyebrow_layer_alpha = self.args.create_alpha_block() + self.eyebrow_layer_color_change = self.args.create_color_change_block() + + def forward(self, image: Tensor) -> List[Tensor]: + feature = self.body(image)[0] + + background_layer_alpha = self.background_layer_alpha(feature) + background_layer_color_change = self.background_layer_color_change(feature) + background_layer_1 = apply_color_change(background_layer_alpha, background_layer_color_change, image) + + eyebrow_layer_alpha = self.eyebrow_layer_alpha(feature) + eyebrow_layer_color_change = self.eyebrow_layer_color_change(feature) + eyebrow_layer = apply_color_change(eyebrow_layer_alpha, image, eyebrow_layer_color_change) + + return [ + eyebrow_layer, # 0 + eyebrow_layer_alpha, # 1 + eyebrow_layer_color_change, # 2 + background_layer_1, # 3 + background_layer_alpha, # 4 + background_layer_color_change, # 5 + ] + + EYEBROW_LAYER_INDEX = 0 + EYEBROW_LAYER_ALPHA_INDEX = 1 + EYEBROW_LAYER_COLOR_CHANGE_INDEX = 2 + BACKGROUND_LAYER_INDEX = 3 + BACKGROUND_LAYER_ALPHA_INDEX = 4 + BACKGROUND_LAYER_COLOR_CHANGE_INDEX = 5 + OUTPUT_LENGTH = 6 + + def forward_from_batch(self, batch: List[Tensor]): + return self.forward(batch[0]) + + +class EyebrowDecomposer00Factory(BatchInputModuleFactory): + def __init__(self, args: EyebrowDecomposer00Args): + super().__init__() + self.args = args + + def create(self) -> BatchInputModule: + return EyebrowDecomposer00(self.args) + + +if __name__ == "__main__": + cuda = torch.device('cuda') + args = EyebrowDecomposer00Args( + image_size=128, + image_channels=4, + start_channels=64, + bottleneck_image_size=16, + num_bottleneck_blocks=3, + block_args=BlockArgs( + initialization_method='xavier', + use_spectral_norm=False, + normalization_layer_factory=InstanceNorm2dFactory(), + nonlinearity_factory=ReLUFactory(inplace=True))) + face_morpher = EyebrowDecomposer00(args).to(cuda) + + image = torch.randn(8, 4, 128, 128, device=cuda) + outputs = face_morpher.forward(image) + for i in range(len(outputs)): + print(i, outputs[i].shape) diff --git a/tha2/nn/eyebrow/eyebrow_morphing_combiner_00.py b/tha2/nn/eyebrow/eyebrow_morphing_combiner_00.py new file mode 100644 index 0000000..de68bef --- /dev/null +++ b/tha2/nn/eyebrow/eyebrow_morphing_combiner_00.py @@ -0,0 +1,117 @@ +from typing import List, Optional + +import torch +from torch import Tensor + +from tha2.nn.backbone.poser_encoder_decoder_00 import PoserEncoderDecoder00Args, PoserEncoderDecoder00 +from tha2.nn.util import apply_color_change, apply_grid_change, apply_rgb_change +from tha2.nn.batch_module.batch_input_module import BatchInputModule, BatchInputModuleFactory +from tha2.nn.base.nonlinearity_factory import ReLUFactory +from tha2.nn.base.normalization import InstanceNorm2dFactory +from tha2.nn.base.util import BlockArgs + + +class EyebrowMorphingCombiner00Args(PoserEncoderDecoder00Args): + def __init__(self, + image_size: int = 128, + image_channels: int = 4, + num_pose_params: int = 12, + start_channels: int = 64, + bottleneck_image_size=16, + num_bottleneck_blocks=6, + max_channels: int = 512, + block_args: Optional[BlockArgs] = None): + super().__init__( + image_size, + 2 * image_channels, + image_channels, + num_pose_params, + start_channels, + bottleneck_image_size, + num_bottleneck_blocks, + max_channels, + block_args) + + +class EyebrowMorphingCombiner00(BatchInputModule): + def __init__(self, args: EyebrowMorphingCombiner00Args): + super().__init__() + self.args = args + self.body = PoserEncoderDecoder00(args) + self.morphed_eyebrow_layer_grid_change = self.args.create_grid_change_block() + self.morphed_eyebrow_layer_alpha = self.args.create_alpha_block() + self.morphed_eyebrow_layer_color_change = self.args.create_color_change_block() + self.combine_alpha = self.args.create_alpha_block() + + def forward(self, background_layer: Tensor, eyebrow_layer: Tensor, pose: Tensor) -> List[Tensor]: + combined_image = torch.cat([background_layer, eyebrow_layer], dim=1) + feature = self.body(combined_image, pose)[0] + + morphed_eyebrow_layer_grid_change = self.morphed_eyebrow_layer_grid_change(feature) + morphed_eyebrow_layer_alpha = self.morphed_eyebrow_layer_alpha(feature) + morphed_eyebrow_layer_color_change = self.morphed_eyebrow_layer_color_change(feature) + warped_eyebrow_layer = apply_grid_change(morphed_eyebrow_layer_grid_change, eyebrow_layer) + morphed_eyebrow_layer = apply_color_change( + morphed_eyebrow_layer_alpha, morphed_eyebrow_layer_color_change, warped_eyebrow_layer) + + combine_alpha = self.combine_alpha(feature) + eyebrow_image = apply_rgb_change(combine_alpha, morphed_eyebrow_layer, background_layer) + eyebrow_image_no_combine_alpha = apply_rgb_change( + (morphed_eyebrow_layer[:, 3:4, :, :] + 1.0) / 2.0, morphed_eyebrow_layer, background_layer) + + return [ + eyebrow_image, # 0 + combine_alpha, # 1 + eyebrow_image_no_combine_alpha, # 2 + morphed_eyebrow_layer, # 3 + morphed_eyebrow_layer_alpha, # 4 + morphed_eyebrow_layer_color_change, # 5 + warped_eyebrow_layer, # 6 + morphed_eyebrow_layer_grid_change, # 7 + ] + + EYEBROW_IMAGE_INDEX = 0 + COMBINE_ALPHA_INDEX = 1 + EYEBROW_IMAGE_NO_COMBINE_ALPHA_INDEX = 2 + MORPHED_EYEBROW_LAYER_INDEX = 3 + MORPHED_EYEBROW_LAYER_ALPHA_INDEX = 4 + MORPHED_EYEBROW_LAYER_COLOR_CHANGE_INDEX = 5 + WARPED_EYEBROW_LAYER_INDEX = 6 + MORPHED_EYEBROW_LAYER_GRID_CHANGE_INDEX = 7 + OUTPUT_LENGTH = 8 + + def forward_from_batch(self, batch: List[Tensor]): + return self.forward(batch[0], batch[1], batch[2]) + + +class EyebrowMorphingCombiner00Factory(BatchInputModuleFactory): + def __init__(self, args: EyebrowMorphingCombiner00Args): + super().__init__() + self.args = args + + def create(self) -> BatchInputModule: + return EyebrowMorphingCombiner00(self.args) + + +if __name__ == "__main__": + cuda = torch.device('cuda') + args = EyebrowMorphingCombiner00Args( + image_size=128, + image_channels=4, + num_pose_params=12, + start_channels=64, + bottleneck_image_size=16, + num_bottleneck_blocks=3, + block_args=BlockArgs( + initialization_method='xavier', + use_spectral_norm=False, + normalization_layer_factory=InstanceNorm2dFactory(), + nonlinearity_factory=ReLUFactory(inplace=True))) + face_morpher = EyebrowMorphingCombiner00(args).to(cuda) + + background_layer = torch.randn(8, 4, 128, 128, device=cuda) + eyebrow_layer = torch.randn(8, 4, 128, 128, device=cuda) + pose = torch.randn(8, 12, device=cuda) + outputs = face_morpher.forward(background_layer, eyebrow_layer, pose) + for i in range(len(outputs)): + print(i, outputs[i].shape) diff --git a/tha2/nn/face/__init__.py b/tha2/nn/face/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha2/nn/face/face_morpher_08.py b/tha2/nn/face/face_morpher_08.py new file mode 100644 index 0000000..247a469 --- /dev/null +++ b/tha2/nn/face/face_morpher_08.py @@ -0,0 +1,219 @@ +import math +from typing import List, Optional + +import torch +from torch import Tensor +from torch.nn import ModuleList, Sequential, Sigmoid, Tanh +from torch.nn.functional import affine_grid, grid_sample + +from tha2.nn.batch_module.batch_input_module import BatchInputModule, BatchInputModuleFactory +from tha2.nn.base.conv import create_conv3_block_from_block_args, \ + create_downsample_block_from_block_args, create_upsample_block_from_block_args, create_conv3_from_block_args, \ + create_conv3 +from tha2.nn.base.nonlinearity_factory import LeakyReLUFactory +from tha2.nn.base.normalization import InstanceNorm2dFactory +from tha2.nn.base.resnet_block import ResnetBlock +from tha2.nn.base.util import BlockArgs + + +class FaceMorpher08Args: + def __init__(self, + image_size: int = 256, + image_channels: int = 4, + num_expression_params: int = 67, + start_channels: int = 16, + bottleneck_image_size=4, + num_bottleneck_blocks=3, + max_channels: int = 512, + block_args: Optional[BlockArgs] = None): + self.max_channels = max_channels + self.num_bottleneck_blocks = num_bottleneck_blocks + assert bottleneck_image_size > 1 + self.bottleneck_image_size = bottleneck_image_size + self.start_channels = start_channels + self.image_channels = image_channels + self.num_expression_params = num_expression_params + self.image_size = image_size + + if block_args is None: + self.block_args = BlockArgs( + normalization_layer_factory=InstanceNorm2dFactory(), + nonlinearity_factory=LeakyReLUFactory(negative_slope=0.2, inplace=True)) + else: + self.block_args = block_args + + +class FaceMorpher08(BatchInputModule): + def __init__(self, args: FaceMorpher08Args): + super().__init__() + self.args = args + self.num_levels = int(math.log2(args.image_size // args.bottleneck_image_size)) + 1 + + self.downsample_blocks = ModuleList() + self.downsample_blocks.append( + create_conv3_block_from_block_args( + args.image_channels, + args.start_channels, + args.block_args)) + current_image_size = args.image_size + current_num_channels = args.start_channels + while current_image_size > args.bottleneck_image_size: + next_image_size = current_image_size // 2 + next_num_channels = self.get_num_output_channels_from_image_size(next_image_size) + self.downsample_blocks.append(create_downsample_block_from_block_args( + in_channels=current_num_channels, + out_channels=next_num_channels, + is_output_1x1=False, + block_args=args.block_args)) + current_image_size = next_image_size + current_num_channels = next_num_channels + assert len(self.downsample_blocks) == self.num_levels + + self.bottleneck_blocks = ModuleList() + self.bottleneck_blocks.append(create_conv3_block_from_block_args( + in_channels=current_num_channels + args.num_expression_params, + out_channels=current_num_channels, + block_args=args.block_args)) + for i in range(1, args.num_bottleneck_blocks): + self.bottleneck_blocks.append( + ResnetBlock.create( + num_channels=current_num_channels, + is1x1=False, + block_args=args.block_args)) + + self.upsample_blocks = ModuleList() + while current_image_size < args.image_size: + next_image_size = current_image_size * 2 + next_num_channels = self.get_num_output_channels_from_image_size(next_image_size) + self.upsample_blocks.append(create_upsample_block_from_block_args( + in_channels=current_num_channels, + out_channels=next_num_channels, + block_args=args.block_args)) + current_image_size = next_image_size + current_num_channels = next_num_channels + + self.iris_mouth_grid_change = self.create_grid_change_block() + self.iris_mouth_color_change = self.create_color_change_block() + self.iris_mouth_alpha = self.create_alpha_block() + + self.eye_color_change = self.create_color_change_block() + self.eye_alpha = self.create_alpha_block() + + def create_alpha_block(self): + return Sequential( + create_conv3( + in_channels=self.args.start_channels, + out_channels=1, + bias=True, + initialization_method=self.args.block_args.initialization_method, + use_spectral_norm=False), + Sigmoid()) + + def create_color_change_block(self): + return Sequential( + create_conv3_from_block_args( + in_channels=self.args.start_channels, + out_channels=self.args.image_channels, + bias=True, + block_args=self.args.block_args), + Tanh()) + + def create_grid_change_block(self): + return create_conv3( + in_channels=self.args.start_channels, + out_channels=2, + bias=False, + initialization_method='zero', + use_spectral_norm=False) + + def get_num_output_channels_from_level(self, level: int): + return self.get_num_output_channels_from_image_size(self.args.image_size // (2 ** level)) + + def get_num_output_channels_from_image_size(self, image_size: int): + return min(self.args.start_channels * (self.args.image_size // image_size), self.args.max_channels) + + def forward(self, image: Tensor, pose: Tensor) -> List[Tensor]: + feature = image + for block in self.downsample_blocks: + feature = block(feature) + n, c = pose.shape + pose = pose.view(n, c, 1, 1).repeat(1, 1, self.args.bottleneck_image_size, self.args.bottleneck_image_size) + feature = torch.cat([feature, pose], dim=1) + for block in self.bottleneck_blocks: + feature = block(feature) + for block in self.upsample_blocks: + feature = block(feature) + + iris_mouth_grid_change = self.iris_mouth_grid_change(feature) + iris_mouth_image_0 = self.apply_grid_change(iris_mouth_grid_change.half(), image) + iris_mouth_color_change = self.iris_mouth_color_change(feature) + iris_mouth_alpha = self.iris_mouth_alpha(feature) + iris_mouth_image_1 = self.apply_color_change(iris_mouth_alpha, iris_mouth_color_change, iris_mouth_image_0) + + eye_color_change = self.eye_color_change(feature) + eye_alpha = self.eye_alpha(feature) + output_image = self.apply_color_change(eye_alpha, eye_color_change, iris_mouth_image_1.detach()) + + # return [ + # output_image, #0 + # eye_alpha, #1 + # eye_color_change, #2 + # iris_mouth_image_1, #3 + # iris_mouth_alpha, #4 + # iris_mouth_color_change, #5 + # iris_mouth_image_0, #6 + # ] + return output_image + + def merge_down(self, top_layer: Tensor, bottom_layer: Tensor): + top_layer_rgb = top_layer[:, 0:3, :, :] + top_layer_a = top_layer[:, 3:4, :, :] + return bottom_layer * (1-top_layer_a) + torch.cat([top_layer_rgb * top_layer_a, top_layer_a], dim=1) + + def apply_grid_change(self, grid_change, image: Tensor) -> Tensor: + n, c, h, w = image.shape + device = grid_change.device + grid_change = torch.transpose(grid_change.view(n, 2, h * w), 1, 2).view(n, h, w, 2) + identity = torch.tensor([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], device=device).unsqueeze(0).repeat(n, 1, 1) + base_grid = affine_grid(identity, [n, c, h, w], align_corners=False) + grid = base_grid + grid_change + resampled_image = grid_sample(image, grid.half(), mode='bilinear', padding_mode='border', align_corners=False) + return resampled_image + + def apply_color_change(self, alpha, color_change, image: Tensor) -> Tensor: + return color_change * alpha + image * (1 - alpha) + + def forward_from_batch(self, batch: List[Tensor]): + return self.forward(batch[0], batch[1]) + + +class FaceMorpher08Factory(BatchInputModuleFactory): + def __init__(self, args: FaceMorpher08Args): + super().__init__() + self.args = args + + def create(self) -> BatchInputModule: + return FaceMorpher08(self.args) + + +if __name__ == "__main__": + cuda = torch.device('cuda') + args = FaceMorpher08Args( + image_size=256, + image_channels=4, + num_expression_params=67, + start_channels=32, + bottleneck_image_size=4, + num_bottleneck_blocks=3, + block_args=BlockArgs( + initialization_method='xavier', + use_spectral_norm=False, + normalization_layer_factory=InstanceNorm2dFactory(), + nonlinearity_factory=LeakyReLUFactory(inplace=True, negative_slope=0.2))) + face_morpher = FaceMorpher08(args).to(cuda) + + image = torch.randn(8, 4, 256, 256, device=cuda) + pose = torch.randn(8, 67, device=cuda) + outputs = face_morpher.forward(image, pose) + for i in range(len(outputs)): + print(i, outputs[i].shape) diff --git a/tha2/nn/util.py b/tha2/nn/util.py new file mode 100644 index 0000000..8ac9ee1 --- /dev/null +++ b/tha2/nn/util.py @@ -0,0 +1,25 @@ +import torch +from torch import Tensor +from torch.nn.functional import affine_grid, grid_sample + + +def apply_rgb_change(alpha: Tensor, color_change: Tensor, image: Tensor): + image_rgb = image[:, 0:3, :, :] + color_change_rgb = color_change[:, 0:3, :, :] + output_rgb = color_change_rgb * alpha + image_rgb * (1 - alpha) + return torch.cat([output_rgb, image[:, 3:4, :, :]], dim=1) + + +def apply_grid_change(grid_change, image: Tensor) -> Tensor: + n, c, h, w = image.shape + device = grid_change.device + grid_change = torch.transpose(grid_change.view(n, 2, h * w), 1, 2).view(n, h, w, 2) + identity = torch.tensor([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], device=device).unsqueeze(0).repeat(n, 1, 1) + base_grid = affine_grid(identity, [n, c, h, w], align_corners=False) + grid = base_grid + grid_change + resampled_image = grid_sample(image, grid, mode='bilinear', padding_mode='border', align_corners=False) + return resampled_image + + +def apply_color_change(alpha, color_change, image: Tensor) -> Tensor: + return color_change * alpha + image * (1 - alpha) diff --git a/tha2/poser/__init__.py b/tha2/poser/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha2/poser/general_poser_02.py b/tha2/poser/general_poser_02.py new file mode 100644 index 0000000..27ec6d9 --- /dev/null +++ b/tha2/poser/general_poser_02.py @@ -0,0 +1,72 @@ +from typing import List, Optional, Tuple, Dict, Callable + +import torch +from torch import Tensor + +from tha2.poser.poser import PoseParameterGroup, Poser +from tha2.nn.batch_module.batch_input_module import BatchInputModule +from tha2.compute.cached_computation_func import TensorListCachedComputationFunc + + +class GeneralPoser02(Poser): + def __init__(self, + module_loaders: Dict[str, Callable[[], BatchInputModule]], + device: torch.device, + output_length: int, + pose_parameters: List[PoseParameterGroup], + output_list_func: TensorListCachedComputationFunc, + subrect: Optional[Tuple[Tuple[int, int], Tuple[int, int]]] = None, + default_output_index: int = 0): + self.default_output_index = default_output_index + self.output_list_func = output_list_func + self.subrect = subrect + self.pose_parameters = pose_parameters + self.device = device + self.module_loaders = module_loaders + + self.modules = None + + self.num_parameters = 0 + for pose_parameter in self.pose_parameters: + self.num_parameters += pose_parameter.get_arity() + + self.output_length = output_length + + def get_modules(self): + if self.modules is None: + self.modules = {} + for key in self.module_loaders: + module = self.module_loaders[key]() + self.modules[key] = module + module.to(self.device) + module.train(False) + return self.modules + + def get_pose_parameter_groups(self) -> List[PoseParameterGroup]: + return self.pose_parameters + + def get_num_parameters(self) -> int: + return self.num_parameters + + def pose(self, image: Tensor, pose: Tensor, output_index: Optional[int] = None) -> Tensor: + if output_index is None: + output_index = self.default_output_index + output_list = self.get_posing_outputs(image, pose) + return output_list[output_index] + + def get_posing_outputs(self, image: Tensor, pose: Tensor) -> List[Tensor]: + modules = self.get_modules() + + if len(image.shape) == 3: + image = image.unsqueeze(0) + if len(pose.shape) == 1: + pose = pose.unsqueeze(0) + if self.subrect is not None: + image = image[:, :, self.subrect[0][0]:self.subrect[0][1], self.subrect[1][0]:self.subrect[1][1]] + batch = [image, pose] + + outputs = {} + return self.output_list_func(modules, batch, outputs) + + def get_output_length(self) -> int: + return self.output_length diff --git a/tha2/poser/modes/__init__.py b/tha2/poser/modes/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha2/poser/modes/mode_20.py b/tha2/poser/modes/mode_20.py new file mode 100644 index 0000000..a4fc8ed --- /dev/null +++ b/tha2/poser/modes/mode_20.py @@ -0,0 +1,286 @@ +import math +from enum import Enum +from typing import List, Dict, Optional + +import torch +from torch import Tensor + +from tha2.compute.cached_computation_func import TensorListCachedComputationFunc +from tha2.compute.cached_computation_protocol import CachedComputationProtocol +from tha2.nn.backcomp.tha.combiner import CombinerFactory +from tha2.nn.backcomp.tha.two_algo_face_rotator import TwoAlgoFaceRotatorFactory, TwoAlgoFaceRotator +from tha2.nn.base.nonlinearity_factory import ReLUFactory +from tha2.nn.base.normalization import InstanceNorm2dFactory +from tha2.nn.base.util import BlockArgs +from tha2.nn.batch_module.batch_input_module import BatchInputModule +from tha2.nn.eyebrow.eyebrow_decomposer_00 import EyebrowDecomposer00, \ + EyebrowDecomposer00Factory, EyebrowDecomposer00Args +from tha2.nn.eyebrow.eyebrow_morphing_combiner_00 import \ + EyebrowMorphingCombiner00Factory, EyebrowMorphingCombiner00Args, EyebrowMorphingCombiner00 +from tha2.nn.face.face_morpher_08 import FaceMorpher08Args, FaceMorpher08Factory +from tha2.poser.general_poser_02 import GeneralPoser02 +from tha2.poser.poser import Poser, PoseParameterCategory, PoseParameters +from tha2.util import torch_load + +KEY_EYEBROW_DECOMPOSER = "eyebrow_decomposer" +KEY_EYEBROW_MORPHING_COMBINER = "eyebrow_combiner" +KEY_FACE_MORPHER = "face_morpher" +KEY_FACE_ROTATER = "face_rotator" +KEY_COMBINER = "combiner" + +KEY_EYEBROW_DECOMPOSER_OUTPUT = "eyebrow_decomposer_output" +KEY_EYEBROW_MORPHING_COMBINER_OUTPUT = "eyebrow_combiner_output" +KEY_FACE_MORPHER_OUTPUT = "face_morpher_output" +KEY_FACE_ROTATER_OUTPUT = "face_rotator_output" +KEY_COMBINER_OUTPUT = "combiner_output" +KEY_ALL_OUTPUT = "all_output" + +NUM_EYEBROW_PARAMS = 12 +NUM_FACE_PARAMS = 27 +NUM_ROTATION_PARAMS = 3 + + +class FiveStepPoserComputationProtocol(CachedComputationProtocol): + def __init__(self, eyebrow_morphed_image_index: int): + super().__init__() + self.eyebrow_morphed_image_index = eyebrow_morphed_image_index + self.cached_batch_0 = None + self.cached_eyebrow_decomposer_output = None + + def compute_func(self) -> TensorListCachedComputationFunc: + def func(modules: Dict[str, BatchInputModule], + batch: List[Tensor], + outputs: Dict[str, List[Tensor]]): + new_batch_0 = self.cached_batch_0 is None or torch.max((batch[0] - self.cached_batch_0).abs()).item() > 0 + if not new_batch_0: + outputs[KEY_EYEBROW_DECOMPOSER_OUTPUT] = self.cached_eyebrow_decomposer_output + output = self.get_output(KEY_ALL_OUTPUT, modules, batch, outputs) + if new_batch_0: + self.cached_batch_0 = batch[0] + self.cached_eyebrow_decomposer_output = outputs[KEY_EYEBROW_DECOMPOSER_OUTPUT] + return output + + return func + + def compute_output(self, key: str, modules: Dict[str, BatchInputModule], batch: List[Tensor], + outputs: Dict[str, List[Tensor]]) -> List[Tensor]: + if key == KEY_EYEBROW_DECOMPOSER_OUTPUT: + input_image = batch[0][:, :, 64:192, 64:192] + return modules[KEY_EYEBROW_DECOMPOSER].forward_from_batch([input_image]) + elif key == KEY_EYEBROW_MORPHING_COMBINER_OUTPUT: + eyebrow_decomposer_output = self.get_output(KEY_EYEBROW_DECOMPOSER_OUTPUT, modules, batch, outputs) + background_layer = eyebrow_decomposer_output[EyebrowDecomposer00.BACKGROUND_LAYER_INDEX] + eyebrow_layer = eyebrow_decomposer_output[EyebrowDecomposer00.EYEBROW_LAYER_INDEX] + eyebrow_pose = batch[1][:, :NUM_EYEBROW_PARAMS] + return modules[KEY_EYEBROW_MORPHING_COMBINER].forward_from_batch([ + background_layer, + eyebrow_layer, + eyebrow_pose + ]) + elif key == KEY_FACE_MORPHER_OUTPUT: + eyebrow_morphing_combiner_output = self.get_output( + KEY_EYEBROW_MORPHING_COMBINER_OUTPUT, modules, batch, outputs) + eyebrow_morphed_image = eyebrow_morphing_combiner_output[self.eyebrow_morphed_image_index] + input_image = batch[0][:, :, 32:32 + 192, 32:32 + 192].clone() + input_image[:, :, 32:32 + 128, 32:32 + 128] = eyebrow_morphed_image + face_pose = batch[1][:, NUM_EYEBROW_PARAMS:NUM_EYEBROW_PARAMS + NUM_FACE_PARAMS] + return modules[KEY_FACE_MORPHER].forward_from_batch([ + input_image, + face_pose + ]) + elif key == KEY_FACE_ROTATER_OUTPUT: + face_morpher_output = self.get_output(KEY_FACE_MORPHER_OUTPUT, modules, batch, outputs) + face_morphed_image = face_morpher_output[0] + input_image = batch[0].clone() + input_image[:, :, 32:32 + 192, 32:32 + 192] = face_morphed_image + rotation_pose = batch[1][:, NUM_EYEBROW_PARAMS + NUM_FACE_PARAMS:] + return modules[KEY_FACE_ROTATER].forward_from_batch([ + input_image, + rotation_pose + ]) + elif key == KEY_COMBINER_OUTPUT: + face_rotater_output = self.get_output(KEY_FACE_ROTATER_OUTPUT, modules, batch, outputs) + color_changed_image = face_rotater_output[ + TwoAlgoFaceRotator.COLOR_CHANGED_IMAGE_INDEX] + resampled_image = face_rotater_output[ + TwoAlgoFaceRotator.RESAMPLED_IMAGE_INDEX] + rotation_pose = batch[1][:, NUM_EYEBROW_PARAMS + NUM_FACE_PARAMS:] + return modules[KEY_COMBINER].forward_from_batch([ + color_changed_image, + resampled_image, + rotation_pose + ]) + elif key == KEY_ALL_OUTPUT: + combiner_output = self.get_output(KEY_COMBINER_OUTPUT, modules, batch, outputs) + rotater_output = self.get_output(KEY_FACE_ROTATER_OUTPUT, modules, batch, outputs) + face_morpher_output = self.get_output(KEY_FACE_MORPHER_OUTPUT, modules, batch, outputs) + eyebrow_morphing_combiner_output = self.get_output( + KEY_EYEBROW_MORPHING_COMBINER_OUTPUT, modules, batch, outputs) + eyebrow_decomposer_output = self.get_output(KEY_EYEBROW_DECOMPOSER_OUTPUT, modules, batch, outputs) + output = combiner_output \ + + rotater_output \ + + face_morpher_output \ + + eyebrow_morphing_combiner_output \ + + eyebrow_decomposer_output + return output + else: + raise RuntimeError("Unsupported key: " + key) + + +def load_eyebrow_decomposer(file_name: str): + factory = EyebrowDecomposer00Factory( + EyebrowDecomposer00Args( + image_size=128, + image_channels=4, + start_channels=64, + bottleneck_image_size=16, + num_bottleneck_blocks=6, + max_channels=512, + block_args=BlockArgs( + initialization_method='he', + use_spectral_norm=False, + normalization_layer_factory=InstanceNorm2dFactory(), + nonlinearity_factory=ReLUFactory(inplace=True)))) + #print("Loading the eyebrow decomposer ... ", end="") + module = factory.create() + module.load_state_dict(torch_load(file_name)) + #print("DONE!!!") + return module + + +def load_eyebrow_morphing_combiner(file_name: str): + factory = EyebrowMorphingCombiner00Factory( + EyebrowMorphingCombiner00Args( + image_size=128, + image_channels=4, + start_channels=64, + num_pose_params=12, + bottleneck_image_size=16, + num_bottleneck_blocks=6, + max_channels=512, + block_args=BlockArgs( + initialization_method='he', + use_spectral_norm=False, + normalization_layer_factory=InstanceNorm2dFactory(), + nonlinearity_factory=ReLUFactory(inplace=True)))) + #print("Loading the eyebrow morphing conbiner ... ", end="") + module = factory.create() + module.load_state_dict(torch_load(file_name)) + #print("DONE!!!") + return module + + +def load_face_morpher(file_name: str): + factory = FaceMorpher08Factory( + FaceMorpher08Args( + image_size=192, + image_channels=4, + num_expression_params=27, + start_channels=64, + bottleneck_image_size=24, + num_bottleneck_blocks=6, + max_channels=512, + block_args=BlockArgs( + initialization_method='he', + use_spectral_norm=False, + normalization_layer_factory=InstanceNorm2dFactory(), + nonlinearity_factory=ReLUFactory(inplace=False)))) + #print("Loading the face morpher ... ", end="") + module = factory.create() + module.load_state_dict(torch_load(file_name)) + #print("DONE") + return module + + +def load_face_rotater(file_name: str): + #print("Loading the face rotater ... ", end="") + module = TwoAlgoFaceRotatorFactory().create() + module.load_state_dict(torch_load(file_name)) + #print("DONE!!!") + return module + + +def load_combiner(file_name: str): + #print("Loading the combiner ... ", end="") + module = CombinerFactory().create() + module.load_state_dict(torch_load(file_name)) + #print("DONE!!!") + return module + + +def get_pose_parameters(): + return PoseParameters.Builder() \ + .add_parameter_group("eyebrow_troubled", PoseParameterCategory.EYEBROW, arity=2) \ + .add_parameter_group("eyebrow_angry", PoseParameterCategory.EYEBROW, arity=2) \ + .add_parameter_group("eyebrow_lowered", PoseParameterCategory.EYEBROW, arity=2) \ + .add_parameter_group("eyebrow_raised", PoseParameterCategory.EYEBROW, arity=2) \ + .add_parameter_group("eyebrow_happy", PoseParameterCategory.EYEBROW, arity=2) \ + .add_parameter_group("eyebrow_serious", PoseParameterCategory.EYEBROW, arity=2) \ + .add_parameter_group("eye_wink", PoseParameterCategory.EYE, arity=2) \ + .add_parameter_group("eye_happy_wink", PoseParameterCategory.EYE, arity=2) \ + .add_parameter_group("eye_surprised", PoseParameterCategory.EYE, arity=2) \ + .add_parameter_group("eye_relaxed", PoseParameterCategory.EYE, arity=2) \ + .add_parameter_group("eye_unimpressed", PoseParameterCategory.EYE, arity=2) \ + .add_parameter_group("eye_raised_lower_eyelid", PoseParameterCategory.EYE, arity=2) \ + .add_parameter_group("iris_small", PoseParameterCategory.IRIS_MORPH, arity=2) \ + .add_parameter_group("mouth_aaa", PoseParameterCategory.MOUTH, arity=1, default_value=1.0) \ + .add_parameter_group("mouth_iii", PoseParameterCategory.MOUTH, arity=1) \ + .add_parameter_group("mouth_uuu", PoseParameterCategory.MOUTH, arity=1) \ + .add_parameter_group("mouth_eee", PoseParameterCategory.MOUTH, arity=1) \ + .add_parameter_group("mouth_ooo", PoseParameterCategory.MOUTH, arity=1) \ + .add_parameter_group("mouth_delta", PoseParameterCategory.MOUTH, arity=1) \ + .add_parameter_group("mouth_lowered_corner", PoseParameterCategory.MOUTH, arity=2) \ + .add_parameter_group("mouth_raised_corner", PoseParameterCategory.MOUTH, arity=2) \ + .add_parameter_group("mouth_smirk", PoseParameterCategory.MOUTH, arity=1) \ + .add_parameter_group("iris_rotation_x", PoseParameterCategory.IRIS_ROTATION, arity=1, range=(-1.0, 1.0)) \ + .add_parameter_group("iris_rotation_y", PoseParameterCategory.IRIS_ROTATION, arity=1, range=(-1.0, 1.0)) \ + .add_parameter_group("head_x", PoseParameterCategory.FACE_ROTATION, arity=1, range=(-1.0, 1.0)) \ + .add_parameter_group("head_y", PoseParameterCategory.FACE_ROTATION, arity=1, range=(-1.0, 1.0)) \ + .add_parameter_group("neck_z", PoseParameterCategory.FACE_ROTATION, arity=1, range=(-1.0, 1.0)) \ + .build() + + +def create_poser( + device: torch.device, + module_file_names: Optional[Dict[str, str]] = None, + eyebrow_morphed_image_index: int = EyebrowMorphingCombiner00.EYEBROW_IMAGE_NO_COMBINE_ALPHA_INDEX, + default_output_index: int = 0) -> Poser: + dir = "data" + if module_file_names is None: + module_file_names = {} + if KEY_EYEBROW_DECOMPOSER not in module_file_names: + file_name = dir + "/eyebrow_decomposer.pt" + module_file_names[KEY_EYEBROW_DECOMPOSER] = file_name + if KEY_EYEBROW_MORPHING_COMBINER not in module_file_names: + file_name = dir + "/eyebrow_morphing_combiner.pt" + module_file_names[KEY_EYEBROW_MORPHING_COMBINER] = file_name + if KEY_FACE_MORPHER not in module_file_names: + file_name = dir + "/face_morpher.pt" + module_file_names[KEY_FACE_MORPHER] = file_name + if KEY_FACE_ROTATER not in module_file_names: + file_name = dir + "/two_algo_face_rotator.pt" + module_file_names[KEY_FACE_ROTATER] = file_name + if KEY_COMBINER not in module_file_names: + file_name = dir + "/combiner.pt" + module_file_names[KEY_COMBINER] = file_name + + loaders = { + KEY_EYEBROW_DECOMPOSER: + lambda: load_eyebrow_decomposer(module_file_names[KEY_EYEBROW_DECOMPOSER]), + KEY_EYEBROW_MORPHING_COMBINER: + lambda: load_eyebrow_morphing_combiner(module_file_names[KEY_EYEBROW_MORPHING_COMBINER]), + KEY_FACE_MORPHER: + lambda: load_face_morpher(module_file_names[KEY_FACE_MORPHER]), + KEY_FACE_ROTATER: + lambda: load_face_rotater(module_file_names[KEY_FACE_ROTATER]), + KEY_COMBINER: + lambda: load_combiner(module_file_names[KEY_COMBINER]), + } + return GeneralPoser02( + module_loaders=loaders, + pose_parameters=get_pose_parameters().get_pose_parameter_groups(), + output_list_func=FiveStepPoserComputationProtocol(eyebrow_morphed_image_index).compute_func(), + subrect=None, + device=device, + output_length=31, + default_output_index=default_output_index) \ No newline at end of file diff --git a/tha2/poser/modes/mode_20_wx.py b/tha2/poser/modes/mode_20_wx.py new file mode 100644 index 0000000..a1fa232 --- /dev/null +++ b/tha2/poser/modes/mode_20_wx.py @@ -0,0 +1,387 @@ +import math +from enum import Enum +from typing import List, Dict, Optional + +import numpy +import scipy.optimize +import wx + +from tha2.mocap.ifacialmocap_constants import * +from tha2.mocap.ifacialmocap_pose_converter import IFacialMocapPoseConverter +from tha2.poser.modes.mode_20 import get_pose_parameters + + +def clamp(x, min_value, max_value): + return max(min_value, min(max_value, x)) + + +class EyebrowDownMode(Enum): + TROUBLED = 1 + ANGRY = 2 + LOWERED = 3 + SERIOUS = 4 + + +class WinkMode(Enum): + NORMAL = 1 + RELAXED = 2 + + +class IFacialMocapPoseConverter20Args: + def __init__(self, + lower_smile_threshold: float = 0.4, + upper_smile_threshold: float = 0.6, + eyebrow_down_mode: EyebrowDownMode = EyebrowDownMode.ANGRY, + wink_mode: WinkMode = WinkMode.NORMAL, + eye_surprised_max_value: float = 0.5, + eye_wink_max_value: float = 0.8, + eyebrow_down_max_value: float = 0.4, + cheek_squint_min_value: float = 0.1, + cheek_squint_max_value: float = 0.7, + eye_rotation_factor: float = 1.0 / 0.75, + jaw_open_min_value: float = 0.1, + jaw_open_max_value: float = 0.4, + mouth_frown_max_value: float = 0.6, + mouth_funnel_min_value: float = 0.25, + mouth_funnel_max_value: float = 0.5, + iris_small_left=0.0, + iris_small_right=0.0): + self.iris_small_right = iris_small_left + self.iris_small_left = iris_small_right + self.wink_mode = wink_mode + self.mouth_funnel_max_value = mouth_funnel_max_value + self.mouth_funnel_min_value = mouth_funnel_min_value + self.mouth_frown_max_value = mouth_frown_max_value + self.jaw_open_max_value = jaw_open_max_value + self.jaw_open_min_value = jaw_open_min_value + self.eye_rotation_factor = eye_rotation_factor + self.cheek_squint_max_value = cheek_squint_max_value + self.cheek_squint_min_value = cheek_squint_min_value + self.eyebrow_down_max_value = eyebrow_down_max_value + self.eye_blink_max_value = eye_wink_max_value + self.eye_wide_max_value = eye_surprised_max_value + self.eyebrow_down_mode = eyebrow_down_mode + self.lower_smile_threshold = lower_smile_threshold + self.upper_smile_threshold = upper_smile_threshold + + +class IFacialMocapPoseConverter20(IFacialMocapPoseConverter): + def __init__(self, args: Optional[IFacialMocapPoseConverter20Args] = None): + super().__init__() + if args is None: + args = IFacialMocapPoseConverter20Args() + self.args = args + pose_parameters = get_pose_parameters() + self.pose_size = 42 + + self.eyebrow_troubled_left_index = pose_parameters.get_parameter_index("eyebrow_troubled_left") + self.eyebrow_troubled_right_index = pose_parameters.get_parameter_index("eyebrow_troubled_right") + self.eyebrow_angry_left_index = pose_parameters.get_parameter_index("eyebrow_angry_left") + self.eyebrow_angry_right_index = pose_parameters.get_parameter_index("eyebrow_angry_right") + self.eyebrow_happy_left_index = pose_parameters.get_parameter_index("eyebrow_happy_left") + self.eyebrow_happy_right_index = pose_parameters.get_parameter_index("eyebrow_happy_right") + self.eyebrow_raised_left_index = pose_parameters.get_parameter_index("eyebrow_raised_left") + self.eyebrow_raised_right_index = pose_parameters.get_parameter_index("eyebrow_raised_right") + self.eyebrow_lowered_left_index = pose_parameters.get_parameter_index("eyebrow_lowered_left") + self.eyebrow_lowered_right_index = pose_parameters.get_parameter_index("eyebrow_lowered_right") + self.eyebrow_serious_left_index = pose_parameters.get_parameter_index("eyebrow_serious_left") + self.eyebrow_serious_right_index = pose_parameters.get_parameter_index("eyebrow_serious_right") + + self.eye_surprised_left_index = pose_parameters.get_parameter_index("eye_surprised_left") + self.eye_surprised_right_index = pose_parameters.get_parameter_index("eye_surprised_right") + self.eye_wink_left_index = pose_parameters.get_parameter_index("eye_wink_left") + self.eye_wink_right_index = pose_parameters.get_parameter_index("eye_wink_right") + self.eye_happy_wink_left_index = pose_parameters.get_parameter_index("eye_happy_wink_left") + self.eye_happy_wink_right_index = pose_parameters.get_parameter_index("eye_happy_wink_right") + self.eye_relaxed_left_index = pose_parameters.get_parameter_index("eye_relaxed_left") + self.eye_relaxed_right_index = pose_parameters.get_parameter_index("eye_relaxed_right") + self.eye_raised_lower_eyelid_left_index = pose_parameters.get_parameter_index("eye_raised_lower_eyelid_left") + self.eye_raised_lower_eyelid_right_index = pose_parameters.get_parameter_index("eye_raised_lower_eyelid_right") + + self.iris_small_left_index = pose_parameters.get_parameter_index("iris_small_left") + self.iris_small_right_index = pose_parameters.get_parameter_index("iris_small_right") + + self.iris_rotation_x_index = pose_parameters.get_parameter_index("iris_rotation_x") + self.iris_rotation_y_index = pose_parameters.get_parameter_index("iris_rotation_y") + + self.head_x_index = pose_parameters.get_parameter_index("head_x") + self.head_y_index = pose_parameters.get_parameter_index("head_y") + self.neck_z_index = pose_parameters.get_parameter_index("neck_z") + + self.mouth_aaa_index = pose_parameters.get_parameter_index("mouth_aaa") + self.mouth_iii_index = pose_parameters.get_parameter_index("mouth_iii") + self.mouth_uuu_index = pose_parameters.get_parameter_index("mouth_uuu") + self.mouth_eee_index = pose_parameters.get_parameter_index("mouth_eee") + self.mouth_ooo_index = pose_parameters.get_parameter_index("mouth_ooo") + + self.mouth_lowered_corner_left_index = pose_parameters.get_parameter_index("mouth_lowered_corner_left") + self.mouth_lowered_corner_right_index = pose_parameters.get_parameter_index("mouth_lowered_corner_right") + self.mouth_raised_corner_left_index = pose_parameters.get_parameter_index("mouth_raised_corner_left") + self.mouth_raised_corner_right_index = pose_parameters.get_parameter_index("mouth_raised_corner_right") + + def convert(self, ifacialmocap_pose: Dict[str, float]) -> List[float]: + pose = [0.0 for i in range(self.pose_size)] + + smile_value = \ + (ifacialmocap_pose[MOUTH_SMILE_LEFT] + ifacialmocap_pose[MOUTH_SMILE_RIGHT]) / 2.0 \ + + ifacialmocap_pose[MOUTH_SHRUG_UPPER] + smile_degree = clamp((smile_value - self.args.lower_smile_threshold) / ( + self.args.upper_smile_threshold - self.args.lower_smile_threshold), 0.0, 1.0) + + # Eyebrow + if True: + brow_inner_up = ifacialmocap_pose[BROW_INNER_UP] + brow_outer_up_right = ifacialmocap_pose[BROW_OUTER_UP_RIGHT] + brow_outer_up_left = ifacialmocap_pose[BROW_OUTER_UP_LEFT] + + brow_up_left = clamp(brow_inner_up + brow_outer_up_left, 0.0, 1.0) + brow_up_right = clamp(brow_inner_up + brow_outer_up_right, 0.0, 1.0) + pose[self.eyebrow_raised_left_index] = brow_up_left + pose[self.eyebrow_raised_right_index] = brow_up_right + + brow_down_left = (1.0 - smile_degree) \ + * clamp(ifacialmocap_pose[BROW_DOWN_LEFT] / self.args.eyebrow_down_max_value, 0.0, 1.0) + brow_down_right = (1.0 - smile_degree) \ + * clamp(ifacialmocap_pose[BROW_DOWN_RIGHT] / self.args.eyebrow_down_max_value, 0.0, 1.0) + if self.args.eyebrow_down_mode == EyebrowDownMode.TROUBLED: + pose[self.eyebrow_troubled_left_index] = brow_down_left + pose[self.eyebrow_troubled_right_index] = brow_down_right + elif self.args.eyebrow_down_mode == EyebrowDownMode.ANGRY: + pose[self.eyebrow_angry_left_index] = brow_down_left + pose[self.eyebrow_angry_right_index] = brow_down_right + elif self.args.eyebrow_down_mode == EyebrowDownMode.LOWERED: + pose[self.eyebrow_lowered_left_index] = brow_down_left + pose[self.eyebrow_lowered_right_index] = brow_down_right + elif self.args.eyebrow_down_mode == EyebrowDownMode.SERIOUS: + pose[self.eyebrow_serious_left_index] = brow_down_left + pose[self.eyebrow_serious_right_index] = brow_down_right + + brow_happy_value = clamp(smile_value, 0.0, 1.0) * smile_degree + pose[self.eyebrow_happy_left_index] = brow_happy_value + pose[self.eyebrow_happy_right_index] = brow_happy_value + + # Eye + if True: + # Surprised + pose[self.eye_surprised_left_index] = clamp( + ifacialmocap_pose[EYE_WIDE_LEFT] / self.args.eye_wide_max_value, 0.0, 1.0) + pose[self.eye_surprised_right_index] = clamp( + ifacialmocap_pose[EYE_WIDE_RIGHT] / self.args.eye_wide_max_value, 0.0, 1.0) + + # Wink + if self.args.wink_mode == WinkMode.NORMAL: + wink_left_index = self.eye_wink_left_index + wink_right_index = self.eye_wink_right_index + else: + wink_left_index = self.eye_relaxed_left_index + wink_right_index = self.eye_relaxed_right_index + pose[wink_left_index] = (1.0 - smile_degree) * clamp( + ifacialmocap_pose[EYE_BLINK_LEFT] / self.args.eye_blink_max_value, 0.0, 1.0) + pose[wink_right_index] = (1.0 - smile_degree) * clamp( + ifacialmocap_pose[EYE_BLINK_RIGHT] / self.args.eye_blink_max_value, 0.0, 1.0) + pose[self.eye_happy_wink_left_index] = smile_degree * clamp( + ifacialmocap_pose[EYE_BLINK_LEFT] / self.args.eye_blink_max_value, 0.0, 1.0) + pose[self.eye_happy_wink_right_index] = smile_degree * clamp( + ifacialmocap_pose[EYE_BLINK_RIGHT] / self.args.eye_blink_max_value, 0.0, 1.0) + + # Lower eyelid + cheek_squint_denom = self.args.cheek_squint_max_value - self.args.cheek_squint_min_value + pose[self.eye_raised_lower_eyelid_left_index] = \ + clamp( + (ifacialmocap_pose[CHEEK_SQUINT_LEFT] - self.args.cheek_squint_min_value) / cheek_squint_denom, + 0.0, 1.0) + pose[self.eye_raised_lower_eyelid_right_index] = \ + clamp( + (ifacialmocap_pose[CHEEK_SQUINT_RIGHT] - self.args.cheek_squint_min_value) / cheek_squint_denom, + 0.0, 1.0) + + # Iris rotation + if True: + eye_rotation_y = (ifacialmocap_pose[EYE_LOOK_IN_LEFT] - + ifacialmocap_pose[EYE_LOOK_OUT_LEFT] - + ifacialmocap_pose[EYE_LOOK_IN_RIGHT] + + ifacialmocap_pose[EYE_LOOK_OUT_RIGHT]) / 2.0 * self.args.eye_rotation_factor + pose[self.iris_rotation_y_index] = clamp(eye_rotation_y, -1.0, 1.0) + + eye_rotation_x = (ifacialmocap_pose[EYE_LOOK_UP_LEFT] + + ifacialmocap_pose[EYE_LOOK_UP_RIGHT] + - ifacialmocap_pose[EYE_LOOK_DOWN_RIGHT] + + ifacialmocap_pose[EYE_LOOK_DOWN_LEFT]) / 2.0 * self.args.eye_rotation_factor + pose[self.iris_rotation_x_index] = clamp(eye_rotation_x, -1.0, 1.0) + + # Iris size + if True: + pose[self.iris_small_left_index] = self.args.iris_small_left + pose[self.iris_small_right_index] = self.args.iris_small_right + + # Head rotation + if True: + pose[self.head_x_index] = clamp(-ifacialmocap_pose[HEAD_BONE_X] * 180.0 / math.pi, -15.0, 15.0) / 15.0 + pose[self.head_y_index] = clamp(-ifacialmocap_pose[HEAD_BONE_Y] * 180.0 / math.pi, -15.0, 15.0) / 15.0 + pose[self.neck_z_index] = clamp(ifacialmocap_pose[HEAD_BONE_Z] * 180.0 / math.pi, -15.0, 15.0) / 15.0 + + # Mouth + if True: + # mouth_open = clamp((ifacialmocap_pose[JAW_OPEN] - 0.10) / 0.7, 0.0, 1.0) + jaw_open_denom = self.args.jaw_open_max_value - self.args.jaw_open_min_value + mouth_open = clamp((ifacialmocap_pose[JAW_OPEN] - self.args.jaw_open_min_value) / jaw_open_denom, 0.0, 1.0) + pose[self.mouth_aaa_index] = mouth_open + pose[self.mouth_raised_corner_left_index] = clamp(smile_value, 0.0, 1.0) + pose[self.mouth_raised_corner_right_index] = clamp(smile_value, 0.0, 1.0) + + is_mouth_open = mouth_open > 0.0 + if not is_mouth_open: + mouth_frown_value = clamp( + (ifacialmocap_pose[MOUTH_FROWN_LEFT] + ifacialmocap_pose[ + MOUTH_FROWN_RIGHT]) / self.args.mouth_frown_max_value, 0.0, 1.0) + pose[self.mouth_lowered_corner_left_index] = mouth_frown_value + pose[self.mouth_lowered_corner_right_index] = mouth_frown_value + else: + mouth_lower_down = clamp( + ifacialmocap_pose[MOUTH_LOWER_DOWN_LEFT] + ifacialmocap_pose[MOUTH_LOWER_DOWN_RIGHT], 0.0, 1.0) + mouth_funnel = ifacialmocap_pose[MOUTH_FUNNEL] + mouth_pucker = ifacialmocap_pose[MOUTH_PUCKER] + + mouth_point = [mouth_open, mouth_lower_down, mouth_funnel, mouth_pucker] + + aaa_point = [1.0, 1.0, 0.0, 0.0] + iii_point = [0.0, 1.0, 0.0, 0.0] + uuu_point = [0.5, 0.3, 0.25, 0.75] + ooo_point = [1.0, 0.5, 0.5, 0.4] + + decomp = numpy.array([0, 0, 0, 0]) + M = numpy.array([ + aaa_point, + iii_point, + uuu_point, + ooo_point + ]) + + def loss(decomp): + return numpy.linalg.norm(numpy.matmul(decomp, M) - mouth_point) \ + + 0.01 * numpy.linalg.norm(decomp, ord=1) + + opt_result = scipy.optimize.minimize(loss, decomp, + bounds=[(0.0, 1.0), (0.0, 1.0), (0.0, 1.0), (0.0, 1.0)]) + decomp = opt_result["x"] + restricted_decomp = [decomp.item(0), decomp.item(1), decomp.item(2), decomp.item(3)] + # restricted_decomp = restrict_to_two_morphs( + # [decomp.item(0), decomp.item(1), decomp.item(2), decomp.item(3)]) + pose[self.mouth_aaa_index] = restricted_decomp[0] + pose[self.mouth_iii_index] = restricted_decomp[1] + mouth_funnel_denom = self.args.mouth_funnel_max_value - self.args.mouth_funnel_min_value + ooo_alpha = clamp((mouth_funnel - self.args.mouth_funnel_min_value) / mouth_funnel_denom, 0.0, 1.0) + uo_value = clamp(restricted_decomp[2] + restricted_decomp[3], 0.0, 1.0) + pose[self.mouth_uuu_index] = uo_value * (1.0 - ooo_alpha) + pose[self.mouth_ooo_index] = uo_value * ooo_alpha + # pose[28] = restricted_decomp[2] + # pose[30] = restricted_decomp[3] + + return pose + + def init_pose_converter_panel(self, parent): + self.panel = wx.Panel(parent, style=wx.SIMPLE_BORDER) + self.panel_sizer = wx.BoxSizer(wx.VERTICAL) + self.panel.SetSizer(self.panel_sizer) + self.panel.SetAutoLayout(1) + parent.GetSizer().Add(self.panel, 0, wx.EXPAND) + + if True: + eyebrow_down_mode_text = wx.StaticText(self.panel, label=" --- Eyebrow Down Mode --- ", + style=wx.ALIGN_CENTER) + self.panel_sizer.Add(eyebrow_down_mode_text, 0, wx.EXPAND) + + self.eyebrow_down_mode_choice = wx.Choice( + self.panel, + choices=[ + "ANGRY", + "TROUBLED", + "SERIOUS", + "LOWERED", + ]) + self.eyebrow_down_mode_choice.SetSelection(0) + self.panel_sizer.Add(self.eyebrow_down_mode_choice, 0, wx.EXPAND) + self.eyebrow_down_mode_choice.Bind(wx.EVT_CHOICE, self.change_eyebrow_down_mode) + + separator = wx.StaticLine(self.panel, -1, size=(256, 5)) + self.panel_sizer.Add(separator, 0, wx.EXPAND) + + if True: + wink_mode_text = wx.StaticText(self.panel, label=" --- Wink Mode --- ", style=wx.ALIGN_CENTER) + self.panel_sizer.Add(wink_mode_text, 0, wx.EXPAND) + + self.wink_mode_choice = wx.Choice( + self.panel, + choices=[ + "NORMAL", + "RELAXED", + ]) + self.wink_mode_choice.SetSelection(0) + self.panel_sizer.Add(self.wink_mode_choice, 0, wx.EXPAND) + self.wink_mode_choice.Bind(wx.EVT_CHOICE, self.change_wink_mode) + + separator = wx.StaticLine(self.panel, -1, size=(256, 5)) + self.panel_sizer.Add(separator, 0, wx.EXPAND) + + if True: + iris_size_text = wx.StaticText(self.panel, label=" --- Iris Size --- ", style=wx.ALIGN_CENTER) + self.panel_sizer.Add(iris_size_text, 0, wx.EXPAND) + + self.iris_left_slider = wx.Slider(self.panel, minValue=0, maxValue=1000, value=0, style=wx.HORIZONTAL) + self.panel_sizer.Add(self.iris_left_slider, 0, wx.EXPAND) + self.iris_left_slider.Bind(wx.EVT_SLIDER, self.change_iris_size) + + self.iris_right_slider = wx.Slider(self.panel, minValue=0, maxValue=1000, value=0, style=wx.HORIZONTAL) + self.panel_sizer.Add(self.iris_right_slider, 0, wx.EXPAND) + self.iris_right_slider.Bind(wx.EVT_SLIDER, self.change_iris_size) + self.iris_right_slider.Enable(False) + + self.link_left_right_irises = wx.CheckBox( + self.panel, label="Use same value for both sides") + self.link_left_right_irises.SetValue(True) + self.panel_sizer.Add(self.link_left_right_irises, wx.SizerFlags().CenterHorizontal().Border()) + self.link_left_right_irises.Bind(wx.EVT_CHECKBOX, self.link_left_right_irises_clicked) + + self.panel_sizer.Fit(self.panel) + + def change_eyebrow_down_mode(self, event: wx.Event): + selected_index = self.eyebrow_down_mode_choice.GetSelection() + if selected_index == 0: + self.args.eyebrow_down_mode = EyebrowDownMode.ANGRY + elif selected_index == 1: + self.args.eyebrow_down_mode = EyebrowDownMode.TROUBLED + elif selected_index == 2: + self.args.eyebrow_down_mode = EyebrowDownMode.SERIOUS + else: + self.args.eyebrow_down_mode = EyebrowDownMode.LOWERED + + def change_wink_mode(self, event: wx.Event): + selected_index = self.wink_mode_choice.GetSelection() + if selected_index == 0: + self.args.wink_mode = WinkMode.NORMAL + else: + self.args.wink_mode = WinkMode.RELAXED + + def change_iris_size(self, event: wx.Event): + if self.link_left_right_irises.GetValue(): + left_value = self.iris_left_slider.GetValue() + right_value = self.iris_right_slider.GetValue() + if left_value != right_value: + self.iris_right_slider.SetValue(left_value) + self.args.iris_small_left = left_value / 1000.0 + self.args.iris_small_right = left_value / 1000.0 + else: + self.args.iris_small_left = self.iris_left_slider.GetValue() / 1000.0 + self.args.iris_small_right = self.iris_right_slider.GetValue() / 1000.0 + + def link_left_right_irises_clicked(self, event: wx.Event): + if self.link_left_right_irises.GetValue(): + self.iris_right_slider.Enable(False) + else: + self.iris_right_slider.Enable(True) + self.change_iris_size(event) + + +def create_ifacialmocap_pose_converter( + args: Optional[IFacialMocapPoseConverter20Args] = None) -> IFacialMocapPoseConverter: + return IFacialMocapPoseConverter20(args) diff --git a/tha2/poser/poser.py b/tha2/poser/poser.py new file mode 100644 index 0000000..9cb440f --- /dev/null +++ b/tha2/poser/poser.py @@ -0,0 +1,148 @@ +from abc import ABC, abstractmethod +from enum import Enum +from typing import Tuple, List, Optional + +from torch import Tensor + + +class PoseParameterCategory(Enum): + EYEBROW = 1 + EYE = 2 + IRIS_MORPH = 3 + IRIS_ROTATION = 4 + MOUTH = 5 + FACE_ROTATION = 6 + + +class PoseParameterGroup: + def __init__(self, + group_name: str, + parameter_index: int, + category: PoseParameterCategory, + arity: int = 1, + discrete: bool = False, + default_value: float = 0.0, + range: Optional[Tuple[float, float]] = None): + assert arity == 1 or arity == 2 + if range is None: + range = (0.0, 1.0) + if arity == 1: + parameter_names = [group_name] + else: + parameter_names = [group_name + "_left", group_name + "_right"] + assert len(parameter_names) == arity + + self.parameter_names = parameter_names + self.range = range + self.default_value = default_value + self.discrete = discrete + self.arity = arity + self.category = category + self.parameter_index = parameter_index + self.group_name = group_name + + def get_arity(self) -> int: + return self.arity + + def get_group_name(self) -> str: + return self.group_name + + def get_parameter_names(self) -> List[str]: + return self.parameter_names + + def is_discrete(self) -> bool: + return self.discrete + + def get_range(self) -> Tuple[float, float]: + return self.range + + def get_default_value(self): + return self.default_value + + def get_parameter_index(self): + return self.parameter_index + + def get_category(self) -> PoseParameterCategory: + return self.category + + +class PoseParameters: + def __init__(self, pose_parameter_groups: List[PoseParameterGroup]): + self.pose_parameter_groups = pose_parameter_groups + + def get_parameter_index(self, name: str) -> int: + index = 0 + for parameter_group in self.pose_parameter_groups: + for param_name in parameter_group.parameter_names: + if name == param_name: + return index + index += 1 + raise RuntimeError("Cannot find parameter with name %s" % name) + + def get_parameter_name(self, index: int) -> str: + assert index >= 0 and index < self.get_parameter_count() + + for group in self.pose_parameter_groups: + if index < group.get_arity(): + return group.get_parameter_names()[index] + index -= group.arity + + raise RuntimeError("Something is wrong here!!!") + + def get_pose_parameter_groups(self): + return self.pose_parameter_groups + + def get_parameter_count(self): + count = 0 + for group in self.pose_parameter_groups: + count += group.arity + return count + + class Builder: + def __init__(self): + self.index = 0 + self.pose_parameter_groups = [] + + def add_parameter_group(self, + group_name: str, + category: PoseParameterCategory, + arity: int = 1, + discrete: bool = False, + default_value: float = 0.0, + range: Optional[Tuple[float, float]] = None): + self.pose_parameter_groups.append( + PoseParameterGroup( + group_name, + self.index, + category, + arity, + discrete, + default_value, + range)) + self.index += arity + return self + + def build(self) -> 'PoseParameters': + return PoseParameters(self.pose_parameter_groups) + + +class Poser(ABC): + @abstractmethod + def get_output_length(self) -> int: + pass + + @abstractmethod + def get_pose_parameter_groups(self) -> List[PoseParameterGroup]: + pass + + @abstractmethod + def get_num_parameters(self) -> int: + pass + + @abstractmethod + def pose(self, image: Tensor, pose: Tensor, output_index: int = 0) -> Tensor: + pass + + @abstractmethod + def get_posing_outputs(self, image: Tensor, pose: Tensor) -> List[Tensor]: + pass diff --git a/tha2/util.py b/tha2/util.py new file mode 100644 index 0000000..f9dc305 --- /dev/null +++ b/tha2/util.py @@ -0,0 +1,213 @@ +import math +import os +from typing import List + +import PIL.Image +import numpy +import torch +from matplotlib import cm +from torch import Tensor + + +def is_power2(x): + return x != 0 and ((x & (x - 1)) == 0) + + +def torch_save(content, file_name): + os.makedirs(os.path.dirname(file_name), exist_ok=True) + with open(file_name, 'wb') as f: + torch.save(content, f) + + +def torch_load(file_name): + with open(file_name, 'rb') as f: + return torch.load(f) + + +def srgb_to_linear(x): + x = numpy.clip(x, 0.0, 1.0) + return numpy.where(x <= 0.04045, x / 12.92, ((x + 0.055) / 1.055) ** 2.4) + + +def linear_to_srgb(x): + x = numpy.clip(x, 0.0, 1.0) + return numpy.where(x <= 0.003130804953560372, x * 12.92, 1.055 * (x ** (1.0 / 2.4)) - 0.055) + + +def image_linear_to_srgb(image): + assert image.shape[2] == 3 or image.shape[2] == 4 + if image.shape[2] == 3: + return linear_to_srgb(image) + else: + height, width, _ = image.shape + rgb_image = linear_to_srgb(image[:, :, 0:3]) + a_image = image[:, :, 3:4] + return numpy.concatenate((rgb_image, a_image), axis=2) + + +def image_srgb_to_linear(image): + assert image.shape[2] == 3 or image.shape[2] == 4 + if image.shape[2] == 3: + return srgb_to_linear(image) + else: + height, width, _ = image.shape + rgb_image = srgb_to_linear(image[:, :, 0:3]) + a_image = image[:, :, 3:4] + return numpy.concatenate((rgb_image, a_image), axis=2) + + +def save_rng_state(file_name): + rng_state = torch.get_rng_state() + torch_save(rng_state, file_name) + + +def load_rng_state(file_name): + rng_state = torch_load(file_name) + torch.set_rng_state(rng_state) + + +def optimizer_to_device(optim, device): + for state in optim.state.values(): + for k, v in state.items(): + if isinstance(v, torch.Tensor): + state[k] = v.to(device) + + +def grid_change_to_numpy_image(torch_image, num_channels=3): + height = torch_image.shape[1] + width = torch_image.shape[2] + size_image = (torch_image[0, :, :] ** 2 + torch_image[1, :, :] ** 2).sqrt().view(height, width, 1).numpy() + hsv = cm.get_cmap('hsv') + angle_image = hsv(((torch.atan2( + torch_image[0, :, :].view(height * width), + torch_image[1, :, :].view(height * width)).view(height, width) + math.pi) / (2 * math.pi)).numpy()) * 3 + numpy_image = size_image * angle_image[:, :, 0:3] + rgb_image = linear_to_srgb(numpy_image) + if num_channels == 3: + return rgb_image + elif num_channels == 4: + return numpy.concatenate([rgb_image, numpy.ones_like(size_image)], axis=2) + else: + raise RuntimeError("Unsupported num_channels: " + str(num_channels)) + + +def rgb_to_numpy_image(torch_image: Tensor, min_pixel_value=-1.0, max_pixel_value=1.0): + assert torch_image.dim() == 3 + assert torch_image.shape[0] == 3 + height = torch_image.shape[1] + width = torch_image.shape[2] + + reshaped_image = torch_image.numpy().reshape(3, height * width).transpose().reshape(height, width, 3) + numpy_image = (reshaped_image - min_pixel_value) / (max_pixel_value - min_pixel_value) + return linear_to_srgb(numpy_image) + + +def rgba_to_numpy_image_greenscreen(torch_image: Tensor, + min_pixel_value=-1.0, + max_pixel_value=1.0, + include_alpha=False): + height = torch_image.shape[1] + width = torch_image.shape[2] + + numpy_image = (torch_image.numpy().reshape(4, height * width).transpose().reshape(height, width, + 4) - min_pixel_value) \ + / (max_pixel_value - min_pixel_value) + rgb_image = linear_to_srgb(numpy_image[:, :, 0:3]) + a_image = numpy_image[:, :, 3] + rgb_image[:, :, 0:3] = rgb_image[:, :, 0:3] * a_image.reshape(a_image.shape[0], a_image.shape[1], 1) + rgb_image[:, :, 1] = rgb_image[:, :, 1] + (1 - a_image) + + if not include_alpha: + return rgb_image + else: + return numpy.concatenate((rgb_image, numpy.ones_like(numpy_image[:, :, 3:4])), axis=2) + + +def rgba_to_numpy_image(torch_image: Tensor, min_pixel_value=-1.0, max_pixel_value=1.0): + assert torch_image.dim() == 3 + assert torch_image.shape[0] == 4 + height = torch_image.shape[1] + width = torch_image.shape[2] + + reshaped_image = torch_image.numpy().reshape(4, height * width).transpose().reshape(height, width, 4) + numpy_image = (reshaped_image - min_pixel_value) / (max_pixel_value - min_pixel_value) + rgb_image = linear_to_srgb(numpy_image[:, :, 0:3]) + a_image = numpy_image[:, :, 3] + rgba_image = numpy.concatenate((rgb_image, a_image.reshape(height, width, 1)), axis=2) + return rgba_image + + +def extract_PIL_image_from_filelike(file): + return PIL.Image.open(file) + + +def extract_pytorch_image_from_filelike(file, scale=2.0, offset=-1.0): + return extract_pytorch_image_from_PIL_image(extract_PIL_image_from_filelike(file), scale, offset) + + +def extract_pytorch_image_from_PIL_image(pil_image, scale=2.0, offset=-1.0): + image = extract_numpy_image_from_PIL_image(pil_image) + h, w, c = image.shape + image = image.reshape(h * w, c) + for pixel in image: + if pixel[3] == 0.0: + pixel[0:3] = 0.0 + image = image.transpose().reshape(c, h, w) + torch_image = torch.from_numpy(image).float() * scale + offset + return torch_image + + +def extract_numpy_image_from_filelike(file): + return extract_numpy_image_from_PIL_image(extract_PIL_image_from_filelike(file)) + + +def extract_numpy_image_from_PIL_image(pil_image): + image = numpy.asarray(pil_image) / 255.0 + image[:, :, 0:3] = srgb_to_linear(image[:, :, 0:3]) + return image + + +def create_parent_dir(file_name): + os.makedirs(os.path.dirname(file_name), exist_ok=True) + + +def run_command(command_parts: List[str]): + command = " ".join(command_parts) + os.system(command) + + +def resize_PIL_image(pil_image, size=(256, 256)): + w, h = pil_image.size + d = min(w, h) + r = ((w - d) // 2, (h - d) // 2, (w + d) // 2, (h + d) // 2) + return pil_image.resize(size, resample=PIL.Image.LANCZOS, box=r) + + +def save_pytorch_image(image, file_name): + if image.shape[0] == 1: + image = image.squeeze() + if image.shape[0] == 4: + numpy_image = rgba_to_numpy_image(image.detach().cpu()) + pil_image = PIL.Image.fromarray(numpy.uint8(numpy.rint(numpy_image * 255.0)), mode='RGBA') + else: + numpy_image = rgb_to_numpy_image(image.detach().cpu()) + pil_image = PIL.Image.fromarray(numpy.uint8(numpy.rint(numpy_image * 255.0)), mode='RGB') + os.makedirs(os.path.dirname(file_name), exist_ok=True) + pil_image.save(file_name) + + +def convert_output_image_from_torch_to_numpy(output_image): + if output_image.shape[2] == 2: + h, w, c = output_image.shape + output_image = torch.transpose(output_image.reshape(h * w, c), 0, 1).reshape(c, h, w) + if output_image.shape[0] == 4: + numpy_image = rgba_to_numpy_image(output_image) + elif output_image.shape[0] == 1: + c, h, w = output_image.shape + alpha_image = torch.cat([output_image.repeat(3, 1, 1) * 2.0 - 1.0, torch.ones(1, h, w)], dim=0) + numpy_image = rgba_to_numpy_image(alpha_image) + elif output_image.shape[0] == 2: + numpy_image = grid_change_to_numpy_image(output_image, num_channels=4) + else: + raise RuntimeError("Unsupported # image channels: %d" % output_image.shape[0]) + return numpy_image