fukasawah/preprocess_image.py

## preprocess_image.py
'''
動画用の前処理生成器。
動画に対して加工する場合、連続した画像に対して同じ変換をかけたくなる場合に使う。

クラスはNumpyとPillowに依存している。

Example
--------

```
import cv2
from preprocess_image import PreprocessImageGenerator

cap = cv2.VideoCapture(0) # カメラを使う

pig = PreprocessImageGenerator()
processor = pig.generate_image_processor()

while True:
    ret, frame = cap.read()

    if not ret:
        break

    # OpenCVはBGRなのでRGBに変換する
    img = frame[:,:,::-1]
    processed_img = processor(img) # 画像の前処理関数を実行する。
    processed_img = processed_img[:,:,::-1] # cv2.imshowはBGRなので変換する

    cv2.imshow('Original-Image', frame)
    cv2.imshow('Processed-Image', processed_img)

    key = cv2.waitKey(args.wait_key)

    if key == 27: # ESC = 終了
        break
    if key == 0x20: # space = 変換方法をランダムで変える
        processor = pig.generate_image_processor()
```

コマンドライン実行
-------------------
OpenCV(cv2)が必要。入力はOpenCVのVideoCaptureがサポートしていれば動作する。

# 入力にカメラを使う(カメラIDを指定する)
python preprocess_image.py 0

# 入力に動画ファイルを与える
python preprocess_image.py movie-1.mp4 movie-1.mp4

# 入力に静止画像の連番ファイルを与える（0001.jpg, 0002.jpg... がある）
python preprocess_image.py "%04d.jpg"

'''

# 動画用の前処理生成器


import numpy as np
from PIL import Image, ImageEnhance, ImageFilter
import sys

