/
sp_demo.py
101 lines (91 loc) · 3.4 KB
/
sp_demo.py
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
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
# -*- coding: utf-8 -*-
#%% Demo for single-pixel imaging using Hadamard patterns as sensing basis
# Loads an object, generates the sensing patterns (Hadamard),
# does the measurements (with and without noise), and recovers the image.
#%% Import libraries
from typing import Tuple
from matplotlib import pyplot as plt
import numpy as np
# Generate hadamard matrices
from scipy.linalg import hadamard
# Working with images
from skimage.transform import resize
from PIL import Image
#%% Define useful functions
def noisify(signal: np.ndarray, end_snr: float) -> Tuple[np.ndarray, np.ndarray]:
'''
Add white gaussian noise to a signal so it ends with end_snr
signal-to-noise ratio.
More info at: https://en.wikipedia.org/wiki/Signal-to-noise_ratio#Decibels
Parameters
----------
signal : np.ndarray
Input signal
end_snr : float
Desired SNR.
Returns
-------
noisy_signal : np.ndarray
Signal with added noise.
noise : np.ndarray
Noise added to the original signal.
'''
# Calculate signal average (~power)
signal_avg = np.mean(signal)
# Convert to dB
signal_avg_db = 10 * np.log10(signal_avg)
# Calculate noise power in dB, using objective SNR
# SNR = P_signal(in dB) - P_noise(in dB)
noise_avg_db = signal_avg_db - end_snr
# Convert to noise average (~power)
noise_avg = 10**(noise_avg_db/10)
# Build noise with desired power
noise = np.random.normal(0, np.sqrt(noise_avg), signal.size)
# Add noise to signal
noisy_signal = signal + noise
return noisy_signal, noise
# %% Load image & generate object
# load image, convert to grayscale, store as array (matrix)
ground_obj = np.asarray(Image.open("./objects/ghost.png").convert('L'))
# Choose new size (don't go higher than 128x128 or Hadamard will kill you)
px = 32
# Resize image to smaller size for simulation
test_obj = resize(ground_obj,(px,px))
#%% Generate measurement patterns using a Hadamard matrix.
## Each row is a 2D pattern (after reshaping)
H = hadamard(px**2) # Complete Hadamard matrix (+1s and -1s)
Hplus = (H + 1) / 2 # H+ Hadamard matrix (1s and 0s)
Hminus = (1 - H) / 2 # H- Hadamard matrix (0s and 1s)
#%% Generate measurements simulating H+ and H-
# (as if patterns were generated on a DMD)
Mplus = Hplus @ test_obj.flatten() # Project H+ patterns, store intensity
Mminus = Hminus @ test_obj.flatten() # Project H- patterns, store intensity
M = Mplus - Mminus # Substract values (H+ - H-)
# Generate measurements from a noisy object
desired_SNR = 20 # SNR desired for the object
# Add noise to the true coefficients
Mplus_noisy, _ = noisify(Mplus, desired_SNR) # Project H+
Mminus_noisy, _ = noisify(Mminus, desired_SNR) # Project H-
M_noisy = Mplus_noisy - Mminus_noisy # Substract values (H+ - H-)
#%% Recover objects
# Inversion from measurements (solve the eq. systems with/without noise)
recovery = np.linalg.solve(H, M)
recovery_noise = np.linalg.solve(H, M_noisy)
# Reshape from vector into image
recovery = recovery.reshape((px, px))
recovery_noise = recovery_noise.reshape((px, px))
#%% Show the results
# Show recovery without noise
plt.figure()
plt.imshow(recovery, cmap = 'hot')
plt.axis('on')
plt.title('Recovered image without noise')
plt.colorbar()
plt.show()
# Show recovery with noise
plt.figure()
plt.imshow(recovery_noise, cmap = 'hot')
plt.axis('on')
plt.title('Recovered image with noise')
plt.colorbar()
plt.show()