まぐらぼ

Unity/Android、Microsoft系のWPFをやってます。

jetbot livedemo

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

import math
from IPython.display import display
import traitlets
#from jetbot import Camera, bgr8_to_jpeg
#from jetbot import Robot

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

angle = 0.0
angle_last = 0.0

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, ...]


def execute(image):
global angle, angle_last
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

# この値からPID+BIASを足し引きしてモーター値とする
speed_slider_value = 0.2#speed_gain_slider_value

steering_gain_slider_value = 0.2 # 左右の係数
steering_dgain_slider_value = 0.1 # 加速度にかける係数
steering_bias_slider_value = 0 # 最後にタス

try :
angle = np.arctan2(x, y)
deg = math.degrees(angle)

pid = angle * steering_gain_slider_value \
+ (angle - angle_last)* steering_dgain_slider_value

angle_last = angle

steering_slider_value = pid + steering_bias_slider_value

left_value = max(min(speed_slider_value + steering_slider_value, 1.0), 0.0)
robot.left_motor.value = left_value
robot.right_motor.value = max(min(speed_slider_value - steering_slider_value, 1.0), 0.0)
except Exception as e:
print(e)

def main():

#image = PIL.Image.open('dataset_xy/xy_061_054_58.jpg')
image = PIL.Image.open('dataset_xy/xy_054_068.jpg')

execute(image)

if __name__ == "__main__":
main()