Mercurial > repos > PredictPDPS / file comparison
comparison: src/PET/ImGenerate.jl
src/PET/ImGenerate.jl
- changeset 8
- e4ad8f7ce671
- child 14
- c286925c0f35
equal
deleted
inserted
replaced
| 7:4bbb246f4cac | 8:e4ad8f7ce671 |
|---|---|
| 1 ################### | |
| 2 # Image generation | |
| 3 ################### | |
| 4 | |
| 5 module ImGenerate | |
| 6 | |
| 7 using ColorTypes: Gray | |
| 8 import TestImages | |
| 9 # We don't really *directly* depend on QuartzImageIO. The import here is | |
| 10 # merely a workaround to suppress warnings when loading TestImages. | |
| 11 # Something is broken in those packages. | |
| 12 import QuartzImageIO | |
| 13 | |
| 14 using AlgTools.Util | |
| 15 using AlgTools.Comms | |
| 16 using ImageTools.Translate | |
| 17 | |
| 18 using ..OpticalFlow: Image, DisplacementConstant, DisplacementFull | |
| 19 using ..Radon | |
| 20 | |
| 21 # Added for reproducibility | |
| 22 import StableRNGs: StableRNG, Random | |
| 23 const rng = StableRNG(9182737465) | |
| 24 | |
| 25 # Added for PET | |
| 26 import PoissonRandom: pois_rand | |
| 27 import Random: shuffle | |
| 28 import Images: center, warp | |
| 29 import CoordinateTransformations: recenter | |
| 30 import Rotations: RotMatrix | |
| 31 import Interpolations: Flat | |
| 32 | |
| 33 | |
| 34 ############## | |
| 35 # Our exports | |
| 36 ############## | |
| 37 | |
| 38 export ImGen, | |
| 39 OnlineData, | |
| 40 imgen_square, | |
| 41 imgen_shake, | |
| 42 PetOnlineData, | |
| 43 imgen_shepplogan_radon | |
| 44 | |
| 45 ################## | |
| 46 # Data structures | |
| 47 ################## | |
| 48 | |
| 49 struct ImGen | |
| 50 f :: Function | |
| 51 dim :: Tuple{Int64,Int64} | |
| 52 Λ :: Float64 | |
| 53 dynrange :: Float64 | |
| 54 name :: String | |
| 55 end | |
| 56 | |
| 57 struct OnlineData{DisplacementT} | |
| 58 b_true :: Image | |
| 59 b_noisy :: Image | |
| 60 v :: DisplacementT | |
| 61 v_true :: DisplacementT | |
| 62 v_cumul_true :: DisplacementT | |
| 63 end | |
| 64 | |
| 65 struct PetOnlineData{DisplacementT} | |
| 66 b_true :: Image | |
| 67 sinogram_true :: Image | |
| 68 sinogram_noisy :: Image | |
| 69 v :: DisplacementT | |
| 70 v_true :: DisplacementT | |
| 71 v_cumul_true :: DisplacementT | |
| 72 theta :: DisplacementT # theta = thetaknown, theta_cumul | |
| 73 S :: Image | |
| 74 end | |
| 75 | |
| 76 ################### | |
| 77 # Shake generation | |
| 78 ################### | |
| 79 | |
| 80 function make_const_v(displ, sz) | |
| 81 v = zeros(2, sz...) | |
| 82 v[1, :, :] .= displ[1] | |
| 83 v[2, :, :] .= displ[2] | |
| 84 return v | |
| 85 end | |
| 86 | |
| 87 function shake(params) | |
| 88 if !haskey(params, :shaketype) || params.shaketype == :gaussian | |
| 89 return () -> params.shake.*randn(rng,2) | |
| 90 elseif params.shaketype == :disk | |
| 91 return () -> begin | |
| 92 θ = 2π*rand(rng,Float64) | |
| 93 r = params.shake*√(rand(rng,Float64)) | |
| 94 return [r*cos(θ), r*sin(θ)] | |
| 95 end | |
| 96 elseif params.shaketype == :circle | |
| 97 return () -> begin | |
| 98 θ = 2π*rand(rng,Float64) | |
| 99 r = params.shake | |
| 100 return [r*cos(θ), r*sin(θ)] | |
| 101 end | |
| 102 else | |
| 103 error("Unknown shaketype $(params.shaketype)") | |
| 104 end | |
| 105 end | |
| 106 | |
| 107 pixelwise = (shakefn, sz) -> () -> make_const_u(shakefn(), sz) | |
| 108 | |
| 109 | |
| 110 function rotatebytheta(params) | |
| 111 r = params.rotation_factor*(2*rand(rng)-1) | |
| 112 return r | |
| 113 end | |
| 114 | |
| 115 function generate_radonmask(params) | |
| 116 imdim = params.radondims | |
| 117 sino_sparse = params.sino_sparsity | |
| 118 numone = Int64(round(sino_sparse*imdim[1]*imdim[2])) | |
| 119 numzero = imdim[1]*imdim[2]-numone | |
| 120 A = shuffle(rng,reshape([ones(numone); zeros(numzero)],(imdim[1],imdim[2]))) | |
| 121 return A | |
| 122 end | |
| 123 | |
| 124 ################ | |
| 125 # Moving square | |
| 126 ################ | |
| 127 | |
| 128 function generate_square(sz, | |
| 129 :: Type{DisplacementT}, | |
| 130 datachannel :: Channel{OnlineData{DisplacementT}}, | |
| 131 params) where DisplacementT | |
| 132 | |
| 133 if false | |
| 134 v₀ = make_const_v(0.1.*(-1, 1), sz) | |
| 135 nextv = () -> v₀ | |
| 136 elseif DisplacementT == DisplacementFull | |
| 137 nextv = pixelwise(shake(params), sz) | |
| 138 elseif DisplacementT == DisplacementConstant | |
| 139 nextv = shake(params) | |
| 140 else | |
| 141 @error "Invalid DisplacementT" | |
| 142 end | |
| 143 | |
| 144 # Constant linear displacement everywhere has Jacobian determinant one | |
| 145 # (modulo the boundaries which we ignore here) | |
| 146 m = round(Int, sz[1]/5) | |
| 147 b_orig = zeros(sz...) | |
| 148 b_orig[sz[1].-(2*m:3*m), 2*m:3*m] .= 1 | |
| 149 | |
| 150 v_true = nextv() | |
| 151 v_cumul = copy(v_true) | |
| 152 | |
| 153 while true | |
| 154 # Flow original data and add noise | |
| 155 b_true = zeros(sz...) | |
| 156 translate_image!(b_true, b_orig, v_cumul; threads=true) | |
| 157 b = b_true .+ params.noise_level.*randn(rng,sz...) | |
| 158 v = v_true.*(1.0 .+ params.shake_noise_level.*randn(rng,size(v_true)...)) | |
| 159 # Pass true data to iteration routine | |
| 160 data = OnlineData{DisplacementT}(b_true, b, v, v_true, v_cumul) | |
| 161 if !put_unless_closed!(datachannel, data) | |
| 162 return | |
| 163 end | |
| 164 # Next step shake | |
| 165 v_true = nextv() | |
| 166 v_cumul .+= v_true | |
| 167 end | |
| 168 end | |
| 169 | |
| 170 function imgen_square(sz) | |
| 171 return ImGen(curry(generate_square, sz), sz, 1, 1, "square$(sz[1])x$(sz[2])") | |
| 172 end | |
| 173 | |
| 174 ################ | |
| 175 # Shake a photo | |
| 176 ################ | |
| 177 | |
| 178 function generate_shake_image(im, sz, | |
| 179 :: Type{DisplacementConstant}, | |
| 180 datachannel :: Channel{OnlineData{DisplacementConstant}}, | |
| 181 params :: NamedTuple) | |
| 182 | |
| 183 # Restart the seed to enable comparison across predictors | |
| 184 Random.seed!(rng,9182737465) | |
| 185 | |
| 186 nextv = shake(params) | |
| 187 v_true = nextv() | |
| 188 v_cumul = copy(v_true) | |
| 189 | |
| 190 while true | |
| 191 # Extract subwindow of original image and add noise | |
| 192 b_true = zeros(sz...) | |
| 193 extract_subimage!(b_true, im, v_cumul; threads=true) | |
| 194 b = b_true .+ params.noise_level.*randn(rng,sz...) | |
| 195 v = v_true.*(1.0 .+ params.shake_noise_level.*randn(rng,size(v_true)...)) | |
| 196 # Pass data to iteration routine | |
| 197 data = OnlineData{DisplacementConstant}(b_true, b, v, v_true, v_cumul) | |
| 198 if !put_unless_closed!(datachannel, data) | |
| 199 return | |
| 200 end | |
| 201 # Next step shake | |
| 202 v_true = nextv() | |
| 203 v_cumul .+= v_true | |
| 204 end | |
| 205 end | |
| 206 | |
| 207 function imgen_shake(imname, sz) | |
| 208 im = Float64.(Gray.(TestImages.testimage(imname))) | |
| 209 dynrange = maximum(im) | |
| 210 return ImGen(curry(generate_shake_image, im, sz), sz, 1, dynrange, | |
| 211 "$(imname)$(sz[1])x$(sz[2])") | |
| 212 end | |
| 213 | |
| 214 | |
| 215 | |
| 216 ######################################################################## | |
| 217 # PETscan | |
| 218 ######################################################################## | |
| 219 function generate_radon(im, sz, | |
| 220 :: Type{DisplacementConstant}, | |
| 221 datachannel :: Channel{PetOnlineData{DisplacementConstant}}, | |
| 222 params :: NamedTuple) | |
| 223 | |
| 224 # Restart the seed to enable comparison across predictors | |
| 225 Random.seed!(rng,9182737465) | |
| 226 | |
| 227 nextv = shake(params) | |
| 228 v_true = nextv() | |
| 229 v_cumul = copy(v_true) | |
| 230 | |
| 231 S_true = generate_radonmask(params) | |
| 232 | |
| 233 theta_true = rotatebytheta(params) | |
| 234 theta_cumul = copy(theta_true) | |
| 235 | |
| 236 while true | |
| 237 # Define the transformation matrix | |
| 238 center_point = (sz[1]/2 + v_true[1], sz[2]/2 + v_true[2]) | |
| 239 tform = recenter(RotMatrix(theta_cumul), center_point) | |
| 240 | |
| 241 # Apply the transformation to the image using warp | |
| 242 b_true = copy(warp(im, tform, axes(im), fillvalue=Flat())) | |
| 243 | |
| 244 v = v_true.*(1.0 .+ params.shake_noise_level.*randn(rng,size(v_true)...)) | |
| 245 theta = theta_true*(1.0 + params.rotation_noise_level.*randn(rng)) | |
| 246 | |
| 247 # Generate the true and noisy sinograms | |
| 248 sinogram_true = zeros(params.radondims...) | |
| 249 sinogram_true .*= params.scale | |
| 250 radon!(sinogram_true, b_true) | |
| 251 sinogram_noisy = copy(sinogram_true) | |
| 252 | |
| 253 for i=1:params.radondims[1], j=1:params.radondims[2] | |
| 254 sinogram_noisy[i, j] += pois_rand(rng,params.noise_level) | |
| 255 end | |
| 256 | |
| 257 # Pass data to iteration routine | |
| 258 data = PetOnlineData{DisplacementConstant}(b_true, sinogram_true, sinogram_noisy, v, v_true, v_cumul, [theta, theta_cumul], S_true) | |
| 259 if !put_unless_closed!(datachannel, data) | |
| 260 return | |
| 261 end | |
| 262 | |
| 263 # Next step shake | |
| 264 v_true = nextv() | |
| 265 v_cumul .+= v_true | |
| 266 # Next theta | |
| 267 theta_true = rotatebytheta(params) | |
| 268 theta_cumul += theta_true | |
| 269 # Next sinogram mask | |
| 270 S_true = generate_radonmask(params) | |
| 271 end | |
| 272 end | |
| 273 | |
| 274 | |
| 275 function imgen_shepplogan_radon(origsize,sz) | |
| 276 im = convert(Array{Float64},TestImages.shepp_logan(origsize, highContrast=true)) | |
| 277 dynrange = maximum(im) | |
| 278 return ImGen(curry(generate_radon, im, sz), sz, 1, dynrange, "shepplogan$(sz[1])x$(sz[2])") | |
| 279 end | |
| 280 | |
| 281 | |
| 282 end # Module |
