Fix PyTorch Hub export inference shapes (#6949)
Browse filesMay resolve https://github.com/ultralytics/yolov5/issues/6947
- models/common.py +3 -4
models/common.py
CHANGED
@@ -544,10 +544,9 @@ class AutoShape(nn.Module):
|
|
544 |
g = (size / max(s)) # gain
|
545 |
shape1.append([y * g for y in s])
|
546 |
imgs[i] = im if im.data.contiguous else np.ascontiguousarray(im) # update
|
547 |
-
shape1 = [make_divisible(x, self.stride) for x in np.
|
548 |
-
x = [letterbox(im, new_shape=shape1
|
549 |
-
x = np.
|
550 |
-
x = np.ascontiguousarray(x.transpose((0, 3, 1, 2))) # BHWC to BCHW
|
551 |
x = torch.from_numpy(x).to(p.device).type_as(p) / 255 # uint8 to fp16/32
|
552 |
t.append(time_sync())
|
553 |
|
|
|
544 |
g = (size / max(s)) # gain
|
545 |
shape1.append([y * g for y in s])
|
546 |
imgs[i] = im if im.data.contiguous else np.ascontiguousarray(im) # update
|
547 |
+
shape1 = [make_divisible(x, self.stride) if self.pt else size for x in np.array(shape1).max(0)] # inf shape
|
548 |
+
x = [letterbox(im, new_shape=shape1, auto=False)[0] for im in imgs] # pad
|
549 |
+
x = np.ascontiguousarray(np.array(x).transpose((0, 3, 1, 2))) # stack and BHWC to BCHW
|
|
|
550 |
x = torch.from_numpy(x).to(p.device).type_as(p) / 255 # uint8 to fp16/32
|
551 |
t.append(time_sync())
|
552 |
|