text:
# -*- coding: utf-8 -*-

# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
# holder of all proprietary rights on this computer program.
# You can only use this computer program if you have closed
# a license agreement with MPG or you get the right to use the computer
# program from someone who is authorized to grant you that right.
# Any use of the computer program without a valid license is prohibited and
# liable to prosecution.
#
# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
# for Intelligent Systems. All rights reserved.
#
# Contact: ps-license@tuebingen.mpg.de

from pytorch3d.structures import Meshes
from lib.pymaf.utils.imutils import uncrop
from lib.common.render_utils import Pytorch3dRasterizer, face_vertices
from pytorch3d.renderer.mesh import rasterize_meshes
from pdb import set_trace
import numpy as np
import cv2
import torch
import torchvision
import trimesh
from pytorch3d.io import load_obj
import os
from termcolor import colored
import os.path as osp
from scipy.spatial import cKDTree
from kaolin.ops.mesh import check_sign
from kaolin.metrics.trianglemesh import point_to_mesh_distance

from pytorch3d.loss import (
    mesh_edge_loss,
    mesh_laplacian_smoothing,
    mesh_normal_consistency,
)

from PIL import Image, ImageFont, ImageDraw

def get_mask(tensor, dim):

    mask = torch.abs(tensor).sum(dim=dim, keepdims=True) > 0.0
    mask = mask.type_as(tensor)

    return mask

def blend_rgb_norm(rgb, norm, mask):

    # [0,0,0] or [127,127,127] should be marked as mask
    final = rgb * (1-mask) + norm * (mask)

    return final.astype(np.uint8)

def unwrap(image, data):

    img_uncrop = uncrop(np.array(Image.fromarray(image).resize(data['uncrop_param']['box_shape'][:2])),
                        data['uncrop_param']['center'],
                        data['uncrop_param']['scale'],
                        data['uncrop_param']['crop_shape'])

    img_orig = cv2.warpAffine(img_uncrop,
                              np.linalg.inv(data['uncrop_param']['M'])[:2, :],
                              data['uncrop_param']['ori_shape'][::-1][1:],
                              flags=cv2.INTER_CUBIC)

    return img_orig

# Losses to smooth / regularize the mesh shape
def update_mesh_shape_prior_losses(mesh, losses):

    # and (b) the edge length of the predicted mesh
    losses["edge"]['value'] = mesh_edge_loss(mesh)
    # mesh normal consistency
    losses["normal"]['value'] = mesh_normal_consistency(mesh)
    # mesh laplacian smoothing
    losses["laplacian"]['value'] = mesh_laplacian_smoothing(mesh,
                                                            method="uniform")

def rename(old_dict, old_name, new_name):
    new_dict = {}
    for key, value in zip(old_dict.keys(), old_dict.values()):
        new_key = key if key != old_name else new_name
        new_dict[new_key] = old_dict[key]
    return new_dict

def load_checkpoint(model, cfg):

    model_dict = model.state_dict()
    main_dict = {}
    normal_dict = {}

    device = torch.device(f"cuda:{cfg['test_gpus'][0]}")

    if os.path.exists(cfg.resume_path) and cfg.resume_path.endswith("ckpt"):
        main_dict = torch.load(cfg.resume_path,
                               map_location=device)['state_dict']

        main_dict = {
            k: v
            for k, v in main_dict.items()
            if k in model_dict and v.shape == model_dict[k].shape and (
                'reconEngine' not in k) and ("normal_filter" not in k) and (
                    'voxelization' not in k)
        }
        print(colored(f"Resume MLP weights from {cfg.resume_path}", 'green'))

    if os.path.exists(cfg.normal_path) and cfg.normal_path.endswith("ckpt"):
        normal_dict = torch.load(cfg.normal_path,
                                 map_location=device)['state_dict']

        for key in normal_dict.keys():
            normal_dict = rename(normal_dict, key,
                                 key.replace("netG", "netG.normal_filter"))

        normal_dict = {
            k: v
            for k, v in normal_dict.items()
            if k in model_dict and v.shape == model_dict[k].shape
        }
        print(colored(f"Resume normal model from {cfg.normal_path}", 'green'))

    model_dict.update(main_dict)
    model_dict.update(normal_dict)
    model.load_state_dict(model_dict)

    model.netG = model.netG.to(device)
    model.reconEngine = model.reconEngine.to(device)

    model.netG.training = False
    model.netG.eval()

    del main_dict
    del normal_dict
    del model_dict

    return model

def read_smpl_constants(folder):
    """Load smpl vertex code"""
    smpl_vtx_std = np.loadtxt(os.path.join(folder, 'vertices.txt'))
    min_x = np.min(smpl_vtx_std[:, 0])
    max_x = np.max(smpl_vtx_std[:, 0])
    min_y = np.min(smpl_vtx_std[:, 1])
    max_y = np.max(smpl_vtx_std[:, 1])
    min_z = np.min(smpl_vtx_std[:, 2])
    max_z = np.max(smpl_vtx_std[:, 2])

    smpl_vtx_std[:, 0] = (smpl_vtx_std[:, 0] - min_x) / (max_x - min_x)
    smpl_vtx_std[:, 1] = (smpl_vtx_std[:, 1] - min_y) / (max_y - min_y)
    smpl_vtx_std[:, 2] = (smpl_vtx_std[:, 2] - min_z) / (max_z - min_z)
    smpl_vertex_code = np.float32(np.copy(smpl_vtx_std))
    """Load smpl faces & tetrahedrons"""
    smpl_faces = np.loadtxt(os.path.join(folder, 'faces.txt'),
                            dtype=np.int32) - 1
    smpl_face_code = (smpl_vertex_code[smpl_faces[:, 0]] +
                      smpl_vertex_code[smpl_faces[:, 1]] +
                      smpl_vertex_code[smpl_faces[:, 2]]) / 3.0
    smpl_tetras = np.loadtxt(os.path.join(folder, 'tetrahedrons.txt'),
                             dtype=np.int32) - 1

    return smpl_vertex_code, smpl_face_code, smpl_faces, smpl_tetras

def feat_select(feat, select):

    # feat [B, featx2, N]
    # select [B, 1, N]
    # return [B, feat, N]

    dim = feat.shape[1] // 2
    idx = torch.tile((1-select), (1, dim, 1))*dim + \
        torch.arange(0, dim).unsqueeze(0).unsqueeze(2).type_as(select)
    feat_select = torch.gather(feat, 1, idx.long())

    return feat_select

