amplify.py

import cv2
import numpy as np
import math
import random
import os
import json

# 定義
targetPath = '/tmp/DataSet'
manifest = 'output.manifest'

# 彩度
def saturation(src, saturation):
    # 一旦、BGRをHSVに変換して彩度を変換する
    img = cv2.cvtColor(src,cv2.COLOR_BGR2HSV)
    img[:,:,(1)] = img[:,:,(1)] * saturation 
    img = cv2.cvtColor(img,cv2.COLOR_HSV2BGR) 
    return img

# 明度
def brightness(src, brightness):
    # 一旦、BGRをHSVに変換して明度を変換する
    img = cv2.cvtColor(src,cv2.COLOR_BGR2HSV) 
    img[:,:,(2)] = img[:,:,(2)] * brightness 
    img = cv2.cvtColor(img,cv2.COLOR_HSV2BGR) 
    return img

# コントラスト
def contrast(src, alpha):
    # 各色相をalphaで演算する
    img = alpha * src
    return np.clip(img, 0, 255).astype(np.uint8)

# モザイク
def mosaic(src, ratio):
    # ratio倍でリサイズして戻す
    img = cv2.resize(src, None, fx=ratio, fy=ratio, interpolation=cv2.INTER_NEAREST)
    return cv2.resize(img, src.shape[:2][::-1], interpolation=cv2.INTER_NEAREST)

# ガウスノイズ
def gaussian(src, sigma):
    # ランダム値で生成した画像と合成する
    row,col,ch = src.shape
    mean = 0
    gauss = np.random.normal(mean, sigma, (row, col, ch)).astype('u1')
    gauss = gauss.reshape(row, col, ch)
    return src + gauss

# ごま塩ノイズ
def noise(src, amount):
    # ランダム値で生成したノイズと合成する
    img = src.copy()
    num_pepper = np.ceil(amount* src.size * (0.5))
    coords = [np.random.randint(0, i-1 , int(num_pepper)) for i in src.shape]
    img[coords[:-1]] = (0,0,0)
    return img

# 1件のデータを表現するクラス
class Data():
    def __init__(self, src):
        # プロジェクト名の取得
        for key in src.keys():
            index = key.rfind("-metadata")
            if(index!=-1):
                self.projectName = key[0:index]

        # メタデータの取得
        self.metadata = src[self.projectName + '-metadata']
        class_map = self.metadata["class-map"]

        # 画像名の取得
        sourceRef = src["source-ref"]
        imgFileName = os.path.basename(sourceRef)
        self.s3Path = sourceRef.replace(imgFileName,"")
        self.baseName = imgFileName.split('.')[0]
        self.ext = imgFileName.split('.')[1]
        # 画像サイズの取得
        project = src[self.projectName]
        image_size = project["image_size"]
        self.img_width = image_size[0]["width"]
        self.img_height = image_size[0]["height"]
        
        self.annotations = []
        # アノテーションの取得
        for annotation in project["annotations"]:
            class_id = annotation["class_id"]
            top = annotation["top"]
            left = annotation["left"]
            width = annotation["width"]
            height = annotation["height"]

            self.annotations.append({
                "class_id": class_id,
                "label": class_map[str(class_id)],
                "width": width,
                "top": top,
                "height": height,
                "left": left
            })
    def dumps(self):
        dstJson = {}
        dstJson["source-ref"] = "{}{}.{}".format(self.s3Path, self.baseName, self.ext)
        dstJson[self.projectName] = {}
        dstJson[self.projectName]["image_size"] = []
        dstJson[self.projectName]["image_size"].append({
            "width": self.img_width,
            "height": self.img_height,
            "depth": 3
        })
        dstJson["boxlabel"]["annotations"] = []
        for annotation in self.annotations:
            dstJson["boxlabel"]["annotations"].append({
                "class_id": annotation["class_id"],
                "top": annotation["top"],
                "left":annotation["left"],
                "width": annotation["width"],
                "height": annotation["height"]
            })
        dstJson[self.projectName + '-metadata'] = self.metadata
        return json.dumps(dstJson)
    

# 全てのJSONデータを読み込む
def getDataList(inputPath, manifest):
    dataList = []
    with open("{}/{}".format(inputPath, manifest), 'r') as f:
        srcList = f.read().split('\n')
        for src in srcList:
            if(src != ''):
                json_src = json.loads(src)
                dataList.append(Data(json.loads(src)))
    return dataList

def main():

    # 変換リスト
    convertList = [
        {"function": saturation, "param": 0.2}, # 彩度
        {"function": saturation, "param": 3.0}, 
        {"function": brightness, "param": 0.6}, # 明度
        {"function": contrast,   "param": 0.8}, # コントラスト
        {"function": mosaic,     "param": 0.05},# モザイク
        {"function": mosaic,     "param": 0.1},
        {"function": mosaic,     "param": 0.2},
        {"function": gaussian,   "param": 20.0},# ガウスノイズ
        {"function": noise,      "param": 0.1}, # ごま塩ノイズ
    ]

    dataList = getDataList(targetPath, manifest)
    print("全データ: {}件 ".format(len(dataList)))
    
    outputManifest = ''
    for data in dataList:
        outputManifest += data.dumps() + '\n'
        srcImage = "{}/{}.{}".format(targetPath, data.baseName, data.ext)
        orgBaseName = data.baseName
        for i,convert in enumerate(convertList):

            baseName = "{}-{}".format(orgBaseName, str(i+1).zfill(3))
            dstImage = "{}/{}.{}".format(targetPath, baseName, data.ext)
            data.baseName = baseName
            outputManifest += data.dumps() + '\n'

            img = cv2.imread(srcImage)
            img = convert["function"](img, convert["param"])
            cv2.imwrite(dstImage,img)

    with open("{}/{}".format(targetPath, manifest), mode='w') as f:
        f.write(outputManifest)

    print("増幅後データ: {}件 ".format(len(outputManifest.split('\n'))))
    
main()