training_data_fsai.py

import sys, os
import math
import itertools
import json
import time

#from jinja2 import Template

from vehicle import Driver

sensorMax = 1000

robot = Driver()

refresh = 1

devices = [ {"camera": "zedcam train mid meta camera", "range": "zedcam train mid meta range finder"},
            {"camera": "zedcam train low meta camera", "range": "zedcam train low meta range finder"}, 
            {"camera": "zedcam train high meta camera", "range": "zedcam train high meta range finder"},
            {"camera": "zedcam train vhigh meta camera", "range": "zedcam train vhigh meta range finder"} ]

#devices = [ {"camera": "zedcam train mid meta camera", "range": "zedcam train mid meta range finder"} ]

# === cameras ===
for d in devices:
    d["camera"] = robot.getCamera( d["camera"] )
    d["camera"].enable( refresh )
    d["camera"].recognitionEnable( refresh )
    
    d["range"] = robot.getRangeFinder( d["range"])
    d["range"].enable( refresh )
    
# === inertial ===    
inertial = robot.getInertialUnit("inertial")    
inertial.enable(refresh)
    
targets = ( ( 0.0,  0.0,  0.0), \
            ( 2.5,  0.0,  6.3), \
            ( 8.5,  0.0,  8.5), \
            (14.5,  0.0,  6.3), \
            (17.3,  0.0,  0.0), \
            (14.5,  0.0, -6.3), \
            ( 8.5,  0.0, -8.5), \
            ( 2.5,  0.0, -6.3), \
            ( 0.0,  0.0,  0.0), \
            (- 2.5,  0.0,  6.3), \
            (- 8.5,  0.0,  8.5), \
            (-14.5,  0.0,  6.3), \
            (-17.3,  0.0,  0.0), \
            (-14.5,  0.0, -6.3), \
            (- 8.5,  0.0, -8.5), \
            (- 2.5,  0.0, -6.3), \
             )
targetId = 0

folder = "trainingdata"
maxImages = 50

def local_point( point, origin, heading ):
    relative = [ p-o for p, o in zip(point,origin) ]
    return relative[0]*math.cos(-heading) + relative[2]*math.sin(-heading), \
           relative[1], \
           relative[2]*math.cos(-heading) - relative[0]*math.sin(-heading)


def save_annotated_image( camera, rangeFinder, position, heading ):
    filename = str(int(time.time()*1000))
    camera.saveImage( os.path.join( folder, "images", filename+".png" ), 100 )
    rangeFinder.saveImage( os.path.join( folder, "depth", filename+".jpg" ), 100 )

    positions = ( ( *object.get_position_on_image(), *object.get_size_on_image() )
                  for object in camera.getRecognitionObjects()
                  if "cone" in object.get_model().decode() )
    filtered = ( (x, y, w, h) for x, y, w, h in positions
                  if x > 15 and x < rangeFinder.getWidth()-15 )
                  

    
    width = rangeFinder.getWidth()
    height = rangeFinder.getHeight()
    scaled = ( (x/width, y/height, w/width, h/height) for x, y, w, h in filtered )
    
    with open( os.path.join( folder, "labels", filename+".txt" ), "w" ) as f:
        for object in scaled:
            print( "0 {} {} {} {}".format( *object ), file=f )
    
targetId = 0
deviceId = 0
while robot.step() != -1:     
    # get next device
    device = devices[ deviceId % len(devices) ]
    deviceId += 1
        
    # get current position
    position = robot.getSelf().getPosition()
    heading = inertial.getRollPitchYaw()[2]
    if heading == float('nan'): continue

    # get next target
    try: target = targets[targetId]
    except IndexError: break

    # convert to local coordinate system and calc required heading    
    localTarget = local_point( target, position, heading )
   
    print( "target> {:.3f} {:.3f} {:.3f}".format(*target) )
       
    print( "local > {:.3f} {:.3f} {:.3f}".format(*localTarget) )
 
    localHeading = math.atan2( localTarget[0], localTarget[2] )
    
    while localHeading > math.pi: localHeading -= math.pi*2
    while localHeading < -math.pi: localHeading += math.pi*2
    
    print( "head  > {:.3f}".format(math.degrees(localHeading)) )
    
    localHeading = min( math.radians(30), max( math.radians(-30), localHeading ) )  
    distance = math.sqrt( abs(localTarget[0]**2) + abs(localTarget[2]**2) )

    print( "dist  > {:.3f}".format(distance) )

    # move to next target if in range    
    if distance < 1:
        targetId += 1
       
    robot.setSteeringAngle( -localHeading )
    robot.setCruisingSpeed( 12 )

    #save_annotated_image( device["camera"], device["range"], position, heading )
    
        
sys.exit()