def get_visibility(xy, z, faces):
    """get the visibility of vertices

    Args:
        xy (torch.tensor): [N,2]
        z (torch.tensor): [N,1]
        faces (torch.tensor): [N,3]
        size (int): resolution of rendered image
    """

    xyz = torch.cat((xy, -z), dim=1)
    xyz = (xyz + 1.0) / 2.0
    faces = faces.long()

    rasterizer = Pytorch3dRasterizer(image_size=2**12)
    meshes_screen = Meshes(verts=xyz[None, ...], faces=faces[None, ...])
    raster_settings = rasterizer.raster_settings

    pix_to_face, zbuf, bary_coords, dists = rasterize_meshes(
        meshes_screen,
        image_size=raster_settings.image_size,
        blur_radius=raster_settings.blur_radius,
        faces_per_pixel=raster_settings.faces_per_pixel,
        bin_size=raster_settings.bin_size,
        max_faces_per_bin=raster_settings.max_faces_per_bin,
        perspective_correct=raster_settings.perspective_correct,
        cull_backfaces=raster_settings.cull_backfaces,
    )

    vis_vertices_id = torch.unique(faces[torch.unique(pix_to_face), :])
    vis_mask = torch.zeros(size=(z.shape[0], 1))
    vis_mask[vis_vertices_id] = 1.0

    # print("------------------------\n")
    # print(f"keep points : {vis_mask.sum()/len(vis_mask)}")

    return vis_mask

def barycentric_coordinates_of_projection(points, vertices):
    ''' https://github.com/MPI-IS/mesh/blob/master/mesh/geometry/barycentric_coordinates_of_projection.py
    '''
    """Given a point, gives projected coords of that point to a triangle
    in barycentric coordinates.
    See
        **Heidrich**, Computing the Barycentric Coordinates of a Projected Point, JGT 05
        at http://www.cs.ubc.ca/~heidrich/Papers/JGT.05.pdf

    :param p: point to project. [B, 3]
    :param v0: first vertex of triangles. [B, 3]
    :returns: barycentric coordinates of ``p``'s projection in triangle defined by ``q``, ``u``, ``v``
            vectorized so ``p``, ``q``, ``u``, ``v`` can all be ``3xN``
    """
    #(p, q, u, v)
    v0, v1, v2 = vertices[:, 0], vertices[:, 1], vertices[:, 2]
    p = points

    q = v0
    u = v1 - v0
    v = v2 - v0
    n = torch.cross(u, v)
    s = torch.sum(n * n, dim=1)
    # If the triangle edges are collinear, cross-product is zero,
    # which makes "s" 0, which gives us divide by zero. So we
    # make the arbitrary choice to set s to epsv (=numpy.spacing(1)),
    # the closest thing to zero
    s[s == 0] = 1e-6
    oneOver4ASquared = 1.0 / s
    w = p - q
    b2 = torch.sum(torch.cross(u, w) * n, dim=1) * oneOver4ASquared
    b1 = torch.sum(torch.cross(w, v) * n, dim=1) * oneOver4ASquared
    weights = torch.stack((1 - b1 - b2, b1, b2), dim=-1)
    # check barycenric weights
    # p_n = v0*weights[:,0:1] + v1*weights[:,1:2] + v2*weights[:,2:3]
    return weights

def cal_sdf_batch(verts, faces, cmaps, vis, points):

    # verts [B, N_vert, 3]
    # faces [B, N_face, 3]
    # triangles [B, N_face, 3, 3]
    # points [B, N_point, 3]
    # cmaps [B, N_vert, 3]

    Bsize = points.shape[0]

    normals = Meshes(verts, faces).verts_normals_padded()

    triangles = face_vertices(verts, faces)
    normals = face_vertices(normals, faces)
    cmaps = face_vertices(cmaps, faces)
    vis = face_vertices(vis, faces)

    residues, pts_ind, _ = point_to_mesh_distance(points, triangles)
    closest_triangles = torch.gather(
        triangles, 1, pts_ind[:, :, None, None].expand(-1, -1, 3, 3)).view(-1, 3, 3)
    closest_normals = torch.gather(
        normals, 1, pts_ind[:, :, None, None].expand(-1, -1, 3, 3)).view(-1, 3, 3)
    closest_cmaps = torch.gather(
        cmaps, 1, pts_ind[:, :, None, None].expand(-1, -1, 3, 3)).view(-1, 3, 3)
    closest_vis = torch.gather(
        vis, 1, pts_ind[:, :, None, None].expand(-1, -1, 3, 1)).view(-1, 3, 1)
    bary_weights = barycentric_coordinates_of_projection(
        points.view(-1, 3), closest_triangles)

    pts_cmap = (closest_cmaps*bary_weights[:, :, None]).sum(1).unsqueeze(0)
    pts_vis = (closest_vis*bary_weights[:,
               :, None]).sum(1).unsqueeze(0).ge(1e-1)
    pts_norm = (closest_normals*bary_weights[:, :, None]).sum(
        1).unsqueeze(0) * torch.tensor([-1.0, 1.0, -1.0]).type_as(normals)
    pts_dist = torch.sqrt(residues) / torch.sqrt(torch.tensor(3))

    pts_signs = 2.0 * (check_sign(verts, faces[0], points).float() - 0.5)
    pts_sdf = (pts_dist * pts_signs).unsqueeze(-1)

    return pts_sdf.view(Bsize, -1, 1), pts_norm.view(Bsize, -1, 3), pts_cmap.view(Bsize, -1, 3), pts_vis.view(Bsize, -1, 1)

def orthogonal(points, calibrations, transforms=None):
    '''
    Compute the orthogonal projections of 3D points into the image plane by given projection matrix
    :param points: [B, 3, N] Tensor of 3D points
    :param calibrations: [B, 3, 4] Tensor of projection matrix
    :param transforms: [B, 2, 3] Tensor of image transform matrix
    :return: xyz: [B, 3, N] Tensor of xyz coordinates in the image plane
    '''
    rot = calibrations[:, :3, :3]
    trans = calibrations[:, :3, 3:4]
    pts = torch.baddbmm(trans, rot, points)  # [B, 3, N]
    if transforms is not None:
        scale = transforms[:2, :2]
        shift = transforms[:2, 2:3]
        pts[:, :2, :] = torch.baddbmm(shift, scale, pts[:, :2, :])
    return pts

def projection(points, calib, format='numpy'):
    if format == 'tensor':
        return torch.mm(calib[:3, :3], points.T).T + calib[:3, 3]
    else:
        return np.matmul(calib[:3, :3], points.T).T + calib[:3, 3]

def load_calib(calib_path):
    calib_data = np.loadtxt(calib_path, dtype=float)
    extrinsic = calib_data[:4, :4]
    intrinsic = calib_data[4:8, :4]
    calib_mat = np.matmul(intrinsic, extrinsic)
    calib_mat = torch.from_numpy(calib_mat).float()
    return calib_mat

