Spaces:
Runtime error
Runtime error
File size: 1,781 Bytes
882f6e2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 |
"""
Copyright (c) Meta Platforms, Inc. and affiliates.
All rights reserved.
This source code is licensed under the license found in the
LICENSE file in the root directory of this source tree.
"""
from typing import Any, Dict
import numpy as np
import torch as th
import torch.nn as nn
import torch.nn.functional as F
def impaint_batch(value: th.Tensor, dst_ij: th.Tensor, src_ij: th.Tensor) -> th.Tensor:
assert len(value.shape) == 4, "expecting a 4D tensor"
preds = value[:]
preds[:, :, dst_ij[:, 0], dst_ij[:, 1]] = value[:, :, src_ij[:, 0], src_ij[:, 1]]
return preds
def resample_tex(tex: th.Tensor, uvs: th.Tensor, weights: th.Tensor) -> th.Tensor:
B = tex.shape[0]
grid = 2.0 * (uvs[np.newaxis].expand(B, -1, -1, -1) - 0.5)
tex_resampled = F.grid_sample(tex, grid, align_corners=False, padding_mode="border")
return (1.0 - weights) * tex + weights * tex_resampled
class SeamSampler(nn.Module):
def __init__(self, seamless_data: Dict[str, Any]) -> None:
super().__init__()
self.register_buffer("dst_ij", seamless_data["dst_ij"])
self.register_buffer("src_ij", seamless_data["src_ij"])
self.register_buffer("uvs", seamless_data["uvs"])
self.register_buffer("weights", seamless_data["weights"])
def impaint(self, value: th.Tensor) -> th.Tensor:
return impaint_batch(value, self.dst_ij, self.src_ij)
def resample(self, tex: th.Tensor) -> th.Tensor:
return resample_tex(tex, self.uvs, self.weights)
def resample_border_only(self, tex: th.Tensor) -> th.Tensor:
tex = resample_tex(tex, self.uvs, self.weights)
return tex
def forward(self, tex: th.Tensor) -> th.Tensor:
x = self.impaint(tex)
x = self.resample(x)
return x
|