diff --git a/.github/test_examples.py b/.github/test_examples.py index 9fd4fe1..6544494 100644 --- a/.github/test_examples.py +++ b/.github/test_examples.py @@ -40,4 +40,4 @@ def run(example, *args, cwd=ROOT): saved = cv2.imread(str(frame)) assert saved is not None and saved.shape == image.shape -assert 'Samoyed:' in run('04_classify.py') +assert 'bow tie:' in run('04_classify.py').lower() diff --git a/.gitignore b/.gitignore index 019644b..2191aa4 100644 --- a/.gitignore +++ b/.gitignore @@ -7,7 +7,6 @@ __pycache__/ /models/ /frames/ /bus.jpg -/zidane.jpg /result.jpg /boxes.jpg /masks.jpg diff --git a/README.md b/README.md index cf730b9..ff48c78 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ # chestnut -Run chat, object detection, segmentation, and image classification with tinygrad on your PC or Chestnut. +Run chat, object detection, segmentation, and image classification with tinygrad on your PC or chestnut. ## Setup @@ -42,10 +42,12 @@ python examples/02_vision.py zidane.jpg --model segment --output masks.jpg Run YOLO26 detection or segmentation on a webcam. Use `--frames 100` for more. `--source` also accepts a video path or stream URL. +Add `--preview` for a live window that runs until you press Esc. ```sh python examples/03_camera.py --source 0 python examples/03_camera.py --source 0 --model segment +python examples/03_camera.py --source 0 --preview ``` ### Image classification @@ -54,14 +56,14 @@ Print the five most likely labels with ResNet18. ```sh python examples/04_classify.py -python examples/04_classify.py photo.jpg +python examples/04_classify.py zidane.jpg ``` -## Run on Chestnut +## Run on chestnut -Plug in the 12V power and connect the USB3 cable from Chestnut's USB3.2 port to your PC or comma. +Plug in the 12V power and connect the USB3 cable from chestnut's USB3.2 port to your PC or comma. -![Chestnut connections](chestnut.png) +![chestnut connections](chestnut.jpg) In the activated environment, check the connection: @@ -69,9 +71,9 @@ In the activated environment, check the connection: python tools/usb.py ``` -Expected: `Chestnut GPU check passed.` +Expected: `chestnut GPU check passed.` -Prefix any example command with `DEV=USB+AMD:LLVM` to run it on Chestnut's GPU: +Prefix any example command with `DEV=USB+AMD:LLVM` to run it on chestnut's GPU: ```sh DEV=USB+AMD:LLVM python examples/01_chat.py @@ -82,21 +84,26 @@ DEV=USB+AMD:LLVM python examples/03_camera.py --source 0 ## Comma camera On a comma device, run setup and activate the environment as above. Requires openpilot at `/data/openpilot`. -`--source comma` selects its camera. Add `--model segment` for segmentation. +Select the road, driver, or wide road camera with `--source comma:road`, `comma:driver`, or `comma:wide`. +Add `--model segment` to any camera command for segmentation. ```sh # Comma CPU -python examples/03_camera.py --source comma - -# Chestnut GPU connected to comma -DEV=USB+AMD:LLVM python examples/03_camera.py --source comma +python examples/03_camera.py --source comma:road +python examples/03_camera.py --source comma:driver +python examples/03_camera.py --source comma:wide + +# chestnut GPU connected to comma +DEV=USB+AMD:LLVM python examples/03_camera.py --source comma:road +DEV=USB+AMD:LLVM python examples/03_camera.py --source comma:driver +DEV=USB+AMD:LLVM python examples/03_camera.py --source comma:wide ``` ## Performance PC CPU: Threadripper PRO 5945WX. Comma CPU: Qualcomm SDM845. -| Model | PC CPU | Comma CPU | Chestnut GPU | +| Model | PC CPU | Comma CPU | chestnut GPU | | --- | ---: | ---: | ---: | | YOLO26n | 526.07 ms | 1846.42 ms | 7.02 ms | | YOLO26n-seg | 704.56 ms | 2349.27 ms | 7.98 ms | diff --git a/chestnut.jpg b/chestnut.jpg new file mode 100644 index 0000000..d28afb7 Binary files /dev/null and b/chestnut.jpg differ diff --git a/chestnut.png b/chestnut.png deleted file mode 100644 index f6b7902..0000000 Binary files a/chestnut.png and /dev/null differ diff --git a/examples/03_camera.py b/examples/03_camera.py index 47c77e4..478ec02 100644 --- a/examples/03_camera.py +++ b/examples/03_camera.py @@ -7,7 +7,7 @@ from vision import Vision -def comma_frames(): +def comma_frames(stream): checkout = Path('/data/openpilot') sys.path.append(str(checkout)) from msgq.visionipc import VisionIpcClient @@ -16,7 +16,7 @@ def comma_frames(): process = subprocess.Popen([str(checkout / 'openpilot/system/camerad/camerad')], cwd=checkout, stdout=subprocess.DEVNULL) try: - client = VisionIpcClient('camerad', 0, True) + client = VisionIpcClient('camerad', stream, True) client.connect(True) while True: buf = client.recv() @@ -45,19 +45,26 @@ def video_frames(source): if __name__ == '__main__': parser = argparse.ArgumentParser(description='YOLO26 on a webcam, video, or comma camera.') - parser.add_argument('--source', default='0', help='Webcam number, video/URL, or comma') + parser.add_argument('--source', default='0', help='Webcam, video/URL, or comma:road, comma:driver, comma:wide') parser.add_argument('--model', choices=['yolo', 'segment'], default='yolo') - parser.add_argument('--frames', type=int, default=10) + parser.add_argument('--frames', type=int, default=10, help='Frames to save without preview') + parser.add_argument('--preview', action='store_true', help='Show a live preview window') args = parser.parse_args() model = Vision(args.model) Path('frames').mkdir(exist_ok=True) - stream = comma_frames() if args.source == 'comma' else video_frames(args.source) + comma_streams = {'comma': 0, 'comma:road': 0, 'comma:wide': 1, 'comma:driver': 2} + stream = comma_frames(comma_streams[args.source]) if args.source in comma_streams else video_frames(args.source) print('Saving to frames/.', flush=True) try: for i, frame in enumerate(stream): - model(frame).save(f'frames/{i:05d}.jpg') + result = model(frame) + result.save(f'frames/{i:05d}.jpg') + if args.preview: + cv2.imshow('chestnut', result.plot()) + if cv2.waitKey(1) == 27: break print(f'Frame {i + 1}', flush=True) - if i + 1 >= args.frames: break + if not args.preview and i + 1 >= args.frames: break finally: stream.close() + if args.preview: cv2.destroyAllWindows() diff --git a/examples/04_classify.py b/examples/04_classify.py index c28fb22..2aed187 100644 --- a/examples/04_classify.py +++ b/examples/04_classify.py @@ -9,13 +9,12 @@ from PIL import Image from torchvision.models import ResNet18_Weights from tinygrad import Tensor -from tinygrad.helpers import fetch from tinygrad.nn.onnx import OnnxRunner parser = argparse.ArgumentParser(description='Classify an image with ResNet18.') parser.add_argument('image', nargs='?') args = parser.parse_args() -image = Image.open(args.image or fetch('https://raw.githubusercontent.com/pytorch/hub/master/images/dog.jpg')).convert('RGB') +image = Image.open(args.image or ROOT / 'zidane.jpg').convert('RGB') weights = ResNet18_Weights.DEFAULT inputs = weights.transforms()(image).unsqueeze(0).numpy() model = OnnxRunner(str(ROOT / 'models/resnet18.onnx')) diff --git a/setup.sh b/setup.sh index 458c533..4e91e46 100755 --- a/setup.sh +++ b/setup.sh @@ -82,9 +82,4 @@ fi uv sync --locked --python 3.12 .venv/bin/python tools/setup.py .venv/bin/python tools/export.py -if [ ! -f zidane.jpg ]; then - download https://ultralytics.com/images/zidane.jpg .cache/zidane.jpg - mv .cache/zidane.jpg zidane.jpg -fi - echo 'Ready. Activate with: source .venv/bin/activate' diff --git a/zidane.jpg b/zidane.jpg new file mode 100644 index 0000000..eeab1cd Binary files /dev/null and b/zidane.jpg differ