You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

worker.py 1.7KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556
  1. """ Wokers definition """
  2. # TODO Implement this with a queue and mutliprocessingt
  3. import threading
  4. from config import Config as Conf
  5. from transform.gan.mask import CorrectToMask, MaskrefToMaskdet, MaskfinToNude
  6. from transform.opencv.correct import DressToCorrect, ColorTransfer
  7. from transform.opencv.mask import MaskToMaskref, MaskdetToMaskfin
  8. from transform.opencv.resize import ImageToResized, ImageToCrop, ImageToOverlay, ImageToResizedCrop, ImageToRescale
  9. workers = {
  10. "gan": {
  11. CorrectToMask: [],
  12. MaskrefToMaskdet: [],
  13. MaskfinToNude: [],
  14. "sem": threading.Semaphore(1)
  15. },
  16. "opencv": {
  17. DressToCorrect: [],
  18. MaskToMaskref: [],
  19. ImageToResized: [],
  20. ImageToCrop: [],
  21. ImageToOverlay: [],
  22. ImageToResizedCrop: [],
  23. ImageToRescale: [],
  24. ColorTransfer: [],
  25. MaskdetToMaskfin: [],
  26. "sem": threading.Semaphore(Conf.args['n_cores'])
  27. }
  28. }
  29. select_sem = threading.Semaphore(1)
  30. def run_worker(klass, image_step, config=None):
  31. r = None
  32. for k in ("gan", "opencv"):
  33. if workers.get(k).get(klass) is not None:
  34. Conf.log.debug("wk {}".format(workers.get(k).get(klass)))
  35. workers.get(k).get("sem").acquire()
  36. select_sem.acquire()
  37. if len(workers.get(k).get(klass)) == 0:
  38. w = klass()
  39. else:
  40. w = workers.get(k).get(klass).pop(0)
  41. select_sem.release()
  42. r = w.run(*[image_step[i] for i in w.input_index], config=config)
  43. select_sem.acquire()
  44. workers.get(k).get(klass).append(w)
  45. select_sem.release()
  46. workers.get(k).get("sem").release()
  47. return r