comma body does a loop around the office
#!/usr/bin/env python3
import os, sys, time, json
from pathlib import Path
import numpy as np
np.set_printoptions(suppress=True)

# openpilot imports
import cereal.messaging as messaging
from openpilot.common.realtime import Ratekeeper
from openpilot.system.hardware import PC

# device gets images not through the compressor
if not PC:
  os.environ["IMAGE"] = "2"
  os.environ["FLOAT16"] = "1"
  import cv2
  from cereal.visionipc import VisionIpcClient, VisionStreamType

# tinygrad!
import tinygrad
from tinygrad.jit import TinyJit
from tinygrad.tensor import Tensor
from tinygrad.nn.state import load_state_dict, safe_load
from tinygrad.helpers import Timing, dtypes, getenv

# repo local imports
from video import live_decode_frames
from train import TinyNet

# this is the yolo foundation model
def get_foundation():
  # add tinygrad and tinygrad examples to python path
  if PC:
    sys.path.append(str(Path(tinygrad.__path__[0]).parent))
    sys.path.append(str(Path(tinygrad.__path__[0]).parent / "examples"))
  else:
    sys.path.append("/data/openpilot/tinygrad_repo/examples")
    sys.path.append("/data/openpilot/tinygrad_repo")

  from yolov8 import YOLOv8, get_variant_multiples
  yolo_variant = "n"
  depth, width, ratio = get_variant_multiples(yolo_variant)
  yolo_infer = YOLOv8(w=width, r=ratio, d=depth, num_classes=80)

  from extra.utils import download_file
  from tinygrad.nn.state import safe_load, load_state_dict
  weights_location = Path("/tmp") / f'yolov8{yolo_variant}.safetensors'
  download_file(f'https://gitlab.com/r3sist/yolov8_weights/-/raw/master/yolov8{yolo_variant}.safetensors', weights_location)
  state_dict = safe_load(weights_location)
  load_state_dict(yolo_infer, state_dict)
  def foundation(imgs:Tensor):
    x = yolo_infer.net(imgs.permute(0,3,1,2).float() / 255)
    x = yolo_infer.fpn(*x)
    return x[2]
  return foundation

def get_tinynet():
  net = TinyNet()
  load_state_dict(net, safe_load("tinynet.safetensors"))
  return net

def get_pred():
  # net runner
  Tensor.no_grad = True
  Tensor.training = False
  foundation = get_foundation()
  net = get_tinynet()
  @TinyJit
  def pred(x):
    return net(foundation(x)).exp()[0].realize()

  # warm up the net runner
  with Timing("building 1: "): pred(Tensor.rand(1,480,640,3, dtype=dtypes.uint8))
  with Timing("building 2: "): pred(Tensor.rand(1,480,640,3, dtype=dtypes.uint8))
  with Timing("testing rt: "): pred(Tensor.rand(1,480,640,3, dtype=dtypes.uint8)).numpy()
  return pred

