Introduce new ReconstructionLoss class

face_swapper/src/models/loss.py

@@ -141,6 +141,28 @@ class FaceSwapperLoss:
 		return translation, scale, rotation
 
 
+class ReconstructionLoss(torch.nn.Module):
+	def __init__(self) -> None:
+		super(ReconstructionLoss, self).__init__()
+
+	def calc(self, source_tensor : Tensor, target_tensor : Tensor, output_tensor : Tensor) -> Tensor:
+		batch_size = CONFIG.getint('training.loader', 'batch_size')
+
+		loss_tensor = torch.pow(output_tensor - target_tensor, 2).reshape(batch_size, -1)
+		loss_tensor = torch.mean(loss_tensor, dim = 1) * 0.5
+
+		if torch.equal(source_tensor, target_tensor):
+			loss_tensor = torch.sum(loss_tensor * torch.tensor(0)) / (torch.tensor(0).sum() + 1e-4)
+		else:
+			loss_tensor = torch.sum(loss_tensor * torch.tensor(1)) / (torch.tensor(1).sum() + 1e-4)
+
+		data_range = float(torch.max(output_tensor) - torch.min(output_tensor))
+		similarity = 1 - ssim(output_tensor, target_tensor, data_range = data_range).mean()
+
+		loss_tensor = (loss_tensor + similarity) * 0.5
+		return loss_tensor
+
+
 class IdentityLoss(torch.nn.Module):
 	def __init__(self) -> None:
 		super(IdentityLoss, self).__init__()
@@ -148,8 +170,8 @@ class IdentityLoss(torch.nn.Module):
 		self.embedder = torch.jit.load(embedder_path, map_location = 'cpu') # type:ignore[no-untyped-call]
 		self.embedder.eval()
 
-	def calc_loss(self, source_tensor : Tensor, output_tensor : Tensor) -> Tensor:
+	def calc(self, source_tensor : Tensor, output_tensor : Tensor) -> Tensor:
 		output_embedding = calc_embedding(self.embedder, output_tensor, (30, 0, 10, 10))
 		source_embedding = calc_embedding(self.embedder, source_tensor, (30, 0, 10, 10))
-		loss = (1 - torch.cosine_similarity(source_embedding, output_embedding)).mean()
-		return loss
+		loss_tensor = (1 - torch.cosine_similarity(source_embedding, output_embedding)).mean()
+		return loss_tensor