まぐらぼ

日々の雑記を書いています。

jetbot train.py

import torchvision
import torch

model = torchvision.models.resnet18(pretrained=False)
model.fc = torch.nn.Linear(512, 2)

model.load_state_dict(torch.load('best_steering_model_xy.pth'))
device = torch.device('cuda')
model = model.to(device)
model = model.eval().half()


import torchvision.transforms as transforms
import torch.nn.functional as F
import cv2
import PIL.Image
import numpy as np

mean = torch.Tensor([0.485, 0.456, 0.406]).cuda().half()
std = torch.Tensor([0.229, 0.224, 0.225]).cuda().half()

def preprocess(image):
image = PIL.Image.fromarray(image)
image = transforms.functional.to_tensor(image).to(device).half()
image.sub_(mean[:, None, None]).div_(std[:, None, None])
return image[None, ...]



from IPython.display import display
import ipywidgets
import traitlets
from jetbot import Camera, bgr8_to_jpeg

camera = Camera()

image_widget = ipywidgets.Image()

traitlets.dlink*1
display(x_slider, steering_slider)

angle = 0.0
angle_last = 0.0

def execute(change):
global angle, angle_last
image = change['new']
xy = model(preprocess(image)).detach().float().cpu().numpy().flatten()
x = xy[0]
y = (0.5 - xy[1]) / 2.0

x_slider.value = x
y_slider.value = y

speed_slider.value = speed_gain_slider.value

angle = np.arctan2(x, y)
pid = angle * steering_gain_slider.value + (angle - angle_last) * steering_dgain_slider.value
angle_last = angle

steering_slider.value = pid + steering_bias_slider.value

robot.left_motor.value = max(min(speed_slider.value + steering_slider.value, 1.0), 0.0)
robot.right_motor.value = max(min(speed_slider.value - steering_slider.value, 1.0), 0.0)

execute({'new': camera.value})


camera.observe(execute, names='value')

*1:camera, 'value'), (image_widget, 'value'), transform=bgr8_to_jpeg) display(image_widget) from jetbot import Robot robot = Robot() speed_gain_slider = ipywidgets.FloatSlider(min=0.0, max=1.0, step=0.01, description='speed gain') steering_gain_slider = ipywidgets.FloatSlider(min=0.0, max=1.0, step=0.01, value=0.2, description='steering gain') steering_dgain_slider = ipywidgets.FloatSlider(min=0.0, max=0.5, step=0.001, value=0.0, description='steering kd') steering_bias_slider = ipywidgets.FloatSlider(min=-0.3, max=0.3, step=0.01, value=0.0, description='steering bias') display(speed_gain_slider, steering_gain_slider, steering_dgain_slider, steering_bias_slider) x_slider = ipywidgets.FloatSlider(min=-1.0, max=1.0, description='x') y_slider = ipywidgets.FloatSlider(min=0, max=1.0, orientation='vertical', description='y') steering_slider = ipywidgets.FloatSlider(min=-1.0, max=1.0, description='steering') speed_slider = ipywidgets.FloatSlider(min=0, max=1.0, orientation='vertical', description='speed') display(ipywidgets.HBox([y_slider, speed_slider]