Square root image filter tool in Python

Question

I have this short Python script for making images brighter:

import skimage.io
import math
import sys


def square_root_filter(input_image_file_name, output_image_file_name):
        image_data = skimage.io.imread(input_image_file_name)
        image_data_height = len(image_data)
        image_data_width  = len(image_data[0])

        for row_index in range(image_data_height):
                for column_index in range(image_data_width):
                        channel = image_data[row_index][column_index]

                        r = float(channel[0])
                        g = float(channel[1])
                        b = float(channel[2])

                        r /= 255.0
                        g /= 255.0
                        b /= 255.0

                        r = math.sqrt(r)
                        g = math.sqrt(g)
                        b = math.sqrt(b)

                        channel[0] = int(255 * r)
                        channel[1] = int(255 * g)
                        channel[2] = int(255 * b)

        skimage.io.imsave(output_image_file_name, image_data)


def main():
        square_root_filter(sys.argv[1], sys.argv[2])


if __name__ == "__main__":
        main()

Output

(I tested it only with Python 3.5. Also, it will take some time to do its work.)

I am not a professional Python programmer, so please tell me anything that will make the code better.

Answer 1

PEP 8 specifies four spaces per level of indentation. Since whitespace matters in Python, this is a pretty strong convention.

You should round() the results to the nearest integer rather than truncating them towards 0.

Each of the three channels is treated identically and independently, so you shouldn't write the same code three times.

def square_root_filter(input_image_file_name, output_image_file_name):
    image_data = skimage.io.imread(input_image_file_name)
    for row_data in image_data:
        for pixel in row_data:
            for channel in range(3):
                pixel[channel] = round(255 * math.sqrt(pixel[channel] / 255))
    skimage.io.imsave(output_image_file_name, image_data)

If you take advantage of the fact that channels is a NumPy array, you can vectorize the calculation.

import numpy as np
import skimage.io

def square_root_filter(input_image_file_name, output_image_file_name):
    image_data = skimage.io.imread(input_image_file_name)
    for row_data in image_data:
        for pixel in row_data:
            pixel[:] = np.rint(255 * (pixel / 255) ** 0.5)
    skimage.io.imsave(output_image_file_name, image_data)