text:
# Copyright (c) Facebook, Inc. and its affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

"""
This code generate shortest-paths for each question.
Paths start from and end at best-view points.

Sampled paths which are list of [{positions, actions, best_ious, ordered}] are saved
as "cache/shortest-paths-mt/qid.json".

{'positions': positions, 'actions': actions, 'best_iou': {obj1_name: obj1_iou, obj2_name: obj2_iou}, 'ordered': ordered}
"""
import argparse
import random
import time
import numpy as np
import os, sys, json
import os.path as osp
from tqdm import tqdm

from queue import Queue
from threading import Thread, Lock

this_dir = osp.dirname(__file__)
sys.path.insert(0, osp.join(this_dir, '../pyutils/House3D'))
from House3D import objrender, Environment, load_config
from House3D.core import local_create_house

sys.path.insert(0, osp.join(this_dir, '../pyutils'))
from house3d import House3DUtils

def get_5points_for_3objects(h3d, obj1_id, obj2_id, obj3_id, args):
  """
  Return 5 grid points (source_point, obj1_point, obj2_point, obj3_point, end_point)
  mentioned by the question, e.g., ``is <obj1> closer to <obj2> than <obj3>?''
  """
  # three objects
  house_id = h3d.env.house.house['id']
  obj1, obj2, obj3 = h3d.objects[obj1_id], h3d.objects[obj2_id], h3d.objects[obj3_id]

  h3d.set_target_object(obj1)
  obj1_conn_map = h3d.env.house.connMapDict[obj1_id][0]
  obj1_best_view_info = json.load(open(osp.join(args.target_obj_best_view_dir, house_id+'_'+obj1['id']+'.json')))  # {points, ious}
  obj1_points, obj1_ious = obj1_best_view_info['points'], obj1_best_view_info['ious']

  h3d.set_target_object(obj2)
  obj2_conn_map = h3d.env.house.connMapDict[obj2_id][0]
  obj2_best_view_info = json.load(open(osp.join(args.target_obj_best_view_dir, house_id+'_'+obj2['id']+'.json')))  # {points, ious}
  obj2_points, obj2_ious = obj2_best_view_info['points'], obj2_best_view_info['ious']

  h3d.set_target_object(obj3)
  obj3_conn_map = h3d.env.house.connMapDict[obj3_id][0]
  obj3_best_view_info = json.load(open(osp.join(args.target_obj_best_view_dir, house_id+'_'+obj3['id']+'.json')))  # {points, ious}
  obj3_points, obj3_ious = obj3_best_view_info['points'], obj3_best_view_info['ious']

  # source point
  if args.source_mode == 'nearby':
    # source point (p \in [target_dist, source_dist] of obj1)
    source_point_cands = np.argwhere((obj1_conn_map > args.min_dist_thresh) & (obj1_conn_map <= args.source_dist_thresh) )
  elif args.source_mode == 'random':
    source_point_cands = np.argwhere((obj1_conn_map > 0) & (obj2_conn_map > 0) & (obj3_conn_map > 0))
  else:
    raise NotImplementedError
  source_point_idx = np.random.choice(source_point_cands.shape[0])
  source_point = (source_point_cands[source_point_idx][0], source_point_cands[source_point_idx][1], np.random.choice(h3d.angles))

  # object 1 point
  obj1_point_cands = [(obj1_points[ix], iou) for ix, iou in enumerate(obj1_ious) if iou > args.good_iou]
  if len(obj1_point_cands) == 0:
    obj1_point_cands = [(obj1_points[0], obj1_ious[0])] if obj1_ious[0] > args.min_iou else []
  assert len(obj1_point_cands) > 0, 'best iou is below %.2f' % args.min_iou
  obj1_point, obj1_iou = obj1_point_cands[np.random.choice(len(obj1_point_cands))]  # (x, y, z, yaw) in coord
  x, z = h3d.env.house.to_grid(obj1_point[0], obj1_point[2])
  obj1_point = (x, z, obj1_point[-1])

  # object 2 point
  obj2_point_cands = [(obj2_points[ix], iou) for ix, iou in enumerate(obj2_ious) if iou > args.good_iou]
  if len(obj2_point_cands) == 0:
    obj2_point_cands = [(obj2_points[0], obj2_ious[0])] if obj2_ious[0] > args.min_iou else []
  assert len(obj2_point_cands) > 0, 'best iou is below %.2f' % args.min_iou
  obj2_point, obj2_iou = obj2_point_cands[np.random.choice(len(obj2_point_cands))]  # (x, y, z, yaw) in coord
  x, z = h3d.env.house.to_grid(obj2_point[0], obj2_point[2])
  obj2_point = (x, z, obj2_point[-1])

  # object 3 point
  obj3_point_cands = [(obj3_points[ix], iou) for ix, iou in enumerate(obj3_ious) if iou > args.good_iou]
  if len(obj3_point_cands) == 0:
    obj3_point_cands = [(obj3_points[0], obj3_ious[0])] if obj3_ious[0] > args.min_iou else []
  assert len(obj3_point_cands) > 0, 'best iou is below %.2f' % args.min_iou
  obj3_point, obj3_iou = obj3_point_cands[np.random.choice(len(obj3_point_cands))]  # (x, y, z, yaw) in coord
  x, z = h3d.env.house.to_grid(obj3_point[0], obj3_point[2])
  obj3_point = (x, z, obj3_point[-1])

  # end point
  if args.source_mode == 'nearby':
    # end point (p \in [target_dist, source_dist] of obj3)
    end_point_cands = np.argwhere((obj3_conn_map > args.min_dist_thresh) & (obj3_conn_map <= args.end_dist_thresh) )
  elif args.source_mode == 'random':
    end_point_cands = np.argwhere((obj1_conn_map > 0) & (obj2_conn_map > 0) & (obj3_conn_map > 0))
  else:
    raise NotImplementedError
  end_point_idx = np.random.choice(end_point_cands.shape[0])
  end_point = (end_point_cands[end_point_idx][0], end_point_cands[end_point_idx][1], np.random.choice(h3d.angles))

  return source_point, obj1_point, obj2_point, obj3_point, end_point, obj1_iou, obj2_iou, obj3_iou

def get_6points_for_3objects(h3d, room_id, obj1_id, obj2_id, obj3_id, args):
  """
  Return 6 grid points (source_point, room_point, obj1_point, obj2_point, obj3_point, end_point)
  mentioned by the question, e.g., ``is <obj1> closer to <obj2> than <obj3> in <room>?''
  """
  # room
  room = h3d.rooms[room_id]
  h3d.set_target_room(room)
  room_conn_map, _, room_inroomDist, _ = h3d.env.house.connMapDict[room_id]

  # three objects
  house_id = h3d.env.house.house['id']
  obj1, obj2, obj3 = h3d.objects[obj1_id], h3d.objects[obj2_id], h3d.objects[obj3_id]

  h3d.set_target_object(obj1)
  obj1_conn_map = h3d.env.house.connMapDict[obj1_id][0]
  obj1_best_view_info = json.load(open(osp.join(args.target_obj_best_view_dir, house_id+'_'+obj1['id']+'.json')))  # {points, ious}
  obj1_points, obj1_ious = obj1_best_view_info['points'], obj1_best_view_info['ious']

  h3d.set_target_object(obj2)
  obj2_conn_map = h3d.env.house.connMapDict[obj2_id][0]
  obj2_best_view_info = json.load(open(osp.join(args.target_obj_best_view_dir, house_id+'_'+obj2['id']+'.json')))  # {points, ious}
  obj2_points, obj2_ious = obj2_best_view_info['points'], obj2_best_view_info['ious']

  h3d.set_target_object(obj3)
  obj3_conn_map = h3d.env.house.connMapDict[obj3_id][0]
  obj3_best_view_info = json.load(open(osp.join(args.target_obj_best_view_dir, house_id+'_'+obj3['id']+'.json')))  # {points, ious}
  obj3_points, obj3_ious = obj3_best_view_info['points'], obj3_best_view_info['ious']

  # source point
  if args.source_mode == 'nearby':
    # source point (p \in [target_dist, source_dist] to room)
    source_point_cands = np.argwhere((room_conn_map > args.min_dist_thresh) & (room_conn_map <= args.source_dist_thresh) )
  elif args.source_mode == 'random':
    source_point_cands = np.argwhere((obj1_conn_map > 0) & (obj2_conn_map > 0) & (obj3_conn_map > 0))
  else:
    raise NotImplementedError
  source_point_idx = np.random.choice(source_point_cands.shape[0])
  source_point = (source_point_cands[source_point_idx][0], source_point_cands[source_point_idx][1], np.random.choice(h3d.angles))

  # room point
  room_point_cands = np.argwhere((room_inroomDist >= 0) & (room_inroomDist <= args.dist_to_room_center_thresh) )
  room_point_idx = np.random.choice(room_point_cands.shape[0])
  room_point = (room_point_cands[room_point_idx][0], room_point_cands[room_point_idx][1], np.random.choice(h3d.angles))

  # object 1 point
  obj1_point_cands = [(obj1_points[ix], iou) for ix, iou in enumerate(obj1_ious) if iou > args.good_iou]
  if len(obj1_point_cands) == 0:
    obj1_point_cands = [(obj1_points[0], obj1_ious[0])] if obj1_ious[0] > args.min_iou else []
  assert len(obj1_point_cands) > 0, 'best iou is below %.2f' % args.min_iou
  obj1_point, obj1_iou = obj1_point_cands[np.random.choice(len(obj1_point_cands))]  # (x, y, z, yaw) in coord
  x, z = h3d.env.house.to_grid(obj1_point[0], obj1_point[2])
  obj1_point = (x, z, obj1_point[-1])

  # object 2 point
  obj2_point_cands = [(obj2_points[ix], iou) for ix, iou in enumerate(obj2_ious) if iou > args.good_iou]
  if len(obj2_point_cands) == 0:
    obj2_point_cands = [(obj2_points[0], obj2_ious[0])] if obj2_ious[0] > args.min_iou else []
  assert len(obj2_point_cands) > 0, 'best iou is below %.2f' % args.min_iou
  obj2_point, obj2_iou = obj2_point_cands[np.random.choice(len(obj2_point_cands))]  # (x, y, z, yaw) in coord
  x, z = h3d.env.house.to_grid(obj2_point[0], obj2_point[2])
  obj2_point = (x, z, obj2_point[-1])

  # object 3 point
  obj3_point_cands = [(obj3_points[ix], iou) for ix, iou in enumerate(obj3_ious) if iou > args.good_iou]
  if len(obj3_point_cands) == 0:
    obj3_point_cands = [(obj3_points[0], obj3_ious[0])] if obj3_ious[0] > args.min_iou else []
  assert len(obj3_point_cands) > 0, 'best iou is below %.2f' % args.min_iou
  obj3_point, obj3_iou = obj3_point_cands[np.random.choice(len(obj3_point_cands))]  # (x, y, z, yaw) in coord
  x, z = h3d.env.house.to_grid(obj3_point[0], obj3_point[2])
  obj3_point = (x, z, obj3_point[-1])

  # end point
  if args.source_mode == 'nearby':
    # end point (p \in [target_dist, source_dist] of obj3)
    end_point_cands = np.argwhere((obj3_conn_map > args.min_dist_thresh) & (obj3_conn_map <= args.end_dist_thresh) )
  elif args.source_mode == 'random':
    end_point_cands = np.argwhere((obj1_conn_map > 0) & (obj2_conn_map > 0) & (obj3_conn_map > 0))
  else:
    raise NotImplementedError
  end_point_idx = np.random.choice(end_point_cands.shape[0])
  end_point = (end_point_cands[end_point_idx][0], end_point_cands[end_point_idx][1], np.random.choice(h3d.angles))

  return source_point, room_point, obj1_point, obj2_point, obj3_point, end_point, obj1_iou, obj2_iou, obj3_iou

