GunwooHan committed
1 parent
ac9ad50ef2
commit
e1e25ae969
53 files changed
+3899
No files matched your search
@@ -2,6 +2,7 @@ pretrained
|
||||
character
|
||||
src
|
||||
dst
|
||||
.idea
|
||||
|
||||
# Byte-compiled / optimized / DLL files
|
||||
__pycache__/
|
||||
|
||||
Whitespace-only changes.
Whitespace-only changes.
@@ -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()
|
||||
@@ -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()
|
||||
Whitespace-only changes.
@@ -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]]
|
||||
@@ -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
|
||||
Whitespace-only changes.
@@ -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"
|
||||
@@ -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
|
||||
Whitespace-only changes.
Whitespace-only changes.
@@ -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)
|
||||
@@ -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
|
||||
Whitespace-only changes.
Whitespace-only changes.
@@ -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))
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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)))))
|
||||
@@ -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
|
||||
Whitespace-only changes.
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
Whitespace-only changes.
@@ -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)
|
||||
@@ -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
|
||||
@@ -0,0 +1,9 @@
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
from torch.nn import Module
|
||||
|
||||
|
||||
class ModuleFactory(ABC):
|
||||
@abstractmethod
|
||||
def create(self) -> Module:
|
||||
pass
|
||||
@@ -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
|
||||
@@ -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()
|
||||
@@ -0,0 +1,9 @@
|
||||
from torch.nn import Module
|
||||
|
||||
|
||||
class PassThrough(Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def forward(self, x):
|
||||
return x
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
Whitespace-only changes.
@@ -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]
|
||||
@@ -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
|
||||
Whitespace-only changes.
@@ -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)
|
||||
@@ -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)
|
||||
Whitespace-only changes.
@@ -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)
|
||||
@@ -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)
|
||||
Whitespace-only changes.
@@ -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
|
||||
Whitespace-only changes.
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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
|
||||
+213
@@ -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
|
||||
Reference in new issue
Block a user