diff --git a/modules/async_worker.py b/modules/async_worker.py index 908cc8c..fd785f0 100644 --- a/modules/async_worker.py +++ b/modules/async_worker.py @@ -22,6 +22,7 @@ def worker(): import traceback import math import numpy as np + import cv2 import torch import time import shared @@ -79,16 +80,20 @@ def worker(): return def build_image_wall(async_task): - results = async_task.results + results = [] - if len(results) < 2: + if len(async_task.results) < 2: return - for img in results: + for img in async_task.results: + if isinstance(img, str) and os.path.exists(img): + img = cv2.imread(img) + img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) if not isinstance(img, np.ndarray): return if img.ndim != 3: return + results.append(img) H, W, C = results[0].shape