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
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