Source code for pysteg.features.wam

#! /usr/bin/env python
##  $Id$
## -*- coding: utf-8 -*-

"""
The WAM features.  Only WAM based on NxN LAW-ML is provided.
An implementation of the original WAM features of Fridrich et al.
can be found in pysteg.analysis.wfeatures.

This module follows the sample code provided in my book.

:Module:    pysteg.features.wam
:Date:      $Date$
:Revision:  $Revision$
:Copyright: (C) 2012: Hans Georg Schaathun <georg@schaathun.net> 
"""

print "[pysteg.features.wam] $Id$"

# *****************************
# pysteg.features.wam
# *****************************
#
# .. module:: pysteg.features.wam
#   

import pywt
import numpy as np

def wam9(C,N=3,nmom=9):
   "Calculate 9 WAM features from the wavelet component C."
   M0 = C - lawmlN(C,N)
   R = np.mean(M0)
   M = M0.flatten().astype(float) - R
   return [R] + [ np.sum(M**k)/M.size for k in range(2,nmom+1) ]

def wamNxN(I,N=3,nmom=9):
    "Calculate the WAM(N)-27 from the pixmap array I."
    if I == None: return wamNxNnames(N,nmom)
    (L,H) = pywt.dwt2(I,"db4")
    return reduce( list.__add__, [ wam9(C,N,nmom) for C in H ] )

def subm(X,N,i,j):
   "Return an NxN submatrix of X, centered at element (i,j)."
   N /= 2
   (m0,m1) = X.shape
   (i0,j0) = (max(0,i-N), max(0,j-N))
   (i1,j1) = (min(m0,i+N)+1, min(m1,j+N)+1)
   return X[i0:i1,j0:j1]

def lawmlN(X,N=3,sig0=0.5):
   "Return the denoised image using NxN LAW-ML."
   (m0,m1) = X.shape
   S = np.zeros((m0,m1))
   for x in xrange(m0):
      for y in xrange(m1):
	 sub = subm(X,N,x,y)
         S[x,y] = max(0, np.var(sub) - sig0 )
   return X*S/(S+sig0)

def wamNxNnames(N=3,nmom=9):
   return [ "WAM%sx%s(%s%s)" % (N,N,d,i+1) for d in ("H","V","D")
      for i in range(nmom) ]

if __name__ == "__main__":
   print wamNxNnames()