mirror of
https://github.com/CopterExpress/clever-show.git
synced 2026-05-26 07:07:58 +00:00
Tools: add change_landing script
This commit is contained in:
Arthur Golubtsov
250
tools/change_landing.py
Normal file
250
tools/change_landing.py
Normal file
@@ -0,0 +1,250 @@
|
||||
import argparse
|
||||
import os
|
||||
import csv
|
||||
import glob
|
||||
import copy
|
||||
import math
|
||||
import numpy as np
|
||||
import logging
|
||||
|
||||
def generate_line(pos1, pos2, speed, frame_delay = 0.1, start_frame = 0, color = [0, 0, 0], yaw = 0):
|
||||
dist = np.linalg.norm(pos1-pos2)
|
||||
delta = pos2 - pos1
|
||||
frames_count = math.ceil(dist/(speed*frame_delay))
|
||||
k = 1./frames_count
|
||||
frames = []
|
||||
for i in range(int(frames_count)+1):
|
||||
frame_pos = pos1 + k*i*delta
|
||||
frames.append({
|
||||
'number': start_frame + i,
|
||||
'x': frame_pos[0],
|
||||
'y': frame_pos[1],
|
||||
'z': frame_pos[2],
|
||||
'yaw': yaw,
|
||||
'red': color[0],
|
||||
'green': color[1],
|
||||
'blue': color[2],
|
||||
})
|
||||
return frames
|
||||
|
||||
def generate_positions(start_pos, nx, ny, dx, dy):
|
||||
len_x = (nx - 1)*dx
|
||||
len_y = (ny - 1)*dy
|
||||
x_center = start_pos[0]
|
||||
y_center = start_pos[1]
|
||||
z = start_pos[2]
|
||||
x0 = x_center - len_x/2.
|
||||
y0 = y_center + len_y/2.
|
||||
positions = []
|
||||
for iy in range(ny):
|
||||
for ix in range(nx):
|
||||
positions.append(np.array([x0+dx*ix, y0-dy*iy, z]))
|
||||
return positions
|
||||
|
||||
def parse_positions_file(filename):
|
||||
names = []
|
||||
pos = []
|
||||
try:
|
||||
pos_file = open(filename)
|
||||
except IOError:
|
||||
logging.error("File {} can't be opened".format(filename))
|
||||
else:
|
||||
with pos_file:
|
||||
n_str = pos_file.readline()[:-1].split(' ')
|
||||
nx = int(n_str[0])
|
||||
ny = int(n_str[1])
|
||||
dx = int(n_str[2])
|
||||
dy = int(n_str[3])
|
||||
pos_str = pos_file.readline()[:-1].split(' ')
|
||||
for i in range(3):
|
||||
pos.append(float(pos_str[i]))
|
||||
for lines in pos_file:
|
||||
names.append(pos_file.readline()[:-1])
|
||||
return nx, ny, dx, dy, pos, names
|
||||
|
||||
def cut_to_closest_position(frames, pos, start_frame = -1):
|
||||
if start_frame == -1:
|
||||
return frames
|
||||
else:
|
||||
distances = np.zeros(len(frames)-start_frame)
|
||||
for i in range(start_frame, len(frames)):
|
||||
current_pos = np.array([frames[i]['x'], frames[i]['y'], frames[i]['z']])
|
||||
distances[i-start_frame] = np.linalg.norm(current_pos - pos)
|
||||
closest_index = distances.argmin()
|
||||
cut_frames = frames[:start_frame+closest_index]
|
||||
return cut_frames
|
||||
|
||||
def parse_animation_file(filename):
|
||||
imported_frames = []
|
||||
anim_id = ''
|
||||
try:
|
||||
animation_file = open(filename)
|
||||
except IOError:
|
||||
logging.error("File {} can't be opened".format(filename))
|
||||
else:
|
||||
with animation_file:
|
||||
csv_reader = csv.reader(
|
||||
animation_file, delimiter=',', quotechar='|'
|
||||
)
|
||||
row_0 = csv_reader.next()
|
||||
if len(row_0) == 1:
|
||||
anim_id = row_0[0]
|
||||
logging.debug("Got animation_id: {}".format(anim_id))
|
||||
else:
|
||||
logging.debug("No animation id in file")
|
||||
frame_number, x, y, z, yaw, red, green, blue = row_0
|
||||
imported_frames.append({
|
||||
'number': int(frame_number),
|
||||
'x': float(x),
|
||||
'y': float(y),
|
||||
'z': float(z),
|
||||
'yaw': float(yaw),
|
||||
'red': int(red),
|
||||
'green': int(green),
|
||||
'blue': int(blue),
|
||||
})
|
||||
for row in csv_reader:
|
||||
frame_number, x, y, z, yaw, red, green, blue = row
|
||||
imported_frames.append({
|
||||
'number': int(frame_number),
|
||||
'x': float(x),
|
||||
'y': float(y),
|
||||
'z': float(z),
|
||||
'yaw': float(yaw),
|
||||
'red': int(red),
|
||||
'green': int(green),
|
||||
'blue': int(blue),
|
||||
})
|
||||
return imported_frames, anim_id
|
||||
|
||||
def cut_file(filename, _from, _to, reverse = False):
|
||||
imported_frames = []
|
||||
anim_id = ""
|
||||
|
||||
try:
|
||||
animation_file = open(filename)
|
||||
except IOError:
|
||||
logging.error("File {} can't be opened".format(filepath))
|
||||
else:
|
||||
with animation_file:
|
||||
csv_reader = csv.reader(
|
||||
animation_file, delimiter=',', quotechar='|'
|
||||
)
|
||||
row_0 = csv_reader.next()
|
||||
if len(row_0) == 1:
|
||||
anim_id = row_0[0]
|
||||
logging.debug("Got animation_id: {}".format(anim_id))
|
||||
else:
|
||||
logging.debug("No animation id in file")
|
||||
frame_number, x, y, z, yaw, red, green, blue = row_0
|
||||
imported_frames.append({
|
||||
'number': int(frame_number),
|
||||
'x': float(x),
|
||||
'y': float(y),
|
||||
'z': float(z),
|
||||
'yaw': float(yaw),
|
||||
'red': int(red),
|
||||
'green': int(green),
|
||||
'blue': int(blue),
|
||||
})
|
||||
for row in csv_reader:
|
||||
frame_number, x, y, z, yaw, red, green, blue = row
|
||||
imported_frames.append({
|
||||
'number': int(frame_number),
|
||||
'x': float(x),
|
||||
'y': float(y),
|
||||
'z': float(z),
|
||||
'yaw': float(yaw),
|
||||
'red': int(red),
|
||||
'green': int(green),
|
||||
'blue': int(blue),
|
||||
})
|
||||
|
||||
if _to == 0 or _to >= len(imported_frames):
|
||||
_to = len(imported_frames)-1
|
||||
|
||||
path = '{}/cut_{}_{}'.format(os.path.dirname(filename),_from,_to)
|
||||
|
||||
if reverse:
|
||||
path += '_r'
|
||||
|
||||
csv_file = open(path+'/'+os.path.basename(filename), mode='w+')
|
||||
with csv_file:
|
||||
csv_writer = csv.writer(csv_file, delimiter=',', quotechar='"', quoting=csv.QUOTE_MINIMAL)
|
||||
if anim_id != "":
|
||||
csv_writer.writerow([anim_id])
|
||||
if not reverse:
|
||||
for i in range(_from, _to+1):
|
||||
csv_writer.writerow([imported_frames[i]['number'], imported_frames[i]['x'], imported_frames[i]['y'], imported_frames[i]['z'],
|
||||
imported_frames[i]['yaw'], imported_frames[i]['red'], imported_frames[i]['green'], imported_frames[i]['blue']])
|
||||
else:
|
||||
for i in range(_from):
|
||||
csv_writer.writerow([imported_frames[i]['number'], imported_frames[i]['x'], imported_frames[i]['y'], imported_frames[i]['z'],
|
||||
imported_frames[i]['yaw'], imported_frames[i]['red'], imported_frames[i]['green'], imported_frames[i]['blue']])
|
||||
for i in range(_to, len(imported_frames)):
|
||||
csv_writer.writerow([imported_frames[i]['number'], imported_frames[i]['x'], imported_frames[i]['y'], imported_frames[i]['z'],
|
||||
imported_frames[i]['yaw'], imported_frames[i]['red'], imported_frames[i]['green'], imported_frames[i]['blue']])
|
||||
|
||||
|
||||
print("Successfully created file {}".format(path+'/'+os.path.basename(filename)))
|
||||
|
||||
def change_landing(frames, land_pos, speed, start_frame = -1):
|
||||
cut_frames = cut_to_closest_position(frames, land_pos, start_frame)
|
||||
last_frame_pos = np.array([cut_frames[-1]['x'], cut_frames[-1]['y'], cut_frames[-1]['z']])
|
||||
line = generate_line(last_frame_pos, land_pos, speed, start_frame=len(cut_frames))
|
||||
return cut_frames + line
|
||||
|
||||
def save_frames(frames, animation_id, filename):
|
||||
csv_file = open(filename, mode='w+')
|
||||
with csv_file:
|
||||
csv_writer = csv.writer(csv_file, delimiter=',', quotechar='"', quoting=csv.QUOTE_MINIMAL)
|
||||
if animation_id != "":
|
||||
csv_writer.writerow([animation_id])
|
||||
for frame in frames:
|
||||
csv_writer.writerow([frame['number'], frame['x'], frame['y'], frame['z'], frame['yaw'], frame['red'], frame['green'], frame['blue']])
|
||||
print("Successfully created file {}".format(path+'/'+os.path.basename(filename)))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
parser = argparse.ArgumentParser(description="Change landing positions")
|
||||
parser.add_argument('directory', nargs='?', default='.',
|
||||
help="Directory with animation csv files. Default is '.'")
|
||||
parser.add_argument('-s','--start', type=int, default=-1,
|
||||
help="Start count from this frame, default is -1 (from the end)")
|
||||
parser.add_argument('-p','--posfile', type=str, default="positions.txt",
|
||||
help="Positions file. Default is positions.txt")
|
||||
parser.add_argument('-v','--speed', type=float, default=0.5,
|
||||
help="Copter speed in m/s. Default is 0.5")
|
||||
args = parser.parse_args()
|
||||
|
||||
start_frame = args.start
|
||||
speed = args.speed
|
||||
|
||||
path = '{}/land_{}'.format(args.directory,start_frame)
|
||||
|
||||
if not os.path.exists(path):
|
||||
try:
|
||||
os.mkdir(path)
|
||||
except OSError:
|
||||
print("Creation of the directory %s failed" % path)
|
||||
else:
|
||||
print("Successfully created the directory %s " % path)
|
||||
|
||||
files = [f for f in glob.glob(args.directory + '/*.csv')]
|
||||
files.sort()
|
||||
|
||||
nx, ny, dx, dy, pos, names = parse_positions_file(args.posfile)
|
||||
land_positions = generate_positions(pos, nx, ny, dx, dy)
|
||||
|
||||
for i in range(len(files)):
|
||||
filename = os.path.basename(files[i])
|
||||
frames, animation_id = parse_animation_file(files[i])
|
||||
land_position = land_positions[i]
|
||||
new_frames = change_landing(frames, land_position, speed, start_frame)
|
||||
save_frames(new_frames, animation_id, path+'/'+filename)
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
2
tools/positions.txt
Normal file
2
tools/positions.txt
Normal file
@@ -0,0 +1,2 @@
|
||||
2 6 1 1
|
||||
-0.6 0 1
|
||||
Reference in New Issue
Block a user