def get_4points_for_2objects(h3d, obj1_id, obj2_id, args):
  """
  Return 4 grid points (source_point, obj1_point, obj2_point, end_point) for two unique objects
  mentioned by question, e.g., ``is color of <obj1> same as <obj2>?''
  """
  # two objects
  house_id = h3d.env.house.house['id']
  obj1, obj2 = h3d.objects[obj1_id], h3d.objects[obj2_id]

  h3d.set_target_object(obj1)
  obj1_conn_map = h3d.env.house.connMapDict[obj1_id][0]
  obj1_best_view_info = json.load(open(osp.join(args.target_obj_best_view_dir, house_id+'_'+obj1['id']+'.json')))  # {points, ious}
  obj1_points, obj1_ious = obj1_best_view_info['points'], obj1_best_view_info['ious']

  h3d.set_target_object(obj2)
  obj2_conn_map = h3d.env.house.connMapDict[obj2_id][0]
  obj2_best_view_info = json.load(open(osp.join(args.target_obj_best_view_dir, house_id+'_'+obj2['id']+'.json')))  # {points, ious}
  obj2_points, obj2_ious = obj2_best_view_info['points'], obj2_best_view_info['ious']

  # source point
  if args.source_mode == 'nearby':
    # source target (p \in [target_dist, source_dist] of obj1)
    source_point_cands = np.argwhere((obj1_conn_map > args.min_dist_thresh) & (obj1_conn_map <= args.source_dist_thresh) )
  elif args.source_mode == 'random':
    source_point_cands = np.argwhere((obj1_conn_map > 0 ) & (obj2_conn_map > 0) )  # randomly spawned
  else:
    raise NotImplementedError
  source_point_idx = np.random.choice(source_point_cands.shape[0])
  source_point = (source_point_cands[source_point_idx][0], source_point_cands[source_point_idx][1], np.random.choice(h3d.angles))

  # object 1 point
  obj1_point_cands = [(obj1_points[ix], iou) for ix, iou in enumerate(obj1_ious) if iou > args.good_iou]
  if len(obj1_point_cands) == 0:
    obj1_point_cands = [(obj1_points[0], obj1_ious[0])] if obj1_ious[0] > args.min_iou else []
  assert len(obj1_point_cands) > 0, 'best iou is below %.2f' % args.min_iou
  obj1_point, obj1_iou = obj1_point_cands[np.random.choice(len(obj1_point_cands))]  # (x, y, z, yaw) in coord
  x, z = h3d.env.house.to_grid(obj1_point[0], obj1_point[2])
  obj1_point = (x, z, obj1_point[-1])

  # object 2 point
  obj2_point_cands = [(obj2_points[ix], iou) for ix, iou in enumerate(obj2_ious) if iou > args.good_iou]
  if len(obj2_point_cands) == 0:
    obj2_point_cands = [(obj2_points[0], obj2_ious[0])] if obj2_ious[0] > args.min_iou else []
  assert len(obj2_point_cands) > 0, 'best iou is below %.2f' % args.min_iou
  obj2_point, obj2_iou = obj2_point_cands[np.random.choice(len(obj2_point_cands))]  # (x, y, z, yaw) in coord
  x, z = h3d.env.house.to_grid(obj2_point[0], obj2_point[2])
  obj2_point = (x, z, obj2_point[-1])

  # end point
  if args.source_mode == 'nearby':
    # end target (p \in [target_dist, source_dist] of obj2)
    end_point_cands = np.argwhere((obj2_conn_map > args.min_dist_thresh) & (obj2_conn_map <= args.end_dist_thresh) )
  elif args.source_mode == 'random':
    end_point_cands = np.argwhere((obj1_conn_map > 0 ) & (obj2_conn_map > 0) )  # randomly spawned
  else:
    raise NotImplementedError
  end_point_idx = np.random.choice(end_point_cands.shape[0])
  end_point = (end_point_cands[end_point_idx][0], end_point_cands[end_point_idx][1], np.random.choice(h3d.angles))

  return source_point, obj1_point, obj2_point, end_point, obj1_iou, obj2_iou

def get_5points_for_2objects(h3d, room_id, obj1_id, obj2_id, args):
  """
  Return 5 grid points (source_point, room_point, obj1_point, obj2_point, end_point) for two non-unique objects
  mentioned by question, e.g., ``is color of <obj1> same as <obj2> in <room>?''
  """
  # room
  room = h3d.rooms[room_id]
  h3d.set_target_room(room)
  room_conn_map, _, room_inroomDist, _ = h3d.env.house.connMapDict[room_id]

  # two objects
  house_id = h3d.env.house.house['id']
  obj1, obj2 = h3d.objects[obj1_id], h3d.objects[obj2_id]

  h3d.set_target_object(obj1)
  obj1_conn_map = h3d.env.house.connMapDict[obj1_id][0]
  obj1_best_view_info = json.load(open(osp.join(args.target_obj_best_view_dir, house_id+'_'+obj1['id']+'.json')))  # {points, ious}
  obj1_points, obj1_ious = obj1_best_view_info['points'], obj1_best_view_info['ious']

  h3d.set_target_object(obj2)
  obj2_conn_map = h3d.env.house.connMapDict[obj2_id][0]
  obj2_best_view_info = json.load(open(osp.join(args.target_obj_best_view_dir, house_id+'_'+obj2['id']+'.json')))  # {points, ious}
  obj2_points, obj2_ious = obj2_best_view_info['points'], obj2_best_view_info['ious']

  # source point
  if args.source_mode == 'nearby':
    # source target (p \in [target_dist, source_dist] of obj1)
    source_point_cands = np.argwhere((room_conn_map > args.min_dist_thresh) & (room_conn_map <= args.source_dist_thresh) )
  elif args.source_mode == 'random':
    source_point_cands = np.argwhere((obj1_conn_map > 0 ) & (obj2_conn_map > 0) )  # randomly spawned
  else:
    raise NotImplementedError
  source_point_idx = np.random.choice(source_point_cands.shape[0])
  source_point = (source_point_cands[source_point_idx][0], source_point_cands[source_point_idx][1], np.random.choice(h3d.angles))

  # room point
  room_point_cands = np.argwhere((room_inroomDist >= 0) & (room_inroomDist <= args.dist_to_room_center_thresh))
  room_point_idx = np.random.choice(room_point_cands.shape[0])
  room_point = (room_point_cands[room_point_idx][0], room_point_cands[room_point_idx][1], np.random.choice(h3d.angles))

  # object1 point
  obj1_point_cands = [(obj1_points[ix], iou) for ix, iou in enumerate(obj1_ious) if iou > args.good_iou]
  if len(obj1_point_cands) == 0:
    obj1_point_cands = [(obj1_points[0], obj1_ious[0])] if obj1_ious[0] > args.min_iou else []
  assert len(obj1_point_cands) > 0, 'best iou is below %.2f' % args.min_iou
  obj1_point, obj1_iou = obj1_point_cands[np.random.choice(len(obj1_point_cands))]  # (x, y, z, yaw) in coord
  x, z = h3d.env.house.to_grid(obj1_point[0], obj1_point[2])
  obj1_point = (x, z, obj1_point[-1])

  # object2 point
  obj2_point_cands = [(obj2_points[ix], iou) for ix, iou in enumerate(obj2_ious) if iou > args.good_iou]
  if len(obj2_point_cands) == 0:
    obj2_point_cands = [(obj2_points[0], obj2_ious[0])] if obj2_ious[0] > args.min_iou else []
  assert len(obj2_point_cands) > 0, 'best iou is below %.2f' % args.min_iou
  obj2_point, obj2_iou = obj2_point_cands[np.random.choice(len(obj2_point_cands))]  # (x, y, z, yaw) in coord
  x, z = h3d.env.house.to_grid(obj2_point[0], obj2_point[2])
  obj2_point = (x, z, obj2_point[-1])

  # end point
  if args.source_mode == 'nearby':
    # end target (p \in [target_dist, source_dist] of obj2)
    end_point_cands = np.argwhere((obj2_conn_map > args.min_dist_thresh) & (obj2_conn_map <= args.end_dist_thresh) )
  elif args.source_mode == 'random':
    end_point_cands = np.argwhere((obj1_conn_map > 0 ) & (obj2_conn_map > 0) )  # randomly spawned
  else:
    raise NotImplementedError
  end_point_idx = np.random.choice(end_point_cands.shape[0])
  end_point = (end_point_cands[end_point_idx][0], end_point_cands[end_point_idx][1], np.random.choice(h3d.angles))

  return source_point, room_point, obj1_point, obj2_point, end_point, obj1_iou, obj2_iou

def get_6points_for_2objects(h3d, room1_id, obj1_id, room2_id, obj2_id, args):
  """
  Return 6 grid points (source_point, room1_point, obj1_point, room2_point, obj2_point, end_point) for two objects
  mentioned by question, e.g., ``is color of <obj1> in <room1> same as <obj2> in <room2>?''
  """
  # two rooms
  room1, room2 = h3d.rooms[room1_id], h3d.rooms[room2_id]
  h3d.set_target_room(room1)
  h3d.set_target_room(room2)
  room1_conn_map, _, room1_inroomDist, _ = h3d.env.house.connMapDict[room1_id]
  room2_conn_map, _, room2_inroomDist, _ = h3d.env.house.connMapDict[room2_id]

  # two objects
  house_id = h3d.env.house.house['id']
  obj1, obj2 = h3d.objects[obj1_id], h3d.objects[obj2_id]

  h3d.set_target_object(obj1)
  obj1_conn_map = h3d.env.house.connMapDict[obj1_id][0]
  obj1_best_view_info = json.load(open(osp.join(args.target_obj_best_view_dir, house_id+'_'+obj1['id']+'.json')))  # {points, ious}
  obj1_points, obj1_ious = obj1_best_view_info['points'], obj1_best_view_info['ious']

  h3d.set_target_object(obj2)
  obj2_conn_map = h3d.env.house.connMapDict[obj2_id][0]
  obj2_best_view_info = json.load(open(osp.join(args.target_obj_best_view_dir, house_id+'_'+obj2['id']+'.json')))  # {points, ious}
  obj2_points, obj2_ious = obj2_best_view_info['points'], obj2_best_view_info['ious']

  # source point
  if args.source_mode == 'nearby':
    # source target (p \in [target_dist, source_dist] of obj1)
    source_point_cands = np.argwhere((room1_conn_map > args.min_dist_thresh) & (room1_conn_map <= args.source_dist_thresh) )
  elif args.source_mode == 'random':
    source_point_cands = np.argwhere((obj1_conn_map > 0 ) & (obj2_conn_map > 0) )  # randomly spawned
  else:
    raise NotImplementedError
  source_point_idx = np.random.choice(source_point_cands.shape[0])
  source_point = (source_point_cands[source_point_idx][0], source_point_cands[source_point_idx][1], np.random.choice(h3d.angles))

  # room1 point
  room1_point_cands = np.argwhere((room1_inroomDist >= 0) & (room1_inroomDist <= args.dist_to_room_center_thresh))
  room1_point_idx = np.random.choice(room1_point_cands.shape[0])
  room1_point = (room1_point_cands[room1_point_idx][0], room1_point_cands[room1_point_idx][1], np.random.choice(h3d.angles))

  # object1 point
  obj1_point_cands = [(obj1_points[ix], iou) for ix, iou in enumerate(obj1_ious) if iou > args.good_iou]
  if len(obj1_point_cands) == 0:
    obj1_point_cands = [(obj1_points[0], obj1_ious[0])] if obj1_ious[0] > args.min_iou else []
  assert len(obj1_point_cands) > 0, 'best iou is below %.2f' % args.min_iou
  obj1_point, obj1_iou = obj1_point_cands[np.random.choice(len(obj1_point_cands))]  # (x, y, z, yaw) in coord
  x, z = h3d.env.house.to_grid(obj1_point[0], obj1_point[2])
  obj1_point = (x, z, obj1_point[-1])

  # room2 point
  room2_point_cands = np.argwhere((room2_inroomDist >= 0) & (room2_inroomDist <= args.dist_to_room_center_thresh) )
  room2_point_idx = np.random.choice(room2_point_cands.shape[0])
  room2_point = (room2_point_cands[room2_point_idx][0], room2_point_cands[room2_point_idx][1], np.random.choice(h3d.angles))

  # object2 point
  obj2_point_cands = [(obj2_points[ix], iou) for ix, iou in enumerate(obj2_ious) if iou > args.good_iou]
  if len(obj2_point_cands) == 0:
    obj2_point_cands = [(obj2_points[0], obj2_ious[0])] if obj2_ious[0] > args.min_iou else []
  assert len(obj2_point_cands) > 0, 'best iou is below %.2f' % args.min_iou
  obj2_point, obj2_iou = obj2_point_cands[np.random.choice(len(obj2_point_cands))]  # (x, y, z, yaw) in coord
  x, z = h3d.env.house.to_grid(obj2_point[0], obj2_point[2])
  obj2_point = (x, z, obj2_point[-1])

  # end point
  if args.source_mode == 'nearby':
    # end target (p \in [target_dist, source_dist] of obj2)
    end_point_cands = np.argwhere((obj2_conn_map > args.min_dist_thresh) & (obj2_conn_map <= args.end_dist_thresh) )
  elif args.source_mode == 'random':
    end_point_cands = np.argwhere((obj1_conn_map > 0 ) & (obj2_conn_map > 0) )  # randomly spawned
  else:
    raise NotImplementedError
  end_point_idx = np.random.choice(end_point_cands.shape[0])
  end_point = (end_point_cands[end_point_idx][0], end_point_cands[end_point_idx][1], np.random.choice(h3d.angles))

  return source_point, room1_point, obj1_point, room2_point, obj2_point, end_point, obj1_iou, obj2_iou