def load_obj_mesh_for_Hoppe(mesh_file):
    vertex_data = []
    face_data = []

    if isinstance(mesh_file, str):
        f = open(mesh_file, "r")
    else:
        f = mesh_file
    for line in f:
        if isinstance(line, bytes):
            line = line.decode("utf-8")
        if line.startswith('#'):
            continue
        values = line.split()
        if not values:
            continue

        if values[0] == 'v':
            v = list(map(float, values[1:4]))
            vertex_data.append(v)

        elif values[0] == 'f':
            # quad mesh
            if len(values) > 4:
                f = list(map(lambda x: int(x.split('/')[0]), values[1:4]))
                face_data.append(f)
                f = list(
                    map(lambda x: int(x.split('/')[0]),
                        [values[3], values[4], values[1]]))
                face_data.append(f)
            # tri mesh
            else:
                f = list(map(lambda x: int(x.split('/')[0]), values[1:4]))
                face_data.append(f)

    vertices = np.array(vertex_data)
    faces = np.array(face_data)
    faces[faces > 0] -= 1

    normals, _ = compute_normal(vertices, faces)

    return vertices, normals, faces

def load_obj_mesh_with_color(mesh_file):
    vertex_data = []
    color_data = []
    face_data = []

    if isinstance(mesh_file, str):
        f = open(mesh_file, "r")
    else:
        f = mesh_file
    for line in f:
        if isinstance(line, bytes):
            line = line.decode("utf-8")
        if line.startswith('#'):
            continue
        values = line.split()
        if not values:
            continue

        if values[0] == 'v':
            v = list(map(float, values[1:4]))
            vertex_data.append(v)
            c = list(map(float, values[4:7]))
            color_data.append(c)

        elif values[0] == 'f':
            # quad mesh
            if len(values) > 4:
                f = list(map(lambda x: int(x.split('/')[0]), values[1:4]))
                face_data.append(f)
                f = list(
                    map(lambda x: int(x.split('/')[0]),
                        [values[3], values[4], values[1]]))
                face_data.append(f)
            # tri mesh
            else:
                f = list(map(lambda x: int(x.split('/')[0]), values[1:4]))
                face_data.append(f)

    vertices = np.array(vertex_data)
    colors = np.array(color_data)
    faces = np.array(face_data)
    faces[faces > 0] -= 1

    return vertices, colors, faces

def load_obj_mesh(mesh_file, with_normal=False, with_texture=False):
    vertex_data = []
    norm_data = []
    uv_data = []

    face_data = []
    face_norm_data = []
    face_uv_data = []

    if isinstance(mesh_file, str):
        f = open(mesh_file, "r")
    else:
        f = mesh_file
    for line in f:
        if isinstance(line, bytes):
            line = line.decode("utf-8")
        if line.startswith('#'):
            continue
        values = line.split()
        if not values:
            continue

        if values[0] == 'v':
            v = list(map(float, values[1:4]))
            vertex_data.append(v)
        elif values[0] == 'vn':
            vn = list(map(float, values[1:4]))
            norm_data.append(vn)
        elif values[0] == 'vt':
            vt = list(map(float, values[1:3]))
            uv_data.append(vt)

        elif values[0] == 'f':
            # quad mesh
            if len(values) > 4:
                f = list(map(lambda x: int(x.split('/')[0]), values[1:4]))
                face_data.append(f)
                f = list(
                    map(lambda x: int(x.split('/')[0]),
                        [values[3], values[4], values[1]]))
                face_data.append(f)
            # tri mesh
            else:
                f = list(map(lambda x: int(x.split('/')[0]), values[1:4]))
                face_data.append(f)

            # deal with texture
            if len(values[1].split('/')) >= 2:
                # quad mesh
                if len(values) > 4:
                    f = list(map(lambda x: int(x.split('/')[1]), values[1:4]))
                    face_uv_data.append(f)
                    f = list(
                        map(lambda x: int(x.split('/')[1]),
                            [values[3], values[4], values[1]]))
                    face_uv_data.append(f)
                # tri mesh
                elif len(values[1].split('/')[1]) != 0:
                    f = list(map(lambda x: int(x.split('/')[1]), values[1:4]))
                    face_uv_data.append(f)
            # deal with normal
            if len(values[1].split('/')) == 3:
                # quad mesh
                if len(values) > 4:
                    f = list(map(lambda x: int(x.split('/')[2]), values[1:4]))
                    face_norm_data.append(f)
                    f = list(
                        map(lambda x: int(x.split('/')[2]),
                            [values[3], values[4], values[1]]))
                    face_norm_data.append(f)
                # tri mesh
                elif len(values[1].split('/')[2]) != 0:
                    f = list(map(lambda x: int(x.split('/')[2]), values[1:4]))
                    face_norm_data.append(f)

    vertices = np.array(vertex_data)
    faces = np.array(face_data)
    faces[faces > 0] -= 1

    if with_texture and with_normal:
        uvs = np.array(uv_data)
        face_uvs = np.array(face_uv_data)
        face_uvs[face_uvs > 0] -= 1
        norms = np.array(norm_data)
        if norms.shape[0] == 0:
            norms, _ = compute_normal(vertices, faces)
            face_normals = faces
        else:
            norms = normalize_v3(norms)
            face_normals = np.array(face_norm_data)
            face_normals[face_normals > 0] -= 1
        return vertices, faces, norms, face_normals, uvs, face_uvs

    if with_texture:
        uvs = np.array(uv_data)
        face_uvs = np.array(face_uv_data) - 1
        return vertices, faces, uvs, face_uvs

    if with_normal:
        norms = np.array(norm_data)
        norms = normalize_v3(norms)
        face_normals = np.array(face_norm_data) - 1
        return vertices, faces, norms, face_normals

    return vertices, faces

def normalize_v3(arr):
    ''' Normalize a numpy array of 3 component vectors shape=(n,3) '''
    lens = np.sqrt(arr[:, 0]**2 + arr[:, 1]**2 + arr[:, 2]**2)
    eps = 0.00000001
    lens[lens < eps] = eps
    arr[:, 0] /= lens
    arr[:, 1] /= lens
    arr[:, 2] /= lens
    return arr