class PreprocessImageGenerator:
    '''
    画像の前処理を行う関数を生成するジェネレータ。

    変換処理については generate_image_processor() で得られた関数を使い、関数毎に変換に使うパラメータ、変換の順序がランダムに変化する。
    '''

    def __init__(self,
        saturation=(0.75, 1.25), # 彩度。 1.0で変化しない。0.0～nが指定できる
        brightness=(0.75, 1.25), # 輝度。1.0で変化しない。0.0～nが指定できる
        contrast=(0.75, 1.25), # 0.0～n
        rotate=(-15.0, +15.0), # 回転
        minimal_crop=(0.75, 1.0), #　切り取る領域の大きさ
        skip_generate_processor_rate=0.1, #
        resample_strategy=Image.LANCZOS,
        return_ndarray=True, # true: numpy.ndarray, false: PIL.Image
        debug=False
    ):
        '''
        変換量については、2値タプルになっているものは(最小値, 最大値)の中からランダムで決定する仕組みになっている。

        Noneを与えると変換を行わないようにする。(大量に変換する必要がある場合、不要なものはNoneを与えるとよい。)

        Prameters
        ----------
        saturation : (float, float)
            彩度の変化量の最小値～最大値。1.0は変化しないのと等しい。Noneで処理しない。
            Pillow実装のPIL.ImageEnhance.Colorを使い変換する。
        brightness : (float, float)
            輝度の変化量の最小値～最大値。1.0は変化しないのと等しい。Noneで処理しない。
            Pillow実装のPIL.ImageEnhance.Brightnessを使い変換する。
        contrast : (float, float)
            コントラストの変化量の最小値～最大値。1.0は変化しないのと等しい。Noneで処理しない。
            Pillow実装のPIL.ImageEnhance.Contrastを使い変換する。
        rotate : (float, float)
            コントラストの変化量の最小値～最大値。1.0は変化しないのと等しい。Noneで処理しない。
            Pillow実装のPIL.Image.rotateを使い変換するが、はみ出た部分は元の画像に収まるようリサイズされる。
        minimal_crop : (float, float)
            切り取る割合の変化量の最小値～最大値。1.0で切り取らない。0.01～1.0(1%～100%)の間で指定できる。Noneで処理しない。
            切り取り開始位置はランダムで決定するが、切り取り領域がはみ出すことはない。
            アスペクト比の考慮は入力画像によるため、指定することはできない。(16:9の画像から正方形にcropすることはできない)
        skip_generate_processor_rate : float
            変換を適用しないことを確率的に決める。0.01～1.0の間で決定する。
            0.1であれば、10%の確立である変換処理を行わないようにする。Noneで常に行うようになる。
        resample_strategy : int
            Pillow実装を使ったリサイズを行う際にリサンプリングをどう行うか決定する。
        return_ndarray: bool
            処理した結果をndarrayで返すかPIL.Imageで返すかを決定する。デフォルトはTrueでndarrayで返す。
        debug : bool
            デバッグ時に変換量の数値を出力したり、1つの画像に対する処理時間を出力する。
        '''

        # validation
        assert saturation is None or len(saturation) == 2
        assert saturation is None or saturation[0] > 0.00
        assert saturation is None or saturation[1] > 0.00

        assert brightness is None or len(brightness) == 2
        assert brightness is None or brightness[0] > 0.00
        assert brightness is None or brightness[1] > 0.00

        assert contrast is None or len(contrast) == 2
        assert contrast is None or contrast[0] > 0.00
        assert contrast is None or contrast[1] > 0.00

        assert rotate is None or len(rotate) == 2
        assert rotate is None or -180.0 <= rotate[0] <= 180.0
        assert rotate is None or -180.0 <= rotate[1] <= 180.0

        assert minimal_crop is None or len(minimal_crop) == 2
        assert minimal_crop is None or 0.01 <= minimal_crop[0] <= 1.0
        assert minimal_crop is None or 0.01 <= minimal_crop[1] <= 1.0

        assert skip_generate_processor_rate is None or 0.01 <= skip_generate_processor_rate <= 1.0

        # set
        self.saturation = saturation
        self.contrast = contrast
        self.brightness = brightness
        self.rotate = rotate
        self.minimal_crop = minimal_crop
        self.resample_strategy = resample_strategy

        self.skip_generate_processor_rate = skip_generate_processor_rate
        self.return_ndarray = return_ndarray

        self.debug = debug

        # register generator
        self.process_generator_list = []

        if self.saturation:
            self.process_generator_list.append(self._generate_saturation)
        if self.brightness:
            self.process_generator_list.append(self._generate_brightness)
        if self.contrast:
            self.process_generator_list.append(self._generate_contrast)
        if self.rotate:
            self.process_generator_list.append(self._generate_rotate)
        if self.minimal_crop:
            self.process_generator_list.append(self._generate_crop)


    def _random_range(self, tuple_value):
        assert len(tuple_value) == 2
        min_value, max_value = tuple_value
        return np.random.uniform(min_value, max_value)

    def _generate_brightness(self):
        val = self._random_range(self.brightness)
        self._debug_val_print('brightness', val)
        return lambda img: ImageEnhance.Brightness(img).enhance(val)

    def _generate_contrast(self):
        val = self._random_range(self.contrast)
        self._debug_val_print('contrast', val)
        return lambda img: ImageEnhance.Contrast(img).enhance(val)

    def _generate_saturation(self): # 彩度
        val = self._random_range(self.saturation)
        self._debug_val_print('saturation', val)
        return lambda img: ImageEnhance.Color(img).enhance(val)

    def _generate_rotate(self):
        val = self._random_range(self.rotate)
        self._debug_val_print('rotate', val)
        def _rotate(img):
            original = img.size
            return img.rotate(val).resize(original, self.resample_strategy)
        return _rotate

    def _generate_crop(self):
        ratio_p = self._random_range(self.minimal_crop) # 切り取り領域
        start_x = np.random.random() * (1.0 - ratio_p) # 開始x座標
        start_y = np.random.random() * (1.0 - ratio_p) # 開始y座標
        self._debug_val_print('crop', {"ratio_p": ratio_p, "start_x": start_x, "start_y": start_y})

        def _crop(img):
            width, height = img.size
            left = width * start_x
            top = height * start_y
            right = left +  width * ratio_p
            bottom = top + height * ratio_p
            return img.crop((left, top, right, bottom)).resize((width, height), self.resample_strategy)
        return _crop


    def _debug_val_print(self, name, val):
        if self.debug:
            print("{}={}".format(name, val), file=sys.stderr)


    def generate_image_processor(self):
        if self.skip_generate_processor_rate is None:
            processor_list = [ f() for f in self.process_generator_list]
        else:
            processor_list = [ f() for f in self.process_generator_list if self.skip_generate_processor_rate < np.random.random()]

        np.random.shuffle(processor_list)
        processor_count = len(processor_list)

        def preprocess_image(img):
            if self.debug:
                import datetime
                start = datetime.datetime.now()

            if isinstance(img, np.ndarray):
                img = Image.fromarray(img)
            elif isinstance(img, Image.Image):
                pass
            else:
                raise "unexpected input type: "+ type(img)

            for processor in processor_list:
                img = processor(img)

            if self.return_ndarray:
                img = np.asarray(img)

            if self.debug:
                import datetime
                end = datetime.datetime.now()
                print("process time: {:7.3f} ms ({} processors)".format((end-start).total_seconds() * 1000.0, processor_count), end='\r')

            return img

        return preprocess_image


