training_data_haar.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

names = [ "zedcam train mid left camera",
            "zedcam train low left camera",
            "zedcam train high left camera",
            "zedcam train vhigh left camera" ]
devices = []

for i in range(robot.getNumberOfDevices()):
    device = robot.getDeviceByIndex(i)
    print( i, device.getName(), type(device) )
  
# === cameras ===
for name in names:
    devices.append( robot.getCamera( name ) )
    devices[-1].enable( refresh )
    devices[-1].recognitionEnable( 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 ):
    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 > 50 and x < camera.getWidth()-50 )
                  
    #camera.saveImage( "temp.png", 100 )
    

    #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 x, y, w, h in positions:
            print( f"{x} {y} {w} {h}", file=f )
            
            #print( "{} {} {} {}".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 )
    
        
sys.exit()