def compute_normal(vertices, faces):
    # Create a zeroed array with the same type and shape as our vertices i.e., per vertex normal
    vert_norms = np.zeros(vertices.shape, dtype=vertices.dtype)
    # Create an indexed view into the vertex array using the array of three indices for triangles
    tris = vertices[faces]
    # Calculate the normal for all the triangles, by taking the cross product of the vectors v1-v0, and v2-v0 in each triangle
    face_norms = np.cross(tris[::, 1] - tris[::, 0], tris[::, 2] - tris[::, 0])
    # n is now an array of normals per triangle. The length of each normal is dependent the vertices,
    # we need to normalize these, so that our next step weights each normal equally.
    normalize_v3(face_norms)
    # now we have a normalized array of normals, one per triangle, i.e., per triangle normals.
    # But instead of one per triangle (i.e., flat shading), we add to each vertex in that triangle,
    # the triangles' normal. Multiple triangles would then contribute to every vertex, so we need to normalize again afterwards.
    # The cool part, we can actually add the normals through an indexed view of our (zeroed) per vertex normal array
    vert_norms[faces[:, 0]] += face_norms
    vert_norms[faces[:, 1]] += face_norms
    vert_norms[faces[:, 2]] += face_norms
    normalize_v3(vert_norms)

    return vert_norms, face_norms

def save_obj_mesh(mesh_path, verts, faces):
    file = open(mesh_path, 'w')
    for v in verts:
        file.write('v %.4f %.4f %.4f\n' % (v[0], v[1], v[2]))
    for f in faces:
        f_plus = f + 1
        file.write('f %d %d %d\n' % (f_plus[0], f_plus[1], f_plus[2]))
    file.close()

def save_obj_mesh_with_color(mesh_path, verts, faces, colors):
    file = open(mesh_path, 'w')

    for idx, v in enumerate(verts):
        c = colors[idx]
        file.write('v %.4f %.4f %.4f %.4f %.4f %.4f\n' %
                   (v[0], v[1], v[2], c[0], c[1], c[2]))
    for f in faces:
        f_plus = f + 1
        file.write('f %d %d %d\n' % (f_plus[0], f_plus[1], f_plus[2]))
    file.close()

def calculate_mIoU(outputs, labels):

    SMOOTH = 1e-6

    outputs = outputs.int()
    labels = labels.int()

    intersection = (
        outputs
        & labels).float().sum()  # Will be zero if Truth=0 or Prediction=0
    union = (outputs | labels).float().sum()  # Will be zzero if both are 0

    iou = (intersection + SMOOTH) / (union + SMOOTH
                                     )  # We smooth our devision to avoid 0/0

    thresholded = torch.clamp(
        20 * (iou - 0.5), 0,
        10).ceil() / 10  # This is equal to comparing with thresolds

    return thresholded.mean().detach().cpu().numpy(
    )  # Or thresholded.mean() if you are interested in average across the batch

def mask_filter(mask, number=1000):
    """only keep {number} True items within a mask

    Args:
        mask (bool array): [N, ]
        number (int, optional): total True item. Defaults to 1000.
    """
    true_ids = np.where(mask)[0]
    keep_ids = np.random.choice(true_ids, size=number)
    filter_mask = np.isin(np.arange(len(mask)), keep_ids)

    return filter_mask

def query_mesh(path):

    verts, faces_idx, _ = load_obj(path)

    return verts, faces_idx.verts_idx

def add_alpha(colors, alpha=0.7):

    colors_pad = np.pad(colors, ((0, 0), (0, 1)),
                        mode='constant',
                        constant_values=alpha)

    return colors_pad

def get_optim_grid_image(per_loop_lst, loss=None, nrow=4, type='smpl'):

    font_path = os.path.join(os.path.dirname(__file__), "tbfo.ttf")
    font = ImageFont.truetype(font_path, 30)
    grid_img = torchvision.utils.make_grid(torch.cat(per_loop_lst, dim=0),
                                           nrow=nrow)
    grid_img = Image.fromarray(
        ((grid_img.permute(1, 2, 0).detach().cpu().numpy() + 1.0) * 0.5 *
         255.0).astype(np.uint8))

    # add text
    draw = ImageDraw.Draw(grid_img)
    grid_size = 512
    if loss is not None:
        draw.text((10, 5), f"error: {loss:.3f}", (255, 0, 0), font=font)

    if type == 'smpl':
        for col_id, col_txt in enumerate(
                ['image', 'smpl-norm(render)', 'cloth-norm(pred)', 'diff-norm', 'diff-mask']):
            draw.text((10+(col_id*grid_size), 5),
                      col_txt, (255, 0, 0), font=font)
    elif type == 'cloth':
        for col_id, col_txt in enumerate(
                ['image', 'cloth-norm(recon)', 'cloth-norm(pred)', 'diff-norm']):
            draw.text((10+(col_id*grid_size), 5),
                      col_txt, (255, 0, 0), font=font)
        for col_id, col_txt in enumerate(
                ['0', '90', '180', '270']):
            draw.text((10+(col_id*grid_size), grid_size*2+5),
                      col_txt, (255, 0, 0), font=font)
    else:
        print(f"{type} should be 'smpl' or 'cloth'")

    grid_img = grid_img.resize((grid_img.size[0], grid_img.size[1]),
                               Image.ANTIALIAS)

    return grid_img

def clean_mesh(verts, faces):

    device = verts.device

    mesh_lst = trimesh.Trimesh(verts.detach().cpu().numpy(),
                               faces.detach().cpu().numpy())
    mesh_lst = mesh_lst.split(only_watertight=False)
    comp_num = [mesh.vertices.shape[0] for mesh in mesh_lst]
    mesh_clean = mesh_lst[comp_num.index(max(comp_num))]

    final_verts = torch.as_tensor(mesh_clean.vertices).float().to(device)
    final_faces = torch.as_tensor(mesh_clean.faces).int().to(device)

    return final_verts, final_faces

def merge_mesh(verts_A, faces_A, verts_B, faces_B):

    sep_mesh = trimesh.Trimesh(np.concatenate([verts_A, verts_B], axis=0),
                               np.concatenate(
                                   [faces_A, faces_B + faces_A.max() + 1],
                                   axis=0),
                               maintain_order=True,
                               process=False)

    colors = np.ones_like(sep_mesh.vertices)
    colors[:verts_A.shape[0]] *= np.array([255.0, 0.0, 0.0])
    colors[verts_A.shape[0]:] *= np.array([0.0, 255.0, 0.0])
    sep_mesh.visual.vertex_colors = colors

    # union_mesh = trimesh.boolean.union([trimesh.Trimesh(verts_A, faces_A),
    #                                     trimesh.Trimesh(verts_B, faces_B)], engine='blender')

    return sep_mesh

def mesh_move(mesh_lst, step, scale=1.0):

    trans = np.array([1.0, 0.0, 0.0]) * step

    resize_matrix = trimesh.transformations.scale_and_translate(
        scale=(scale), translate=trans)

    results = []

    for mesh in mesh_lst:
        mesh.apply_transform(resize_matrix)
        results.append(mesh)

    return results

