aiden
/
DragGAN
mirror of https://github.com/XingangPan/DragGAN
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

SOLVED AttributeError: 'FreeTypeFont' object has no attribute 'getsize'

Browse Source
pull/239/head
Mert Cobanov 3 years ago
parent 1a2d37d03d
commit 3623b1ade0
1 changed files with 171 additions and 109 deletions
Show Stats Download Patch File Download Diff File

280
viz/renderer.py
Unescape Escape View File

@ -20,9 +20,10 @@ import torch.nn.functional as F
import matplotlib.cm import matplotlib.cm
import dnnlib import dnnlib
from torch_utils.ops import upfirdn2d from torch_utils.ops import upfirdn2d
import legacy # pylint: disable=import-error import legacy # pylint: disable=import-error
# ----------------------------------------------------------------------------
#----------------------------------------------------------------------------
class CapturedException(Exception): class CapturedException(Exception):
def __init__(self, msg=None): def __init__(self, msg=None):
@ -36,26 +37,35 @@ class CapturedException(Exception):
assert isinstance(msg, str) assert isinstance(msg, str)
super().__init__(msg) super().__init__(msg)
#----------------------------------------------------------------------------
# ----------------------------------------------------------------------------
class CaptureSuccess(Exception): class CaptureSuccess(Exception):
def __init__(self, out): def __init__(self, out):
super().__init__() super().__init__()
self.out = out self.out = out
#----------------------------------------------------------------------------
# ----------------------------------------------------------------------------
def add_watermark_np(input_image_array, watermark_text="AI Generated"): def add_watermark_np(input_image_array, watermark_text="AI Generated"):
image = Image.fromarray(np.uint8(input_image_array)).convert("RGBA") image = Image.fromarray(np.uint8(input_image_array)).convert("RGBA")
# Initialize text image # Initialize text image
txt = Image.new('RGBA', image.size, (255, 255, 255, 0)) txt = Image.new("RGBA", image.size, (255, 255, 255, 0))
font = ImageFont.truetype('arial.ttf', round(25/512*image.size[0])) font = ImageFont.truetype("arial.ttf", round(25 / 512 * image.size[0]))
d = ImageDraw.Draw(txt) d = ImageDraw.Draw(txt)
text_width, text_height = font.getsize(watermark_text) text_width, text_height = [0, 0] # font.getsize(watermark_text)
text_position = (image.size[0] - text_width - 10, image.size[1] - text_height - 10) text_position = (image.size[0] - text_width - 10, image.size[1] - text_height - 10)
text_color = (255, 255, 255, 128) # white color with the alpha channel set to semi-transparent text_color = (
255,
255,
255,
128,
) # white color with the alpha channel set to semi-transparent
# Draw the text onto the text canvas # Draw the text onto the text canvas
d.text(text_position, watermark_text, font=font, fill=text_color) d.text(text_position, watermark_text, font=font, fill=text_color)
@ -65,22 +75,30 @@ def add_watermark_np(input_image_array, watermark_text="AI Generated"):
watermarked_array = np.array(watermarked) watermarked_array = np.array(watermarked)
return watermarked_array return watermarked_array
#----------------------------------------------------------------------------
# ----------------------------------------------------------------------------
class Renderer: class Renderer:
def __init__(self, disable_timing=False): def __init__(self, disable_timing=False):
self._device = torch.device('cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu') self._device = torch.device(
self._dtype = torch.float32 if self._device.type == 'mps' else torch.float64 "cuda"
self._pkl_data = dict() # {pkl: dict | CapturedException, ...} if torch.cuda.is_available()
self._networks = dict() # {cache_key: torch.nn.Module, ...} else "mps"
self._pinned_bufs = dict() # {(shape, dtype): torch.Tensor, ...} if torch.backends.mps.is_available()
self._cmaps = dict() # {name: torch.Tensor, ...} else "cpu"
self._is_timing = False )
self._dtype = torch.float32 if self._device.type == "mps" else torch.float64
self._pkl_data = dict() # {pkl: dict | CapturedException, ...}
self._networks = dict() # {cache_key: torch.nn.Module, ...}
self._pinned_bufs = dict() # {(shape, dtype): torch.Tensor, ...}
self._cmaps = dict() # {name: torch.Tensor, ...}
self._is_timing = False
if not disable_timing: if not disable_timing:
self._start_event = torch.cuda.Event(enable_timing=True) self._start_event = torch.cuda.Event(enable_timing=True)
self._end_event = torch.cuda.Event(enable_timing=True) self._end_event = torch.cuda.Event(enable_timing=True)
self._disable_timing = disable_timing self._disable_timing = disable_timing
self._net_layers = dict() # {cache_key: [dnnlib.EasyDict, ...], ...} self._net_layers = dict() # {cache_key: [dnnlib.EasyDict, ...], ...}
def render(self, **args): def render(self, **args):
if self._disable_timing: if self._disable_timing:
@ -91,21 +109,21 @@ class Renderer:
res = dnnlib.EasyDict() res = dnnlib.EasyDict()
try: try:
init_net = False init_net = False
if not hasattr(self, 'G'): if not hasattr(self, "G"):
init_net = True init_net = True
if hasattr(self, 'pkl'): if hasattr(self, "pkl"):
if self.pkl != args['pkl']: if self.pkl != args["pkl"]:
init_net = True init_net = True
if hasattr(self, 'w_load'): if hasattr(self, "w_load"):
if self.w_load is not args['w_load']: if self.w_load is not args["w_load"]:
init_net = True init_net = True
if hasattr(self, 'w0_seed'): if hasattr(self, "w0_seed"):
if self.w0_seed != args['w0_seed']: if self.w0_seed != args["w0_seed"]:
init_net = True init_net = True
if hasattr(self, 'w_plus'): if hasattr(self, "w_plus"):
if self.w_plus != args['w_plus']: if self.w_plus != args["w_plus"]:
init_net = True init_net = True
if args['reset_w']: if args["reset_w"]:
init_net = True init_net = True
res.init_net = init_net res.init_net = init_net
if init_net: if init_net:
@ -115,12 +133,12 @@ class Renderer:
res.error = CapturedException() res.error = CapturedException()
if not self._disable_timing: if not self._disable_timing:
self._end_event.record(torch.cuda.current_stream(self._device)) self._end_event.record(torch.cuda.current_stream(self._device))
if 'image' in res: if "image" in res:
res.image = self.to_cpu(res.image).detach().numpy() res.image = self.to_cpu(res.image).detach().numpy()
res.image = add_watermark_np(res.image, 'AI Generated') res.image = add_watermark_np(res.image, "AI Generated")
if 'stats' in res: if "stats" in res:
res.stats = self.to_cpu(res.stats).detach().numpy() res.stats = self.to_cpu(res.stats).detach().numpy()
if 'error' in res: if "error" in res:
res.error = str(res.error) res.error = str(res.error)
# if 'stop' in res and res.stop: # if 'stop' in res and res.stop:
@ -133,14 +151,14 @@ class Renderer:
def get_network(self, pkl, key, **tweak_kwargs): def get_network(self, pkl, key, **tweak_kwargs):
data = self._pkl_data.get(pkl, None) data = self._pkl_data.get(pkl, None)
if data is None: if data is None:
print(f'Loading "{pkl}"... ', end='', flush=True) print(f'Loading "{pkl}"... ', end="", flush=True)
try: try:
with dnnlib.util.open_url(pkl, verbose=False) as f: with dnnlib.util.open_url(pkl, verbose=False) as f:
data = legacy.load_network_pkl(f) data = legacy.load_network_pkl(f)
print('Done.') print("Done.")
except: except:
data = CapturedException() data = CapturedException()
print('Failed!') print("Failed!")
self._pkl_data[pkl] = data self._pkl_data[pkl] = data
self._ignore_timing() self._ignore_timing()
if isinstance(data, CapturedException): if isinstance(data, CapturedException):
@ -151,19 +169,26 @@ class Renderer:
net = self._networks.get(cache_key, None) net = self._networks.get(cache_key, None)
if net is None: if net is None:
try: try:
if 'stylegan2' in pkl: if "stylegan2" in pkl:
from training.networks_stylegan2 import Generator from training.networks_stylegan2 import Generator
elif 'stylegan3' in pkl: elif "stylegan3" in pkl:
from training.networks_stylegan3 import Generator from training.networks_stylegan3 import Generator
elif 'stylegan_human' in pkl: elif "stylegan_human" in pkl:
from stylegan_human.training_scripts.sg2.training.networks import Generator from stylegan_human.training_scripts.sg2.training.networks import (
Generator,
)
else: else:
raise NameError('Cannot infer model type from pkl name!') raise NameError("Cannot infer model type from pkl name!")
print(data[key].init_args) print(data[key].init_args)
print(data[key].init_kwargs) print(data[key].init_kwargs)
if 'stylegan_human' in pkl: if "stylegan_human" in pkl:
net = Generator(*data[key].init_args, **data[key].init_kwargs, square=False, padding=True) net = Generator(
*data[key].init_args,
**data[key].init_kwargs,
square=False,
padding=True,
)
else: else:
net = Generator(*data[key].init_args, **data[key].init_kwargs) net = Generator(*data[key].init_args, **data[key].init_kwargs)
net.load_state_dict(data[key].state_dict()) net.load_state_dict(data[key].state_dict())
@ -193,7 +218,7 @@ class Renderer:
def _ignore_timing(self): def _ignore_timing(self):
self._is_timing = False self._is_timing = False
def _apply_cmap(self, x, name='viridis'): def _apply_cmap(self, x, name="viridis"):
cmap = self._cmaps.get(name, None) cmap = self._cmaps.get(name, None)
if cmap is None: if cmap is None:
cmap = matplotlib.cm.get_cmap(name) cmap = matplotlib.cm.get_cmap(name)
@ -205,26 +230,32 @@ class Renderer:
x = torch.nn.functional.embedding(x, cmap) x = torch.nn.functional.embedding(x, cmap)
return x return x
def init_network(self, res, def init_network(
pkl = None, self,
w0_seed = 0, res,
w_load = None, pkl=None,
w_plus = True, w0_seed=0,
noise_mode = 'const', w_load=None,
trunc_psi = 0.7, w_plus=True,
trunc_cutoff = None, noise_mode="const",
input_transform = None, trunc_psi=0.7,
lr = 0.001, trunc_cutoff=None,
**kwargs input_transform=None,
): lr=0.001,
**kwargs,
):
# Dig up network details. # Dig up network details.
self.pkl = pkl self.pkl = pkl
G = self.get_network(pkl, 'G_ema') G = self.get_network(pkl, "G_ema")
self.G = G self.G = G
res.img_resolution = G.img_resolution res.img_resolution = G.img_resolution
res.num_ws = G.num_ws res.num_ws = G.num_ws
res.has_noise = any('noise_const' in name for name, _buf in G.synthesis.named_buffers()) res.has_noise = any(
res.has_input_transform = (hasattr(G.synthesis, 'input') and hasattr(G.synthesis.input, 'transform')) "noise_const" in name for name, _buf in G.synthesis.named_buffers()
)
res.has_input_transform = hasattr(G.synthesis, "input") and hasattr(
G.synthesis.input, "transform"
)
# Set input transform. # Set input transform.
if res.has_input_transform: if res.has_input_transform:
@ -242,11 +273,15 @@ class Renderer:
if self.w_load is None: if self.w_load is None:
# Generate random latents. # Generate random latents.
z = torch.from_numpy(np.random.RandomState(w0_seed).randn(1, 512)).to(self._device, dtype=self._dtype) z = torch.from_numpy(np.random.RandomState(w0_seed).randn(1, 512)).to(
self._device, dtype=self._dtype
)
# Run mapping network. # Run mapping network.
label = torch.zeros([1, G.c_dim], device=self._device) label = torch.zeros([1, G.c_dim], device=self._device)
w = G.mapping(z, label, truncation_psi=trunc_psi, truncation_cutoff=trunc_cutoff) w = G.mapping(
z, label, truncation_psi=trunc_psi, truncation_cutoff=trunc_cutoff
)
else: else:
w = self.w_load.clone().to(self._device) w = self.w_load.clone().to(self._device)
@ -263,42 +298,43 @@ class Renderer:
self.points0_pt = None self.points0_pt = None
def update_lr(self, lr): def update_lr(self, lr):
del self.w_optim del self.w_optim
self.w_optim = torch.optim.Adam([self.w], lr=lr) self.w_optim = torch.optim.Adam([self.w], lr=lr)
print(f'Rebuild optimizer with lr: {lr}') print(f"Rebuild optimizer with lr: {lr}")
print(' Remain feat_refs and points0_pt') print(" Remain feat_refs and points0_pt")
def _render_drag_impl(self, res, def _render_drag_impl(
points = [], self,
targets = [], res,
mask = None, points=[],
lambda_mask = 10, targets=[],
reg = 0, mask=None,
feature_idx = 5, lambda_mask=10,
r1 = 3, reg=0,
r2 = 12, feature_idx=5,
random_seed = 0, r1=3,
noise_mode = 'const', r2=12,
trunc_psi = 0.7, random_seed=0,
force_fp32 = False, noise_mode="const",
layer_name = None, trunc_psi=0.7,
sel_channels = 3, force_fp32=False,
base_channel = 0, layer_name=None,
img_scale_db = 0, sel_channels=3,
img_normalize = False, base_channel=0,
untransform = False, img_scale_db=0,
is_drag = False, img_normalize=False,
reset = False, untransform=False,
to_pil = False, is_drag=False,
**kwargs reset=False,
to_pil=False,
**kwargs,
): ):
G = self.G G = self.G
ws = self.w ws = self.w
if ws.dim() == 2: if ws.dim() == 2:
ws = ws.unsqueeze(1).repeat(1,6,1) ws = ws.unsqueeze(1).repeat(1, 6, 1)
ws = torch.cat([ws[:,:6,:], self.w0[:,6:,:]], dim=1) ws = torch.cat([ws[:, :6, :], self.w0[:, 6:, :]], dim=1)
if hasattr(self, 'points'): if hasattr(self, "points"):
if len(points) != len(self.points): if len(points) != len(self.points):
reset = True reset = True
if reset: if reset:
@ -308,7 +344,14 @@ class Renderer:
# Run synthesis network. # Run synthesis network.
label = torch.zeros([1, G.c_dim], device=self._device) label = torch.zeros([1, G.c_dim], device=self._device)
img, feat = G(ws, label, truncation_psi=trunc_psi, noise_mode=noise_mode, input_is_w=True, return_feature=True) img, feat = G(
ws,
label,
truncation_psi=trunc_psi,
noise_mode=noise_mode,
input_is_w=True,
return_feature=True,
)
h, w = G.img_resolution, G.img_resolution h, w = G.img_resolution, G.img_resolution
@ -316,14 +359,18 @@ class Renderer:
X = torch.linspace(0, h, h) X = torch.linspace(0, h, h)
Y = torch.linspace(0, w, w) Y = torch.linspace(0, w, w)
xx, yy = torch.meshgrid(X, Y) xx, yy = torch.meshgrid(X, Y)
feat_resize = F.interpolate(feat[feature_idx], [h, w], mode='bilinear') feat_resize = F.interpolate(feat[feature_idx], [h, w], mode="bilinear")
if self.feat_refs is None: if self.feat_refs is None:
self.feat0_resize = F.interpolate(feat[feature_idx].detach(), [h, w], mode='bilinear') self.feat0_resize = F.interpolate(
feat[feature_idx].detach(), [h, w], mode="bilinear"
)
self.feat_refs = [] self.feat_refs = []
for point in points: for point in points:
py, px = round(point[0]), round(point[1]) py, px = round(point[0]), round(point[1])
self.feat_refs.append(self.feat0_resize[:,:,py,px]) self.feat_refs.append(self.feat0_resize[:, :, py, px])
self.points0_pt = torch.Tensor(points).unsqueeze(0).to(self._device) # 1, N, 2 self.points0_pt = (
torch.Tensor(points).unsqueeze(0).to(self._device)
) # 1, N, 2
# Point tracking with feature matching # Point tracking with feature matching
with torch.no_grad(): with torch.no_grad():
@ -333,9 +380,11 @@ class Renderer:
down = min(point[0] + r + 1, h) down = min(point[0] + r + 1, h)
left = max(point[1] - r, 0) left = max(point[1] - r, 0)
right = min(point[1] + r + 1, w) right = min(point[1] + r + 1, w)
feat_patch = feat_resize[:,:,up:down,left:right] feat_patch = feat_resize[:, :, up:down, left:right]
L2 = torch.linalg.norm(feat_patch - self.feat_refs[j].reshape(1,-1,1,1), dim=1) L2 = torch.linalg.norm(
_, idx = torch.min(L2.view(1,-1), -1) feat_patch - self.feat_refs[j].reshape(1, -1, 1, 1), dim=1
)
_, idx = torch.min(L2.view(1, -1), -1)
width = right - left width = right - left
point = [idx.item() // width + up, idx.item() % width + left] point = [idx.item() // width + up, idx.item() % width + left]
points[j] = point points[j] = point
@ -346,24 +395,35 @@ class Renderer:
loss_motion = 0 loss_motion = 0
res.stop = True res.stop = True
for j, point in enumerate(points): for j, point in enumerate(points):
direction = torch.Tensor([targets[j][1] - point[1], targets[j][0] - point[0]]) direction = torch.Tensor(
[targets[j][1] - point[1], targets[j][0] - point[0]]
)
if torch.linalg.norm(direction) > max(2 / 512 * h, 2): if torch.linalg.norm(direction) > max(2 / 512 * h, 2):
res.stop = False res.stop = False
if torch.linalg.norm(direction) > 1: if torch.linalg.norm(direction) > 1:
distance = ((xx.to(self._device) - point[0])**2 + (yy.to(self._device) - point[1])**2)**0.5 distance = (
(xx.to(self._device) - point[0]) ** 2
+ (yy.to(self._device) - point[1]) ** 2
) ** 0.5
relis, reljs = torch.where(distance < round(r1 / 512 * h)) relis, reljs = torch.where(distance < round(r1 / 512 * h))
direction = direction / (torch.linalg.norm(direction) + 1e-7) direction = direction / (torch.linalg.norm(direction) + 1e-7)
gridh = (relis-direction[1]) / (h-1) * 2 - 1 gridh = (relis - direction[1]) / (h - 1) * 2 - 1
gridw = (reljs-direction[0]) / (w-1) * 2 - 1 gridw = (reljs - direction[0]) / (w - 1) * 2 - 1
grid = torch.stack([gridw,gridh], dim=-1).unsqueeze(0).unsqueeze(0) grid = torch.stack([gridw, gridh], dim=-1).unsqueeze(0).unsqueeze(0)
target = F.grid_sample(feat_resize.float(), grid, align_corners=True).squeeze(2) target = F.grid_sample(
loss_motion += F.l1_loss(feat_resize[:,:,relis,reljs], target.detach()) feat_resize.float(), grid, align_corners=True
).squeeze(2)
loss_motion += F.l1_loss(
feat_resize[:, :, relis, reljs], target.detach()
)
loss = loss_motion loss = loss_motion
if mask is not None: if mask is not None:
if mask.min() == 0 and mask.max() == 1: if mask.min() == 0 and mask.max() == 1:
mask_usq = mask.to(self._device).unsqueeze(0).unsqueeze(0) mask_usq = mask.to(self._device).unsqueeze(0).unsqueeze(0)
loss_fix = F.l1_loss(feat_resize * mask_usq, self.feat0_resize * mask_usq) loss_fix = F.l1_loss(
feat_resize * mask_usq, self.feat0_resize * mask_usq
)
loss += lambda_mask * loss_fix loss += lambda_mask * loss_fix
loss += reg * F.l1_loss(ws, self.w0) # latent code regularization loss += reg * F.l1_loss(ws, self.w0) # latent code regularization
@ -375,14 +435,16 @@ class Renderer:
# Scale and convert to uint8. # Scale and convert to uint8.
img = img[0] img = img[0]
if img_normalize: if img_normalize:
img = img / img.norm(float('inf'), dim=[1,2], keepdim=True).clip(1e-8, 1e8) img = img / img.norm(float("inf"), dim=[1, 2], keepdim=True).clip(1e-8, 1e8)
img = img * (10 ** (img_scale_db / 20)) img = img * (10 ** (img_scale_db / 20))
img = (img * 127.5 + 128).clamp(0, 255).to(torch.uint8).permute(1, 2, 0) img = (img * 127.5 + 128).clamp(0, 255).to(torch.uint8).permute(1, 2, 0)
if to_pil: if to_pil:
from PIL import Image from PIL import Image
img = img.cpu().numpy() img = img.cpu().numpy()
img = Image.fromarray(img) img = Image.fromarray(img)
res.image = img res.image = img
res.w = ws.detach().cpu().numpy() res.w = ws.detach().cpu().numpy()
#----------------------------------------------------------------------------
# ----------------------------------------------------------------------------

Write Preview
Loading…
Cancel
Save