zjc
/
yolov5
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

autoshape() update for PIL greyscale inputs (#1279) 

* autoshape update for PIL greyscale inputs

* autoshape update for PIL greyscale inputs
This commit is contained in:
Glenn Jocher 2020-11-04 14:20:11 +01:00 committed by GitHub
parent 187f7c2ed1
commit 15a10609fe
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
models/common.py
View File

@ -1,6 +1,7 @@
# This file contains modules common to various models # This file contains modules common to various models
import math import math
import numpy as np import numpy as np
import torch import torch
import torch.nn as nn import torch.nn as nn
@ -144,7 +145,8 @@ class autoShape(nn.Module):
shape0, shape1 = [], [] # image and inference shapes shape0, shape1 = [], [] # image and inference shapes
batch = range(len(x)) # batch size batch = range(len(x)) # batch size
for i in batch: 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 s = x[i].shape[:2] # HWC
shape0.append(s) # image shape shape0.append(s) # image shape
g = (size / max(s)) # gain g = (size / max(s)) # gain