Skip to content

Instantly share code, notes, and snippets.

@surt91
Last active Dec 30, 2015
Embed
What would you like to do?
a simple implementation of a black-and-white png denoiser (inclusive adding the noise) based on the Ising model. You need python3 and pypng (https://pypi.python.org/pypi/pypng/)
#!/usr/bin/env python3
import png
from random import random, sample
class Bitmap():
def __init__(self, filename):
pic = png.Reader(filename)
w,h,pixel,meta = pic.asDirect()
self.w = w
self.h = h
self.n = w*h
self.bmp = self.convertToBitmap(pixel)
def savePNG(self, filename):
outPic = png.Writer(width=self.w, height=self.h)
with open(filename, "wb") as f:
outPic.write(f, self.convertToPng())
def convertToBitmap(self, pixel):
out = []
for row in pixel:
while row:
r,g,b = row[:3]
del row[:3]
out.append(-1 if r==0 and g==0 and b==0 else 1)
return out
def convertToPng(self):
tmpRow = []
for i,p in enumerate(self.bmp):
tmpRow.extend((255,255,255) if p==1 else (0,0,0))
if (i+1)%self.w==0:
yield tmpRow
tmpRow = []
def compareBmp(self, bmp2):
score = 0
for a,b in zip(self.bmp, bmp2.bmp):
if a != b:
score += 1
return (1 - score/(self.n))
class RandomizedBitmap(Bitmap):
@staticmethod
def rand(p):
return random() < p
def flip(self, i):
self.bmp[i] *= -1
def addNoise(self, noise=0.05):
for i in range(self.n):
if self.rand(noise):
self.flip(i)
class MarkovRandomField(RandomizedBitmap):
def __init__(self, filename):
super().__init__(filename)
self.nodes = list(range(self.n))
self.neighbors = {i:self.findNeighbors(i) for i in self.nodes}
def findNeighbors(self, i):
# Boundary
neighbors = {i-1, i+1, i-self.w, i+self.w}
if i%self.w == 0: # left
neighbors.discard(i-1)
if i%self.w == self.w-1: # right
neighbors.discard(i+1)
if i+self.w >= self.n: # upper
neighbors.discard(i+self.w)
if i-self.w < 0: # lower
neighbors.discard(i-self.w)
return neighbors
def energyOfSite(self, i):
h = 0
beta = 1
eta = 2.5
b = 0
E = 0
if h:
E = h*self.bmp[i]
if beta:
for j in self.neighbors[i]:
b += self.bmp[i]*self.bmp[j]
E -= beta * b
if eta:
E -= eta * self.bmp[i]*self.yBmp[i]
return E
def denoise(self):
self.yBmp = list(self.bmp)
somethingChanged = True
while somethingChanged:
somethingChanged = False
for i in sample(self.nodes, len(self.nodes)):
E1 = self.energyOfSite(i)
self.flip(i)
E2 = self.energyOfSite(i)
if E2 - E1 > 0:
self.flip(i)
else:
somethingChanged = True
def denoise_simulated_annealing(self):
self.yBmp = list(self.bmp)
for T in (j / 10 for j in range(31, 0, -5)):
print("T = {}".format(T))
for _ in range(3):
for i in sample(self.nodes, len(self.nodes)):
E1 = self.energyOfSite(i)
self.flip(i)
E2 = self.energyOfSite(i)
if E2 - E1 > 0 and not self.rand(exp(-(E2 - E1)/T)):
self.flip(i)
if __name__ == "__main__":
import sys
if len(sys.argv) > 1:
filename = sys.argv[1]
else:
filename = "dag.png"
tmpFilename = filename.replace(".png", "Noisy.png")
outFilename = filename.replace(".png", "Denoised.png")
original = Bitmap(filename)
bmp = MarkovRandomField(filename)
bmp.addNoise()
print("Similarity: {:.2f}%".format(bmp.compareBmp(original) *100))
bmp.savePNG(tmpFilename)
bmp.denoise()
# bmp.denoise_simulated_annealing()
print("Similarity: {:.2f}%".format(bmp.compareBmp(original) *100))
bmp.savePNG(outFilename)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment