gerador de mapas de diferença ft. multiprocessing

teste.py

#!/usr/bin/env python3
# borges 2021 - nenhum direito reservado
# escrito para o lendário lucas "janela" giannella

import numpy as np
import matplotlib.pyplot as plt
import skimage
from skimage import io, feature
import time
import os
from multiprocessing import Process, Queue

# fica lendo um queue, no qual serão inseridos pares (i, j)
# escreve no outro queue ternas (i, j, val)
class SubGerador(Process):
  def __init__(self, base, sample, margin, pos_in: Queue, vals_out: Queue):
    super().__init__()
    self.base = base
    self.sample = sample
    self.margin = margin
    self.pos_in = pos_in
    self.vals_out = vals_out

  def run(self):
    while (got := self.pos_in.get()) is not None:
      i, j = got
      template = self.base[i:i+self.margin,j:j+self.margin]
      corr = skimage.feature.match_template(self.sample, template)
      ix, jx = np.divmod(np.argmax(corr), corr.shape[0])
      val = np.absolute(i+self.margin-ix) + np.absolute(j+self.margin-jx)
      self.vals_out.put((i+self.margin, j+self.margin, val))
    self.vals_out.put(None)
    return 0


def gera_mapa(fn_base, fn_comp, nome_mapa, nome_sobrepos, step=1, nproc=None):
  base = skimage.io.imread(fn_base, 1)
  sample = skimage.io.imread(fn_comp, 1)
  tmp_size =30
  margin = tmp_size//2
  dispmap = np.zeros(base.shape)
  pos_queue, vals_queue = (Queue(), Queue())
  subprocessos = []
  nproc = nproc or max(os.cpu_count()-1, 1)
  print(f"Iniciando {nproc} subprocessos...")
  for _ in range(nproc):
    sp = SubGerador(base, sample, margin, pos_queue, vals_queue)
    subprocessos.append(sp)
    sp.start()
  print("Enchendo posições...")
  for i in range(0,base.shape[0]-margin,step):
    for j in range(0,base.shape[1]-margin,step):
      pos_queue.put((i, j))
  print("Enviando Nones...")
  for _ in range(nproc):
    pos_queue.put(None)
  print("Aguardando subprocessos...")
  deads = 0
  while deads != nproc:
    got = vals_queue.get()
    if got is None:
      deads += 1
    else:
      i, j, val = got
      dispmap[i, j] = val
  print("Montando figuras...")
  plt.figure()
  plt.imshow(dispmap)
  plt.axis('off')
  #plt.show()
  plt.imsave(nome_mapa, dispmap)
  plt.figure()
  plt.imshow(base, cmap="Greys_r",alpha=0.5)
  plt.imshow(sample, cmap="Greys_r",alpha=0.5)
  plt.imshow(dispmap,alpha=0.5)
  plt.axis('off')
  plt.savefig(nome_sobrepos, bbox_inches="tight")
  #plt.show()
  print("Pronto.")

def main(argc, argv):
  if argc < 5:
    print(
      "Especifique caminho do arquivo base, caminho do arquivo de comparação, "
      "nome do mapa, e nome da sobreposição."
    )
    return 1
  _, fn_base, fn_comp, nome_mapa, nome_sobrepos = argv
  try:
    gera_mapa(fn_base, fn_comp, nome_mapa, nome_sobrepos)
    return 0
  except:
    print("Deu big ruim:")
    raise

if __name__ == "__main__":
  from sys import argv
  exit(main(len(argv), argv))