def get_4points_for_2rooms(h3d, room1_id, room2_id, args):
  """
  Return 4 grid points (source_point, room1_point, room2_point, end_point) for two rooms
  mentioned by question, e.g., ``is room1 bigger than room2?''
  """
  # two rooms
  room1, room2 = h3d.rooms[room1_id], h3d.rooms[room2_id]
  h3d.set_target_room(room1)
  h3d.set_target_room(room2)
  room1_conn_map, _, room1_inroomDist, _ = h3d.env.house.connMapDict[room1_id]
  room2_conn_map, _, room2_inroomDist, _ = h3d.env.house.connMapDict[room2_id]

  if args.source_mode == 'nearby':
    # source target (p \in [target_dist, source_dist] of room1)
    source_point_cands = np.argwhere((room1_conn_map > args.min_dist_thresh) & (room1_conn_map <= args.source_dist_thresh) )
  elif args.source_mode == 'random':
    source_point_cands = np.argwhere((room1_conn_map > 0) & (room2_conn_map > 0) )  # randomly spawned
  else:
    raise NotImplementedError
  source_point_idx = np.random.choice(source_point_cands.shape[0])
  source_point = (source_point_cands[source_point_idx][0], source_point_cands[source_point_idx][1], np.random.choice(h3d.angles))

  # room 1 point (p \in [0, target_dist])
  room1_point_cands = np.argwhere((room1_inroomDist >= 0) & (room1_inroomDist <= args.dist_to_room_center_thresh))
  room1_point_idx = np.random.choice(room1_point_cands.shape[0])
  room1_point = (room1_point_cands[room1_point_idx][0], room1_point_cands[room1_point_idx][1], np.random.choice(h3d.angles))

  # room 2 point (p \in [0, target_dist])
  room2_point_cands = np.argwhere((room2_inroomDist >= 0) & (room2_inroomDist <= args.dist_to_room_center_thresh) )
  room2_point_idx = np.random.choice(room2_point_cands.shape[0])
  room2_point = (room2_point_cands[room2_point_idx][0], room2_point_cands[room2_point_idx][1], np.random.choice(h3d.angles))

  if args.source_mode == 'nearby':
    # end target (p \in [target_dist, source_dist] for room2)
    end_point_cands = np.argwhere((room2_conn_map > args.min_dist_thresh) & (room2_conn_map <= args.end_dist_thresh) )
  elif args.source_mode == 'random':
    end_point_cands = np.argwhere((room1_conn_map > 0) & (room2_conn_map > 0) )  # randomly spawned
  else:
    raise NotImplementedError
  end_point_idx = np.random.choice(end_point_cands.shape[0])
  end_point = (end_point_cands[end_point_idx][0], end_point_cands[end_point_idx][1], np.random.choice(h3d.angles))

  return source_point, room1_point, room2_point, end_point

def path_to_actions(h3d, shortest_path):
  """
  return list of positions and actions, the last of which is always 'STOP'
  """
  act_q, pos_q = [], []
  for i in shortest_path:
    gx, gy = h3d.env.house.to_coor(i[0], i[1])
    yaw = i[2]
    pos_q.append([float(gx), 1.0, float(gy), float(yaw)])

  # action_map = {0: 'FRWD', 1: 'LEFT', 2: 'RGHT', 3: 'STOP'}
  for i in range(len(pos_q)-1):
    if pos_q[i][0] == pos_q[i+1][0] and pos_q[i][2] == pos_q[i+1][2]:
      # turn
      if pos_q[i][3] > 0 and pos_q[i+1][3] == -180:   # 150 -> -180
        act_q.append(2)
      elif pos_q[i][3] == -180 and pos_q[i+1][3] > 0: # -180 -> 150
        act_q.append(1)
      elif pos_q[i+1][3] > pos_q[i][3]:
        act_q.append(2)
      elif pos_q[i+1][3] < pos_q[i][3]:
        act_q.append(1)
      else:
        raise NotImplementedError
    else:
      act_q.append(0)
  act_q.append(3)
  return act_q, pos_q

def sample_paths(h3d, path_points):
  """
  Inputs:
  - path_points: list of grid points
  Return:
  list of sampled path information:
  - positions: list of list of [[(pos_x, 1.0, pos_z, yaw)]], in coordinates (not grids).
  - actions: list of list of [[actions]]
  we only try two version: source_point -> points, source_point -> points[::-1]
  most of the cases are fine, unless we have 3 points we will miss some combinations
  """
  # source --> points
  positions, actions = [], []
  for pt0, pt1 in zip(path_points[:-1], path_points[1:]):
    spath = h3d.compute_shortest_path(pt0, pt1)
    act_q, pos_q = path_to_actions(h3d, spath)
    positions.append(pos_q)
    actions.append(act_q)
  return positions, actions

def main(house_to_qns, args):

  invalid = []

  # data pool
  lock = Lock()
  q = Queue()
  for i, house_id in enumerate(list(house_to_qns.keys())):
    qns = house_to_qns[house_id]
    q.put((i, house_id, qns))

  # worker
  def worker():
    api_thread = objrender.RenderAPIThread(w=224, h=224)
    while True:
      i, house_id, qns = q.get()
      print('Processing house[%s] %s/%s' % (house_id, i+1, len(house_to_qns)))
      env = Environment(api_thread, house_id, cfg, ColideRes=args.colide_resolution)
      h3d = House3DUtils(env, build_graph=False, graph_dir=args.graph_dir,
              target_obj_conn_map_dir=args.target_obj_conn_map_dir)

      # compute shortest path for each qn
      for qn in qns:

        tic = time.time()
        samples = []

        for _ in range(args.num_samples):

          sample = {'shortest_paths': [], 'positions': []}

          if qn['type'] in ['object_color_compare_inroom', 'object_size_compare_inroom']:
            # 2 objects
            assert len(qn['bbox']) == 2
            obj1_id, obj2_id = qn['bbox'][0]['id'], qn['bbox'][1]['id']
            obj1_name, obj2_name = qn['bbox'][0]['name'], qn['bbox'][1]['name']
            assert qn['bbox'][0]['room_id'] == qn['bbox'][1]['room_id']
            room_id = qn['bbox'][0]['room_id']
            if 'in the' in qn['question']:
              # S -> rm -> obj1 -> obj2 -> E
              ntnp = '1rm_2obj_5p'
              try:
                source_point, room_point, obj1_point, obj2_point, end_point, obj1_iou, obj2_iou = get_5points_for_2objects(h3d, room_id, obj1_id, obj2_id, args)
              except:
                print('5 points for [S -> rm -> obj1 -> obj2 -> E not found.')
                continue
              # compute shortest paths
              try:
                ordered = True
                positions, actions = sample_paths(h3d, [source_point, room_point, obj1_point, obj2_point, end_point])
                samples.append({'positions': positions, 'actions': actions, 'best_iou': {obj1_name: obj1_iou, obj2_name: obj2_iou}, 'ordered': ordered})
              except:
                print('shortest path not found for question[%s](%s).' % (qn['id'], qn['type']))
                continue
              # fname
              fname = '.'.join([str(qn['house']), str(room_id), str(obj1_id), str(obj2_id)])
              assert fname == qn['path'], '%s, %s' % (fname, qn['path'])

            else:
              # S -> obj1 -> obj2 -> E
              ntnp = '2obj_4p'
              try:
                source_point, obj1_point, obj2_point, end_point, obj1_iou, obj2_iou = get_4points_for_2objects(h3d, obj1_id, obj2_id, args)
              except:
                print('4 points for [S -> obj1 -> obj2 -> E] not found.')
                continue
              # compute shortest paths
              try:
                ordered = True
                positions, actions = sample_paths(h3d, [source_point, obj1_point, obj2_point, end_point])
                samples.append({'positions': positions, 'actions': actions, 'best_iou': {obj1_name: obj1_iou, obj2_name: obj2_iou}, 'ordered': ordered})
              except:
                print('shortest path not found for question[%s](%s).' % (qn['id'], qn['type']))
                continue
              # fname
              fname = '.'.join([str(qn['house']), str(obj1_id), str(obj2_id)])
              assert fname == qn['path'], '%s, %s' % (fname, qn['path'])

          elif qn['type'] in ['object_color_compare_xroom', 'object_size_compare_xroom']:
            # S -> rm1 -> obj1 -> rm2 -> obj2 -> E
            ntnp = '2rm_2obj_6p'
            assert len(qn['bbox']) == 2
            obj1_id, obj2_id = qn['bbox'][0]['id'], qn['bbox'][1]['id']
            obj1_name, obj2_name = qn['bbox'][0]['name'], qn['bbox'][1]['name']
            room1_id, room2_id = qn['bbox'][0]['room_id'], qn['bbox'][1]['room_id']
            assert room1_id != room2_id
            try:
              source_point, room1_point, obj1_point, room2_point, obj2_point, end_point, obj1_iou, obj2_iou = \
                get_6points_for_2objects(h3d, room1_id, obj1_id, room2_id, obj2_id, args)
            except:
              print('6 points for [S -> rm1 -> obj1 -> rm2 -> obj2 -> E] not found.')
              continue
            # compute shortest paths
            try:
              ordered = True
              positions, actions = sample_paths(h3d, [source_point, room1_point, obj1_point, room2_point, obj2_point, end_point])
              samples.append({'positions': positions, 'actions': actions, 'best_iou': {obj1_name: obj1_iou, obj2_name: obj2_iou}, 'ordered': ordered})
            except:
              print('shortest path not found for question[%s](%s).' % (qn['id'], qn['type']))
              continue
            # fname
            fname = '.'.join([str(qn['house']), str(room1_id), str(obj1_id), str(room2_id), str(obj2_id)])
            assert fname == qn['path'], '%s, %s' % (fname, qn['path'])

          elif qn['type'] == 'object_dist_compare_inroom':
            assert len(qn['bbox']) == 3
            obj1_id, obj2_id, obj3_id = qn['bbox'][0]['id'], qn['bbox'][1]['id'], qn['bbox'][2]['id']
            obj1_name, obj2_name, obj3_name = qn['bbox'][0]['name'], qn['bbox'][1]['name'], qn['bbox'][2]['name']
            room_id = qn['bbox'][0]['room_id']
            assert room_id == qn['bbox'][1]['room_id'] == qn['bbox'][2]['room_id']
            if 'in the ' in qn['question']:
              # S -> rm -> obj1 -> obj2 -> obj3 -> E
              ntnp = '1rm_3obj_6p'
              try:
                source_point, room_point, obj1_point, obj2_point, obj3_point, end_point, obj1_iou, obj2_iou, obj3_iou = \
                  get_6points_for_3objects(h3d, room_id, obj1_id, obj2_id, obj3_id, args)
              except:
                print('6 points for [S -> rm -> obj1 -> obj2 -> obj3 -> E] not found.')
                continue
              # compute shortest paths
              try:
                ordered = True
                positions, actions = sample_paths(h3d, [source_point, room_point, obj1_point, obj2_point, obj3_point, end_point])
                samples.append({'positions': positions, 'actions': actions, 'best_iou': {obj1_name: obj1_iou, obj2_name: obj2_iou, obj3_name: obj3_iou}, 'ordered': ordered})
              except:
                print('shortest path not found for question[%s](%s).' % (qn['id'], qn['type']))
                continue
              # fname
              fname = '.'.join([str(qn['house']), str(room_id), str(obj1_id), str(obj2_id), str(obj3_id)])
              assert fname == qn['path'], '%s, %s' % (fname, qn['path'])

            else:
              # S -> obj1 -> obj2 -> obj3 -> E
              ntnp = '3obj_5p'
              try:
                source_point, obj1_point, obj2_point, obj3_point, end_point, obj1_iou, obj2_iou, obj3_iou = \
                  get_5points_for_3objects(h3d, obj1_id, obj2_id, obj3_id, args)
              except:
                print('5 points for [S -> obj1 -> obj2 -> obj3 -> E] not found.')
                continue
              # compute shortest paths
              try:
                ordered = True
                positions, actions = sample_paths(h3d, [source_point, obj1_point, obj2_point, obj3_point, end_point])
                samples.append({'positions': positions, 'actions': actions, 'best_iou': {obj1_name: obj1_iou, obj2_name: obj2_iou, obj3_name: obj3_iou}, 'ordered': ordered})
              except:
                print('shortest path not found for question[%s](%s).' % (qn['id'], qn['type']))
                continue
              # fname
              fname = '.'.join([str(qn['house']), str(obj1_id), str(obj2_id), str(obj3_id)])
              assert fname == qn['path'], '%s, %s' % (fname, qn['path'])

          elif qn['type'] == 'room_size_compare':
            # S -> room1 -> room2 -> E
            ntnp = '2rm4p'
            # 2 rooms
            assert len(qn['bbox']) == 2
            room1_id, room2_id = qn['bbox'][0]['id'], qn['bbox'][1]['id']
            assert room1_id != room2_id
            # 4 points
            try:
              source_point, room1_point, room2_point, end_point = get_4points_for_2rooms(h3d, room1_id, room2_id, args)
            except:
              print('4 points for [S -> room1 -> room2 -> E] not found.')
              continue
            # compute shortest paths
            try:
              ordered = True
              positions, actions = sample_paths(h3d, [source_point, room1_point, room2_point, end_point])
              samples.append({'positions': positions, 'actions': actions, 'ordered': ordered})
            except:
              print('shortest path not found for question[%s][%s].' % (qn['id'], qn['type']))
              continue
            # fname
            fname = '.'.join([str(qn['house']), str(room1_id), str(room2_id)])
            assert fname == qn['path'], '%s, %s' % (fname, qn['path'])
        # save
        if len(samples) == 0:
          invalid.append(qn['id'])  # do not use += here!
        else:
          print('%s [%s] samples for question[%s] in %.2fs.' % (len(samples), ntnp, qn['id'], time.time()-tic))
          assert fname == qn['path']
          with open(osp.join(args.shortest_path_dir, fname + '.json'), 'w') as f:
            json.dump(samples, f)

      # finished this job
      q.task_done()

  for i in range(args.num_workers):
    t = Thread(target=worker)
    t.daemon = True
    t.start()
  q.join()

  return invalid

