Add embedding, need to complete the module and the RJ sampling part.

NS/sampling/sampler.py

+import numpy as np
+import itertools
+import scipy.linalg as LA
+
+from tqdm.notebook import tqdm
+
+class Sampler():
+	def __init__(self, cover, Q, K, seed, strategy = "4-Lattices"):
+		self.cover = cover #Cover file
+		self.Q = Q #Quality factor
+		self.K = K #Alphabet
+		self.seed = seed #Seed for the RNG
+		#self.cov = cov #Covariance matrix
+		self.strategy = strategy
+
+		self.eps = 10**(-10)
+
+	def process(self, cov, obs, i, j):
+		#For overloading purposes
+		return self
+
+	def getStego(self):
+		if(self.strategy == "4-Lattices"):
+			self.Embedding_4_lattices()
+		elif(self.strategy == "Intra-only"):
+			self.Embedding_intra_only
+		elif(self.strategy == "Indepent"):
+			self.Embedding_independant()
+		
+		return self.stego
+
+	def Embedding_4_lattices(self):
+		np.random.seed(self.seed) #Feed the RNG
+
+		h, w = self.cover.shape
+
+		cover, dct_cover = self.cover, np.zeros((h, w))
+		dct_stego = np.zeros((h, w))
+		#dct_im_cover = raw_to_dct(raw_cover,stego_file[:-4]+'.tiff')
+        
+		M_1 = comp_M(1)
+		M_3 = comp_M(3)
+
+		# Lattices loaders
+		pbar1 = tqdm(total=(h//16-1)*(w//16-1), desc = "Lattices 1")
+		pbar2 = tqdm(total=(h//16-2)*(w//16-2), desc = "Lattices 2")
+		pbar3 = tqdm(total=(h//16-2)*(w//16-2), desc = "Lattices 3")
+		pbar4 = tqdm(total=(h//16-2)*(w//16-2), desc = "Lattices 4")
+
+		#Lattice 1
+		for i, j in itertools.product(range((h//16-1)), range(w//16-1)):
+			obs = dct_stego[i*16+8:i*16+16 , j*16+8:j*16+16].ravel()
+			c = cover[i*16+7:i*16+17 , j*16+7:j*16+17]
+			PP_T = compute_PP_T(c, 1)
+			cov_L1 = compute_cov(M_1, PP_T)
+			dct_stego[i*16+8:i*16+16 , j*16+8:j*16+16 ] = self.process(cov_L1, obs, pos = {'i_left' : i*16+8, 'i_right' : i*16+16, 'j_left' : j*16+8, 'j_right' : j*16+16 })
+
+			pbar1.update(1)
+		
+		#Lattice 2  
+		for i, j in itertools.product(range((h//16-2)), range(w//16-2)):          
+			obs = np.concatenate(((dct_stego[i*16+8:i*16+16 , j*16+8:j*16+16 ]).ravel(),\
+                                (dct_stego[i*16+8:i*16+16 , j*16+24:j*16+32 ]).ravel(),\
+                                (dct_stego[i*16+24:i*16+32 , j*16+8:j*16+16 ]).ravel(),\
+                                (dct_stego[i*16+24:i*16+32 , j*16+24:j*16+32 ]).ravel(),\
+                                (dct_stego[i*16+8:i*16+16 , j*16+8:j*16+16]).ravel()))
+			
+			c = cover[i*16+7:i*16+33 , j*16+7:j*16+33 ]
+			PP_T = compute_PP_T(c, 3)
+			cov_L4 = compute_cov(M_3,PP_T)            
+			cov_L2 = np.block([ [cov_L4[-4*64:, -4*64:], cov_L4[-4*64:, :64]],
+                                [cov_L4[:64, -4*64:],    cov_L4[:64, :64]]])
+			
+			dct_stego[i*16+16:i*16+24 , j*16+16:j*16+24 ] = self.process(cov_L2, obs, pos = {'i_left' : i*16+16, 'i_right' : i*16+24, 'j_left' : j*16+16, 'j_right' : j*16+24 })
+
+			pbar2.update(1)
+		
+		#Lattice 3 
+		for i, j in itertools.product(range((h//16-2)), range(w//16-2)):
+			obs = np.concatenate(((dct_stego[i*16+8:i*16+16 , j*16+24:j*16+32 ]).flatten(),\
+								(dct_stego[i*16+16:i*16+24 , j*16+16:j*16+24 ]).flatten(),\
+								(dct_stego[i*16+16:i*16+24 , j*16+32:j*16+40 ]).flatten(),\
+								(dct_stego[i*16+24:i*16+32 , j*16+24:j*16+32 ]).flatten()))
+
+			c = cover[i*16+7:i*16+33 , j*16+15:j*16+41 ]
+			PP_T = compute_PP_T(c, 3)
+			cov_L4 = compute_cov(M_3, PP_T)            
+			cov_L3 = cov_L4[:5*64,:5*64]
+
+			dct_stego[i*16+16:i*16+24 , j*16+24:j*16+32 ] = self.process(cov_L3, obs, pos = {'i_left' : i*16+16, 'i_right' : i*16+24, 'j_left' : j*16+24, 'j_right' : j*16+32 })
+			pbar3.update(1)
+			
+		#Lattice 4
+		for i, j in itertools.product(range((h//16-2)), range(w//16-2)):
+			obs = np.concatenate(((dct_stego[i*16+16:i*16+24 , j*16+16:j*16+24 ]).flatten(),\
+								(dct_stego[i*16+24:i*16+32 , j*16+8:j*16+16 ]).flatten(),\
+								(dct_stego[i*16+24:i*16+32 , j*16+24:j*16+32 ]).flatten(),\
+								(dct_stego[i*16+32:i*16+40 , j*16+16:j*16+24 ]).flatten(),\
+								(dct_stego[i*16+16:i*16+24 , j*16+8:j*16+16 ]).flatten(),\
+								(dct_stego[i*16+16:i*16+24 , j*16+24:j*16+32 ]).flatten(),\
+								(dct_stego[i*16+32:i*16+40 , j*16+8:j*16+16 ]).flatten(),\
+								(dct_stego[i*16+32:i*16+40 , j*16+24:j*16+32 ]).flatten()))
+
+			c = im_cover[i*16+15:i*16+41 , j*16+7:j*16+33 ]
+			PP_T = compute_PP_T(c, 3)
+			cov_L4 = compute_cov(M_3,PP_T)            
+	 
+			dct_stego[i*16+24:i*16+32 , j*16+16:j*16+24 ] = self.process(cov_L4, obs, pos = {'i_left' : i*16+24, 'i_right' : i*16+32, 'j_left' : j*16+16, 'j_right' : j*16+24 })
+			pbar4.update()
+
+		self.stego = dct_stego + dct_cover
+		return self.stego
+
+	def Embedding_intra_only(self):
+		# Lattice 1 loader
+		pbar1 = tqdm_notebook(total=(h//8)*(w//8), desc = "Intra-only")
+		for i, j in itertools.product(range((h//8)), range(w//8)):
+			pbar1.update(1)
+
+		self.stego = dct_stego
+		return self.stego
+
+	def Embedding_independant(self):
+		# Lattices loaders
+		pbar1 = tqdm_notebook(total=(h//8)*(w//8), desc = "Intra-only")
+		for i_bloc, j_bloc in itertools.product(range((h//8)), range(w//8)):
+			pbar1.update(1)
+
+		self.stego = dct_stego
+		return self.stego
+
+class MVG(Sampler):
+	def __init__(self, cover, Q, K, seed, strategy = "4-Lattices"):
+		Sampler.__init__(self, cover, Q, K, seed, strategy = "4-Lattices")
+
+	def process(self, cov, obs, i = 0, j = 0):
+		cov_11 = cov[:64,:64]
+		cov_12 = cov[:64,64:]
+		cov_21 = cov[64:,:64]
+		cov_22 = cov[64:,64:]
+
+		reg_mat = np.dot(cov_12, np.linalg.pinv(cov_22))
+		cov_schur = cov_11 - np.dot(reg_mat, cov_21)
+
+		mean_vec = np.zeros(cov.shape[0])
+		mean = mean_vec[:64] + np.dot(reg_mat, obs-mean_vec[64:])
+		cov_schur = (cov_schur + cov_schur.T)/2 + self.eps*np.eye(64)
+
+		return np.random.multivariate_normal(mean, cov_schur).reshape( (8, 8) )
+
+class Cholesky(Sampler):
+	def __init__(self, cover, Q, K, seed, strategy = "4-Lattices"):
+		Sampler.__init__(self, cover, Q, K, seed, strategy = "4-Lattices")
+		self.wn = np.random.normal(0, 1, cover.shape) #White Noise
+
+	def process(self, cov, obs, pos):
+		cov_11 = cov[:64,:64]
+		cov_12 = cov[:64,64:]
+		cov_21 = cov[64:,:64]
+		cov_22 = cov[64:,64:]
+
+		reg_mat = np.dot(cov_12, np.linalg.pinv(cov_22))
+		cov_schur = cov_11 - np.dot(reg_mat, cov_21) + self.eps*np.eye(64)
+
+		L = LA.cholesky(cov_schur, lower=True) #Cholesky decomposition
+		D = np.diag(L)*np.eye(64)
+
+		mean_arr = np.dot(L-D, obs[-64:])
+		std_arr = np.diag(D).reshape(64, 1)
+
+		if(obs.shape[0] > 64):
+			print(reg_mat.shape)
+			print(obs[:64].shape) 
+			return (np.dot(L, self.wn[pos['i_left']:pos['i_right'], pos['j_left']:pos['j_right']].ravel()) + np.dot(reg_mat, obs[:64].reshape(64, 1))).reshape((8, 8))
+		else:
+			return np.dot(L, self.wn[pos['i_left']:pos['i_right'], pos['j_left']:pos['j_right']].ravel()).reshape((8, 8))
\ No newline at end of file