if __name__ == '__main__':
    import cv2
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument('--wait-key', help='OpenCV waitKey Time(ms) in image preview (Debug mode Only.)', type=int, default=100)
    parser.add_argument('input_sources', metavar='InputSource', type=str, nargs='+', help='OpenCV2 support input source.')
    parser.add_argument('--input-width', type=int, default=1280)
    parser.add_argument('--input-height', type=int, default=720)

    args = parser.parse_args()

    pig = PreprocessImageGenerator(debug=True)

    for source in args.input_sources:
        if source.isdecimal():
            source = int(source)

        processor = pig.generate_image_processor() # 画像の前処理関数をgenerate_image_processorで作る。
        cap = cv2.VideoCapture(source)

        # カメラ取り込み時の画像サイズを決定する
        cap.set(cv2.CAP_PROP_FRAME_WIDTH, args.input_width)
        cap.set(cv2.CAP_PROP_FRAME_HEIGHT, args.input_height)

        while True:
            ret, frame = cap.read()

            if not ret:
                break

            processed_img = processor(frame[:,:,::-1]) [:,:,::-1] # 画像の前処理関数を実行する。（OpenCVはBGRなのでRGBにしてから変換する）

            cv2.imshow('Original-Image', frame)
            cv2.imshow('Processed-Image', processed_img)
            key = cv2.waitKey(args.wait_key)
            if key == 27: # ESC
                break
            if key == 0x20: # space
                # update processor
                processor = pig.generate_image_processor()
	'''
	動画用の前処理生成器。
	動画に対して加工する場合、連続した画像に対して同じ変換をかけたくなる場合に使う。

	クラスはNumpyとPillowに依存している。

	Example
	--------

	```
	import cv2
	from preprocess_image import PreprocessImageGenerator

	cap = cv2.VideoCapture(0) # カメラを使う

	pig = PreprocessImageGenerator()
	processor = pig.generate_image_processor()

	while True:
	ret, frame = cap.read()

	if not ret:
	break

	# OpenCVはBGRなのでRGBに変換する
	img = frame[:,:,::-1]
	processed_img = processor(img) # 画像の前処理関数を実行する。
	processed_img = processed_img[:,:,::-1] # cv2.imshowはBGRなので変換する

	cv2.imshow('Original-Image', frame)
	cv2.imshow('Processed-Image', processed_img)

	key = cv2.waitKey(args.wait_key)

	if key == 27: # ESC = 終了
	break
	if key == 0x20: # space = 変換方法をランダムで変える
	processor = pig.generate_image_processor()
	```

	コマンドライン実行
	-------------------
	OpenCV(cv2)が必要。入力はOpenCVのVideoCaptureがサポートしていれば動作する。

	# 入力にカメラを使う(カメラIDを指定する)
	python preprocess_image.py 0

	# 入力に動画ファイルを与える
	python preprocess_image.py movie-1.mp4 movie-1.mp4

	# 入力に静止画像の連番ファイルを与える（0001.jpg, 0002.jpg... がある）
	python preprocess_image.py "%04d.jpg"

	'''

	# 動画用の前処理生成器


	import numpy as np
	from PIL import Image, ImageEnhance, ImageFilter
	import sys

	class PreprocessImageGenerator:
	'''
	画像の前処理を行う関数を生成するジェネレータ。

	変換処理については generate_image_processor() で得られた関数を使い、関数毎に変換に使うパラメータ、変換の順序がランダムに変化する。
	'''

	def __init__(self,
	saturation=(0.75, 1.25), # 彩度。 1.0で変化しない。0.0～nが指定できる
	brightness=(0.75, 1.25), # 輝度。1.0で変化しない。0.0～nが指定できる
	contrast=(0.75, 1.25), # 0.0～n
	rotate=(-15.0, +15.0), # 回転
	minimal_crop=(0.75, 1.0), #　切り取る領域の大きさ
	skip_generate_processor_rate=0.1, #
	resample_strategy=Image.LANCZOS,
	return_ndarray=True, # true: numpy.ndarray, false: PIL.Image
	debug=False
	):
	'''
	変換量については、2値タプルになっているものは(最小値, 最大値)の中からランダムで決定する仕組みになっている。

	Noneを与えると変換を行わないようにする。(大量に変換する必要がある場合、不要なものはNoneを与えるとよい。)

	Prameters
	----------
	saturation : (float, float)
	彩度の変化量の最小値～最大値。1.0は変化しないのと等しい。Noneで処理しない。
	Pillow実装のPIL.ImageEnhance.Colorを使い変換する。
	brightness : (float, float)
	輝度の変化量の最小値～最大値。1.0は変化しないのと等しい。Noneで処理しない。
	Pillow実装のPIL.ImageEnhance.Brightnessを使い変換する。
	contrast : (float, float)
	コントラストの変化量の最小値～最大値。1.0は変化しないのと等しい。Noneで処理しない。
	Pillow実装のPIL.ImageEnhance.Contrastを使い変換する。
	rotate : (float, float)
	コントラストの変化量の最小値～最大値。1.0は変化しないのと等しい。Noneで処理しない。
	Pillow実装のPIL.Image.rotateを使い変換するが、はみ出た部分は元の画像に収まるようリサイズされる。
	minimal_crop : (float, float)
	切り取る割合の変化量の最小値～最大値。1.0で切り取らない。0.01～1.0(1%～100%)の間で指定できる。Noneで処理しない。
	切り取り開始位置はランダムで決定するが、切り取り領域がはみ出すことはない。
	アスペクト比の考慮は入力画像によるため、指定することはできない。(16:9の画像から正方形にcropすることはできない)
	skip_generate_processor_rate : float
	変換を適用しないことを確率的に決める。0.01～1.0の間で決定する。
	0.1であれば、10%の確立である変換処理を行わないようにする。Noneで常に行うようになる。
	resample_strategy : int
	Pillow実装を使ったリサイズを行う際にリサンプリングをどう行うか決定する。
	return_ndarray: bool
	処理した結果をndarrayで返すかPIL.Imageで返すかを決定する。デフォルトはTrueでndarrayで返す。
	debug : bool
	デバッグ時に変換量の数値を出力したり、1つの画像に対する処理時間を出力する。
	'''

	# validation
	assert saturation is None or len(saturation) == 2
	assert saturation is None or saturation[0] > 0.00
	assert saturation is None or saturation[1] > 0.00

	assert brightness is None or len(brightness) == 2
	assert brightness is None or brightness[0] > 0.00
	assert brightness is None or brightness[1] > 0.00

	assert contrast is None or len(contrast) == 2
	assert contrast is None or contrast[0] > 0.00
	assert contrast is None or contrast[1] > 0.00

	assert rotate is None or len(rotate) == 2
	assert rotate is None or -180.0 <= rotate[0] <= 180.0
	assert rotate is None or -180.0 <= rotate[1] <= 180.0

	assert minimal_crop is None or len(minimal_crop) == 2
	assert minimal_crop is None or 0.01 <= minimal_crop[0] <= 1.0
	assert minimal_crop is None or 0.01 <= minimal_crop[1] <= 1.0

	assert skip_generate_processor_rate is None or 0.01 <= skip_generate_processor_rate <= 1.0

	# set
	self.saturation = saturation
	self.contrast = contrast
	self.brightness = brightness
	self.rotate = rotate
	self.minimal_crop = minimal_crop
	self.resample_strategy = resample_strategy

	self.skip_generate_processor_rate = skip_generate_processor_rate
	self.return_ndarray = return_ndarray

	self.debug = debug

	# register generator
	self.process_generator_list = []

	if self.saturation:
	self.process_generator_list.append(self._generate_saturation)
	if self.brightness:
	self.process_generator_list.append(self._generate_brightness)
	if self.contrast:
	self.process_generator_list.append(self._generate_contrast)
	if self.rotate:
	self.process_generator_list.append(self._generate_rotate)
	if self.minimal_crop:
	self.process_generator_list.append(self._generate_crop)


	def _random_range(self, tuple_value):
	assert len(tuple_value) == 2
	min_value, max_value = tuple_value
	return np.random.uniform(min_value, max_value)

	def _generate_brightness(self):
	val = self._random_range(self.brightness)
	self._debug_val_print('brightness', val)
	return lambda img: ImageEnhance.Brightness(img).enhance(val)

	def _generate_contrast(self):
	val = self._random_range(self.contrast)
	self._debug_val_print('contrast', val)
	return lambda img: ImageEnhance.Contrast(img).enhance(val)

	def _generate_saturation(self): # 彩度
	val = self._random_range(self.saturation)
	self._debug_val_print('saturation', val)
	return lambda img: ImageEnhance.Color(img).enhance(val)

	def _generate_rotate(self):
	val = self._random_range(self.rotate)
	self._debug_val_print('rotate', val)
	def _rotate(img):
	original = img.size
	return img.rotate(val).resize(original, self.resample_strategy)
	return _rotate

	def _generate_crop(self):
	ratio_p = self._random_range(self.minimal_crop) # 切り取り領域
	start_x = np.random.random() * (1.0 - ratio_p) # 開始x座標
	start_y = np.random.random() * (1.0 - ratio_p) # 開始y座標
	self._debug_val_print('crop', {"ratio_p": ratio_p, "start_x": start_x, "start_y": start_y})

	def _crop(img):
	width, height = img.size
	left = width * start_x
	top = height * start_y
	right = left + width * ratio_p
	bottom = top + height * ratio_p
	return img.crop((left, top, right, bottom)).resize((width, height), self.resample_strategy)
	return _crop


	def _debug_val_print(self, name, val):
	if self.debug:
	print("{}={}".format(name, val), file=sys.stderr)


	def generate_image_processor(self):
	if self.skip_generate_processor_rate is None:
	processor_list = [ f() for f in self.process_generator_list]
	else:
	processor_list = [ f() for f in self.process_generator_list if self.skip_generate_processor_rate < np.random.random()]

	np.random.shuffle(processor_list)
	processor_count = len(processor_list)

	def preprocess_image(img):
	if self.debug:
	import datetime
	start = datetime.datetime.now()

	if isinstance(img, np.ndarray):
	img = Image.fromarray(img)
	elif isinstance(img, Image.Image):
	pass
	else:
	raise "unexpected input type: "+ type(img)

	for processor in processor_list:
	img = processor(img)

	if self.return_ndarray:
	img = np.asarray(img)

	if self.debug:
	import datetime
	end = datetime.datetime.now()
	print("process time: {:7.3f} ms ({} processors)".format((end-start).total_seconds() * 1000.0, processor_count), end='\r')

	return img

	return preprocess_image



	if __name__ == '__main__':
	import cv2
	import argparse

	parser = argparse.ArgumentParser()
	parser.add_argument('--wait-key', help='OpenCV waitKey Time(ms) in image preview (Debug mode Only.)', type=int, default=100)
	parser.add_argument('input_sources', metavar='InputSource', type=str, nargs='+', help='OpenCV2 support input source.')
	parser.add_argument('--input-width', type=int, default=1280)
	parser.add_argument('--input-height', type=int, default=720)

	args = parser.parse_args()

	pig = PreprocessImageGenerator(debug=True)

	for source in args.input_sources:
	if source.isdecimal():
	source = int(source)

	processor = pig.generate_image_processor() # 画像の前処理関数をgenerate_image_processorで作る。
	cap = cv2.VideoCapture(source)

	# カメラ取り込み時の画像サイズを決定する
	cap.set(cv2.CAP_PROP_FRAME_WIDTH, args.input_width)
	cap.set(cv2.CAP_PROP_FRAME_HEIGHT, args.input_height)

	while True:
	ret, frame = cap.read()

	if not ret:
	break

	processed_img = processor(frame[:,:,::-1]) [:,:,::-1] # 画像の前処理関数を実行する。（OpenCVはBGRなのでRGBにしてから変換する）

	cv2.imshow('Original-Image', frame)
	cv2.imshow('Processed-Image', processed_img)
	key = cv2.waitKey(args.wait_key)
	if key == 27: # ESC
	break
	if key == 0x20: # space
	# update processor
	processor = pig.generate_image_processor()