class SMPLX():
    def __init__(self):

        self.current_dir = osp.join(osp.dirname(__file__),
                                    "../../data/smpl_related")

        self.smpl_verts_path = osp.join(self.current_dir,
                                        "smpl_data/smpl_verts.npy")
        self.smplx_verts_path = osp.join(self.current_dir,
                                         "smpl_data/smplx_verts.npy")
        self.faces_path = osp.join(self.current_dir,
                                   "smpl_data/smplx_faces.npy")
        self.cmap_vert_path = osp.join(self.current_dir,
                                       "smpl_data/smplx_cmap.npy")

        self.faces = np.load(self.faces_path)
        self.verts = np.load(self.smplx_verts_path)
        self.smpl_verts = np.load(self.smpl_verts_path)

        self.model_dir = osp.join(self.current_dir, "models")
        self.tedra_dir = osp.join(self.current_dir, "../tedra_data")

    def get_smpl_mat(self, vert_ids):

        mat = torch.as_tensor(np.load(self.cmap_vert_path)).float()
        return mat[vert_ids, :]

    def smpl2smplx(self, vert_ids=None):
        """convert vert_ids in smpl to vert_ids in smplx

        Args:
            vert_ids ([int.array]): [n, knn_num]
        """
        smplx_tree = cKDTree(self.verts, leafsize=1)
        _, ind = smplx_tree.query(self.smpl_verts, k=1)  # ind: [smpl_num, 1]

        if vert_ids is not None:
            smplx_vert_ids = ind[vert_ids]
        else:
            smplx_vert_ids = ind

        return smplx_vert_ids

    def smplx2smpl(self, vert_ids=None):
        """convert vert_ids in smplx to vert_ids in smpl

        Args:
            vert_ids ([int.array]): [n, knn_num]
        """
        smpl_tree = cKDTree(self.smpl_verts, leafsize=1)
        _, ind = smpl_tree.query(self.verts, k=1)  # ind: [smplx_num, 1]
        if vert_ids is not None:
            smpl_vert_ids = ind[vert_ids]
        else:
            smpl_vert_ids = ind

        return smpl_vert_ids

max_stars_repo_name: YuliangXiu/ICON

max_issues_repo_name: YuliangXiu/ICON

max_forks_repo_name: YuliangXiu/ICON

reason: import numpy,from scipy


text:
from core.algorithm import Algorithm
import numpy as np

class PID(Algorithm):
    def __init__(self, temp):
        self.aimingTemp = temp
        self.lastTemp = None
        self.lastError = 0
        self.errorSum = 0
        self.kp = 0.2
        self.ki = 0.1
        self.kd = 0.3

    def __call__(self, cluster, clock, cooling_equip=None):
        machinesInletTemp = []
        for machine in cluster.machines:
            machinesInletTemp.append(machine.inlet_temp)
        mean = np.array(machinesInletTemp).mean()
        errorNow = abs(self.aimingTemp - mean)
        self.errorSum += errorNow
        output = self.kp * errorNow + self.ki * self.errorSum + self.kd * (errorNow - self.lastError)
        self.lastError = errorNow
        return output

ext: py

max_stars_repo_name: kangyifei/CloudSimPy

max_issues_repo_name: kangyifei/CloudSimPy

max_forks_repo_name: kangyifei/CloudSimPy

reason: import numpy


text:
""" Simple module to handle modifying images """
from typing import Tuple
from pathlib import Path
import numpy as np
import imutils
import cv2

DEPTH_GRAY = 1  # the depth of grayscale image
DEPTH_COLOR = 3  # the depth of a color image
DEPTH_COLOR_ALPHA = 4  # the depth of a color and alpha image

# convert a color from a byte depth to the next depth
# key = (from, to), value = the opencv conversion constant
DEPTH_CONVERSIONS = {
    (1, 3): cv2.COLOR_GRAY2RGB,
    (3, 1): cv2.COLOR_RGB2GRAY,
    (1, 4): cv2.COLOR_GRAY2RGBA,
    (3, 4): cv2.COLOR_RGB2RGBA,
    (4, 3): cv2.COLOR_RGBA2RGB,
    (4, 1): cv2.COLOR_RGBA2GRAY
}

# a representation of the opencv interpolation constants
INTERPOLATIONS = {
    'area': cv2.INTER_AREA,
    'nearest': cv2.INTER_NEAREST,
    'lancz': cv2.INTER_LANCZOS4
}

def has_depth(img: np.ndarray) -> bool:
    """ Determines if an image has a depth element """
    return len(img.shape) >= 3

def image_dims(img: np.ndarray) -> Tuple[int, int]:
    """ Gets the dimensions of an image in wxh format """
    if (len(img.shape) == 3 and img.shape[0] < img.shape[-1]) or len(img.shape) == 4:
        return img.shape[2], img.shape[1]
    elif 2 <=len(img.shape) < 4:
        return img.shape[1], img.shape[0]
    else:
        raise RuntimeError('can not get the dimensions of a flat or timed array')

def fix_depth(img: np.ndarray, depth: int=DEPTH_COLOR):
    """ Fixes the image depth if it doesn't match the current depth in bytes

    :param img: the input image
    :param depth: the image depth in bytes
    """
    if len(img.shape) <= 1:  # the array has been flattened
        raise RuntimeError('we cannot assume the dimesions of a flattened or empty array')
    elif len(img.shape) == 2:  # we only have wxh and a grayscale image
        if depth == DEPTH_GRAY:
            img = img.reshape(img.shape + (depth,))
        else:
            img = cv2.cvtColor(img, DEPTH_CONVERSIONS[(DEPTH_GRAY, depth)])
    elif img.shape[-1] != depth:  # this is a multi-page/time image
        if DEPTH_GRAY <= img.shape[-1] < DEPTH_COLOR:  # convert from grayscale to another depth
            if depth == DEPTH_GRAY:
                img = img[:, :, 0].reshape(img.shape[:2] + (DEPTH_GRAY,))
            else:
                img = cv2.cvtColor(img[:, :, 0], DEPTH_CONVERSIONS[(DEPTH_GRAY, depth)])
        elif img.shape[-1] == DEPTH_COLOR and depth != DEPTH_COLOR:
            img = cv2.cvtColor(img, DEPTH_CONVERSIONS[(DEPTH_COLOR, depth)])
            if depth == DEPTH_GRAY:  # opencv will return just a wxh image when in reality we need a wxhxd image
                img = img.reshape(img.shape + (DEPTH_GRAY,))
        else:
            raise RuntimeError('we cannot fix the depth of a multi-page/time array of shape {}'.format(img.shape))
    return img

