Last active
December 30, 2015 06:29
-
-
Save surt91/7789753 to your computer and use it in GitHub Desktop.
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/)
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#!/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