add comma cameras

.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()

.gitignore

@@ -7,7 +7,6 @@ __pycache__/
 /models/
 /frames/
 /bus.jpg
-/zidane.jpg
 /result.jpg
 /boxes.jpg
 /masks.jpg

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,24 +56,24 @@ 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:
 
 ```sh
 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 |

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

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

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'