def fix_simple_depth(img: np.ndarray) -> np.ndarray:
    """ Fixes the simple grayscale depth of the images """
    if not has_depth(img):
        img = img.reshape(img.shape + (DEPTH_GRAY,))
    return img

def resize_aspect(img: np.ndarray, width: int=None, height: int=None, interp=cv2.INTER_AREA) -> np.ndarray:
    """ Resizes an image keeping its aspect ratio

    :param img: the image to resize
    :param width: the width of the new image
    :param height: the height of the new image
    :param interp: the interpolation of the image
    :return: a resized image with the new aspect ratio
    """
    # initialize the dimensions of the image to be resized and
    # grab the image size
    dim = None
    (h, w) = img.shape[:2]

    # if both the width and height are None, then return the
    # original image
    if width is None and height is None:
        return img

    # check to see if the width is None
    if width is None:
        # calculate the ratio of the height and construct the
        # dimensions
        r = height / float(h)
        dim = (int(w * r), height)

    # otherwise, the height is None
    else:
        # calculate the ratio of the width and construct the
        # dimensions
        r = width / float(w)
        dim = (width, int(h * r))

    # resize the image
    resized = cv2.resize(img, dim, interpolation=interp)

    # return the resized image
    return resized

def resize_image(img: np.ndarray, dims: Tuple[int, int], interp: str='nearest', aspect=None, fix_depth=True) -> np.ndarray:
    """ Resizes an image to the specified dimensions

    :param img: the image to resize
    :param dims: the dimensions of the image wxh
    :param interp: the interpolation of the image (default: 'nearest')
    :param aspect: keep the aspect ratio (default: None)
    :param fix_depth: fix the image to have a depth if it doesn't have it (default: True)
    :return: the resized image
    """
    if dims[0] != img.shape[1] or dims[1] != img.shape[0]:
        try:
            if (None in dims or aspect == 'aspect') and aspect != 'none':  # we want to keep some aspect ratio
                img = imutils.resize(img, dims[0], dims[1], INTERPOLATIONS[interp])
            elif aspect == 'contain':  # we want to contain the image inside the dims bounding box
                width, height = dims[:2]
                if img.shape[0] >= img.shape[1]:  # adjust by height
                    width = None
                else:
                    height = None

                # get the initial image with the fixed depth
                cp_img = fix_simple_depth(resize_aspect(img, width, height, INTERPOLATIONS[interp]))

                # paste the resized image into the container if the dimensions don't match
                if dims[0] != img.shape[1] or dims[1] != img.shape[0]:
                    cpy = np.zeros((dims[1], dims[0], img.shape[-1] if has_depth(img) else 1), img.dtype)
                    cpy[0:cp_img.shape[0], 0:cp_img.shape[1]] = cp_img  # place the resized image into the black background image
                    cp_img = cpy

                # replace the original image with the contained one
                img = cp_img
            else:
                img = cv2.resize(img, dims[:2], INTERPOLATIONS[interp])  # distortion is allowed
        except KeyError:
            raise ValueError('invalid interpolation method {}. Please use one of the following {}'.format(interp, list(INTERPOLATIONS.keys())))

        if not has_depth(img) and fix_depth:  # fix the grayscale depth if applicable
            img = img.reshape(img.shape + (DEPTH_GRAY,))
    return img

def resize_stack(stack: np.ndarray, dims: Tuple[int, int], interp: str='nearest', aspect=None, fix_depth=True) -> np.ndarray:
    """ Resizes a stack of images to the specified dimensions

    :param img: the image stack to resize
    :param dims: the dimensions of the image wxh
    :param interp: the interpolation of the image (default: 'nearest')
    :param aspect: keep the aspect ratio (default: None)
    :param fix_depth: fix the image to have a depth if it doesn't have it (default: True)
    :return: the resized image stack
    """
    if len(stack.shape) != 4:
        raise RuntimeError('invalid image stack shape of {}'.format(stack.shape))

    # make sure the sizes are correct if not we have to make a copy of the array and loop through to resize the stack
    if dims[0] != stack.shape[2] or dims[1] != stack.shape[1]:
        cp_stack = np.zeros((stack.shape[0],) + (dims[1], dims[0]) + (stack.shape[-1],), stack.dtype)  # this will assume the same dtype and depth as the current stack

        # loop through each layer of the stack
        for i in range(stack.shape[0]):
            cp_stack[i] = resize_image(stack[i], dims, interp=interp, aspect=aspect, fix_depth=True)

        # replace the current stack
        stack = cp_stack
    return stack

def resize_image_nparray(img: np.ndarray, dims: Tuple[int, int], interp: str='nearest', aspect=None, fix_depth=True):
    if img is None:
        return img
    elif (len(img.shape) == 3 and img.shape[0] < img.shape[-1]) or len(img.shape) == 4:  # most likely this image is stacked if we have fewer elements in our first shape then the last
        is_stack = True  # we can assume the array has probably already been fixed
    else:
        is_stack = False  # let's assume this array is just a single image

    ref_fix_shape = False  # flag to determine if we need to fix the last depth shape in case of grayscale images
    if is_stack:
        if len(img.shape) == 3:
            img = img.reshape(img.shape + (1,))
            ref_fix_shape = True
        img = resize_stack(img, dims, interp, aspect, fix_depth)
        if ref_fix_shape:
            return img.reshape(img.shape[:3])
        return img

    if len(img.shape) == 2:
        img = img.reshape(img.shape + (1,))
        ref_fix_shape = True
    img = resize_image(img, dims, interp, aspect, fix_depth)
    if ref_fix_shape:
        return img.reshape(img.shape[:2])
    return img

max_stars_repo_name: smerkousdavid/membrane-analysis

max_issues_repo_name: smerkousdavid/membrane-analysis

max_forks_repo_name: smerkousdavid/membrane-analysis


text:
"""
Example dynamic_extension_tracking.py
Author: Joshua A. Marshall <joshua.marshall@queensu.ca>
GitHub: https://github.com/botprof/agv-examples
"""

# %%
# SIMULATION SETUP

import numpy as np
import matplotlib.pyplot as plt
from mobotpy.models import DiffDrive
from mobotpy.integration import rk_four
from scipy import signal

# Set the simulation time [s] and the sample period [s]
SIM_TIME = 15.0
T = 0.04

# Create an array of time values [s]
t = np.arange(0.0, SIM_TIME, T)
N = np.size(t)

# %%
# COMPUTE THE REFERENCE TRAJECTORY

# Radius of the circle [m]
R = 10

# Angular rate [rad/s] at which to traverse the circle
OMEGA = 0.1