if __name__ == '__main__':

  parser = argparse.ArgumentParser()
  parser.add_argument('--question_json', default='cache/question-gen-outputs/questions_pruned_mt_with_conn_program.json')
  parser.add_argument('--house3d_metadata_dir', default='pyutils/House3D/House3D/metadata')
  parser.add_argument('--suncg_data_dir', default='data/SUNCGdata')
  parser.add_argument('--graph_dir', default='cache/3d-graphs', help='directory for saving graphs')
  parser.add_argument('--target_obj_conn_map_dir', default='cache/target-obj-conn-maps')
  parser.add_argument('--target_obj_best_view_dir', default='cache/target-obj-bestview-pos')
  parser.add_argument('--colide_resolution', default=500, type=int, help='house grid resolution')
  parser.add_argument('--seed', default=24, type=int, help='random seed')
  parser.add_argument('--source_mode', default='nearby', type=str, help='nearby or random')
  parser.add_argument('--source_dist_thresh', default=30, type=int, help='max steps to source, only useful for source_mode == nearby')
  parser.add_argument('--min_dist_thresh', default=5, type=int, help='max steps to target')
  parser.add_argument('--dist_to_room_center_thresh', default=5, type=int, help='max steps to room center')
  parser.add_argument('--end_dist_thresh', default=20, type=int, help='max steps to end, only useful for source_mode == nearby')
  parser.add_argument('--good_iou', default=0.25, type=float, help='a good iou for best-view point')
  parser.add_argument('--min_iou', default=0.10, type=float, help='a minimum iou for best-view point')
  parser.add_argument('--num_samples', default=3, type=int, help='number of shortest paths per question')
  parser.add_argument('--shortest_path_dir', default='cache/shortest-paths-mt', help='directory saving shortest paths for every question')
  parser.add_argument('--num_workers', default=2, type=int, help='number of workers')
  args = parser.parse_args()

  # set random seed
  random.seed(args.seed)
  np.random.seed(args.seed)

  # config
  cfg = {}
  cfg['colorFile'] = osp.join(args.house3d_metadata_dir, 'colormap_fine.csv')
  cfg['roomTargetFile'] = osp.join(args.house3d_metadata_dir, 'room_target_object_map.csv')
  cfg['modelCategoryFile'] = osp.join(args.house3d_metadata_dir, 'ModelCategoryMapping.csv')
  cfg['prefix'] = osp.join(args.suncg_data_dir, 'house')
  for d in cfg.values():
    assert osp.exists(d), d

  # update shortest_path_dir
  shortest_path_dir = osp.join(args.shortest_path_dir, args.source_mode+'_source_best_view')
  if not osp.isdir(shortest_path_dir):
    os.makedirs(shortest_path_dir)
  args.shortest_path_dir = shortest_path_dir

  # load questions = [{bbox, id, house, question, answer, type}]
  questions = json.load(open(args.question_json, 'r'))
  Questions = {qn['id']: qn for qn in questions}

  # how many unique paths are there for input questions
  visited_paths = {}
  for qn in questions:
    path = qn['path']
    visited_paths[path] = False
  print('There are %s unique paths for %s questions.' % (len(visited_paths), len(questions)))

  # merge questions that share same targets
  merged_questions = []
  for qn in questions:
    path = qn['path']
    if visited_paths[path] is False:
      merged_questions.append(qn)
      visited_paths[path] = True
  print('Correspondingly, we prepared %s merged_questions.' % len(merged_questions))

  # main
  house_to_qns = {}
  merged_questions = merged_questions
  for qn in merged_questions:
    # check if we've already computed this path
    path = qn['path']
    path_file = osp.join(args.shortest_path_dir, path+'.json')
    if osp.exists(path_file):
      print('%s computed already.' % path_file)
      continue
    house_id = qn['house']
    if house_id not in house_to_qns: house_to_qns[house_id] = []
    house_to_qns[house_id].append(qn)

  invalid = main(house_to_qns, args)

  # invalid
  invalid_json = osp.join(args.shortest_path_dir, args.source_mode+'_source_best_view_invalid.json')
  with open(invalid_json, 'w') as f:
    json.dump(invalid, f)
  print('%s invalid questions found, saved in %s' % (len(invalid), invalid_json))
  print('Done.')

max_stars_repo_name: sawravchy/MT-EQA

max_issues_repo_name: sawravchy/MT-EQA

max_forks_repo_name: sawravchy/MT-EQA


text:
#!/usr/bin/env python
# -*- coding: utf-8 -*-

# 3rd party imports
import numpy as np

from scipy import interpolate

__author__ = "Louis Richard"
__email__ = "louisr@irfu.se"
__copyright__ = "Copyright 2020-2021"
__license__ = "MIT"
__version__ = "2.3.7"
__status__ = "Prototype"

def _get_vol_ten(r_xyz, time):
    if len(time) == 1:
        time = np.array([time, time, time, time])

    tckr_x, tckr_y, tckr_z = [[], [], []]

    for i in range(4):
        tckr_x.append(interpolate.interp1d(r_xyz[i].time.data,
                                           r_xyz[i].data[:, 0]))
        tckr_y.append(interpolate.interp1d(r_xyz[i].time.data,
                                           r_xyz[i].data[:, 1]))
        tckr_z.append(interpolate.interp1d(r_xyz[i].time.data,
                                           r_xyz[i].data[:, 2]))

        r_xyz[i] = np.array([tckr_x[i](time[0]), tckr_y[i](time[0]),
                             tckr_z[i](time[0])])

    # Volumetric tensor with SC1 as center.
    dr_mat = (np.vstack(r_xyz[1:]) - np.tile(r_xyz[0], (3, 1))).T

    return dr_mat

def c_4_v(r_xyz, time):
    r"""Calculates velocity or time shift of discontinuity as in [6]_.

    Parameters
    ----------
    r_xyz : list
        Time series of the positions of the 4 spacecraft.
    time : list
        Crossing times or time and velocity.

    Returns
    -------
    out : ndarray
        Discontinuity velocity or time shift with respect to mms1.

    References
    ----------
    .. [6]	Vogt, J., Haaland, S., and Paschmann, G. (2011) Accuracy
            of multi-point boundary crossing time analysis, Ann.
            Geophys., 29, 2239-2252, doi :
            https://doi.org/10.5194/angeo-29-2239-2011

    """

    if isinstance(time, np.ndarray) and time.dtype == np.datetime64:
        flag = "v_from_t"

        time = time.view("i8") * 1e-9
    elif time[1] > 299792.458:
        flag = "v_from_t"
    else:
        flag = "dt_from_v"

    if flag.lower() == "v_from_t":
        # Time input, velocity output
        dr_mat = _get_vol_ten(r_xyz, time)
        tau = np.array(time[1:]) - time[0]
        slowness = np.linalg.solve(dr_mat, tau)

        # "1/v vector"
        out = slowness / np.linalg.norm(slowness) ** 2

    elif flag.lower() == "dt_from_v":
        # Time and velocity input, time output
        time_center = time[0]  # center time
        velocity = np.array(time[1:])  # Input velocity
        slowness = velocity / np.linalg.norm(velocity) ** 2

        dr_mat = _get_vol_ten(r_xyz, time_center)

        delta_t = np.matmul(dr_mat, slowness)
        out = np.hstack([0, delta_t])

    else:
        raise ValueError("Invalid flag")

    return out

ext: py

max_stars_repo_name: ablotekar/irfu-python

max_issues_repo_name: ablotekar/irfu-python

max_forks_repo_name: ablotekar/irfu-python

reason: import numpy,from scipy


text:
"""
Utility functions for parsing data files. Especially designed for the CSVs produced
by trial_util.py
"""
import csv
import datetime
import os

import numpy as np

from common import render_exception

def lookup_data_file(data_prefix, filename):
    full_name = os.path.join(data_prefix, filename)
    if not os.path.exists(full_name):
        raise Exception('Could not find "{}"'.format(filename))
    return full_name

def compute_summary_stats(data, trait_name, trait_values, average_key='time', is_numeric=False):
    """
    Expects the data to be a list of dicts.
    For each entry r in the data and each value v in trait_values,
    this function assembles the values of all fields
    r[trait_key] where r[trait_name] == v.
    Returns the mean, median, and std dev of all values in a
    dict with fields "mean", "median", and "std"

    is_numeric: Whether the trait is numeric or not (expects string by default)
    """
    vals = []
    for value in trait_values:
        def filter_func(r):
            if is_numeric:
                return int(r[trait_name]) == value
            return r[trait_name] == value
        vals += list(map(lambda r: float(r[average_key]),
                         filter(filter_func, data)))
    return {
        'mean': np.mean(vals),
        'median': np.median(vals),
        'std': np.std(vals)
    }

def summarize_over_reps(data, num_reps):
    return compute_summary_stats(data, 'rep', range(num_reps), is_numeric=True)

def obtain_data_rows(data_dir, framework, task_name, parameter_names, params_to_match):
    """
    Returns all data rows from the given framework from the
    given task where the specified parameters match.

    params_to_match as a dictionary {param names => value to match}
    """
    filename = lookup_data_file(data_dir, '{}-{}.csv'.format(framework, task_name))
    with open(filename, newline='') as csvfile:
        # even though it seems redundant, parameter names does
        # need to be a separate arg because *order matters*
        # whereas it doesn't in a dict
        fieldnames = parameter_names + ['rep', 'run', 'time']
        reader = csv.DictReader(csvfile, fieldnames)

        def filter_func(row):
            for (name, value) in params_to_match.items():
                comp = value
                if not isinstance(value, str):
                    comp = str(value)
                if row[name] != comp:
                    return False
            return True
        return list(filter(filter_func, reader))

def trials_stat_summary(data_dir, framework, task_name,
                        num_reps, parameter_names, params_to_match):
    """
    Returns a full summary of statistics on the specified framework
    and task across all reps where the specified parameters match.

    Returns (summary, success, message)
    """
    try:
        data = obtain_data_rows(data_dir, framework, task_name, parameter_names, params_to_match)
        summary = summarize_over_reps(data, num_reps)
        return (summary, True, 'success')
    except Exception as e:
        return (-1, False,
                'Encountered exception on {}, {} using params {}:\n{}'.format(
                    framework, task_name, params_to_match,
                    render_exception(e)))

def add_detailed_summary(report, detailed_summary, *fields):
    """
    Nasty hack provided for including a more detailed statistical
    summary in analysis files. The main reason this is here is
    because old reports contained only the mean and the graphing,
    etc., made assumptions about the layout of records.
    Eventually the old records should be migrated.
    """
    current = report
    all_fields = ['detailed', *fields]
    for field in all_fields[:-1]:
        if field not in current:
            current[field] = {}
        current = current[field]
    current[all_fields[-1]] = detailed_summary

max_stars_repo_name: carol-hsu/relay-bench

max_issues_repo_name: carol-hsu/relay-bench

max_forks_repo_name: carol-hsu/relay-bench

reason: import numpy


