xdobro23/YOLOv8-16bit
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

changed data/loaders and engine/predictor so they accept 1 channel images for predictions

Browse Source
This commit is contained in:
Clea Parcerisas
2023-08-24 14:42:29 +02:00
parent 95e2c2424a
commit 82b97d5d4c
2 changed files with 3 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
ultralytics/engine/predictor.py
View File

@ -122,7 +122,8 @@ class BasePredictor:
not_tensor = not isinstance(im, torch.Tensor)
if not_tensor:
im = np.stack(self.pre_transform(im))
im = im[..., ::-1].transpose((0, 3, 1, 2)) # BGR to RGB, BHWC to BCHW, (n, 3, h, w)
im = np.expand_dims(im, -1)
im = im[..., ::].transpose((0, 3, 1, 2)) # BGR to RGB, BHWC to BCHW, (n, 3, h, w)
im = np.ascontiguousarray(im) # contiguous
im = torch.from_numpy(im)