# Precompute the desired trajectory
x_d = np.zeros((3, N))
u_d = np.zeros((2, N))
xi_d = np.zeros((4, N))
ddz_d = np.zeros((2, N))
for k in range(0, N):
    x_d[0, k] = R * np.sin(OMEGA * t[k])
    x_d[1, k] = R * (1 - np.cos(OMEGA * t[k]))
    x_d[2, k] = OMEGA * t[k]
    u_d[0, k] = R * OMEGA
    u_d[1, k] = OMEGA

# Precompute the extended system reference trajectory
for k in range(0, N):
    xi_d[0, k] = x_d[0, k]
    xi_d[1, k] = x_d[1, k]
    xi_d[2, k] = u_d[0, k] * np.cos(x_d[2, k])
    xi_d[3, k] = u_d[0, k] * np.sin(x_d[2, k])

# Precompute the extended system reference acceleration
for k in range(0, N):
    ddz_d[0, k] = -u_d[0, k] * u_d[1, k] * np.sin(x_d[2, k])
    ddz_d[1, k] = u_d[0, k] * u_d[1, k] * np.cos(x_d[2, k])

# %%
# VEHICLE SETUP

# Set the track length of the vehicle [m]
ELL = 1.0

# Create a vehicle object of type DiffDrive
vehicle = DiffDrive(ELL)

# %%
# SIMULATE THE CLOSED-LOOP SYSTEM

# Initial conditions
x_init = np.zeros(3)
x_init[0] = 0.0
x_init[1] = 3.0
x_init[2] = 0.0

# Setup some arrays
x = np.zeros((3, N))
xi = np.zeros((4, N))
u = np.zeros((2, N))
x[:, 0] = x_init

# Set the initial speed [m/s] to be non-zero to avoid singularity
w = np.zeros(2)
u_unicycle = np.zeros(2)
u_unicycle[0] = u_d[0, 0]

# Initial extended state
xi[0, 0] = x_init[0]
xi[1, 0] = x_init[1]
xi[2, 0] = u_d[0, 0] * np.cos(x_init[2])
xi[3, 0] = u_d[0, 0] * np.sin(x_init[2])

# Defined feedback linearized state matrices
A = np.array([[0, 0, 1, 0], [0, 0, 0, 1], [0, 0, 0, 0], [0, 0, 0, 0]])
B = np.array([[0, 0], [0, 0], [1, 0], [0, 1]])

# Choose pole locations for closed-loop linear system
p = np.array([-1.0, -2.0, -2.5, -1.5])
K = signal.place_poles(A, B, p)

for k in range(1, N):

    # Simulate the vehicle motion
    x[:, k] = rk_four(vehicle.f, x[:, k - 1], u[:, k - 1], T)

    # Update the extended system states
    xi[0, k] = x[0, k]
    xi[1, k] = x[1, k]
    xi[2, k] = u_unicycle[0] * np.cos(x[2, k])
    xi[3, k] = u_unicycle[0] * np.sin(x[2, k])

    # Compute the extended linear system input control signals
    eta = K.gain_matrix @ (xi_d[:, k - 1] - xi[:, k - 1]) + ddz_d[:, k - 1]

    # Compute the new (unicycle) vehicle inputs
    B_inv = np.array(
        [
            [np.cos(x[2, k - 1]), np.sin(x[2, k - 1])],
            [-np.sin(x[2, k - 1]) / u_unicycle[0], np.cos(x[2, k - 1]) / u_unicycle[0]],
        ]
    )
    w = B_inv @ eta
    u_unicycle[0] = u_unicycle[0] + T * w[0]
    u_unicycle[1] = w[1]

    # Convert unicycle inputs to differential drive wheel speeds
    u[:, k] = vehicle.uni2diff(u_unicycle)

# %%
# MAKE PLOTS

# Change some plot settings (optional)
plt.rc("text", usetex=True)
plt.rc("text.latex", preamble=r"\usepackage{cmbright,amsmath,bm}")
plt.rc("savefig", format="pdf")
plt.rc("savefig", bbox="tight")

# Plot the states as a function of time
fig1 = plt.figure(1)
fig1.set_figheight(6.4)
ax1a = plt.subplot(411)
plt.plot(t, x_d[0, :], "C1--")
plt.plot(t, x[0, :], "C0")
plt.grid(color="0.95")
plt.ylabel(r"$x$ [m]")
plt.setp(ax1a, xticklabels=[])
plt.legend(["Desired", "Actual"])
ax1b = plt.subplot(412)
plt.plot(t, x_d[1, :], "C1--")
plt.plot(t, x[1, :], "C0")
plt.grid(color="0.95")
plt.ylabel(r"$y$ [m]")
plt.setp(ax1b, xticklabels=[])
ax1c = plt.subplot(413)
plt.plot(t, x_d[2, :] * 180.0 / np.pi, "C1--")
plt.plot(t, x[2, :] * 180.0 / np.pi, "C0")
plt.grid(color="0.95")
plt.ylabel(r"$\theta$ [deg]")
plt.setp(ax1c, xticklabels=[])
ax1d = plt.subplot(414)
plt.step(t, u[0, :], "C2", where="post", label="$v_L$")
plt.step(t, u[1, :], "C3", where="post", label="$v_R$")
plt.grid(color="0.95")
plt.ylabel(r"$\bm{u}$ [m/s]")
plt.xlabel(r"$t$ [s]")
plt.legend()

# Save the plot
plt.savefig("../agv-book/figs/ch4/dynamic_extension_tracking_fig1.pdf")

# Plot the position of the vehicle in the plane
fig2 = plt.figure(2)
plt.plot(x_d[0, :], x_d[1, :], "C1--", label="Desired")
plt.plot(x[0, :], x[1, :], "C0", label="Actual")
plt.axis("equal")
X_L, Y_L, X_R, Y_R, X_B, Y_B, X_C, Y_C = vehicle.draw(x[0, 0], x[1, 0], x[2, 0])
plt.fill(X_L, Y_L, "k")
plt.fill(X_R, Y_R, "k")
plt.fill(X_C, Y_C, "k")
plt.fill(X_B, Y_B, "C2", alpha=0.5, label="Start")
X_L, Y_L, X_R, Y_R, X_B, Y_B, X_C, Y_C = vehicle.draw(
    x[0, N - 1], x[1, N - 1], x[2, N - 1]
)
plt.fill(X_L, Y_L, "k")
plt.fill(X_R, Y_R, "k")
plt.fill(X_C, Y_C, "k")
plt.fill(X_B, Y_B, "C3", alpha=0.5, label="End")
plt.xlabel(r"$x$ [m]")
plt.ylabel(r"$y$ [m]")
plt.legend()

# Save the plot
plt.savefig("../agv-book/figs/ch4/dynamic_extension_tracking_fig2.pdf")

