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