WavMark / models /my_model_v7_recover.py
my
Add application file
32ca76b
raw
history blame
3.54 kB
import pdb
import torch.optim
import torch.nn as nn
from models.hinet import Hinet
# from utils.attacks import attack_layer, mp3_attack_v2, butterworth_attack
import numpy as np
import random
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
class Model(nn.Module):
def __init__(self, num_point, num_bit, n_fft, hop_length, use_recover_layer, num_layers):
super(Model, self).__init__()
self.hinet = Hinet(num_layers=num_layers)
self.watermark_fc = torch.nn.Linear(num_bit, num_point)
self.watermark_fc_back = torch.nn.Linear(num_point, num_bit)
self.n_fft = n_fft
self.hop_length = hop_length
self.dropout1 = torch.nn.Dropout()
self.identity = torch.nn.Identity()
self.recover_layer = SameSizeConv2d(2, 2)
self.use_recover_layer = use_recover_layer
def stft(self, data):
window = torch.hann_window(self.n_fft).to(data.device)
tmp = torch.stft(data, n_fft=self.n_fft, hop_length=self.hop_length, window=window, return_complex=False)
# [1, 501, 41, 2]
return tmp
def istft(self, signal_wmd_fft):
window = torch.hann_window(self.n_fft).to(signal_wmd_fft.device)
# Changed in version 2.0: Real datatype inputs are no longer supported. Input must now have a complex datatype, as returned by stft(..., return_complex=True).
return torch.istft(signal_wmd_fft, n_fft=self.n_fft, hop_length=self.hop_length, window=window,
return_complex=False)
def encode(self, signal, message, need_fft=False):
# 1.信号执行fft
signal_fft = self.stft(signal)
# import pdb
# pdb.set_trace()
# (batch,freq_bins,time_frames,2)
# 2.Message执行fft
message_expand = self.watermark_fc(message)
message_fft = self.stft(message_expand)
# 3.encode
signal_wmd_fft, msg_remain = self.enc_dec(signal_fft, message_fft, rev=False)
# (batch,freq_bins,time_frames,2)
signal_wmd = self.istft(signal_wmd_fft)
if need_fft:
return signal_wmd, signal_fft, message_fft
return signal_wmd
def decode(self, signal):
signal_fft = self.stft(signal)
if self.use_recover_layer:
signal_fft = self.recover_layer(signal_fft)
watermark_fft = signal_fft
# watermark_fft = torch.randn(signal_fft.shape).cuda()
_, message_restored_fft = self.enc_dec(signal_fft, watermark_fft, rev=True)
message_restored_expanded = self.istft(message_restored_fft)
message_restored_float = self.watermark_fc_back(message_restored_expanded).clamp(-1, 1)
return message_restored_float
def enc_dec(self, signal, watermark, rev):
signal = signal.permute(0, 3, 2, 1)
# [4, 2, 41, 501]
watermark = watermark.permute(0, 3, 2, 1)
# pdb.set_trace()
signal2, watermark2 = self.hinet(signal, watermark, rev)
return signal2.permute(0, 3, 2, 1), watermark2.permute(0, 3, 2, 1)
class SameSizeConv2d(nn.Module):
def __init__(self, in_channels, out_channels):
super(SameSizeConv2d, self).__init__()
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
def forward(self, x):
# (batch,501,41,2]
x1 = x.permute(0, 3, 1, 2)
# (batch,2,501,41]
x2 = self.conv(x1)
# (batch,2,501,41]
x3 = x2.permute(0, 2, 3, 1)
# (batch,501,41,2]
return x3