if __name__ == "__main__":
  # controlling the body
  # REQUIRED on PC: run bridge on device back to YOUR ip
  # /data/openpilot/cereal/messaging/bridge 192.168.60.251 customReservedRawData1
  pm = messaging.PubMaster(['customReservedRawData1'])
  def control(x, y):
    dat = messaging.new_message()
    dat.customReservedRawData1 = json.dumps({'x': x, 'y': y}).encode()
    pm.send('customReservedRawData1', dat)
  print("pm connected")

  # getting frames
  if PC:
    # TODO: replace with DEVICE ip
    dcamData = messaging.sub_sock("driverEncodeData", None, addr="192.168.63.69", conflate=False)
  else:
    vipc_client = VisionIpcClient("camerad", VisionStreamType.VISION_STREAM_DRIVER, True)
    while not vipc_client.connect(False): time.sleep(0.1)
  print("vipc connected")

  # get the net runner
  pred = get_pred()

  if PC:
    seen_iframe = False
    from openpilot.common.window import Window
    win = Window(640, 480)

  # run at max 5hz
  rk = Ratekeeper(5)
  while 1:
    # get frame
    if PC:
      frm = messaging.drain_sock(dcamData, wait_for_one=True)
      if not seen_iframe: frm = frm[-20:]
      frms, seen_iframe = live_decode_frames(frm, seen_iframe)
      if not frms: continue
      img = frms[-1]
    else:
      yuv_img_raw = vipc_client.recv()
      if yuv_img_raw is None or not yuv_img_raw.data.any(): continue
      imgff = yuv_img_raw.data.reshape(-1, vipc_client.stride)
      imgff = imgff[8:8+vipc_client.height * 3 // 2, :vipc_client.width]
      img = cv2.cvtColor(imgff, cv2.COLOR_YUV2RGB_NV12)
      img = cv2.resize(img, (640, 480))

    # run the model
    probs = pred(Tensor(img).reshape(1, 480, 640, 3)).numpy()

    # control policy, turn harder if we are more confident
    if probs[0] > 0.99: choice = 0
    elif probs[2] > 0.99: choice = 4
    elif probs[1] > 0.99: choice = 2
    elif probs[0] > probs[2]: choice = 1
    else: choice = 3

    # go mostly straight, go faster if confident, maybe even go reverse
    x = -0.5 - (probs[1]*0.5) + (probs[3]*1.5)
    # minus y is right, plus y is left
    y = [-0.7, -0.2, 0, 0.2, 0.7][choice]
    if not getenv("NODRIVE"): control(x, y)
    print(f"{x:5.2f} {y:5.2f}", probs)

    # draw frame and run at 5hz
    if PC: win.draw(frms[-1])
    rk.keep_time()

control.py
geohot/body_loop
#!/usr/bin/env python3
import os
from tinygrad.tensor import Tensor
from tinygrad.nn.state import safe_save
from control import get_foundation
from video import frames_from_file
from tinygrad.tensor import Tensor
Tensor.training = False
Tensor.no_grad = True
import numpy as np
from tqdm import tqdm

def preprocess(foundation, fn, fn_out, flip=False):
  if not fn.endswith(".hevc"): return
  if 'online' in fn: return
  print("preprocessing", fn, fn_out)
  if os.path.isfile(fn_out): return
  frms = []
  xs = []
  for frm in tqdm(frames_from_file(fn)):
    frms.append(frm)
    if len(frms) == 256:
      data = Tensor(np.concatenate([x.reshape(1,480,640,3) for x in frms]))
      if flip: data = data[:, :, ::-1]
      xs.append(foundation(data).realize())
      frms = []
  if frms:
    data = Tensor(np.concatenate([x.reshape(1,480,640,3) for x in frms]))
    if flip: data = data[:, :, ::-1]
    xs.append(foundation(data).realize())
  big_x = Tensor.cat(*xs, dim=0)
  print(big_x.shape)
  return safe_save({"x": big_x}, fn_out)

if __name__ == "__main__":
  foundation = get_foundation()
  for fn in os.listdir("data"):
    preprocess(foundation, "data/"+fn, "data/"+fn.replace(".hevc", ".safetensors"))
    preprocess(foundation, "data/"+fn, "data/"+fn.replace(".hevc", "_flip.safetensors"), True)


preprocess.py
geohot/body_loop
#!/usr/bin/env python3
import sys
from control import get_pred
from video import frames_from_file
from tinygrad.jit import TinyJit
from tinygrad.tensor import Tensor
import numpy as np

if __name__ == "__main__":
  test_data = "data/straight_2_dcamera.hevc" if len(sys.argv) == 1 else sys.argv[1]
  pred = get_pred()

  all_probs = []
  for i,frm in enumerate(frames_from_file(test_data)):
    probs = pred(Tensor(frm.reshape(1, 480, 640, 3))).numpy().copy()
    print(i, probs)
    all_probs.append(probs)

  import matplotlib.pyplot as plt
  plt.figure(figsize=(16,16))
  plt.plot(all_probs)
  plt.legend(["left", "straight", "right", "reverse"])
  plt.show()

test_online.py
geohot/body_loop
#!/usr/bin/env python3
import random
from tinygrad.nn.state import safe_load
from tinygrad.tensor import Tensor
from tinygrad.nn import Conv2d, Linear
from tinygrad.nn.state import get_parameters, get_state_dict, safe_save
from tinygrad.nn.optim import AdamW
from tqdm import trange
from tinygrad.jit import TinyJit
import numpy as np

train_set = [
  ("left_1_dcamera", 0),
  ("left_3_dcamera", 0),
  ("right_1_dcamera", 2),
  ("right_3_dcamera", 2),
  ("straight_1_dcamera", 1),
  ("ljust_1_dcamera", 0),
  ("rjust_1_dcamera", 2),
  ("left_day_1_dcamera", 0),
  ("right_day_1_dcamera", 2),
  ("left_day_2_dcamera", 0),
  ("right_day_2_dcamera", 2),
  ("left_day_3_dcamera", 0),
  ("right_day_3_dcamera", 2),
  ("left_day_4_dcamera", 0),
  ("right_day_4_dcamera", 2),
  ("extra_left_1_dcamera", 0),
  ("extra_right_1_dcamera", 2),
  ("extra_left_2_dcamera", 0),
  ("extra_right_2_dcamera", 2),
  ("extra_straight_2_dcamera", 1),
  ("extra_left_3_dcamera", 0),
  ("extra_right_3_dcamera", 2),
  ("messup_left_1_dcamera", 0),
  ("messup_left_2_dcamera", 0),
  ("messup_left_3_dcamera", 0),
  ("messup_left_4_dcamera", 0),
  ("reverse_1_dcamera", 3),
  ("reverse_2_dcamera", 3),
  ("reverse_3_dcamera", 3),
  ("messup_right_1_dcamera", 2),
  ("messup_right_2_dcamera", 2),
  ("messup_right_3_dcamera", 2),
  ("messup_right_4_dcamera", 2),
  ("messup_right_5_dcamera", 2),
  ("messup_right_6_dcamera", 2),
  ("messup_right_7_dcamera", 2),
  ("messup_right_8_dcamera", 2),
  ("messup_right_9_dcamera", 2),
  ("messup_reverse_1_dcamera", 3),
  ("messup_reverse_2_dcamera", 3),
  ("messup_reverse_3_dcamera", 3),
  ("messup_straight_1_dcamera", 1),
]

test_set = [
  ("left_2_dcamera", 0),
  ("straight_2_dcamera", 1),
  ("right_2_dcamera", 2),
]

# this is the net used after the yolo foundation model
class TinyNet:
  def __init__(self):
    self.c1 = Conv2d(256,64,3)
    self.c2 = Conv2d(64,8,3)
    self.l = Linear(1408,4)
  def __call__(self, x):
    x = self.c1(x).gelu()
    x = self.c2(x).gelu().dropout(0.85)
    x = x.reshape(x.shape[0], -1)
    x = self.l(x)
    return x.log_softmax()

# TODO: fix dropout in the JIT
#@TinyJit
def train_step(x,y):
  Tensor.training = True
  optim.lr *= 0.999
  out = net(x)
  loss = out.sparse_categorical_crossentropy(y)
  optim.zero_grad()
  loss.backward()
  optim.step()
  cat = out.argmax(axis=-1)
  accuracy = (cat == y).mean()
  return loss.realize(), accuracy.realize()

@TinyJit
def test_step(tx,ty):
  Tensor.training = False
  out = net(tx)
  loss = out.sparse_categorical_crossentropy(ty)
  cat = out.argmax(axis=-1)
  return loss.realize(), (cat == ty).mean().realize()

def get_minibatch(sets, bs=32):
  xs, ys = [], []
  for _ in range(bs):
    src, val = random.choice(sets)
    sel = random.randint(0, src.shape[0]-1)
    xs.append(src[sel:sel+1])
    ys.append(val)
  return Tensor(np.concatenate(xs, axis=0)), Tensor(np.array(ys))

# add the flips to the training set with the inverse option
def get_flips(x):
  ret = []
  for t,y in x:
    if y == 2: y = 0
    elif y == 0: y = 2
    ret.append((t+"_flip", y))
  return ret

if __name__ == "__main__":
  train_set += get_flips(train_set)
  test_set += get_flips(test_set)

  train_sets = [(safe_load(f"data/{fn}.safetensors")["x"].numpy(),y) for fn,y in train_set]
  test_sets = [(safe_load(f"data/{fn}.safetensors")["x"].numpy(),y) for fn,y in test_set]

  # get test set
  tx,ty = get_minibatch(test_sets, 1024)

  Tensor.no_grad = False
  Tensor.training = True
  net = TinyNet()
  optim = AdamW(get_parameters(net))

  acc, tacc, losses, tlosses = [], [], [], []
  for i in (t:=trange(600)):
    if i%10 == 0: test_loss, test_accuracy = test_step(tx, ty)
    x,y = get_minibatch(train_sets, 64)
    loss, accuracy = train_step(x, y)
    losses.append(float(loss.numpy()))
    tlosses.append(float(test_loss.numpy()))
    acc.append(float(accuracy.numpy()))
    tacc.append(float(test_accuracy.numpy()))
    t.set_description("loss %.2f accuracy %.3f test_accuracy %.3f test_loss %.4f lr: %f" % (loss.numpy(), accuracy.numpy(), test_accuracy.numpy(), test_loss.numpy(), optim.lr.numpy()))

  safe_save(get_state_dict(net), "tinynet.safetensors")

  import matplotlib.pyplot as plt
  plt.plot(losses)
  plt.plot(tlosses)
  plt.plot(acc)
  plt.plot(tacc)
  plt.show()

train.py
geohot/body_loop
import av
fmt = av.video.format.VideoFormat('rgb24')

# pass in msgs to decode them to frames
codec = None
def live_decode_frames(msgs, seen_iframe=False):
  global codec
  if not seen_iframe: codec = av.CodecContext.create("hevc", "r")
  imgs = []
  V4L2_BUF_FLAG_KEYFRAME = 8
  for evt in msgs:
    evta = getattr(evt, evt.which())
    if not seen_iframe and not (evta.idx.flags & V4L2_BUF_FLAG_KEYFRAME): continue
    if not seen_iframe:
      codec.decode(av.packet.Packet(evta.header))
      seen_iframe = True
    frames = codec.decode(av.packet.Packet(evta.data))
    if not len(frames): continue
    imgs.append(frames[0].to_ndarray(format=fmt))
  return imgs, seen_iframe

# get frames from dcamera files
def frames_from_file(fn):
  container = av.open(fn)
  for frame in container.decode(video=0): yield frame.to_ndarray(format=fmt)

video.py
geohot/body_loop