From 15a10609fe0e0f1fd3a03b11bf631a053cf001ca Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Wed, 4 Nov 2020 14:20:11 +0100 Subject: [PATCH] autoshape() update for PIL greyscale inputs (#1279) * autoshape update for PIL greyscale inputs * autoshape update for PIL greyscale inputs --- models/common.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/models/common.py b/models/common.py index 022ad00..b48ad48 100644 --- a/models/common.py +++ b/models/common.py @@ -1,6 +1,7 @@ # This file contains modules common to various models import math + import numpy as np import torch import torch.nn as nn @@ -144,7 +145,8 @@ class autoShape(nn.Module): shape0, shape1 = [], [] # image and inference shapes batch = range(len(x)) # batch size for i in batch: - x[i] = np.array(x[i])[:, :, :3] # up to 3 channels if png + x[i] = np.array(x[i]) # to numpy + x[i] = x[i][:, :, :3] if x[i].ndim == 3 else np.tile(x[i][:, :, None], 3) # enforce 3ch input s = x[i].shape[:2] # HWC shape0.append(s) # image shape g = (size / max(s)) # gain