images.py

############################################################################
###    image processing
############################################################################

# importing out own Matrix class.
from matrix import *
import random

def add_gauss(mat, sigma=10):
  ''' Generates Gaussian noise with mean 0 and SD sigma.
      Adds indep. noise to pixel, keeping values in 0..255'''
  n, m = mat.dim()
  new = Matrix(n,m)
  for i in range (n):
      for j in range (m):
          noise = round(random.gauss(0, sigma))
          if noise > 0:
              new[i, j] = min(mat[i, j] + noise, 255)
          elif noise < 0:
              new[i, j] = max(mat[i, j] + noise, 0)
          elif noise == 0:
              new[i, j] = mat[i, j]

  return new

def add_SP(mat, p=0.01):
  ''' Generates salt and pepper noise: Each pixel is "hit" indep.
      with prob. p. If hit, it has fifty fifty chance of becoming
      white or black. '''
  n, m = mat.dim()
  new = Matrix(n, m)
  for i in range (n):
      for j in range (m):
          rand = random.randint(0, 20000)
          if rand < p * 20000:
              if rand % 2 == 1:
                  new[i, j] = 0
              elif rand % 2 == 0:
                  new[i,j ] = 255
          else: 
              new[i, j] = mat[i, j]
  return new

def average(lst):
  return round(sum(lst) / len(lst))

def median(lst):
  sort_lst = sorted(lst)
  l = len(sort_lst)
  if l % 2 == 1:  # odd number of elements. well defined median
    return sort_lst[l // 2]
  else:  # even number of elements. average of middle two
    return average(sort_lst[l//2 - 1: l//2 + 1])

## Local methods
def local_means(A, k=1):
  n, m = A.dim()
  res = A[:, :]  # brand new copy of A
  for i in range(k, n - k):
    for j in range(k, m - k):
      neighborhood = A[i-k:i+k+1, j-k:j+k+1]
      res[i, j] = average(list(neighborhood))
  return res

def local_medians(A, k=1):
  n, m = A.dim()
  res = A[:, :]  # brand new copy of A

  for i in range(k, n-k):
    for j in range(k, m-k):
      neighborhood = A[i-k:i+k+1, j-k:j+k+1]
      res[i, j] = median(list(neighborhood))
  return res

def join_h(mat1, mat2):
    """ joins two matrices, side by side with some separation """
    n1, m1 = mat1.dim()
    n2, m2 = mat2.dim()
    m = m1 + m2 + 10
    n = max(n1, n2)
    # New matrix with white pixels.
    new = Matrix(n, m, val=255)

    for i in range(n1):
        for j in range(m1):
            new[i, j] = mat1[i,j]
            
    for i in range(n2):
        for j in range(m2):
            new[i, j + m1 + 10] = mat2[i, j]

    return new

############################################################################
###    convert .jpg/.bmp <--> .bitmap
###    .bitmap format will represent BW images readable by Matrix class
############################################################################

try:
    # These 2 modules can be downloaded from
    # http://www.lfd.uci.edu/~gohlke/pythonlibs/
    import numpy            #numerical python module. Used here for arrays
    from PIL import Image   #PIL is Python's Image Library
except ImportError:
    print("Missing numpy and/or PIL modules.")

def image2bitmap(file):
    ''' convert .jpg/.bmp file to format .bitmap
        which can be loaded into class Matrix using Matrix.load() ''' 
    image = Image.open(file).convert("L") # converts to 8 bit black and white image
    im = numpy.asarray(image)             # stores image in an array
    im.setflags(write=1)                  # makes array writeable
    n,m = im.shape[0], im.shape[1]        # image dimensions

    new_file_name = file.split(".")[0] + ".bitmap"
    new_file = open(new_file_name, 'w')

    print(n,m, file=new_file)
    for i in range(n):
        for j in range(m):
            print(int(im[i,j]), end=" ", file = new_file)
        print("", file=new_file)#newline
    new_file.close()

def bitmap2image(file):
    ''' convert format .bitmap image to .jpg/.bmp '''
    im = Matrix.load(file)
    n,m = im.dim()
    ar = numpy.zeros((n,m)) #an nXm array of zeros
    for i in range(n):
        for j in range(m):
            ar[i,j] = im[i,j]
    image = Image.fromarray(numpy.uint8(ar))
    image.save(file.split(".")[0] + ".jpg")