text:
"""
Copyright 2021 Google LLC

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

     https://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

"""Unit tests for RobotComponent and RobotGroupConfig."""

from typing import Any

from absl.testing import absltest
import numpy as np

from adept_envs.components.robot.robot import ControlMode, RobotComponent
from adept_envs.utils.testing.mock_sim_scene import MockSimScene

class RobotComponentTest(absltest.TestCase):
    """Unit test class for RobotComponent."""

    def test_get_state(self):
        """Tests querying the state of multiple groups."""
        sim_scene = MockSimScene(nq=10)  # type: Any
        sim_scene.data.qpos[:] = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
        sim_scene.data.qvel[:] = [11, 12, 13, 14, 15, 16, 17, 18, 19, 20]
        robot = RobotComponent(
            sim_scene,
            groups={
                'a': {
                    'qpos_indices': [0, 1, 3, 5],
                },
                'b': {
                    'qpos_indices': [2, 6],
                    'qvel_indices': [7, 8, 9],
                },
            })
        a_state, b_state = robot.get_state(['a', 'b'])
        np.testing.assert_array_equal(a_state.qpos, [1, 2, 4, 6])
        np.testing.assert_array_equal(a_state.qvel, [11, 12, 14, 16])
        np.testing.assert_array_equal(b_state.qpos, [3, 7])
        np.testing.assert_array_equal(b_state.qvel, [18, 19, 20])

    def test_step(self):
        """Tests stepping with an action for multiple groups."""
        sim_scene = MockSimScene(nq=10, ctrl_range=[-1, 1])  # type: Any
        robot = RobotComponent(
            sim_scene,
            groups={
                'a': {
                    'qpos_indices': [0, 1, 3, 5],
                },
                'b': {
                    'actuator_indices': [7, 8, 9],
                },
            })
        robot.step({
            'a': np.array([.2, .4, .6, .8]),
            'b': np.array([.1, .3, .5])
        })
        np.testing.assert_allclose(sim_scene.data.ctrl,
                                   [.2, .4, 0, .6, 0, .8, 0, .1, .3, .5])

        self.assertEqual(sim_scene.advance.call_count, 1)

    def test_step_denormalize(self):
        """Tests denormalizing the actions to the sim control range."""
        sim_scene = MockSimScene(nq=5, ctrl_range=[0, 10])  # type: Any
        robot = RobotComponent(
            sim_scene, groups={'a': {
                'qpos_indices': [0, 1, 2, 3, 4],
            }})
        robot.step({
            'a': np.array([-1, 1, -0.5, 0.5, 0]),
        })
        np.testing.assert_allclose(sim_scene.data.ctrl, [0, 10, 2.5, 7.5, 5])

    def test_step_position_control_bounds(self):
        """Tests action clamping when doing position control."""
        sim_scene = MockSimScene(nq=5, ctrl_range=[-1, 1])  # type: Any
        sim_scene.data.qpos[:] = [-0.4, -0.2, 0, 0.2, 0.4]
        robot = RobotComponent(
            sim_scene,
            groups={
                'a': {
                    'qpos_indices': [0, 1, 2, 3, 4],
                    'qpos_range': [(-0.5, 0.5)] * 5,
                    'qvel_range': [(-0.2, 0.2)] * 5,
                }
            })
        robot.step({'a': np.array([-1, -1, 0.2, 1, 1])})
        np.testing.assert_allclose(sim_scene.data.ctrl,
                                   [-0.5, -0.4, 0.1, 0.4, 0.5])

    def test_step_velocity_control_bounds(self):
        """Tests action clamping when doing velocity control."""
        sim_scene = MockSimScene(nq=3, ctrl_range=[-10, 10])  # type: Any
        robot = RobotComponent(
            sim_scene,
            groups={
                'a': {
                    'control_mode': ControlMode.JOINT_VELOCITY,
                    'qpos_indices': [0, 1, 2],
                    'qvel_range': [(-2, 2), (-1, 5), (-7, -4)],
                }
            })
        robot.step({'a': np.array([-0.5, 1, -1])})
        np.testing.assert_allclose(sim_scene.data.ctrl, [-1, 5, -7])

if __name__ == '__main__':
    absltest.main()

max_issues_repo_name: google-research/DBAP-simulation


text:
"""Tests for my attempt at calculating Accuracy Score"""

import unittest

import numpy as np

from metrics import accuracy

class AccuracyScoreTest(unittest.TestCase):
    """Tests for my attempt at calculating Accuracy Score"""

    def test_accuracy_score(self):
        """Test with some trivial but real data"""

        # Yuan 2004 Iris
        conf_matrix = np.array([
            [0, 0, 50],
            [3, 47, 0],
            [36, 14, 0]
            ])

        self.assertAlmostEqual(
            accuracy.from_matrix(conf_matrix),
            0.886667,
            places=6)

        # Yuan 2004 Wine
        conf_matrix = np.array([
            [31, 1, 27],
            [7, 64, 0],
            [11, 37, 0]
            ])

        self.assertAlmostEqual(
            accuracy.from_matrix(conf_matrix),
            0.685393,
            places=6)

max_stars_repo_name: simonharris/pykmeans

max_issues_repo_name: simonharris/pycluster

max_forks_repo_name: simonharris/pycluster


text:
from pytomo.inversion.cmc import ConstrainedMonteCarlo, InputFile
from pytomo.work.jp.params import get_model, get_dataset
from pytomo.inversion.inversionresult import InversionResult
from dsmpy.seismicmodel import SeismicModel
from dsmpy.modelparameters import ModelParameters, ParameterType
import numpy as np
import time
import matplotlib.pyplot as plt
from dsmpy.event import Event
from dsmpy.station import Station
from dsmpy.utils.cmtcatalog import read_catalog
from dsmpy.dataset import Dataset
from dsmpy.dsm import PyDSMInput, compute, compute_models_parallel
from dsmpy.windowmaker import WindowMaker
from dsmpy.component import Component
from mpi4py import MPI
import sys

if __name__ == '__main__':
    comm = MPI.COMM_WORLD
    n_cores = comm.Get_size()
    rank = comm.Get_rank()

    input_file = InputFile(sys.argv[1])
    params = input_file.read()

    tlen = params['tlen']
    nspc = params['nspc']
    sampling_hz = params['sampling_hz']
    verbose = params['verbose']
    result_path = params['result_path']
    n_mod = params['n_mod']
    n_block = params['n_block']
    mode = params['mode']
    freq = params['freq']
    freq2 = params['freq2']
    filter_type = params['filter_type']

    n_pass = n_mod // n_block
    assert n_mod % n_block == 0
    if verbose > 1:
        print('n_pass={}'.format(n_pass))

    if rank == 0:
        n_upper_mantle = 20
        n_mtz = 10
        n_lower_mantle = 12
        types = [ParameterType.VSH]
        g = 1.
        l = 1.
        seed = 42

        model_ref, model_params = get_model(
            n_upper_mantle, n_mtz, n_lower_mantle, types, verbose=verbose)

        cov = ConstrainedMonteCarlo.smooth_damp_cov(
            model_params._n_nodes, g, l)

        cmc = ConstrainedMonteCarlo(
            model_ref, model_params, cov,
            mesh_type='boxcar', seed=seed)
        models = cmc.sample_models(n_mod)
    else:
        model_params = None
        models = None
        model_ref = None

    dataset = get_dataset(tlen, nspc, sampling_hz, mode=mode)

    if filter_type is not None:
        dataset.filter(freq, freq2, filter_type)

    if rank == 0:
        windows_S = WindowMaker.windows_from_dataset(
            dataset, 'prem', ['s', 'S'],
            [Component.T], t_before=30., t_after=50.)
        windows_P = WindowMaker.windows_from_dataset(
            dataset, 'prem', ['p', 'P'],
            [Component.Z], t_before=30., t_after=50.)
        windows = windows_S #+ windows_P
    else:
        windows = None

    ipass = 0
    misfit_dict = None
    if rank == 0:
        start_time = time.time_ns()
    while ipass < n_pass:
        if rank == 0:
            start = ipass * n_block
            end = start + n_block
            current_models = models[start:end]
        else:
            current_models = None
        outputs = compute_models_parallel(
            dataset, current_models, tlen, nspc, sampling_hz,
            comm, mode=mode)

        if rank == 0:
            if filter_type is not None:
                for imod in range(len(outputs)):
                    for iev in range(len(outputs[0])):
                        outputs[imod][iev].filter(
                            freq, freq2, filter_type)

            if misfit_dict is None:
                misfit_dict = process_outputs(
                    outputs, dataset, current_models, windows)
                result = InversionResult(
                    dataset, current_models, windows, misfit_dict)
            else:
                result.add_result(current_models, misfit_dict)

        ipass += 1
        comm.Barrier()

    if rank == 0:
        end_time = time.time_ns()
        result.save(result_path)
        if verbose > 0:
            print('Models and misfits computation done in {} s'
                .format((end_time-start_time) * 1e-9))
            print('Results saved to \'{}\''.format(result_path))

        # fig, ax = dataset.plot_event(
        #     0, windows, component=Component.T,
        #     align_zero=True, color='black')
        # cycler = plt.rcParams['axes.prop_cycle']
        # for imod, sty in enumerate(cycler[:n_block]):
        #     _, ax = outputs[imod][0].plot_component(
        #         Component.T, windows, ax=ax, align_zero=True, **sty)
        # plt.show()

        # fig, ax = result.plot_models(types=[ParameterType.VSH])
        # ax.set(ylim=[model_params._radii[0]-100, 6371])
        # plt.show()

max_stars_repo_name: afeborgeaud/pytomo

max_issues_repo_name: afeborgeaud/pytomo

max_forks_repo_name: afeborgeaud/pytomo


text:
import numpy as np
import tensorflow as tf
from tqdm.auto import tqdm

def get_pod_bases(U, eps=1e-10, eps_init_step=None):
    if eps_init_step is not None:
        eps_init_step = 1e-8
        print("Performing initial time-trajectory POD")
        # Number of DOFs
        n_h = U.shape[0]
        # Number of snapshots n_s x Number of space nodes (n_x * n_y * ...)
        n_s = U.shape[-1]
        T_raw = []

        # Init at the max it can be, n_t
        n_L_init = U.shape[1]
        T = np.zeros((n_h, n_L_init, n_s))
        for k in range(n_s):
            T_k = perform_pod(U[:, :, k], eps_init_step, verbose=False)
            if T_k.shape[1] < n_L_init:
                n_L_init = T_k.shape[1]
            T[:, :n_L_init, k] = T_k[:, :n_L_init]

        # Cropping the results accordingly and stacking
        T = T[:, :n_L_init, :]
        # Reshaping the 3d-mat into a 2d-mat
        U = np.reshape(T, (n_h, n_s*n_L_init))

    print("Performing SVD...")
    return perform_pod(U, eps)

def perform_pod(U, eps, verbose=True):
    # Number of DOFs
    n_h = U.shape[0]

    # Number of snapshots n_s x Number of time steps
    n_st = U.shape[1]

    # SVD algoritm call
    W, D, ZT = np.linalg.svd(U, full_matrices=False)

    # Getting MATLAB-like orientation
    Z = ZT.T

    # Storing eigenvalues and their sum
    lambdas = D**2
    sum_lambdas = np.sum(lambdas)

    # Finding n_L
    n_L = 0
    sum_lambdas_trunc = 0.
    for i in range(n_st):
        sum_lambdas_trunc += lambdas[i]
        n_L += 1
        if sum_lambdas_trunc/sum_lambdas >= (1 - eps):
            break

    # Truncating according to n_L
    lambdas_trunc = lambdas[0:n_L]

    if verbose:
        print("Contructing the reduced bases V...")
    V = np.zeros((n_h, n_L))
    for i in tqdm(range(n_L), disable=(not verbose)):
        V[:, i] = U.dot(Z[:, i]) / np.sqrt(lambdas_trunc[i])

    return V

max_stars_repo_name: BJUT-SIPL/POD-NN

max_issues_repo_name: BJUT-SIPL/POD-NN

max_forks_repo_name: BJUT-SIPL/POD-NN


text:
import numpy as np
import matplotlib.pyplot as plt
from math import sin,cos,pi,sqrt
import pandas as pd

class A:
	def __init__(self,nodes,angles,rays):
		self.n = nodes
		self.m = angles
		self.nr = rays
		self.A = np.zeros([n*m,n*n])

		#make grid
		self.bounds = np.linspace(-n/2,n/2,n+1)

		self.grd_x,self.grd_y = np.meshgrid(self.bounds,self.bounds)

		length = len(self.bounds)
		self.arr_x =[]
		self.arr_y = []
		y_0 = np.array(range(n))
		x_0 = np.arange(0,n*n,n)
		for i in range(length-1):
			y_vec = y_0+i*n
			x_vec = x_0 + i
			self.arr_x.append(list(x_vec))
			self.arr_y.append(list(y_vec))

	def discretize_rays(self,theta):
		rotation_M = [[sin(theta),cos(theta)],[-sin(theta),cos(theta)]]
		X = np.asarray(self.grd_x).ravel()
		Y = np.asarray(self.grd_y).ravel()
		M = [list(X),list(Y)]
		MR = np.matmul(rotation_M,M)

		lower = min(MR[1])
		upper = max(MR[1])

		rays = np.linspace(lower,upper,self.nr)#[1:self.nr+1]

		discretized = np.linspace(sqrt(2)*min(MR[0]),sqrt(2)*max(MR[0]),self.n**3)#[1:self.n**3-1]
		return(rays,discretized)

	def pointToBox(self,x,y,theta,verbose=False):
		th = 2*pi-theta

		#invert rotation by rotating by 2pi-theta
		irot_M = [[cos(th),sin(th)],[-sin(th),cos(th)]]

		X,Y = np.matmul(irot_M,[x,y])

		X_ind = np.where(X <= self.bounds)[0]
		Y_ind = np.where(Y <= self.bounds)[0]

		if verbose:
			print(X_ind)
			print(Y_ind)

		if len(X_ind)==0 or len(Y_ind)==0:
			box = None
		else:
			index_X = min(X_ind)
			index_Y = min(Y_ind)
			if (index_X == 0 and not X==min(self.bounds)) or (index_Y==0 and not Y == min(self.bounds)):  #out of bounds of grid
				box = None
			else:
				X_list = self.arr_x[index_X-1]
				Y_list = self.arr_y[index_Y-1]

				box = list(set(X_list)&set(Y_list))[0]

		return(box)

	def scan_data(self):

		angles = np.linspace(0,pi,self.m+1)[0:self.m]
		i=0
		rays,discretized = self.discretize_rays(0)

		j=1

		for theta in angles:
			print("Run "+str(j)+ " of "+str(self.m))
			j+=1
			for ray in rays:
				boxes = [0]*(n*n)
				for x in discretized:
					box = self.pointToBox(x,ray,theta)
					try:
						boxes[box]+=1
					except(TypeError):
						#print([x,ray])
						#self.pointToBox(x,ray,theta,verbose=True)
						pass
				maximum = max(boxes)
				self.A[i,:] = [a/maximum for a in boxes]
				i+=1
		print(self.A)

		df = pd.DataFrame(self.A)
		df.to_csv("A.csv",header=False,index=False)

if __name__ == "__main__":
	n = 62
	m = 66
	nr = 60

	scan = A(n,m,nr)
	scan.scan_data()

max_stars_repo_name: ShadowWarden/leastsquaresinversion

max_issues_repo_name: ShadowWarden/leastsquaresinversion

max_forks_repo_name: ShadowWarden/leastsquaresinversion


text:
import os
import copy
import torch
import imageio
from typing import Sequence
import matplotlib.pyplot as plt
from torchvision.utils import make_grid, save_image
import numpy as np

from pytorch_lightning.callbacks import Callback
from pytorch_lightning.utilities import rank_zero_only
from pytorch_lightning import LightningModule, Trainer

from models.backprojection_utils import backproject
from models.model_utils import collapse_trajectory_dim, expand_trajectory_dim, resize_trajectory
from utils.utils import exclusive_mean

class GSNVizCallback(Callback):
    def __init__(self, log_dir, voxel_res, voxel_size) -> None:
        super().__init__()
        self.viz_dir = os.path.join(log_dir, 'viz')
        if not os.path.exists(self.viz_dir):
            os.makedirs(self.viz_dir)
        self.voxel_res = voxel_res
        self.voxel_size = voxel_size

    @rank_zero_only
    def on_validation_batch_end(
        self,
        trainer: Trainer,
        pl_module: LightningModule,
        outputs: Sequence,
        batch: Sequence,
        batch_idx: int,
        dataloader_idx: int,
    ) -> None:

        if batch_idx == 0:
            T = 8

            filename = 'real_samples_epoch_{:06d}.png'.format(trainer.current_epoch)
            sample_static(self.viz_dir, filename, collapse_trajectory_dim(batch['rgb'].clone()), nrow=T)

            filename = 'real_floorplan_epoch_{:06d}.png'.format(trainer.current_epoch)
            sample_floorplans(self.viz_dir, filename, self.voxel_res, self.voxel_size, batch)

            filename = 'real_trajectories_epoch_{:06d}.png'.format(trainer.current_epoch)
            sample_trajectories(self.viz_dir, filename, self.voxel_res, self.voxel_size, batch['Rt'].clone())

            del batch['Rt']  # remove Rt so that trajectory sampler fills it in
            for k in batch.keys():
                batch[k] = batch[k].cuda()
            with torch.no_grad():
                z = torch.rand(batch['K'].shape[0], pl_module.z_dim, device=batch['K'].device)
                y_hat_rgb, y_hat_depth, Rt, K = pl_module(z, batch)

            filename = 'fake_samples_epoch_{:06d}.png'.format(trainer.current_epoch)
            sample_static(self.viz_dir, filename, collapse_trajectory_dim(y_hat_rgb), nrow=T)

            fake_batch = {'rgb': y_hat_rgb, 'depth': y_hat_depth, 'Rt': Rt, 'K': K}
            filename = 'fake_floorplan_epoch_{:06d}.png'.format(trainer.current_epoch)
            sample_floorplans(self.viz_dir, filename, self.voxel_res, self.voxel_size, fake_batch)

            filename = 'fake_trajectories_epoch_{:06d}.png'.format(trainer.current_epoch)
            sample_trajectories(self.viz_dir, filename, self.voxel_res, self.voxel_size, Rt)

def sample_static(viz_dir, filename, rgb, nrow=4):
    # save trajectories like a film strip
    filepath = os.path.join(viz_dir, filename)
    save_image(rgb, fp=filepath, nrow=nrow, range=(0, 1))

def sample_video(sample_dir, filename, rgb, n_trajectories=64, frame_duration=0.2):
    video_filepath = os.path.join(sample_dir, 'video')
    if not os.path.exists(video_filepath):
        os.mkdir(video_filepath)
    filepath = os.path.join(video_filepath, filename)

    frames = []
    for i in range(rgb.shape[1]):
        grid = make_grid(rgb[:n_trajectories, i], nrow=int(np.sqrt(n_trajectories)))
        grid = grid.permute(1, 2, 0) * 255
        grid = grid.cpu().numpy().astype(np.uint8)

        frames.append(grid)
    imageio.mimsave(filepath, frames, duration=frame_duration)

def sample_video_single(sample_dir, filename, rgb_in, rgb_pred, frame_duration=0.2):
    video_filepath = os.path.join(sample_dir, 'video')
    if not os.path.exists(video_filepath):
        os.mkdir(video_filepath)
    filepath = os.path.join(video_filepath, filename)

    frames = []
    for frame_in, frame_pred in zip(rgb_in, rgb_pred):
        grid = make_grid([frame_in, frame_pred], nrow=2)
        grid = grid.permute(1, 2, 0) * 255
        grid = grid.cpu().numpy().astype(np.uint8)
        frames.append(grid)

    imageio.mimsave(filepath, frames, duration=frame_duration)

def get_floorplans(data, voxel_res, voxel_size, floorplan_res, batch_size=8):
    B, T, C, H, W = data['rgb'].shape

    volume_length = voxel_res * voxel_size  # recalculate original length of volume
    voxel_size = volume_length / floorplan_res  # scale to new resolution

    if data['depth'].shape[3] != H:
        data['depth'] = resize_trajectory(x=data['depth'], size=H)

    floorplans = []
    for i in range(0, B, batch_size):
        with torch.no_grad():
            volume = backproject(
                voxel_dim=(floorplan_res, floorplan_res, floorplan_res),
                voxel_size=voxel_size,  # should be roughly equivalent to (self.opt.far * 2 / vd)
                world_center=(0, 0, 0),
                Rt=collapse_trajectory_dim(data['Rt'][i : i + batch_size]),
                K=collapse_trajectory_dim(data['K'][i : i + batch_size]),
                features=collapse_trajectory_dim(data['rgb'][i : i + batch_size]),
                depth=collapse_trajectory_dim(data['depth'][i : i + batch_size]),
            )

            volume = expand_trajectory_dim(volume, T=T)
            volume = exclusive_mean(volume, dim=1)  # merge along trajectory dim

            # remove the top half of the scene (i.e the ceiling)
            height = volume.shape[3]
            volume = volume[:, :, :, : height // 2, :]

            # take the first nonzero pixel, as if looking down from a bird's-eye view
            depth_idx = torch.argmax((volume > 0).float(), dim=3, keepdim=True)
            floorplan = torch.gather(volume, 3, depth_idx).squeeze(3)
            floorplans.append(floorplan.cpu())

    floorplans = torch.cat(floorplans, dim=0)
    floorplans = floorplans.permute(0, 1, 3, 2)  # swap height and width dimensions
    return floorplans

def sample_floorplans(viz_dir, filename, voxel_res, voxel_size, data, floorplan_res=200, n_trajectories=8):
    filepath = os.path.join(viz_dir, filename)

    data = copy.deepcopy(data)  # copy so we don't mess up the original dictionary
    data['rgb'] = data['rgb'][:n_trajectories]
    data['depth'] = data['depth'][:n_trajectories]
    data['Rt'] = data['Rt'][:n_trajectories]
    data['K'] = data['K'][:n_trajectories]

    floorplans = get_floorplans(data, voxel_res=voxel_res, voxel_size=voxel_size, floorplan_res=floorplan_res)

    save_image(floorplans, fp=filepath, nrow=2, range=(0, 1), pad_value=1)

def sample_trajectories(viz_dir, filename, voxel_res, voxel_size, Rts):
    filepath = os.path.join(viz_dir, filename)

    xyz = Rts.inverse()[:, :, 0:3, 3].cpu().numpy()
    for i in range(len(xyz)):
        plt.plot(xyz[i, :, 0], xyz[i, :, 2], c='blue', alpha=0.1, linewidth=3)
        plt.scatter(xyz[i, :, 0], xyz[i, :, 2], c='blue', alpha=0.1, linewidth=3)
    plt.xlabel('X')
    plt.ylabel('Z')
    extents = voxel_res * voxel_size / 2
    plt.xlim(-extents, extents)
    plt.ylim(-extents, extents)
    plt.savefig(filepath)
    plt.close()

max_stars_repo_name: tjtanaa/ml-gsn

max_issues_repo_name: tjtanaa/ml-gsn

max_forks_repo_name: tjtanaa/ml-gsn


text:
import argparse
import os
import shutil

import h5py
# import librosa
import numpy as np
from tqdm import tqdm
from scipy.signal import decimate
from scipy import interpolate
from scipy import signal
from scipy.signal import butter, lfilter, freqz, wiener
from scipy.io import wavfile
import multiprocessing as mp
import utils
from facealigner import FaceAligner, crop_image
from helper import shape_to_np
import cv2
import dlib
import math
from multiprocessing import Process, Queue, Manager
import subprocess
import random as rn
from random import randint, shuffle
import time
from collections import deque
from scipy.interpolate import CubicSpline
#-----------------------------------------#
#           Reproducible results          #
#-----------------------------------------#
os.environ['PYTHONHASHSEED'] = '123'
np.random.seed(123)
rn.seed(123)
# torch.manual_seed(999)
#-----------------------------------------#
#----------------PARSER------------------#
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("-i", "--input-folder", type=str, help='Path to folder that contains video files')
parser.add_argument("-p", "--pred-path", type=str, help="Predictor Path", default='data/dlib_data/shape_predictor_68_face_landmarks.dat')
parser.add_argument("-ti", "--temp-img", type=str, help="Template face image path", default='data_prep/template_face.jpg')
parser.add_argument("-fs", "--fs", type=int, help='sampling rate', default=8000)
parser.add_argument("-nw", "--nw", type=int, help='num workers', default=1)
parser.add_argument("-m", "--mode", type=int, help='Mode 0: Aligns all frames, Mode 1: Uses first frame to align all frames', default=1)
parser.add_argument("-d", "--debug", type=bool, help='Writes videos for debug purposes', default=False)
parser.add_argument("-o", "--out-path", type=str, help='Output path')
args = parser.parse_args()

in_path = args.input_folder
fs = args.fs
out_path = args.out_path
TI_path = args.temp_img
num_processes = args.nw
print("num_processes: ", num_processes)
predictor_path = args.pred_path

if not os.path.exists(out_path):
    os.makedirs(out_path)
else:
    shutil.rmtree(out_path)
    os.mkdir(out_path)

class TemplateProcessor():
    def __init__(self, path):
        template_I = cv2.imread(path)
        detector = dlib.get_frontal_face_detector()
        predictor = dlib.shape_predictor(predictor_path)
        fa = FaceAligner(predictor, desiredLeftEye=(0.25, 0.25), desiredFaceWidth=128)
        gray = cv2.cvtColor(template_I, cv2.COLOR_BGR2GRAY)

        dets = detector(template_I, 1)
        for k, d in enumerate(dets):
            shape = predictor(gray, d)

        template_I, scale = fa.align(template_I, gray, d, shape, None)

        gray = cv2.cvtColor(template_I, cv2.COLOR_BGR2GRAY)

        dets = detector(template_I, 1)
        for k, d in enumerate(dets):
            shape = predictor(gray, d)

        template_I = cv2.cvtColor(template_I, cv2.COLOR_BGR2RGB)
        utils.easy_show(template_I, 'template_face.png')

        scl = 0.6

        shape = shape_to_np(shape)
        shape -= np.tile(np.min(shape, axis=0), [68, 1])
        shape = ((128*(shape/np.tile(np.max(shape, axis=0), [68, 1])))*scl) + 128*((1-scl)/2)

        self.points = shape

    def getShape(self):
        return self.points

class Extractor():
    def __init__(self, in_path, mean_shape):
        self.dataQueue = Manager().Queue()
        self.fileList = []
        self.mean_shape = mean_shape
        for root, dirnames, filenames in os.walk(in_path):
            for filename in filenames:
                if os.path.splitext(filename)[1] == '.mp4' or os.path.splitext(filename)[1] == '.mpg' or os.path.splitext(filename)[1] == '.mov' or os.path.splitext(filename)[1] == '.flv':
                    self.fileList.append((root, filename))

        self.fileList = self.fileList[:]
        print(self.fileList)
        # exit()

    def processSample(self, process_id):
            import librosa
            n = len(self.fileList)
            increment = n // num_processes
            if process_id == num_processes-1:
                sampleList = list(range(process_id*increment, n))
            else:
                sampleList = list(range(process_id*increment, (process_id+1)*increment))
            print(process_id, len(sampleList), n,  process_id*increment, (process_id+1)*increment)

            # exit()

            detector = dlib.get_frontal_face_detector()
            predictor = dlib.shape_predictor(predictor_path)
            desired_dim = 128
            fa = FaceAligner(predictor, desiredLeftEye=(0.25, 0.25), desiredFaceWidth=desired_dim)

            shuffle(sampleList)

            for j in tqdm(sampleList):
                print(j)
                root, filename = self.fileList[j]

                y, sr = librosa.load(os.path.join(root, filename), sr=fs)

                try:
                    y = y-np.mean(y)
                    speech = y/np.max(np.abs(y))
                    if np.isnan(speech).any():
                        print('NaN encountered! Skipping file...')
                        continue
                except:
                    print('Exception! Skipping file...')
                    continue

                frame_list = []
                lmarks_list = []
                try:
                    cap = cv2.VideoCapture(os.path.join(root, filename))
                    length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
                    print("Length: ", length)
                    dset = h5py.File(os.path.join(out_path, os.path.splitext(filename)[0]) + '.hdf5', 'w')
                    frame_cnt = 0
                    scale = None
                    video = []
                    tmp_lmarks = []
                    for frame_cnt in range(length):
                        ret, frame = cap.read()
                        if frame is None:
                            break
                        video.append(frame)
                        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

                        dets = detector(frame, 1)
                        for k, d in enumerate(dets):
                            tmp_lmarks.append(shape_to_np(predictor(gray, d)))

                    lmarks = np.array(tmp_lmarks)

                    if args.mode == 1:
                        fa.get_tform(video[0], lmarks[0, ...], self.mean_shape, scale)

                    for frame_cnt in range(length):
                        # frame, scale = fa.align_box(video[frame_cnt], lmarks[frame_cnt, ...], self.mean_shape, scale)
                        # frame, scale = fa.align_three_points(frame, np.average(np.array(buffer), axis=0, weights=[x/sum(list(range(buffersize))) for x in range(buffersize)]), self.mean_shape, scale)
                        if args.mode == 1:
                            frame, scale = fa.apply_tform(video[frame_cnt])
                        else:
                            frame, scale = fa.align_three_points(video[frame_cnt], lmarks[frame_cnt, ...], self.mean_shape, scale)

                        frame_list.append(frame)

                        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
                        dets = detector(frame, 1)
                        for k, d in enumerate(dets):
                            shape = shape_to_np(predictor(gray, d))
                        lmarks_list.append(shape)

                        frame_cnt+=1

                    dset.create_dataset('speech', data=speech)
                    dset.create_dataset('video', data=np.array(frame_list))
                    dset.create_dataset('lmarks', data=np.array(lmarks_list))
                    self.dataQueue.put(('dummy', None))
                except:
                    os.remove(os.path.join(out_path, os.path.splitext(filename)[0]) + '.hdf5')
                    print('Exception! Deleting file...', os.path.join(out_path, os.path.splitext(filename)[0]) + '.hdf5')
                    continue
                # utils.write_video_cv(np.array(frame_list), speech, 8000, '', '{}_test.mp5'.format(j), 25.0)
                # if j > 9:
                #     exit()
                if args.debug:
                    utils.write_video_cv(np.array(frame_list), speech, 8000, out_path, os.path.join(os.path.splitext(filename)[0] + '.mp4'), 25.0)

            self.dataQueue.put(('end', process_id))
            print('Thread Ended #', process_id)

    def writeToFile(self):
        cnt = 0
        threadStatus = [0] * num_processes
        pbar = tqdm(total = len(self.fileList))
        while True:
            if all(threadStatus):
                break
            data = self.dataQueue.get()
            if isinstance(data[0], str) and data[0] == 'end':
                print('End ', data[1])
                threadStatus[data[1]] = 1
                continue
            # speech_dset.create_dataset(data[0]+'_'+data[1], data=data[3])
            # image_seq_dset.create_dataset(data[0]+'_'+data[1], data=data[2])
            cnt += 1
            pbar.update(1)
            # utils.write_video(data[2], data[3], 8000, out_path, data[0]+'_'+data[1], 29.97)
        print('Main Ended.')
        # dset.close()
        pbar.close()

if __name__ == '__main__':
    TI_process = TemplateProcessor(TI_path)
    mean_shape = TI_process.getShape()
    del TI_process

    cnt = 0
    ext = Extractor(in_path, mean_shape)

    if num_processes < 2:
        ext.processSample(0)
    else:
        processes = []
        for i in range(num_processes):
            processes.append(Process(target=ext.processSample, args=(i, )))

        for i, p in enumerate(processes):
            # p.daemon = True
            p.start()
            print('Process #', i)

        p = Process(target=ext.writeToFile, args=())
        p.start()
        p.join()
        print('Main joined.')

        for i, p in enumerate(processes):
            p.join()
            print('Joined #', i)

max_stars_repo_name: fjbriones/emotalkingface

max_issues_repo_name: ramizf/emotalkingface

max_forks_repo_name: ramizf/emotalkingface


text:
# python imports
import signal, time, os, sys, random
import pickle as pkl
import argparse as ap
from itertools import chain
from datetime import date
# numpy imports
import numpy as np
# torch imports
import torch
import torch.autograd as ta
import torch.nn.functional as F
import torch.nn as nn
from torch.utils import data
# code imports
import process_data as pd
import process_sentences as ps
import code.dataset as dt
import code.embedding_modules as em
import code.pooling_modules as pm
import code.classifier_modules as cm

models_path = os.path.join('.', 'models')
results_path = os.path.join('.', 'results')
batch_size = 32
num_workers = 16
max_epochs = None

dataset_params = { 'batch_size' : batch_size,
                   'shuffle' : True,
                   'num_workers' : num_workers }

model_params = { 'device' : torch.device('cpu'),
                 'sparse_grads' : True,
                 'label_hashing' : False,
                 'type_smoothing' : False,
                 'mask_rate' : 0.01,
                 'default_pc' : True,
                 'lr_decay' : False,
                 'att_aggr' : True,
                 'mha_heads' : 2,
                 'edge_type' : None,
                 'dep_gate' : True,
                 # learning rate
                 'lr' : 0.001,
                 'depth_cap' : None,
                 'aggr_type' : 'sum',
                 'dep_match_type' : 'label',
                 'dep_depth' : 0,
                 'dropout' : None,#0.2,
                 # node info
                 'node_ct' : 30000,
                 'node_emb_dim' : None,
                 'node_state_dim' : None,
                 'lstm_state_dim' : None,
                 'pretrained_emb_dim' : 50,
                 # edge info
                 'edge_ct' : 150,
                 'edge_emb_dim' : None,
                 'edge_state_dim' : None,
                 # module used to get initial node embeddings
                 'init_node_embedder' : em.MPNN,
                 # if init node embedder is DagLSTM
                 'init_node_embedder_acc_dir' : em.ACCNN.down_acc,
                 # if init node embedder is GCN or MPNN
                 'num_rounds' : 0,
                 # pooling type used for graph embedding
                 'pooling_module' : pm.SimpleMaxPool,
                 # to use leaf or root pooling
                 'pooling_dir' : em.ACCNN.up_acc }

if __name__ == '__main__':

    parser = ap.ArgumentParser(description='Train formula classifier module')
    parser.add_argument('--dataset', help='Dataset to train model for')
    parser.add_argument('--device', help='Device to use')
    parser.add_argument('--node_embedder', help='Node embedding module')
    parser.add_argument('--pool_type', help='Graph pooling type')
    parser.add_argument('--acc_dir', help='Node embedder update direction (if applicable, e.g., for DAG LSTM).')
    parser.add_argument('--pool_dir', help='Pooling direction (if applicable, e.g., for DAG LSTM)')
    parser.add_argument('--depth_cap', type=int, help='Maximum depth of formulas')
    parser.add_argument('--num_epochs', type=int, help='Epochs over training data')
    parser.add_argument('--num_rounds', type=int,
                        help='Number of rounds of updates (if applicable, e.g., for MPNN or GCN)')
    parser.add_argument('--sparse_grads', help='Use sparse grads')
    parser.add_argument('--type_smoothing',
                        help='Use type smoothing for unknown symbols')
    parser.add_argument('--default_pc', help='Dataset default training method')
    parser.add_argument('--aggr_type', help='Aggregation method')
    parser.add_argument('--save_model', help='Save model or not')
    parser.add_argument('--edge_type', help='Edge labeling strategy')
    parser.add_argument('--val_is_test', help='Use test set as validation')
    parser.add_argument('--small_network', help='Halve all dimensionalities')
    parser.add_argument('--dep_match_type', help='Dependent embedding match type')
    parser.add_argument('--dep_depth', help='Dependent embedding match depth', type=int)
    parser.add_argument('--batch_size', help='Batch size', type=int)
    parser.add_argument('--mha_heads', help='Number of multi-headed attention heads', type=int)

    args = parser.parse_args()

    val_is_test = False
    if args.val_is_test:
        assert args.val_is_test in ['True', 'False'], 'Validation as test should be True or False'
        val_is_test = args.val_is_test == 'True'

    if args.dep_match_type:
        v_m_types = ['label', 'alpha', 'all', 'type', 'leaf', 'iso',
                     'depth', 'depth_typed', 'leaf_label']
        assert args.dep_match_type in v_m_types, 'Dependent match type must be one of ' + ', '.join(v_m_types)
        model_params['dep_match_type'] = args.dep_match_type

    if args.dep_depth is not None:
        assert args.dep_depth >= 0, 'Non-negative values for dep_depth required'
        model_params['dep_depth'] = args.dep_depth

    if args.mha_heads is not None:
        assert args.mha_heads >= 1, 'Positive values for mha_heads required'
        model_params['mha_heads'] = args.mha_heads

    small_network = False
    if args.small_network:
        assert args.small_network in ['True', 'False'], 'Smallify network should be True or False'
        small_network = args.small_network == 'True'

    save_model = True
    if args.save_model:
        assert args.save_model in ['True', 'False'], 'Save model should be either True or False'
        save_model = args.save_model == 'True'

    if args.acc_dir:
        assert args.acc_dir in em.ACCNN.acc_dirs, 'Embedding direction must be one of: ' + \
            ', '.join(em.ACCNN.acc_dirs)
        model_params['init_node_embedder_acc_dir'] = args.acc_dir

    if args.edge_type:
        valid_edge_types = ['typed_ord', 'untyped_ord', 'typed_unord', 'untyped_unord', 'None']
        assert args.edge_type in valid_edge_types, 'Embedding direction must be one of: ' + \
            ', '.join(valid_edge_types)
        model_params['edge_type'] = None if args.edge_type == 'None' else args.edge_type

    if args.pool_dir:
        good_pool_dirs = [em.ACCNN.up_acc, em.ACCNN.down_acc]
        assert args.pool_dir in good_pool_dirs, \
            'Pooling direction must be one of: ' + ', '.join(good_pool_dirs)
        model_params['pooling_dir'] = args.pool_dir

    if args.aggr_type:
        good_aggr_types = ['mean', 'sum']
        assert args.aggr_type in good_aggr_types, \
            'Aggregation method must be one of: ' + ', '.join(good_aggr_types)
        model_params['aggr_type'] = args.aggr_type

    if args.sparse_grads:
        assert args.sparse_grads in ['True', 'False'], \
            'Sparse grads either True or False...'
        model_params['sparse_grads'] = args.sparse_grads == 'True'

    if args.type_smoothing:
        assert args.type_smoothing in ['True', 'False'], \
            'Type smoothing either True or False...'
        model_params['type_smoothing'] = args.type_smoothing == 'True'

    if args.device:
        model_params['device'] = torch.device(args.device)

    if args.pool_type:
        avail_pts = ['MaxPool', 'DepDagLSTM', 'DagLSTM']
        assert args.pool_type in avail_pts, \
            'Node pooling module not recognized. Option must be one of ' + \
            ', '.join(avail_pts)
        p_d = { 'MaxPool' : pm.SimpleMaxPool, 'DepDagLSTM' : pm.DepDagLSTMPool,
                'DagLSTM' : pm.DagLSTMPool }
        model_params['pooling_module'] = p_d[args.pool_type]

    if args.node_embedder:
        avail_ets = ['MPNN', 'DagLSTM', 'BidirDagLSTM', 'GCN']
        assert args.node_embedder in avail_ets, \
            'Node embedding module not recognized. Option must be one of ' + \
            ', '.join(avail_ets)
        ne_d = { 'MPNN' : em.MPNN, 'DagLSTM' : em.DagLSTM,
                 'BidirDagLSTM' : em.BidirDagLSTM, 'GCN' : em.GCN }
        model_params['init_node_embedder'] = ne_d[args.node_embedder]

    if args.num_rounds is not None:
        assert args.num_rounds >= 0, 'Round count must be at least 0...'
        model_params['num_rounds'] = args.num_rounds

    model_params['depth_cap'] = None
    if args.depth_cap is not None:
        assert args.depth_cap >= 0, 'Maximum formula depth must be at least 0...'
        model_params['depth_cap'] = args.depth_cap

    if args.default_pc:
        assert args.default_pc in ['True', 'False'], '--default_pc must be True or False'
        model_params['default_pc'] = args.default_pc == 'True'

    dataset = args.dataset
    assert dataset in ['holstep', 'mizar', 'mizar_cnf', 'scitail'], 'Unknown dataset...'
    model_params['pretrained_embs'] = dataset == 'scitail' and pd.emb_info_loc
    if dataset == 'holstep':
        div_emb_sz = 4 if small_network else 1
        tr_dobj = pd.hol_tr_dobj_loc
        val_dobj = pd.hol_val_dobj_loc
        te_dobj = pd.hol_te_dobj_loc
        tr_collator = dt.HolstepCollator(model_params['depth_cap'], model_params['default_pc'],
                                         model_params['edge_type'])
        te_collator = dt.HolstepCollator(model_params['depth_cap'], model_params['default_pc'],
                                         model_params['edge_type'])
        max_epochs = 5
        model_params['node_emb_dim'] = int(128 / div_emb_sz)
        model_params['node_state_dim'] = int(256 / div_emb_sz)
        model_params['lstm_state_dim'] = int(256 / div_emb_sz)
        # edge info
        model_params['edge_emb_dim'] = int(32 / div_emb_sz)
        model_params['edge_state_dim'] = int(64 / div_emb_sz)
    elif dataset == 'scitail':
        div_emb_sz = 2 if small_network else 1
        tr_dobj = pd.sci_tr_dobj_loc
        val_dobj = pd.sci_val_dobj_loc
        te_dobj = pd.sci_te_dobj_loc
        tr_collator = dt.HolstepCollator(model_params['depth_cap'], model_params['default_pc'],
                                         model_params['edge_type'])
        te_collator = dt.HolstepCollator(model_params['depth_cap'], model_params['default_pc'],
                                         model_params['edge_type'])
        max_epochs = 10
        model_params['node_emb_dim'] = int(16 / div_emb_sz)
        #model_params['node_state_dim'] = int(32 / div_emb_sz)
        #model_params['lstm_state_dim'] = int(32 / div_emb_sz)
        model_params['node_state_dim'] = model_params['pretrained_emb_dim']
        model_params['lstm_state_dim'] = model_params['pretrained_emb_dim']
        # edge info
        model_params['edge_emb_dim'] = int(4 / div_emb_sz)
        model_params['edge_state_dim'] = int(8 / div_emb_sz)
    elif dataset in ['mizar', 'mizar_cnf']:
        #if small_network: div_emb_sz = 2 if model_params['default_pc'] else 4
        if small_network: div_emb_sz = 2
        else: div_emb_sz = 1
        if dataset == 'mizar':
            tr_dobj = pd.miz_tr_dobj_loc
            val_dobj = pd.miz_val_dobj_loc
            te_dobj = pd.miz_te_dobj_loc
        elif dataset == 'mizar_cnf':
            tr_dobj = pd.miz_cnf_tr_dobj_loc
            val_dobj = pd.miz_cnf_val_dobj_loc
            te_dobj = pd.miz_cnf_te_dobj_loc
        tr_collator = dt.MizarCollator(model_params['depth_cap'], model_params['default_pc'],
                                       model_params['edge_type'])
        te_collator = dt.MizarCollator(model_params['depth_cap'], model_params['default_pc'],
                                       model_params['edge_type'])
        model_params['node_emb_dim'] = int(64 / div_emb_sz)
        model_params['node_state_dim'] = int(128 / div_emb_sz)
        model_params['lstm_state_dim'] = int(128 / div_emb_sz)
        # edge info
        model_params['edge_emb_dim'] = int(16 / div_emb_sz)
        model_params['edge_state_dim'] = int(32 / div_emb_sz)
        # each example in a batch is actually the complete set of positive and
        # negative premises, so there's actually around 100ish graphs in a batch
        batch_size = 16
        dataset_params['batch_size'] = batch_size
        max_epochs = 30 if model_params['default_pc'] else 5
    dataset_params['num_workers'] = batch_size

    if args.batch_size is not None:
        assert args.batch_size > 0, 'Positive values for batch size required'
        batch_size = args.batch_size
        model_params['batch_size'] = args.batch_size

    if args.num_epochs is not None:
        assert args.num_epochs >= 1, 'Epochs count must be at least 1...'
        max_epochs = args.num_epochs

    def get_tr_val_generators():
        # get training data
        training_set = pkl.load(open(tr_dobj, 'rb'))
        training_generator = data.DataLoader(training_set, collate_fn=tr_collator,
                                             **dataset_params)
        # get validation data
        validation_set = pkl.load(open(val_dobj, 'rb'))
        validation_generator = data.DataLoader(validation_set, collate_fn=te_collator,
                                               **dataset_params)
        if val_is_test:
            training_generator = chain(training_generator, validation_generator)
            validation_set = pkl.load(open(te_dobj, 'rb'))
            validation_generator = data.DataLoader(validation_set,
                                                   collate_fn=collator,
                                                   **dataset_params)
        return training_set,training_generator,validation_set,validation_generator

    model = cm.FormulaRelevanceClassifier(**model_params)
    today = date.today()
    model_param_str = '_'.join([model_params[x].__name__
                                for x in ['init_node_embedder',
                                          'pooling_module']]) + '_' + \
                                          '_'.join([str(model_params['num_rounds']),
                                                    model_params['aggr_type'],
                                                    str(model_params['edge_type']),
                                                    model_params['dep_match_type'],
                                                    str(model_params['dep_depth']),
                                                    str(model_params['default_pc']),
                                                    today.strftime('%Y%m%d')])

    # actually train model
    best_performance, time_info, loss_info, acc_info, track_ct = -1, [], [], [], 500
    gap_filler = '\n' + '='.join(['' for _ in range(30)]) + '\n'
    for epoch in range(max_epochs):
        if model_params['lr_decay'] and epoch > 0:
            for p_group in model.sparse_optimizer.param_groups:
                p_group['lr'] = p_group['lr'] / 3.0
            for p_group in model.dense_optimizer.param_groups:
                p_group['lr'] = p_group['lr'] / 3.0

        ( training_set, training_generator,
          validation_set, validation_generator ) = get_tr_val_generators()

        print(gap_filler)
        print('Starting training with model params ' + model_param_str + '...')
        print(gap_filler)

        tr_ex_ct, prev_pc, acc_tr_time = 0, None, 0

        tr_acc_nums, tr_loss_nums = [], []

        # training
        model.train()
        for batch in training_generator:

            tr_ex_ct += batch_size

            if (not model_params['default_pc']) and dataset == 'mizar':
                random.shuffle(batch)
                pseudo_batches = [batch[i : i + batch_size]
                                  for i in range(0, len(batch), batch_size)]
                if len(pseudo_batches[-1]) < batch_size / 2:
                    up_to = pseudo_batches[:-1]
                    up_to[-1].extend(pseudo_batches[-1])
                    pseudo_batches = up_to
            else:
                pseudo_batches = [batch]

            for pb in pseudo_batches:
                st_time = time.time()
                b_loss, b_acc = model.train_classifier(pb)
                assert not np.isnan(b_loss), 'Loss is nan...'

                tr_time = time.time() - st_time
                acc_tr_time += tr_time

                time_info.append(tr_time)
                loss_info.append(b_loss)
                acc_info.append(b_acc)
                if len(time_info) > track_ct: time_info.pop(0)
                if len(loss_info) > track_ct: loss_info.pop(0)
                if len(acc_info) > track_ct: acc_info.pop(0)
                if epoch == 0:
                    tr_acc_nums.append(b_acc)
                    tr_loss_nums.append(b_loss)
                print('At epoch: ' + str(epoch))
                print('Average training loss: ' + str(np.mean(loss_info)))
                print('Average training accuracy: ' + str(np.mean(acc_info)))
                print('Average training time: ' + \
                      str(np.mean(time_info) / len(pb)))
                print('Training completion: ' + \
                      str(round(tr_ex_ct/len(training_set.data_ids)*100, 2)) + '%')
                print()

        # validation
        model.eval()
        val_batch_ct, val_total_ct, corr_ct, acc_val_time = 0, 0, 0, 0
        for batch in validation_generator:
            val_batch_ct += batch_size

            if (not model_params['default_pc']) and dataset == 'mizar':
                random.shuffle(batch)
                pseudo_batches = [batch[i : i + batch_size]
                                  for i in range(0, len(batch), batch_size)]
            else:
                pseudo_batches = [batch]

            for pb in pseudo_batches:
                with torch.no_grad():
                    st_time = time.time()
                    outputs, labels, parse_failures = model.run_classifier(pb)
                    val_time = time.time() - st_time
                acc_val_time += val_time
                for output, label in zip(outputs, labels):
                    pred = 1 if output >= 0.5 else 0
                    if pred == label: corr_ct += 1
                    val_total_ct += 1
                val_total_ct += len(parse_failures)
                print('Validation completion: ' + \
                      str(round(val_batch_ct/len(validation_set.data_ids) * 100, 2))+'%')
                print('Current accuracy: ' + str(corr_ct / val_total_ct))
                print()

        epoch_performance = corr_ct / val_total_ct

        print(gap_filler)
        print(epoch_performance)
        print(gap_filler)

        results_file = dataset + '_' + model_param_str + '_val_results.csv'
        with open(os.path.join(results_path, results_file), 'a') as f:
            f.write(','.join([str(epoch), str(corr_ct), str(val_batch_ct),
                              str(val_total_ct),
                              str(acc_tr_time), str(acc_val_time),
                              str(epoch_performance)]) + '\n')

        if epoch_performance > best_performance and save_model:
            best_spec = dataset + '_best_' + model_param_str
            pkl.dump(model, open(os.path.join(models_path, best_spec), 'wb'))
            best_performance = epoch_performance

        latest_spec = dataset + '_epoch_' + str(epoch) + '_' + model_param_str
        pkl.dump(model, open(os.path.join(models_path, latest_spec), 'wb'))

        if epoch == 0:
            tr_perf_file = dataset + '_' + model_param_str + '_tr_curve.csv'
            with open(os.path.join(results_path, tr_perf_file), 'w') as f:
                f.write(','.join([str(round(num, 5)) for num in tr_acc_nums]))
                f.write('\n')
                f.write(','.join([str(round(num, 5)) for num in tr_loss_nums]))
                f.write('\n')

max_stars_repo_name: IBM/LogicalFormulaEmbedder

max_issues_repo_name: IBM/LogicalFormulaEmbedder

max_forks_repo_name: IBM/LogicalFormulaEmbedder