# Show all the plots to the screen
plt.show()

# %%
# MAKE AN ANIMATION

# Create and save the animation
ani = vehicle.animate_trajectory(
    x, x_d, T, True, "../agv-book/gifs/ch4/dynamic_extension_tracking.gif"
)

# Show the movie to the screen
plt.show()

# # Show animation in HTML output if you are using IPython or Jupyter notebooks
# plt.rc('animation', html='jshtml')
# display(ani)
# plt.close()

max_stars_repo_name: botprof/agv-examples

max_issues_repo_name: botprof/agv-examples

max_forks_repo_name: botprof/agv-examples


text:
# -*- coding: utf-8 -*-
# =============================================================================
# MIT License
#
# Copyright (c) 2018 Charles Jekel
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
# =============================================================================

#   Define a function that combines two vander matricies for use with least
#   Square polynomial regression fitting for three dimensional data
#   Useage A = doubleVander(X,Y,degree)
#   Where X, Y are your vectors
#   Degree = the order of the desired polynomial
#   A is the resulting double vander matrix

import numpy as np

def doubleVander(x,y,Degree):
    N = Degree + 1
    Ax = np.vander(x,Degree+1)
    Ay = np.vander(y,Degree+1)
    Aout = np.zeros((len(x),N*N))
    count = 0
    for i in range(0,N):
        for j in range(0,N):
            Aout[:,count] = Ax[:,i]*Ay[:,j]
            count = count +1
    return Aout;


text:
#   Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
#
# 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
#
#     http://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.

import sys
import os
import traceback
import logging
import json

import numpy as np
import paddle

from text2sql.utils import nn_utils

def collate_batch_data_v2(origin_batch, config):
    """format origin batch data for model forward
    """
    TOKEN_IDS = []
    SENT_IDS = []
    INPUT_MASK = []
    POSITION_IDS = []
    TASK_IDS = []

    QUESTION_TOKENS_INDEX = []
    TABLE_INDEX = []
    COLUMN_INDEX = []
    VALUE_INDEX = []
    RELATION_MATRIXES = []

    lst_orig_inputs = []
    lst_orig_labels = []
    for orig_input, orig_label in origin_batch:
        if orig_input.value_indexes[-1] > 510:
            logging.warning('sequence is too long: %d. question is %s',
                            orig_input.value_indexes[-1] + 2,
                            orig_input.question)
            continue
        lst_orig_inputs.append(orig_input)
        lst_orig_labels.append(orig_label)

        TOKEN_IDS.append(orig_input.token_ids)
        SENT_IDS.append(orig_input.sent_ids)

        # orig_input.span_lens[0] 即 question 包含 [cls], [sep] 的长度
        QUESTION_TOKENS_INDEX.append(
            list(range(1, orig_input.column_indexes[0] - 1)))
        TABLE_INDEX.append(orig_input.table_indexes)
        COLUMN_INDEX.append(orig_input.column_indexes)
        VALUE_INDEX.append(orig_input.value_indexes)

        relations = orig_input.relations
        RELATION_MATRIXES.append(
            np.pad(relations, (0, config.max_seq_len - relations.shape[0])))

    TOKEN_IDS = nn_utils.pad_sequences(TOKEN_IDS, max_len=config.max_seq_len)
    SENT_IDS = nn_utils.pad_sequences(SENT_IDS, max_len=config.max_seq_len)

    QUESTION_TOKENS_INDEX = nn_utils.pad_sequences(
        QUESTION_TOKENS_INDEX, max_len=config.max_question_len)
    TABLE_INDEX = nn_utils.pad_sequences(TABLE_INDEX,
                                         max_len=config.max_table_num)
    COLUMN_INDEX = nn_utils.pad_sequences(COLUMN_INDEX,
                                          max_len=config.max_column_num)
    VALUE_INDEX = nn_utils.pad_sequences(VALUE_INDEX,
                                         max_len=config.max_column_num * 2)

    inputs = {
        'src_ids': TOKEN_IDS,
        'sent_ids': SENT_IDS,
        'question_tokens_index': QUESTION_TOKENS_INDEX,
        'table_indexes': TABLE_INDEX,
        'column_indexes': COLUMN_INDEX,
        'value_indexes': VALUE_INDEX,
        'orig_inputs': lst_orig_inputs,
    }
    RELATION_MATRIXES = np.array(RELATION_MATRIXES).astype(np.int64)
    inputs["relations"] = RELATION_MATRIXES

    for key, value in inputs.items():
        if key in ('orig_inputs', ):
            continue
        inputs[key] = paddle.to_tensor(value)
    return (inputs, lst_orig_labels)

class DataLoader(object):
    """Data Loader for train, test and inference"""

    def __init__(self,
                 config,
                 dataset,
                 batch_size=1,
                 collate_fn=collate_batch_data_v2,
                 shuffle=False,
                 drop_last=False,
                 use_data_parallel=False,
                 use_multiprocess=False):
        super(DataLoader, self).__init__()
        assert batch_size > 0, "batch_size must be an interger that > 0"

        self.config = config
        self._dataset = dataset
        self._batch_size = batch_size
        self._collate_fn = collate_fn
        self._shuffle = shuffle
        self._drop_last = drop_last
        self._use_data_parallel = use_data_parallel
        self._use_multiprocess = use_multiprocess

        self.dataloader = paddle.io.DataLoader.from_generator(
            capacity=1000, return_list=True, use_multiprocess=use_multiprocess)
        self.dataloader.set_batch_generator(self.create_generator())
        if use_data_parallel:
            self.dataloader = paddle.distributed_batch_reader(self.dataloader)

    def __call__(self):
        """call"""
        return self.create_generator()()

    def create_generator(self):
        """Returns a generator, each iteration returns a batch of data"""

        def _reader():
            range_fn = np.random.permutation if self._shuffle else np.arange
            batch = []
            for iid in range_fn(len(self._dataset)):
                batch.append(self._dataset[iid])
                if len(batch) == self._batch_size:
                    outputs = self._collate_fn(batch, self.config.model)
                    batch = []
                    if len(outputs[1]) == 0:
                        continue
                    yield outputs

            if len(batch) > 0 and not self._drop_last:
                yield self._collate_fn(batch, self.config.model)

        return _reader

    @property
    def name(self):
        """read property of name"""
        return self._dataset.name

if __name__ == "__main__":
    """run simple tests"""
    if len(sys.argv) != 5:
        print("usage: %s schema content data grammar_file" % (sys.argv[0]))
        sys.exit(1)

max_stars_repo_name: mukaiu/PaddleNLP

max_issues_repo_name: mukaiu/PaddleNLP

max_forks_repo_name: mukaiu/PaddleNLP