How to parallelize CNN evaluation of images with PyTorch
I have a fairly long script that uses a CNN to classify people from drone footage as either human or non-human. The general flow is as such: (1) Create a video object and extract "captures" from it based on intervals of seconds specified. (2) Instantiate the Model class, loading a PyTorch CNN. (3) For each capture, break into much smaller (and overlapping) images which the CNN will classify as human or non-human. (4) Create a list of coordinates, which satisfy previous step. (5) Draw red squares surrounding coordinates, and save marked up image. (6) Repeat process for every capture.

There's a bottleneck at the image level. Crops are being produced and evaluated by the CNN sequentially. I would love to parallelize this process but it goes beyond my current knowledge/experience level. 

Any recommendations? The script is below for reference.

import torch, torchvision
from torchvision import datasets, models, transforms
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import time
from torchsummary import summary

import numpy as np
import matplotlib.pyplot as plt
import os

from PIL import Image
import shutil
from PIL import Image, ImageDraw
import random
import cv2

class Model():

    def __init__(self,model):

        self.idx_to_class = {1:'No human',0:'Human'}
        self.image_transforms = { 
            'test': transforms.Compose([
                transforms.Resize(size=256),
                transforms.CenterCrop(size=224),
                transforms.ToTensor(),
                transforms.Normalize([0.485, 0.456, 0.406],
                                     [0.229, 0.224, 0.225])
            ])
        }
        self.model = torch.load(model)

    def predict(self, test_image_name):
        transform = self.image_transforms['test']
#         test_image = Image.open(test_image_name)
        test_image_tensor = transform(test_image_name)

        if torch.cuda.is_available():
            test_image_tensor = test_image_tensor.view(1, 3, 224, 224).cuda()
        else:
            test_image_tensor = test_image_tensor.view(1, 3, 224, 224)

        with torch.no_grad():
            self.model.eval()
            out = self.model(test_image_tensor)
            ps = torch.exp(out)
            topk, topclass = ps.topk(1, dim=1)
            return topclass.cpu().numpy()[0][0]


class Image_classifier():

    def __init__(self,image,folder,positive_location,model):
        self.name = image
        self.alias = image.split('.')[0]
        self.folder = folder
        self.src = Image.open(f"{self.folder}/{self.name}")
        self.width = self.src.size[0]
        self.square_size = int(self.width/25)
        self.max_down = int(self.src.height/self.square_size) * self.square_size - self.square_size
        self.max_right = int(self.src.width/self.square_size) * self.square_size - self.square_size
        self.offset = int(self.square_size/3)
        self.positive_location = positive_location 
        self.model = model


    def window_coordinates(self):

        def right_pass(y):
            x_coords = [x for x in range(0,self.max_right,self.offset)]
            y_coords = [y for x in range(0,self.max_right,self.offset)]
            return [(x,y,x+self.square_size,y+self.square_size) for x,y in zip(x_coords,y_coords)]

        #v_pass = np.vectorize(right_pass)

        y_values = [y for y in range(0,self.max_down,self.offset)]
        coordinates = [right_pass(y) for y in y_values]
        self.coordinates = [item for sublist in coordinates for item in sublist]

    def predict_coord(self,coord):
        sample = self.src.crop(coord)
        return self.model.predict(sample)


    def parse_coordinates(self):

        new_coords = [coord if self.predict_coord(coord) == 0 else 0 for coord in self.coordinates]
        while 0 in new_coords:
            new_coords.remove(0)
        self.coordinates = new_coords

    def select_squares(self):
        self.window_coordinates()
        self.parse_coordinates()
        self.drawable = ImageDraw.Draw(self.src)
        for coord in self.coordinates:
            self.drawable.rectangle(list(coord), fill = None, outline = 'red')
        self.src.save(f"{self.positive_location}/{self.alias}.jpg")      


class Video_classifier():
    def __init__(self,video,root,seconds):
        self.video = video
        self.alias = self.video.split('.')[0]
        self.root = root
        self.seconds = seconds
        self.model = Model(model='/home/team4/output/_model_142.pt')

        self.folder = f"{self.root}/{self.alias}"
        if os.path.exists(self.folder):
            shutil.rmtree(self.folder)
        os.mkdir(self.folder)

        self.positive_location = f"{self.root}/{self.alias}/positives"
        if os.path.exists(self.positive_location):
            shutil.rmtree(self.positive_location)
        os.mkdir(self.positive_location)

    def get_frames(self):
        import cv2
        cam = cv2.VideoCapture(f"{self.root}/{self.video}")

        (major_ver, minor_ver, subminor_ver) = (cv2.__version__).split('.')
        if int(major_ver)  &lt; 3 :
            fps = round(video.get(cv2.cv.CV_CAP_PROP_FPS))
        else:
            fps = round(cam.get(cv2.CAP_PROP_FPS))

        current_frame = 0
        while(True):
            ret,frame = cam.read()
            if ret: 
                if current_frame % (self.seconds*fps) == 0:
                    f_name = f"{self.alias}_{current_frame}.jpg"
                    cv2.imwrite(f_name, frame)
                    shutil.move(f_name, self.folder)

                current_frame += 1
            else:
                break

        cam.release()
        cv2.destroyAllWindows()

    def read_dir(self):
        self.files = [f for f in os.listdir(self.folder) if os.path.isfile(os.path.join(self.folder, f))]

    def classify_frames(self):
        self.get_frames()
        self.read_dir()
        for file in self.files:
            image = Image_classifier(image=file, folder=self.folder, positive_location=self.positive_location,
                                     model=self.model)
            image.select_squares()

test = Video_classifier(video='refugee_test.mp4',root='/home/team4/Untitled Folder 1', seconds=10)
test.classify_frames()



I apologize for the lack of comments, doc-strings, etc. It's the in the works. 

So you are gonna want to learn how to use the library multiprocessing. There are no doubt many ways to go about this as that library is extensive. Here is how I would go about it assuming you need to keep track of the different processes to put the image back together correctly. This is more complicated than just creating a process pool but allows you to keep tack of all processes. 

from multiprocessing import Process, Pipe

#keeps track of connections
conList = []
for i in range(numberOfProcessesNeeded):
   #creates connection
   recv, send = Pipe()

   #gets process ready with function to complete and arguments
   pid = Process(target=imageProcessingFunction, args=(i,command,send))

   #starts process
   pid.start()

   #keeps track of process
   conList.append([i,recv])


Now one of the arguments to that function must be 'send' from the pipe creation so that when the process ends it can send information back.

In the Image processing function that you want to use 'send.send(DataToReturn)' instead of 'return(DataToReturn)'. 

In the main function to collect the data do this.

data = []
for i in conList:
    #this will wait until each process has finished and then collect the data
    out = i[1].recv()
    data.append(out)


After this the array data will have all of your images in the order you called them to be processed.

Hopefully that wasn't to confusing but this is a tricky subject. More info can be found here:
https://docs.python.org/3.4/library/multiprocessing.html
