マルチプロセス処理のサンプル

multiprocess.py

import pandas as pd
import cv2
from tqdm import tqdm
import numpy as np
from scipy.stats import multivariate_normal
import matplotlib.pyplot as plt
import time
import pathlib
from multiprocessing import Process, Queue
import multiprocessing

def runner(id,q,s,e,csv):

    csvData = {
        "x":list(),
        "y":list()
    }

    for i in tqdm(range(s, e), position=id, desc=str(s)+", "+str(e)):
        img = cv2.imread(csv["x"][i])
        h, w, c = img.shape

        y = np.loadtxt(csv["y"][i], delimiter=',')
        p_x = y[::2]
        p_y = y[1::2]

        plt.clf()
        data = dict()
        for j in range(len(p_x)):
            pos = np.dstack(np.mgrid[0:h:1, 0:w:1])
            rv = multivariate_normal(mean=[p_y[j], p_x[j]], cov=100)
            mask = rv.pdf(pos).astype(np.float64)
            mask = cv2.normalize(mask, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)
            mask = cv2.resize(mask,(128,128))

            data["x_" + str(j)] = np.ravel(mask)

        path = pathlib.Path(csv["y"][i])
        maskPath = "./y_mask/" + path.name

        df = pd.DataFrame.from_dict(data)
        #df = df.T
        df.to_csv(maskPath, index=False,header=False)

        csvData["x"].append(csv["x"][i])
        csvData["y"].append(maskPath)

    q.put(csvData)


if __name__ == '__main__':
    csv = pd.read_csv("./train.csv")

    worker = 10
    #size = len(csv["x"])
    size = 100
    index = list(range(0, size+1, int(size/worker)))

    q = multiprocessing.Manager().Queue()
    taskList = []

    for i in range(len(index)-1):
        s = index[i]
        e = index[i+1]
        p = Process(target=runner, args=[i,q, s, e, csv])
        p.start()
        taskList.append(p)

    for i in range(len(taskList)):
        taskList[i].join()

    data = {
        "x":list(),
        "y":list()
    }
    for i in range(len(taskList)):
        tmp_df = q.get()
        data["x"].extend(tmp_df["x"])
        data["y"].extend(tmp_df["y"])

    df = pd.DataFrame.from_dict(data)
    df.to_csv("./dataset.csv", index=False)

    totalSize = len(df.index)
    trainSize = int(float(totalSize) * 0.7)
    testSize = totalSize - trainSize

    print(trainSize, testSize)

    train_df = df[df.index < trainSize]
    test_df = df[df.index >= trainSize]
    train_df.to_csv("./trainV3.csv", index=False)
    test_df.to_csv("./testV3.csv", index=False)

    print(train_df)
    print(test_df)

    print("Finish !")