github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-09 15:08:02 -05:00
Code Issues Packages Projects Releases Wiki Activity

fix yolov8 example (#5003)

Browse Source 
it was creating Tensor from a list of numpy arrays, which is not supported after moving creating from a list not using numpy.
This commit is contained in:
chenyu
2024-06-16 20:47:29 -04:00
committed by GitHub
parent 8bc0cbf67b
commit c52352bd9a
1 changed files with 5 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
examples/yolov8.py
View File

@@ -13,7 +13,7 @@ from tinygrad.nn.state import safe_load, load_state_dict
#The upsampling class has been taken from this pull request https://github.com/tinygrad/tinygrad/pull/784 by dc-dc-dc. Now 2(?) models use upsampling. (retinet and this)
#Pre processing image functions.
def compute_transform(image, new_shape=(640, 640), auto=False, scaleFill=False, scaleup=True, stride=32):
def compute_transform(image, new_shape=(640, 640), auto=False, scaleFill=False, scaleup=True, stride=32) -> Tensor:
shape = image.shape[:2] # current shape [height, width]
new_shape = (new_shape, new_shape) if isinstance(new_shape, int) else new_shape
r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
@@ -28,15 +28,15 @@ def compute_transform(image, new_shape=(640, 640), auto=False, scaleFill=False,
top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
image = cv2.copyMakeBorder(image, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(114, 114, 114))
return image
return Tensor(image)
def preprocess(im, imgsz=640, model_stride=32, model_pt=True):
same_shapes = all(x.shape == im[0].shape for x in im)
auto = same_shapes and model_pt
im = Tensor([compute_transform(x, new_shape=imgsz, auto=auto, stride=model_stride) for x in im])
im = Tensor.stack(*im) if im.shape[0] > 1 else im
im = [compute_transform(x, new_shape=imgsz, auto=auto, stride=model_stride) for x in im]
im = Tensor.stack(*im) if len(im) > 1 else im[0].unsqueeze(0)
im = im[..., ::-1].permute(0, 3, 1, 2) # BGR to RGB, BHWC to BCHW, (n, 3, h, w)
im /= 255 # 0 - 255 to 0.0 - 1.0
im = im / 255.0 # 0 - 255 to 0.0 - 1.0
return im
# Post Processing functions