text:
"""
Copyright 2016 Max Planck Society, Federica Bogo, Angjoo Kanazawa. All rights reserved.
This software is provided for research purposes only.
By using this software you agree to the terms of the SMPLify license here:
     http://smplify.is.tue.mpg.de/license

About this Script:
============
This is a demo version of the algorithm implemented in the paper,
which fits the SMPL body model to the image given the joint detections.
The code is organized to be run on the LSP dataset.
See README to see how to download images and the detected joints.
"""

from os.path import join, exists, abspath, dirname
from os import makedirs
import logging
import cPickle as pickle
from time import time
from glob import glob
import argparse

import cv2
import numpy as np
import chumpy as ch
import pyopenpose as op

from multiprocessing import Process

from opendr.camera import ProjectPoints
from lib.robustifiers import GMOf
from smpl_webuser.serialization import load_model
from smpl_webuser.lbs import global_rigid_transformation
from smpl_webuser.verts import verts_decorated
from lib.sphere_collisions import SphereCollisions
from lib.max_mixture_prior import MaxMixtureCompletePrior
from render_model import render_model

# Mapping from LSP joints to SMPL joints.
# 0 Right ankle  8
# 1 Right knee   5
# 2 Right hip    2
# 3 Left hip     1
# 4 Left knee    4
# 5 Left ankle   7
# 6 Right wrist  21
# 7 Right elbow  19
# 8 Right shoulder 17
# 9 Left shoulder  16
# 10 Left elbow    18
# 11 Left wrist    20
# 12 Neck           -
# 13 Head top       added
POSE = ["R ankle", "R knee", "R hip", "L hip", "L knee", "L ankle", "R wrist", "R elbow", "R shoulder", "L shoulder", "L elbow", "L wrist", "Neck", "head"]

# --------------------Camera estimation --------------------
def guess_init(model, focal_length, j2d, init_pose):
    """Initialize the camera translation via triangle similarity, by using the torso joints        .
    :param model: SMPL model
    :param focal_length: camera focal length (kept fixed)
    :param j2d: 14x2 array of CNN joints
    :param init_pose: 72D vector of pose parameters used for initialization (kept fixed)
    :returns: 3D vector corresponding to the estimated camera translation
    """
    cids = np.arange(0, 12)
    # map from LSP to SMPL joints
    j2d_here = j2d[cids]
    smpl_ids = [8, 5, 2, 1, 4, 7, 21, 19, 17, 16, 18, 20]

    opt_pose = ch.array(init_pose)
    (_, A_global) = global_rigid_transformation(
        opt_pose, model.J, model.kintree_table, xp=ch)
    Jtr = ch.vstack([g[:3, 3] for g in A_global])
    Jtr = Jtr[smpl_ids].r

    # 9 is L shoulder, 3 is L hip
    # 8 is R shoulder, 2 is R hip
    diff3d = np.array([Jtr[9] - Jtr[3], Jtr[8] - Jtr[2]])
    mean_height3d = np.mean(np.sqrt(np.sum(diff3d**2, axis=1)))

    diff2d = np.array([j2d_here[9] - j2d_here[3], j2d_here[8] - j2d_here[2]])
    mean_height2d = np.mean(np.sqrt(np.sum(diff2d**2, axis=1)))

    est_d = focal_length * (mean_height3d / mean_height2d)
    # just set the z value
    init_t = np.array([0., 0., est_d])
    return init_t

def initialize_camera(model,
                      j2d,
                      img,
                      init_pose,
                      flength=5000.,
                      pix_thsh=25.):
    """Initialize camera translation and body orientation
    :param model: SMPL model
    :param j2d: 14x2 array of CNN joints
    :param img: h x w x 3 image
    :param init_pose: 72D vector of pose parameters used for initialization
    :param flength: camera focal length (kept fixed)
    :param pix_thsh: threshold (in pixel), if the distance between shoulder joints in 2D
                     is lower than pix_thsh, the body orientation as ambiguous (so a fit is run on both
                     the estimated one and its flip)
    :returns: a tuple containing the estimated camera,
              a boolean deciding if both the optimized body orientation and its flip should be considered,
              3D vector for the body orientation
    """
    # optimize camera translation and body orientation based on torso joints
    # LSP torso ids:
    # 2=right hip, 3=left hip, 8=right shoulder, 9=left shoulder
    torso_cids = [2, 3, 8, 9]
    # corresponding SMPL torso ids
    torso_smpl_ids = [2, 1, 17, 16]

    center = np.array([img.shape[1] / 2, img.shape[0] / 2])

    # initialize camera rotation
    rt = ch.zeros(3)
    # initialize camera translation
    print 'initializing translation via similar triangles'
    init_t = guess_init(model, flength, j2d, init_pose)
    t = ch.array(init_t)

    # check how close the shoulder joints are
    try_both_orient = np.linalg.norm(j2d[8] - j2d[9]) < pix_thsh

    opt_pose = ch.array(init_pose)
    (_, A_global) = global_rigid_transformation(
        opt_pose, model.J, model.kintree_table, xp=ch)
    Jtr = ch.vstack([g[:3, 3] for g in A_global])

    # initialize the camera
    cam = ProjectPoints(
        f=np.array([flength, flength]), rt=rt, t=t, k=np.zeros(5), c=center)

    # we are going to project the SMPL joints
    cam.v = Jtr
    on_step = None

    # optimize for camera translation and body orientation
    free_variables = [cam.t, opt_pose[:3]]
    ch.minimize(
        # data term defined over torso joints...
        {'cam': j2d[torso_cids] - cam[torso_smpl_ids],
         # ...plus a regularizer for the camera translation
         'cam_t': 1e2 * (cam.t[2] - init_t[2])},
        x0=free_variables,
        method='dogleg',
        callback=on_step,
        # maxiter: 100, e_3: .0001
        options={'maxiter': 100,
                 'e_3': .0001,
                 # disp set to 1 enables verbose output from the optimizer
                 'disp': 0})

    return (cam, try_both_orient, opt_pose[:3].r)

# --------------------Core optimization --------------------
def optimize_on_joints(j2d,
                       model,
                       cam,
                       img,
                       prior,
                       try_both_orient,
                       body_orient,
                       n_betas=10,
                       regs=None,
                       conf=None):
    """Fit the model to the given set of joints, given the estimated camera
    :param j2d: 14x2 array of CNN joints
    :param model: SMPL model
    :param cam: estimated camera
    :param img: h x w x 3 image
    :param prior: mixture of gaussians pose prior
    :param try_both_orient: boolean, if True both body_orient and its flip are considered for the fit
    :param body_orient: 3D vector, initialization for the body orientation
    :param n_betas: number of shape coefficients considered during optimization
    :param regs: regressors for capsules' axis and radius, if not None enables the interpenetration error term
    :param conf: 14D vector storing the confidence values from the CNN
    :returns: a tuple containing the optimized model, its joints projected on image space, the camera translation
    """
    t0 = time()
    # define the mapping LSP joints -> SMPL joints
    # cids are joints ids for LSP:
    cids = range(12) + [13]
    # joint ids for SMPL
    # SMPL does not have a joint for head, instead we use a vertex for the head
    # and append it later.
    smpl_ids = [8, 5, 2, 1, 4, 7, 21, 19, 17, 16, 18, 20]

    # the vertex id for the joint corresponding to the head
    head_id = 411

    # weights assigned to each joint during optimization;
    # the definition of hips in SMPL and LSP is significantly different so set
    # their weights to zero
    base_weights = np.array(
        [1, 1, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1], dtype=np.float64)

    if try_both_orient:
        flipped_orient = cv2.Rodrigues(body_orient)[0].dot(
            cv2.Rodrigues(np.array([0., np.pi, 0]))[0])
        flipped_orient = cv2.Rodrigues(flipped_orient)[0].ravel()
        orientations = [body_orient, flipped_orient]
    else:
        orientations = [body_orient]

    if try_both_orient:
        # store here the final error for both orientations,
        # and pick the orientation resulting in the lowest error
        errors = []

    svs = []
    cams = []
    for o_id, orient in enumerate(orientations):
        # initialize the shape to the mean shape in the SMPL training set
        betas = ch.zeros(n_betas)

        # initialize the pose by using the optimized body orientation and the
        # pose prior
        init_pose = np.hstack((orient, prior.weights.dot(prior.means)))

        # instantiate the model:
        # verts_decorated allows us to define how many
        # shape coefficients (directions) we want to consider (here, n_betas)
        sv = verts_decorated(
            trans=ch.zeros(3),
            pose=ch.array(init_pose),
            v_template=model.v_template,
            J=model.J_regressor,
            betas=betas,
            shapedirs=model.shapedirs[:, :, :n_betas],
            weights=model.weights,
            kintree_table=model.kintree_table,
            bs_style=model.bs_style,
            f=model.f,
            bs_type=model.bs_type,
            posedirs=model.posedirs)

        # make the SMPL joints depend on betas
        Jdirs = np.dstack([model.J_regressor.dot(model.shapedirs[:, :, i])
                           for i in range(len(betas))])
        J_onbetas = ch.array(Jdirs).dot(betas) + model.J_regressor.dot(
            model.v_template.r)

        # get joint positions as a function of model pose, betas and trans
        (_, A_global) = global_rigid_transformation(
            sv.pose, J_onbetas, model.kintree_table, xp=ch)
        Jtr = ch.vstack([g[:3, 3] for g in A_global]) + sv.trans

        # add the head joint, corresponding to a vertex...
        Jtr = ch.vstack((Jtr, sv[head_id]))

        # ... and add the joint id to the list
        if o_id == 0:
            smpl_ids.append(len(Jtr) - 1)

        # update the weights using confidence values
        weights = base_weights * conf[
            cids] if conf is not None else base_weights

        # project SMPL joints on the image plane using the estimated camera
        cam.v = Jtr

        # data term: distance between observed and estimated joints in 2D
        obj_j2d = lambda w, sigma: (
            w * weights.reshape((-1, 1)) * GMOf((j2d[cids] - cam[smpl_ids]), sigma))

        # mixture of gaussians pose prior
        pprior = lambda w: w * prior(sv.pose)
        # joint angles pose prior, defined over a subset of pose parameters:
        # 55: left elbow,  90deg bend at -np.pi/2
        # 58: right elbow, 90deg bend at np.pi/2
        # 12: left knee,   90deg bend at np.pi/2
        # 15: right knee,  90deg bend at np.pi/2
        alpha = 10
        my_exp = lambda x: alpha * ch.exp(x)
        obj_angle = lambda w: w * ch.concatenate([my_exp(sv.pose[55]), my_exp(-sv.pose[
                                                 58]), my_exp(-sv.pose[12]), my_exp(-sv.pose[15])])

        on_step = None

        if regs is not None:
            # interpenetration term
            sp = SphereCollisions(
                pose=sv.pose, betas=sv.betas, model=model, regs=regs)
            sp.no_hands = True
        # weight configuration used in the paper, with joints + confidence values from the CNN
        # (all the weights used in the code were obtained via grid search, see the paper for more details)
        # the first list contains the weights for the pose priors,
        # the second list contains the weights for the shape prior
        opt_weights = zip([4.04 * 1e2, 4.04 * 1e2, 57.4, 4.78],
                          [1e2, 5 * 1e1, 1e1, .5 * 1e1])

        # run the optimization in 4 stages, progressively decreasing the
        # weights for the priors
        for stage, (w, wbetas) in enumerate(opt_weights):
            print 'stage', stage
            objs = {}

            objs['j2d'] = obj_j2d(1., 100)

            objs['pose'] = pprior(w)

            objs['pose_exp'] = obj_angle(0.317 * w)

            objs['betas'] = wbetas * betas

            if regs is not None:
                objs['sph_coll'] = 1e3 * sp

            # HYUNHO
            start = time()
            # maxiter: 100, e_3: .0001
            ch.minimize(objs,
                x0=[sv.betas, sv.pose],
                method='dogleg',
                callback=on_step,
                options={'maxiter': 100,
                         'e_3': .0001,
                         'disp': 0})
            print time()-start

        t1 = time()
        print 'elapsed ', str(t1 - t0)
        if try_both_orient:
            errors.append((objs['j2d'].r**2).sum())
        svs.append(sv)
        cams.append(cam)

    if try_both_orient and errors[0] > errors[1]:
        choose_id = 1
    else:
        choose_id = 0

    return (svs[choose_id], cams[choose_id].r, cams[choose_id].t.r)

def run_single_fit(img,
                   j2d,
                   conf,
                   model,
                   regs=None,
                   n_betas=10,
                   flength=5000.,
                   pix_thsh=25.,
                   scale_factor=1,
                   do_degrees=None):
    """Run the fit for one specific image.
    :param img: h x w x 3 image
    :param j2d: 14x2 array of CNN joints
    :param conf: 14D vector storing the confidence values from the CNN
    :param model: SMPL model
    :param regs: regressors for capsules' axis and radius, if not None enables the interpenetration error term
    :param n_betas: number of shape coefficients considered during optimization
    :param flength: camera focal length (kept fixed during optimization)
    :param pix_thsh: threshold (in pixel), if the distance between shoulder joints in 2D
                     is lower than pix_thsh, the body orientation as ambiguous (so a fit is run on both
                     the estimated one and its flip)
    :param scale_factor: int, rescale the image (for LSP, slightly greater images -- 2x -- help obtain better fits)
    :param do_degrees: list of degrees in azimuth to render the final fit when saving results
    :returns: a tuple containing camera/model parameters and images with rendered fits
    """
    if do_degrees is None:
        do_degrees = []

    # create the pose prior (GMM over CMU)
    prior = MaxMixtureCompletePrior(n_gaussians=8).get_gmm_prior()
    # get the mean pose as our initial pose
    init_pose = np.hstack((np.zeros(3), prior.weights.dot(prior.means)))

    if scale_factor != 1:
        img = cv2.resize(img, (img.shape[1] * scale_factor,
                               img.shape[0] * scale_factor))
        j2d[:, 0] *= scale_factor
        j2d[:, 1] *= scale_factor

    # estimate the camera parameters
    (cam, try_both_orient, body_orient) = initialize_camera(
        model,
        j2d,
        img,
        init_pose,
        flength=flength,
        pix_thsh=pix_thsh)

    # fit
    (sv, opt_j2d, t) = optimize_on_joints(
        j2d,
        model,
        cam,
        img,
        prior,
        try_both_orient,
        body_orient,
        n_betas=n_betas,
        conf=conf,
        regs=regs, )

    h = img.shape[0]
    w = img.shape[1]
    dist = np.abs(cam.t.r[2] - np.mean(sv.r, axis=0)[2])

    images = []
    orig_v = sv.r

    for deg in do_degrees:
        if deg != 0:
            aroundy = cv2.Rodrigues(np.array([0, np.radians(deg), 0]))[0]
            center = orig_v.mean(axis=0)
            new_v = np.dot((orig_v - center), aroundy)
            verts = new_v + center
        else:
            verts = orig_v

        # now render
        im = (render_model(
            verts, model.f, w, h, cam, far=20 + dist) * 255.).astype('uint8')
        images.append(im)

    # return fit parameters
    params = {'cam_t': cam.t.r,
              'f': cam.f.r,
              'pose': sv.pose.r,
              'betas': sv.betas.r,
              'conf': conf}

    return params, images

def openPose(img, gamma=1.0, ui=False):
    params = dict()
    model_folder = ('' if ui else '../') + 'models'
    params["model_folder"] = model_folder
    params["number_people_max"] = 1;

    opWrapper = op.WrapperPython()
    opWrapper.configure(params)
    opWrapper.start()

    datum = op.Datum()

    # hyunho gamma
    if gamma != 1.0:
        img2 = adjust_gamma(img, gamma=gamma)
        datum.cvInputData = img2
    else:
        datum.cvInputData = img

    opWrapper.emplaceAndPop([datum])

    mydata = datum.poseKeypoints[0]
    joints = np.zeros((14,2))
    conf = np.zeros(14)

    joints[0] = mydata[11, :2]
    joints[1] = mydata[10, :2]
    joints[2] = mydata[9, :2]
    joints[3] = mydata[12, :2]
    joints[4] = mydata[13, :2]
    joints[5] = mydata[14, :2]
    joints[6] = mydata[4, :2]
    joints[7] = mydata[3, :2]
    joints[8] = mydata[2, :2]
    joints[9] = mydata[5, :2]
    joints[10] = mydata[6, :2]
    joints[11] = mydata[7, :2]
    joints[12] = mydata[1, :2]
    joints[13] = mydata[0, :2]

    conf[0] = mydata[11, 2]
    conf[1] = mydata[10, 2]
    conf[2] = mydata[9, 2]
    conf[3] = mydata[12, 2]
    conf[4] = mydata[13, 2]
    conf[5] = mydata[14, 2]
    conf[6] = mydata[4, 2]
    conf[7] = mydata[3, 2]
    conf[8] = mydata[2, 2]
    conf[9] = mydata[5, 2]
    conf[10] = mydata[6, 2]
    conf[11] = mydata[7, 2]
    conf[12] = mydata[1, 2]
    conf[13] = mydata[0, 2]

    # head joint estimate
    # joints[13] += 0.5 * (joints[13] - joints[12])

    return joints, conf

def smplify(img, out_path, out_video_dir, op_joints, op_confs, cnt, genders, model_female, model_male, sph_regs_female, sph_regs_male, n_betas, flength, pix_thsh, do_degrees):
        print out_path
        if exists(out_path):
            out_path = out_path[:-4] + '_new.pkl'

        print 'Fitting 3D body saving to', out_path

        if img.ndim == 2:
            print "The image is grayscale!"
            img = np.dstack((img, img, img))

        # openpose
        joints, conf = op_joints[cnt], op_confs[cnt]

        # decide gender
        gender = 'male' if int(genders[cnt]) == 0 else 'female'
        if gender == 'female':
            model = model_female
            sph_regs = sph_regs_female
        elif gender == 'male':
            model = model_male
            sph_regs = sph_regs_male

        # female
        model = model_female
        sph_regs = sph_regs_female

        # run
        params, vis = run_single_fit(
            img,
            joints,
            conf,
            model,
            regs=sph_regs,
            n_betas=n_betas,
            flength=flength,
            pix_thsh=pix_thsh,
            scale_factor=1,
            do_degrees=do_degrees)

        with open(out_path, 'w') as outf:
            params['op_joints'] = joints
            pickle.dump(params, outf)

        # This only saves the first rendering.
        if do_degrees is not None:
            cv2.imwrite(out_path.replace('.pkl', '.png'), vis[0])
#            out.write(vis[0])

def getFrames(temps, frame_cnt, ratio):
    frames = []
    cnt = 0
    if ratio > 1.0:
        for i, temp in enumerate(temps):
            if cnt < int(i * ratio):
                for _ in range(cnt, int(i * ratio)):
                    frames.append(temps[i - 1])
                    cnt += 1

                    if cnt == frame_cnt:
                        break

            if cnt == frame_cnt:
                break

            frames.append(temp)
            cnt += 1
    elif ratio < 1.0:
        for i, temp in enumerate(temps):
            if cnt > int(i * ratio):
                continue

            frames.append(temp)
            cnt += 1

            if cnt == frame_cnt:
                break
    else:
        frames = temps[:frame_cnt]
    return frames

def mmain(video="False",
         ui=False,
         female=True,
         num=-1,
         gamma=1.0):
    n_betas=10
    flength=5000.
    pix_thsh=25.

    # Set up paths & load models.
    # Assumes 'models' in the 'code/' directory where this file is in.
    MODEL_DIR = join(abspath(dirname(__file__)), 'models')

    # Model paths:
    MODEL_NEUTRAL_PATH = join(MODEL_DIR, 'basicModel_neutral_lbs_10_207_0_v1.0.0.pkl')
    MODEL_FEMALE_PATH = join(MODEL_DIR, 'basicModel_f_lbs_10_207_0_v1.0.0.pkl')
    MODEL_MALE_PATH = join(MODEL_DIR, 'basicmodel_m_lbs_10_207_0_v1.0.0.pkl')

    # paths to the npz files storing the regressors for capsules
    SPH_REGS_NEUTRAL_PATH = join(MODEL_DIR, 'regressors_locked_normalized_hybrid.npz')
    SPH_REGS_FEMALE_PATH = join(MODEL_DIR, 'regressors_locked_normalized_female.npz')
    SPH_REGS_MALE_PATH = join(MODEL_DIR, 'regressors_locked_normalized_male.npz')

    img_dir = ('' if ui else '../') + 'image'
    data_dir = ('opmlify/' if ui else '') + 'results/lsp'
    out_dir = ('opmlify/' if ui else '') + 'result/'
    csv_dir = ('opmlify/' if ui else '') + 'lsp/lsp_gender.csv'
    est_dir = ('opmlify/' if ui else '') + 'lsp/est_joints.npz'

    if not exists(out_dir):
        makedirs(out_dir)

    # Render degrees: List of degrees in azimuth to render the final fit.
    # Note that rendering many views can take a while.
    do_degrees = [0.]

    sph_regs = None
    # File storing information about gender in LSP
    with open(csv_dir) as f:
        genders = f.readlines()
    model_female = load_model(MODEL_FEMALE_PATH)
    model_male = load_model(MODEL_MALE_PATH)
    sph_regs_male = np.load(SPH_REGS_MALE_PATH)
    sph_regs_female = np.load(SPH_REGS_FEMALE_PATH)

    # Load joints
    est = np.load(est_dir)['est_joints']

    # hyunho
    # Load video
    if video != "False":
        out_video_dir = out_dir + video[:-4]

        if not exists (out_video_dir):
            makedirs (out_video_dir)

        # first, openpose
        cap = cv2.VideoCapture(join(img_dir, video))

        fps = float(cap.get(cv2.CAP_PROP_FPS))
        video_h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        video_w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        frame_cnt = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        ratio = 10.0 / fps

        # only 60 sec
        if frame_cnt > 60.0 * fps:
            frame_cnt = int(60.0 * fps)

        if ratio > 1.0:
            frame_cnt = int(frame_cnt * ratio)

        # get captures
        temps = []
        cnt = 0
        while cap.isOpened():
            _, img = cap.read()
            temps.append(img)

            cnt += 1
            if cnt == frame_cnt:
                break
        cap.release()

        # get frames
        frames = getFrames(temps, frame_cnt, ratio)

        if not exists(out_video_dir + '/img'):
            makedirs(out_video_dir + '/img')

        for cnt, img in enumerate(frames):
            if cnt <= num:
                continue
            out_path = '%s/%04d.png' % (out_video_dir + '/img', cnt)
            cv2.imwrite(out_path, img)
        return

        # make openpose pkl
        op_pickle_path = join(out_video_dir, 'hyunho.pkl')
        if not exists(op_pickle_path):
            temp_joints = []
            temp_confs = []

            # get openpose result
            for i, frame in enumerate(frames):
                print i
                temp_joint, temp_conf = openPose(frame, gamma, ui)
                temp_joints.append(temp_joint)
                temp_confs.append(temp_conf)

            # save to 'hyunho.pkl'
            with open(op_pickle_path, 'w') as outf:
                temp_params = dict()
                temp_params['op_joints'] = temp_joints
                temp_params['conf'] = temp_confs
                pickle.dump(temp_params, outf)

        # open openpose pkl
        with open(op_pickle_path, "rb") as f:
            op_datas = pickle.load(f)
            op_joints = op_datas['op_joints']
            op_confs = op_datas['conf']

        # view
        op_view = False
        if op_view:
            for cnt, img in enumerate(frames):
                for i, j in enumerate(op_joints[cnt]):
                    if op_confs[cnt][i] < 0.33:
                        temp_color = (0, 0, 255)
                    elif op_confs[cnt][i] > 0.66:
                        temp_color = (0, 255, 0)
                    else:
                        temp_color = (255, 0, 0)

                    cv2.putText(img, str(int(op_confs[cnt][i] * 100)), (int(j[0]), int(j[1])), cv2.FONT_HERSHEY_SCRIPT_SIMPLEX, 0.5, temp_color)
                    cv2.circle(img, (int(j[0]), int(j[1])), 10, temp_color)

                cv2.imshow("img", img)
                k = cv2.waitKey(0)
                if k == 27:
                    return
                elif k == ord('q'):
                    break

        frame_cnt = len(op_joints)

        # openpose pose filter
        op_filter = True
        changed = []
        if op_filter:
            for cnt, img in enumerate(frames):
                for num, _ in enumerate(op_joints[cnt]):
                    joint = op_joints[cnt][num]
                    conf = op_confs[cnt][num]

                    if conf < 0.33:
                        finish_l = finish_r = False

                        # set range
                        for i in range(1, 3):
                            if not finish_l and 0 <= cnt - i:
                                if 0.66 < op_confs[cnt - i][num]:
                                    conf_l = op_confs[cnt - i][num]
                                    joint_l = op_joints[cnt - i][num]
                                    index_l = cnt - i
                                    finish_l = True

                            if not finish_r and cnt + i < frame_cnt:
                                if 0.66 < op_confs[cnt + i][num]:
                                    conf_r = op_confs[cnt + i][num]
                                    joint_r = op_joints[cnt + i][num]
                                    index_r = cnt + i
                                    finish_r = True

                        if finish_l and finish_r:
                            if cnt not in changed:
                                changed.append(cnt)

                            conf_gradient = (conf_r - conf_l) / (index_r - index_l)
                            conf_bias = conf_l - conf_gradient * index_l
                            joint_gradient = (joint_r - joint_l) / (index_r - index_l)
                            joint_bias = joint_l - joint_gradient * index_l

                            for j in range(index_l + 1, index_r):
                                op_confs[j][num] = - (conf_gradient * j + conf_bias)
                                op_joints[j][num] = joint_gradient * j + joint_bias

	print changed

        # smplify
        max_proc = 5
        procs = []
        for cnt, img in enumerate(frames):
            if cnt <= num:
                continue

            print "video frame num: " + str(cnt)
            out_path = '%s/%04d.pkl' % (out_video_dir, cnt)

            if not exists(out_path) or (cnt in changed and not exists(out_path[:-4] + '_new.pkl')):
                if len(procs) < max_proc:
                    procs.append(Process(target=smplify, args=(img, out_path, out_video_dir, op_joints, op_confs, cnt, genders, model_female, model_male, sph_regs_female, sph_regs_male, n_betas, flength, pix_thsh, do_degrees)))

            if len(procs) == max_proc:
                for proc in procs:
                    proc.start()
                for proc in procs:
                    proc.join()

                procs[:] = []
                cap.release()

                return cnt, False

        if len(procs) > 0:
            for proc in procs:
                proc.start()
            for proc in procs:
                proc.join()
            procs[:] = []

        return cnt, True

def adjust_gamma(image, gamma=1.0):
    image = image / 255.0
    image = cv2.pow(image, 1.0 / gamma)
    return np.uint8(image * 255)

if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO)

    parser = argparse.ArgumentParser(description='run SMPLify on LSP dataset')
    parser.add_argument(
        '--video',
        default="False",
        help="Input Video Path")
    args = parser.parse_args()

    mmain(args.video)

max_stars_repo_path: etc code/opmlify2.py

max_issues_repo_path: etc code/opmlify2.py

max_forks_repo_path: etc code/opmlify2.py


text:
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
import importlib
import os
import json
import logging
import pickle
from unittest import mock

import pytest
import numpy as np
import pandas as pd
import sklearn.datasets as datasets
import yaml

import mlflow.pyfunc as pyfunc
import mlflow.pytorch
import mlflow.pyfunc.scoring_server as pyfunc_scoring_server
from mlflow import tracking
from mlflow.exceptions import MlflowException
from mlflow.models import Model, infer_signature
from mlflow.models.utils import _read_example
from mlflow.pytorch import pickle_module as mlflow_pytorch_pickle_module
from mlflow.store.artifact.s3_artifact_repo import S3ArtifactRepository
from mlflow.tracking.artifact_utils import _download_artifact_from_uri
from mlflow.utils.environment import _mlflow_conda_env
from mlflow.utils.file_utils import TempDir
from mlflow.utils.model_utils import _get_flavor_configuration
from mlflow.tracking._model_registry import DEFAULT_AWAIT_MAX_SLEEP_SECONDS

_logger = logging.getLogger(__name__)

# This test suite is included as a code dependency when testing PyTorch model scoring in new
# processes and docker containers. In these environments, the `tests` module is not available.
# Therefore, we attempt to import from `tests` and gracefully emit a warning if it's unavailable.
try:
    from tests.helper_functions import pyfunc_serve_and_score_model
    from tests.helper_functions import score_model_in_sagemaker_docker_container
    from tests.helper_functions import set_boto_credentials  # pylint: disable=unused-import
    from tests.helper_functions import mock_s3_bucket  # pylint: disable=unused-import
except ImportError:
    _logger.warning(
        "Failed to import test helper functions. Tests depending on these functions may fail!"
    )

@pytest.fixture(scope="module")
def data():
    iris = datasets.load_iris()
    data = pd.DataFrame(
        data=np.c_[iris["data"], iris["target"]], columns=iris["feature_names"] + ["target"]
    )
    y = data["target"]
    x = data.drop("target", axis=1)
    return x, y

def get_dataset(data):
    x, y = data
    dataset = [(xi.astype(np.float32), yi.astype(np.float32)) for xi, yi in zip(x.values, y.values)]
    return dataset

def train_model(model, data):
    dataset = get_dataset(data)
    criterion = nn.MSELoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
    batch_size = 16
    num_workers = 4
    dataloader = DataLoader(
        dataset, batch_size=batch_size, num_workers=num_workers, shuffle=True, drop_last=False
    )

    model.train()
    for _ in range(5):
        for batch in dataloader:
            optimizer.zero_grad()
            batch_size = batch[0].shape[0]
            y_pred = model(batch[0]).squeeze(dim=1)
            loss = criterion(y_pred, batch[1])
            loss.backward()
            optimizer.step()

@pytest.fixture(scope="module")
def sequential_model(data):
    model = nn.Sequential(nn.Linear(4, 3), nn.ReLU(), nn.Linear(3, 1),)

    train_model(model=model, data=data)
    return model

def get_subclassed_model_definition():
    """
    Defines a PyTorch model class that inherits from ``torch.nn.Module``. This method can be invoked
    within a pytest fixture to define the model class in the ``__main__`` scope. Alternatively, it
    can be invoked within a module to define the class in the module's scope.
    """

    # pylint: disable=W0223
    class SubclassedModel(torch.nn.Module):
        def __init__(self):
            super(SubclassedModel, self).__init__()
            self.linear = torch.nn.Linear(4, 1)

        def forward(self, x):
            # pylint: disable=arguments-differ
            y_pred = self.linear(x)
            return y_pred

    return SubclassedModel

@pytest.fixture(scope="module")
def main_scoped_subclassed_model(data):
    """
    A custom PyTorch model inheriting from ``torch.nn.Module`` whose class is defined in the
    "__main__" scope.
    """
    model_class = get_subclassed_model_definition()
    model = model_class()
    train_model(model=model, data=data)
    return model

# pylint: disable=W0223
class ModuleScopedSubclassedModel(get_subclassed_model_definition()):
    """
    A custom PyTorch model class defined in the test module scope. This is a subclass of
    ``torch.nn.Module``.
    """

@pytest.fixture(scope="module")
def module_scoped_subclassed_model(data):
    """
    A custom PyTorch model inheriting from ``torch.nn.Module`` whose class is defined in the test
    module scope.
    """
    model = ModuleScopedSubclassedModel()
    train_model(model=model, data=data)
    return model

@pytest.fixture
def model_path(tmpdir):
    return os.path.join(str(tmpdir), "model")

@pytest.fixture
def pytorch_custom_env(tmpdir):
    conda_env = os.path.join(str(tmpdir), "conda_env.yml")
    _mlflow_conda_env(
        conda_env,
        additional_conda_deps=["pytorch", "torchvision", "pytest"],
        additional_conda_channels=["pytorch"],
    )
    return conda_env

def _predict(model, data):
    dataset = get_dataset(data)
    batch_size = 16
    num_workers = 4
    dataloader = DataLoader(
        dataset, batch_size=batch_size, num_workers=num_workers, shuffle=False, drop_last=False
    )
    predictions = np.zeros((len(dataloader.sampler),))
    model.eval()
    with torch.no_grad():
        for i, batch in enumerate(dataloader):
            y_preds = model(batch[0]).squeeze(dim=1).numpy()
            predictions[i * batch_size : (i + 1) * batch_size] = y_preds
    return predictions

@pytest.fixture(scope="module")
def sequential_predicted(sequential_model, data):
    return _predict(sequential_model, data)

@pytest.mark.large
def test_signature_and_examples_are_saved_correctly(sequential_model, data):
    model = sequential_model
    signature_ = infer_signature(*data)
    example_ = data[0].head(3)
    for signature in (None, signature_):
        for example in (None, example_):
            with TempDir() as tmp:
                path = tmp.path("model")
                mlflow.pytorch.save_model(
                    model, path=path, signature=signature, input_example=example
                )
                mlflow_model = Model.load(path)
                assert signature == mlflow_model.signature
                if example is None:
                    assert mlflow_model.saved_input_example_info is None
                else:
                    assert all((_read_example(mlflow_model, path) == example).all())

@pytest.mark.large
def test_log_model(sequential_model, data, sequential_predicted):
    old_uri = tracking.get_tracking_uri()
    # should_start_run tests whether or not calling log_model() automatically starts a run.
    for should_start_run in [False, True]:
        with TempDir(chdr=True, remove_on_exit=True) as tmp:
            try:
                tracking.set_tracking_uri(tmp.path("test"))
                if should_start_run:
                    mlflow.start_run()

                artifact_path = "pytorch"
                mlflow.pytorch.log_model(sequential_model, artifact_path=artifact_path)
                model_uri = "runs:/{run_id}/{artifact_path}".format(
                    run_id=mlflow.active_run().info.run_id, artifact_path=artifact_path
                )

                # Load model
                sequential_model_loaded = mlflow.pytorch.load_model(model_uri=model_uri)

                test_predictions = _predict(sequential_model_loaded, data)
                np.testing.assert_array_equal(test_predictions, sequential_predicted)
            finally:
                mlflow.end_run()
                tracking.set_tracking_uri(old_uri)

def test_log_model_calls_register_model(module_scoped_subclassed_model):
    custom_pickle_module = pickle
    artifact_path = "model"
    register_model_patch = mock.patch("mlflow.register_model")
    with mlflow.start_run(), register_model_patch:
        mlflow.pytorch.log_model(
            artifact_path=artifact_path,
            pytorch_model=module_scoped_subclassed_model,
            conda_env=None,
            pickle_module=custom_pickle_module,
            registered_model_name="AdsModel1",
        )
        model_uri = "runs:/{run_id}/{artifact_path}".format(
            run_id=mlflow.active_run().info.run_id, artifact_path=artifact_path
        )
        mlflow.register_model.assert_called_once_with(
            model_uri, "AdsModel1", await_registration_for=DEFAULT_AWAIT_MAX_SLEEP_SECONDS
        )

def test_log_model_no_registered_model_name(module_scoped_subclassed_model):
    custom_pickle_module = pickle
    artifact_path = "model"
    register_model_patch = mock.patch("mlflow.register_model")
    with mlflow.start_run(), register_model_patch:
        mlflow.pytorch.log_model(
            artifact_path=artifact_path,
            pytorch_model=module_scoped_subclassed_model,
            conda_env=None,
            pickle_module=custom_pickle_module,
        )
        mlflow.register_model.assert_not_called()

@pytest.mark.large
def test_raise_exception(sequential_model):
    with TempDir(chdr=True, remove_on_exit=True) as tmp:
        path = tmp.path("model")
        with pytest.raises(IOError):
            mlflow.pytorch.load_model(path)

        with pytest.raises(TypeError):
            mlflow.pytorch.save_model([1, 2, 3], path)

        mlflow.pytorch.save_model(sequential_model, path)
        with pytest.raises(RuntimeError):
            mlflow.pytorch.save_model(sequential_model, path)

        from mlflow import sklearn
        import sklearn.neighbors as knn

        path = tmp.path("knn.pkl")
        knn = knn.KNeighborsClassifier()
        with open(path, "wb") as f:
            pickle.dump(knn, f)
        path = tmp.path("knn")
        sklearn.save_model(knn, path=path)
        with pytest.raises(MlflowException):
            mlflow.pytorch.load_model(path)

@pytest.mark.large
def test_save_and_load_model(sequential_model, model_path, data, sequential_predicted):
    mlflow.pytorch.save_model(sequential_model, model_path)

    # Loading pytorch model
    sequential_model_loaded = mlflow.pytorch.load_model(model_path)
    np.testing.assert_array_equal(_predict(sequential_model_loaded, data), sequential_predicted)

    # Loading pyfunc model
    pyfunc_loaded = mlflow.pyfunc.load_pyfunc(model_path)
    np.testing.assert_array_almost_equal(
        pyfunc_loaded.predict(data[0]).values[:, 0], sequential_predicted, decimal=4
    )

@pytest.mark.large
def test_load_model_from_remote_uri_succeeds(
    sequential_model, model_path, mock_s3_bucket, data, sequential_predicted
):
    mlflow.pytorch.save_model(sequential_model, model_path)

    artifact_root = "s3://{bucket_name}".format(bucket_name=mock_s3_bucket)
    artifact_path = "model"
    artifact_repo = S3ArtifactRepository(artifact_root)
    artifact_repo.log_artifacts(model_path, artifact_path=artifact_path)

    model_uri = artifact_root + "/" + artifact_path
    sequential_model_loaded = mlflow.pytorch.load_model(model_uri=model_uri)
    np.testing.assert_array_equal(_predict(sequential_model_loaded, data), sequential_predicted)

@pytest.mark.large
def test_model_save_persists_specified_conda_env_in_mlflow_model_directory(
    sequential_model, model_path, pytorch_custom_env
):
    mlflow.pytorch.save_model(
        pytorch_model=sequential_model, path=model_path, conda_env=pytorch_custom_env
    )

    pyfunc_conf = _get_flavor_configuration(model_path=model_path, flavor_name=pyfunc.FLAVOR_NAME)
    saved_conda_env_path = os.path.join(model_path, pyfunc_conf[pyfunc.ENV])
    assert os.path.exists(saved_conda_env_path)
    assert saved_conda_env_path != pytorch_custom_env

    with open(pytorch_custom_env, "r") as f:
        pytorch_custom_env_text = f.read()
    with open(saved_conda_env_path, "r") as f:
        saved_conda_env_text = f.read()
    assert saved_conda_env_text == pytorch_custom_env_text

@pytest.mark.large
def test_model_save_accepts_conda_env_as_dict(sequential_model, model_path):
    conda_env = dict(mlflow.pytorch.get_default_conda_env())
    conda_env["dependencies"].append("pytest")
    mlflow.pytorch.save_model(pytorch_model=sequential_model, path=model_path, conda_env=conda_env)

    pyfunc_conf = _get_flavor_configuration(model_path=model_path, flavor_name=pyfunc.FLAVOR_NAME)
    saved_conda_env_path = os.path.join(model_path, pyfunc_conf[pyfunc.ENV])
    assert os.path.exists(saved_conda_env_path)

    with open(saved_conda_env_path, "r") as f:
        saved_conda_env_parsed = yaml.safe_load(f)
    assert saved_conda_env_parsed == conda_env

@pytest.mark.large
def test_model_log_persists_specified_conda_env_in_mlflow_model_directory(
    sequential_model, pytorch_custom_env
):
    artifact_path = "model"
    with mlflow.start_run():
        mlflow.pytorch.log_model(
            pytorch_model=sequential_model,
            artifact_path=artifact_path,
            conda_env=pytorch_custom_env,
        )
        model_path = _download_artifact_from_uri(
            "runs:/{run_id}/{artifact_path}".format(
                run_id=mlflow.active_run().info.run_id, artifact_path=artifact_path
            )
        )

    pyfunc_conf = _get_flavor_configuration(model_path=model_path, flavor_name=pyfunc.FLAVOR_NAME)
    saved_conda_env_path = os.path.join(model_path, pyfunc_conf[pyfunc.ENV])
    assert os.path.exists(saved_conda_env_path)
    assert saved_conda_env_path != pytorch_custom_env

    with open(pytorch_custom_env, "r") as f:
        pytorch_custom_env_text = f.read()
    with open(saved_conda_env_path, "r") as f:
        saved_conda_env_text = f.read()
    assert saved_conda_env_text == pytorch_custom_env_text

@pytest.mark.large
def test_model_save_without_specified_conda_env_uses_default_env_with_expected_dependencies(
    sequential_model, model_path
):
    mlflow.pytorch.save_model(pytorch_model=sequential_model, path=model_path, conda_env=None)

    pyfunc_conf = _get_flavor_configuration(model_path=model_path, flavor_name=pyfunc.FLAVOR_NAME)
    conda_env_path = os.path.join(model_path, pyfunc_conf[pyfunc.ENV])
    with open(conda_env_path, "r") as f:
        conda_env = yaml.safe_load(f)

    assert conda_env == mlflow.pytorch.get_default_conda_env()

@pytest.mark.large
def test_model_log_without_specified_conda_env_uses_default_env_with_expected_dependencies(
    sequential_model,
):
    artifact_path = "model"
    with mlflow.start_run():
        mlflow.pytorch.log_model(
            pytorch_model=sequential_model, artifact_path=artifact_path, conda_env=None
        )
        model_path = _download_artifact_from_uri(
            "runs:/{run_id}/{artifact_path}".format(
                run_id=mlflow.active_run().info.run_id, artifact_path=artifact_path
            )
        )

    pyfunc_conf = _get_flavor_configuration(model_path=model_path, flavor_name=pyfunc.FLAVOR_NAME)
    conda_env_path = os.path.join(model_path, pyfunc_conf[pyfunc.ENV])
    with open(conda_env_path, "r") as f:
        conda_env = yaml.safe_load(f)

    assert conda_env == mlflow.pytorch.get_default_conda_env()

@pytest.mark.large
def test_load_model_with_differing_pytorch_version_logs_warning(sequential_model, model_path):
    mlflow.pytorch.save_model(pytorch_model=sequential_model, path=model_path)
    saver_pytorch_version = "1.0"
    model_config_path = os.path.join(model_path, "MLmodel")
    model_config = Model.load(model_config_path)
    model_config.flavors[mlflow.pytorch.FLAVOR_NAME]["pytorch_version"] = saver_pytorch_version
    model_config.save(model_config_path)

    log_messages = []

    def custom_warn(message_text, *args, **kwargs):
        log_messages.append(message_text % args % kwargs)

    loader_pytorch_version = "0.8.2"
    with mock.patch("mlflow.pytorch._logger.warning") as warn_mock, mock.patch(
        "torch.__version__"
    ) as torch_version_mock:
        torch_version_mock.__str__ = lambda *args, **kwargs: loader_pytorch_version
        warn_mock.side_effect = custom_warn
        mlflow.pytorch.load_model(model_uri=model_path)

    assert any(
        [
            "does not match installed PyTorch version" in log_message
            and saver_pytorch_version in log_message
            and loader_pytorch_version in log_message
            for log_message in log_messages
        ]
    )

@pytest.mark.large
def test_pyfunc_model_serving_with_module_scoped_subclassed_model_and_default_conda_env(
    module_scoped_subclassed_model, model_path, data
):
    mlflow.pytorch.save_model(
        path=model_path,
        pytorch_model=module_scoped_subclassed_model,
        conda_env=None,
        code_paths=[__file__],
    )

    scoring_response = pyfunc_serve_and_score_model(
        model_uri=model_path,
        data=data[0],
        content_type=pyfunc_scoring_server.CONTENT_TYPE_JSON_SPLIT_ORIENTED,
        extra_args=["--no-conda"],
    )
    assert scoring_response.status_code == 200

    deployed_model_preds = pd.DataFrame(json.loads(scoring_response.content))
    np.testing.assert_array_almost_equal(
        deployed_model_preds.values[:, 0],
        _predict(model=module_scoped_subclassed_model, data=data),
        decimal=4,
    )

def test_save_model_with_wrong_codepaths_fails_corrrectly(
    module_scoped_subclassed_model, model_path, data
):
    # pylint: disable=unused-argument
    with pytest.raises(TypeError) as exc_info:
        mlflow.pytorch.save_model(
            path=model_path,
            pytorch_model=module_scoped_subclassed_model,
            conda_env=None,
            code_paths="some string",
        )
    assert "TypeError: Argument code_paths should be a list, not {}".format(type("")) in str(
        exc_info
    )
    assert not os.path.exists(model_path)

@pytest.mark.large
def test_pyfunc_model_serving_with_main_scoped_subclassed_model_and_custom_pickle_module(
    main_scoped_subclassed_model, model_path, data
):
    mlflow.pytorch.save_model(
        path=model_path,
        pytorch_model=main_scoped_subclassed_model,
        conda_env=None,
        pickle_module=mlflow_pytorch_pickle_module,
    )

    scoring_response = pyfunc_serve_and_score_model(
        model_uri=model_path,
        data=data[0],
        content_type=pyfunc_scoring_server.CONTENT_TYPE_JSON_SPLIT_ORIENTED,
        extra_args=["--no-conda"],
    )
    assert scoring_response.status_code == 200

    deployed_model_preds = pd.DataFrame(json.loads(scoring_response.content))
    np.testing.assert_array_almost_equal(
        deployed_model_preds.values[:, 0],
        _predict(model=main_scoped_subclassed_model, data=data),
        decimal=4,
    )

@pytest.mark.large
def test_load_model_succeeds_with_dependencies_specified_via_code_paths(
    module_scoped_subclassed_model, model_path, data
):
    # Save a PyTorch model whose class is defined in the current test suite. Because the
    # `tests` module is not available when the model is deployed for local scoring, we include
    # the test suite file as a code dependency
    mlflow.pytorch.save_model(
        path=model_path,
        pytorch_model=module_scoped_subclassed_model,
        conda_env=None,
        code_paths=[__file__],
    )

    # Define a custom pyfunc model that loads a PyTorch model artifact using
    # `mlflow.pytorch.load_model`
    class TorchValidatorModel(pyfunc.PythonModel):
        def load_context(self, context):
            # pylint: disable=attribute-defined-outside-init
            self.pytorch_model = mlflow.pytorch.load_model(context.artifacts["pytorch_model"])

        def predict(self, context, model_input):
            with torch.no_grad():
                input_tensor = torch.from_numpy(model_input.values.astype(np.float32))
                output_tensor = self.pytorch_model(input_tensor)
                return pd.DataFrame(output_tensor.numpy())

    pyfunc_artifact_path = "pyfunc_model"
    with mlflow.start_run():
        pyfunc.log_model(
            artifact_path=pyfunc_artifact_path,
            python_model=TorchValidatorModel(),
            artifacts={"pytorch_model": model_path},
        )
        pyfunc_model_path = _download_artifact_from_uri(
            "runs:/{run_id}/{artifact_path}".format(
                run_id=mlflow.active_run().info.run_id, artifact_path=pyfunc_artifact_path
            )
        )

    # Deploy the custom pyfunc model and ensure that it is able to successfully load its
    # constituent PyTorch model via `mlflow.pytorch.load_model`
    scoring_response = pyfunc_serve_and_score_model(
        model_uri=pyfunc_model_path,
        data=data[0],
        content_type=pyfunc_scoring_server.CONTENT_TYPE_JSON_SPLIT_ORIENTED,
        extra_args=["--no-conda"],
    )
    assert scoring_response.status_code == 200

    deployed_model_preds = pd.DataFrame(json.loads(scoring_response.content))
    np.testing.assert_array_almost_equal(
        deployed_model_preds.values[:, 0],
        _predict(model=module_scoped_subclassed_model, data=data),
        decimal=4,
    )

@pytest.mark.large
def test_load_pyfunc_loads_torch_model_using_pickle_module_specified_at_save_time(
    module_scoped_subclassed_model, model_path
):
    custom_pickle_module = pickle

    mlflow.pytorch.save_model(
        path=model_path,
        pytorch_model=module_scoped_subclassed_model,
        conda_env=None,
        pickle_module=custom_pickle_module,
    )

    import_module_fn = importlib.import_module
    imported_modules = []

    def track_module_imports(module_name):
        imported_modules.append(module_name)
        return import_module_fn(module_name)

    with mock.patch("importlib.import_module") as import_mock, mock.patch(
        "torch.load"
    ) as torch_load_mock:
        import_mock.side_effect = track_module_imports
        pyfunc.load_pyfunc(model_path)

    torch_load_mock.assert_called_with(mock.ANY, pickle_module=custom_pickle_module)
    assert custom_pickle_module.__name__ in imported_modules

@pytest.mark.large
def test_load_model_loads_torch_model_using_pickle_module_specified_at_save_time(
    module_scoped_subclassed_model,
):
    custom_pickle_module = pickle

    artifact_path = "pytorch_model"
    with mlflow.start_run():
        mlflow.pytorch.log_model(
            artifact_path=artifact_path,
            pytorch_model=module_scoped_subclassed_model,
            conda_env=None,
            pickle_module=custom_pickle_module,
        )
        model_uri = "runs:/{run_id}/{artifact_path}".format(
            run_id=mlflow.active_run().info.run_id, artifact_path=artifact_path
        )

    import_module_fn = importlib.import_module
    imported_modules = []

    def track_module_imports(module_name):
        imported_modules.append(module_name)
        return import_module_fn(module_name)

    with mock.patch("importlib.import_module") as import_mock, mock.patch(
        "torch.load"
    ) as torch_load_mock:
        import_mock.side_effect = track_module_imports
        pyfunc.load_pyfunc(model_uri=model_uri)

    torch_load_mock.assert_called_with(mock.ANY, pickle_module=custom_pickle_module)
    assert custom_pickle_module.__name__ in imported_modules

@pytest.mark.large
def test_load_pyfunc_succeeds_when_data_is_model_file_instead_of_directory(
    module_scoped_subclassed_model, model_path, data
):
    """
    This test verifies that PyTorch models saved in older versions of MLflow are loaded successfully
    by ``mlflow.pytorch.load_model``. The ``data`` path associated with these older models is
    serialized PyTorch model file, as opposed to the current format: a directory containing a
    serialized model file and pickle module information.
    """
    mlflow.pytorch.save_model(
        path=model_path, pytorch_model=module_scoped_subclassed_model, conda_env=None
    )

    model_conf_path = os.path.join(model_path, "MLmodel")
    model_conf = Model.load(model_conf_path)
    pyfunc_conf = model_conf.flavors.get(pyfunc.FLAVOR_NAME)
    assert pyfunc_conf is not None
    model_data_path = os.path.join(model_path, pyfunc_conf[pyfunc.DATA])
    assert os.path.exists(model_data_path)
    assert mlflow.pytorch._SERIALIZED_TORCH_MODEL_FILE_NAME in os.listdir(model_data_path)
    pyfunc_conf[pyfunc.DATA] = os.path.join(
        model_data_path, mlflow.pytorch._SERIALIZED_TORCH_MODEL_FILE_NAME
    )
    model_conf.save(model_conf_path)

    loaded_pyfunc = pyfunc.load_pyfunc(model_path)

    np.testing.assert_array_almost_equal(
        loaded_pyfunc.predict(data[0]),
        pd.DataFrame(_predict(model=module_scoped_subclassed_model, data=data)),
        decimal=4,
    )

@pytest.mark.large
def test_load_model_succeeds_when_data_is_model_file_instead_of_directory(
    module_scoped_subclassed_model, model_path, data
):
    """
    This test verifies that PyTorch models saved in older versions of MLflow are loaded successfully
    by ``mlflow.pytorch.load_model``. The ``data`` path associated with these older models is
    serialized PyTorch model file, as opposed to the current format: a directory containing a
    serialized model file and pickle module information.
    """
    artifact_path = "pytorch_model"
    with mlflow.start_run():
        mlflow.pytorch.log_model(
            artifact_path=artifact_path,
            pytorch_model=module_scoped_subclassed_model,
            conda_env=None,
        )
        model_path = _download_artifact_from_uri(
            "runs:/{run_id}/{artifact_path}".format(
                run_id=mlflow.active_run().info.run_id, artifact_path=artifact_path
            )
        )

    model_conf_path = os.path.join(model_path, "MLmodel")
    model_conf = Model.load(model_conf_path)
    pyfunc_conf = model_conf.flavors.get(pyfunc.FLAVOR_NAME)
    assert pyfunc_conf is not None
    model_data_path = os.path.join(model_path, pyfunc_conf[pyfunc.DATA])
    assert os.path.exists(model_data_path)
    assert mlflow.pytorch._SERIALIZED_TORCH_MODEL_FILE_NAME in os.listdir(model_data_path)
    pyfunc_conf[pyfunc.DATA] = os.path.join(
        model_data_path, mlflow.pytorch._SERIALIZED_TORCH_MODEL_FILE_NAME
    )
    model_conf.save(model_conf_path)

    loaded_pyfunc = pyfunc.load_pyfunc(model_path)

    np.testing.assert_array_almost_equal(
        loaded_pyfunc.predict(data[0]),
        pd.DataFrame(_predict(model=module_scoped_subclassed_model, data=data)),
        decimal=4,
    )

@pytest.mark.large
def test_load_model_allows_user_to_override_pickle_module_via_keyword_argument(
    module_scoped_subclassed_model, model_path
):
    mlflow.pytorch.save_model(
        path=model_path,
        pytorch_model=module_scoped_subclassed_model,
        conda_env=None,
        pickle_module=pickle,
    )

    mlflow_torch_pickle_load = mlflow_pytorch_pickle_module.load
    pickle_call_results = {
        "mlflow_torch_pickle_load_called": False,
    }

    def validate_mlflow_torch_pickle_load_called(*args, **kwargs):
        pickle_call_results["mlflow_torch_pickle_load_called"] = True
        return mlflow_torch_pickle_load(*args, **kwargs)

    log_messages = []

    def custom_warn(message_text, *args, **kwargs):
        log_messages.append(message_text % args % kwargs)

    with mock.patch(
        "mlflow.pytorch.pickle_module.load"
    ) as mlflow_torch_pickle_load_mock, mock.patch("mlflow.pytorch._logger.warning") as warn_mock:
        mlflow_torch_pickle_load_mock.side_effect = validate_mlflow_torch_pickle_load_called
        warn_mock.side_effect = custom_warn
        mlflow.pytorch.load_model(model_uri=model_path, pickle_module=mlflow_pytorch_pickle_module)

    assert all(pickle_call_results.values())
    assert any(
        [
            "does not match the pickle module that was used to save the model" in log_message
            and pickle.__name__ in log_message
            and mlflow_pytorch_pickle_module.__name__ in log_message
            for log_message in log_messages
        ]
    )

@pytest.mark.large
def test_load_model_raises_exception_when_pickle_module_cannot_be_imported(
    main_scoped_subclassed_model, model_path
):
    mlflow.pytorch.save_model(
        path=model_path, pytorch_model=main_scoped_subclassed_model, conda_env=None
    )

    bad_pickle_module_name = "not.a.real.module"

    pyfunc_conf = _get_flavor_configuration(model_path=model_path, flavor_name=pyfunc.FLAVOR_NAME)
    model_data_path = os.path.join(model_path, pyfunc_conf[pyfunc.DATA])
    assert os.path.exists(model_data_path)
    assert mlflow.pytorch._PICKLE_MODULE_INFO_FILE_NAME in os.listdir(model_data_path)
    with open(
        os.path.join(model_data_path, mlflow.pytorch._PICKLE_MODULE_INFO_FILE_NAME), "w"
    ) as f:
        f.write(bad_pickle_module_name)

    with pytest.raises(MlflowException) as exc_info:
        mlflow.pytorch.load_model(model_uri=model_path)

    assert "Failed to import the pickle module" in str(exc_info)
    assert bad_pickle_module_name in str(exc_info)

@pytest.mark.release
def test_sagemaker_docker_model_scoring_with_sequential_model_and_default_conda_env(
    model, model_path, data, sequential_predicted
):
    mlflow.pytorch.save_model(pytorch_model=model, path=model_path, conda_env=None)

    scoring_response = score_model_in_sagemaker_docker_container(
        model_uri=model_path,
        data=data[0],
        content_type=pyfunc_scoring_server.CONTENT_TYPE_JSON_SPLIT_ORIENTED,
        flavor=mlflow.pyfunc.FLAVOR_NAME,
        activity_polling_timeout_seconds=360,
    )
    deployed_model_preds = pd.DataFrame(json.loads(scoring_response.content))

    np.testing.assert_array_almost_equal(
        deployed_model_preds.values[:, 0], sequential_predicted, decimal=4
    )

ext: py

max_stars_repo_name: thomasfrederikhoeck/mlflow

max_issues_repo_name: thomasfrederikhoeck/mlflow

max_forks_repo_name: thomasfrederikhoeck/mlflow

reason: import numpy


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

# @Time : 2020/10/15 下午12:09
# @Author : TaoWang
# @Description :

from pytorchtools import EarlyStopping
import torch
import torch.autograd as autograd
import torch.nn as nn
import torch.optim as optim
from model import TextCNN
from data import MR_Dataset
import numpy as np
import config as argumentparser
config = argumentparser.ArgumentParser()
config.filters = list(map(int, config.filters.split(",")))

torch.manual_seed(config.seed)

if torch.cuda.is_available():
    torch.cuda.set_device(config.gpu)

def get_test_result(data_iter, data_set):
    # 生成测试结果
    model.eval()
    data_loss = 0
    true_sample_num = 0
    for data, label in data_iter:
        if config.cuda and torch.cuda.is_available():
            data = data.cuda()
            label = label.cuda()
        out = model(data)
        loss = criterion(out, autograd.Variable(label.long()))
        data_loss += loss.data.item()
        true_sample_num += np.sum((torch.argmax(out, 1) == label).cpu().numpy())
    acc = true_sample_num / data_set.__len__()

    return data_loss, acc

acc = 0
for i in range(0, 10):
    # 10-cv
    early_stopping = EarlyStopping(patience=10,
                                   verbose=True,
                                   cv_index=i)

    training_set = MR_Dataset(state="train",
                              k=i,
                              embedding_type=config.embedding_type)

    if config.use_pretrained_embed:
        config.embedding_pretrained = torch.from_numpy(training_set.weight).float()
    else:
        config.embedding_pretrained = False

    config.n_vocab = training_set.n_vocab
    training_iter = torch.utils.data.DataLoader(dataset=training_set,
                                                batch_size=config.batch_size,
                                                shuffle=True,
                                                num_workers=2)
    valid_set = MR_Dataset(state="valid",
                           k=i,
                           embedding_type="no")

    valid_iter = torch.utils.data.DataLoader(dataset=valid_set,
                                             batch_size=config.batch_size,
                                             shuffle=False,
                                             num_workers=2)

    test_set = MR_Dataset(state="test",
                          k=i,
                          embedding_type="no")

    test_iter = torch.utils.data.DataLoader(dataset=test_set,
                                            batch_size=config.batch_size,
                                            shuffle=False,
                                            num_workers=2)
    model = TextCNN(config)
    if config.cuda and torch.cuda.is_available():
        model.cuda()
        config.embedding_pretrained.cuda()

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=config.learning_rate)
    count = 0
    loss_sum = 0
    for epoch in range(config.epoch):
        # 开始训练
        model.train()
        for data, label in training_iter:
            if config.cuda and torch.cuda.is_available():
                data = data.cuda()
                label = label.cuda()
            else:
                data = torch.autograd.Variable(data).long()
            label = torch.autograd.Variable(label).squeeze()
            out = model(data)
            l2_loss = config.l2*torch.sum(torch.pow(list(model.parameters())[1],2))
            loss = criterion(out, autograd.Variable(label.long()))+l2_loss
            loss_sum += loss.data.item()
            count += 1
            if count % 100 == 0:
                print("epoch", epoch, end='  ')
                print("The loss is: %.5f" % (loss_sum / 100))
                loss_sum = 0
                count = 0
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            # save the model in every epoch
        # 一轮训练结束，在验证集测试
        valid_loss,valid_acc = get_test_result(valid_iter,valid_set)
        early_stopping(valid_loss, model)
        print ("The valid acc is: %.5f" % valid_acc)
        if early_stopping.early_stop:
            print("Early stopping")
            break
    # 1 fold训练结果
    model.load_state_dict(torch.load('./checkpoints/checkpoint%d.pt'%i))
    test_loss, test_acc = get_test_result(test_iter, test_set)
    print("The test acc is: %.5f" % test_acc)
    acc += test_acc/10

# 输出10-fold的平均acc
print("The test acc is: %.5f" % acc)


text:
import numpy as np

class Wiggle(object):

    def __init__(self):
        # Initialize the system.
        self.reset()

    def reset(self):
        # Reset the simulation.
        # Spring/piston configuration.
        self.k = 1 # N/meter
        self.l0 = 1
        self.p = 0 # piston position

        # Mass parameters.
        self.mass = 10 * 1e-3 # ten grams (0.01 kg)
        self.gamma = 9e-1 # air resistance
        self.x = 1.00 * self.l0 # mass position
        self.v = 0
        self.a = 0

        # Simulation parameters.
        self.dt = 1e-2
        self.t = 0
        self.X = []
        self.V = []
        self.A = []
        self.P = []
        self.T = []

        # Return system state on a reset.
        return self.state

    @property
    def action_space(self):
        # 1-dimensional action space.
        return np.zeros(1)

    @property
    def observation_space(self):
        # 3-dimensional observation space.
        return np.zeros(4)

    def step(self, action):
        # Action corresponds to position of piston, p:
        self.p = action[0]
        self.P.append(self.p)

        # Move the simulation forward one discrete timestep.
        dt = self.dt
        self.T.append(self.t)
        self.t += dt

        # Compute current acceleration.
        self.A.append(self.a)
        self.a = -(self.x - self.p - self.l0)/self.mass - self.gamma * self.v

        # Compute current velocity.
        self.V.append(self.v)
        self.v += self.a * dt

        # Compute the current position.
        self.X.append(self.x)
        self.x += self.v * dt

        return self.state, self.reward, self.done, self.info

    @property
    def info(self):
        # Return relevant environment information.
        return {}

    @property
    def state(self):
        # Current state vector of the system
        return [self.x, self. v, self.a, self.t]

    @property
    def done(self):
        # We will run the simulation for 10 seconds.
        return self.t > 10

    @property
    def reward(self):
        # Reward specifiction.
        bonus = 0
        if self.t < 5:
            if np.abs(self.x - self.l0)>1e-2:
                bonus = -100
            return -(self.x - self.l0)**2 + bonus
        else:
            if np.abs(self.x - 1.5 * self.l0)>1e-2:
                bonus = -100
            return -(self.x - 1.5 * self.l0)**2 + bonus
        # if self.t < 5:
        #     if np.abs(self.x - self.l0) > 1e-2:
        #         return -100 - np.abs(self.x - self.l0)
        #     else:
        #         return -np.abs(self.x - self.l0)
        # else:
        #     if np.abs(self.x - 1.5 * self.l0) > 1e-2:
        #         return -100 - np.abs(self.x - 1.5 * self.l0)
        #     else:
        #         return -np.abs(self.x - 1.5 * self.l0)

if __name__ == '__main__':
    import pylab as plt
    plt.ion()
    plt.close('all')

    sim = Wiggle()
    dt = sim.dt
    T = 10 # seconds
    p = 0

    for itr in range(int(T/sim.dt)):
        t = itr * dt
        if (t>5):
            while p<=0.5:
                p += 0.01/0.5 * dt
        sim.step([p])

    plt.plot(sim.T, sim.X)
    plt.ylim([0.5, 2.5])
    plt.xlabel('Time (seconds)')
    plt.ylabel('Mass Position (meters)')
    plt.savefig('ringing.png')

max_stars_repo_name: lightscalar/dopamine

max_issues_repo_name: lightscalar/dopamine

max_forks_repo_name: lightscalar/dopamine


text:
from keras import backend as K
from keras.callbacks import ModelCheckpoint, TensorBoard
from keras.layers import Input, Conv2D, Conv2DTranspose
from keras.layers.advanced_activations import LeakyReLU
from keras.losses import mean_squared_error
from keras.models import Model
from keras.regularizers import l2

import matplotlib.pyplot as plt
import numpy as np
from pydub import AudioSegment
from utils.sounds import dct_by_interval, idct_by_interval

def sampling(args):
    z_mean, z_log_var = args
    raise NotImplementedError('Learn how to sample (maybe from Costco lol)')

def conv_bn_lrelu(x, filters, kernel_size, strides):
    """Applies Conv2D, performs batch normalization, and applies Leaky ReLU
    activation to input `x`."""
    raise NotImplementedError('Your batches are not normal!!!')

# -------------------- #
# Data processing here #
# -------------------- #

# hyperparameters
num_epochs = 200
batch_size = 40
sound_shape = (512, 1378, 2)
latent_shape = (1, 1, 64)

# encoder network
input_sound = Input(shape=sound_shape, name='encoder_input')
x = Conv2D(16, kernel_size=(8, 16), strides=(4, 8))

print('Made it here in one piece!!!')

max_stars_repo_name: sagnibak/vaes

max_issues_repo_name: sagnibak/vaes

max_forks_repo_name: sagnibak/vaes


text:
# -----------------------------------------------------------------------------
# Copyright (c) 2009-2016 Nicolas P. Rougier. All rights reserved.
# Distributed under the (new) BSD License.
# -----------------------------------------------------------------------------
import numpy as np
from galaxy_specrend import *
from galaxy_simulation import Galaxy
from glumpy import app, gloo, gl, glm, data

vertex = """
#version 120
uniform mat4  u_model;
uniform mat4  u_view;
uniform mat4  u_projection;
uniform sampler1D u_colormap;

attribute float a_size;
attribute float a_type;
attribute vec2  a_position;
attribute float a_temperature;
attribute float a_brightness;

varying vec3 v_color;
void main (void)
{
    gl_Position = u_projection * u_view * u_model * vec4(a_position,0.0,1.0);
    if (a_size > 2.0)
    {
        gl_PointSize = a_size;
    } else {
        gl_PointSize = 0.0;
    }
    v_color = texture1D(u_colormap, a_temperature).rgb * a_brightness;
    if (a_type == 2)
        v_color *= vec3(2,1,1);
    else if (a_type == 3)
        v_color = vec3(.9);
}
"""

fragment = """
#version 120
uniform sampler2D u_texture;
varying vec3 v_color;
void main()
{
    gl_FragColor = vec4(texture2D(u_texture, gl_PointCoord).r*v_color, 1.0);
}
"""

window = app.Window(width=800, height=800)

@window.event
def on_init():
    gl.glBlendFunc(gl.GL_SRC_ALPHA, gl.GL_ONE)

@window.event
def on_draw(dt):
    window.clear()
    galaxy.update(100000) # in years !
    program['a_size'] = galaxy['size'] * max(window.width/800.0, window.height/800.0)
    program['a_position'] = galaxy['position'] / 13000.0
    program.draw(gl.GL_POINTS)

@window.event
def on_resize(width,height):
    gl.glViewport(0, 0, width, height)
    projection = glm.perspective(45.0, width/float(height), 1.0, 1000.0)
    program['u_projection'] = projection

galaxy = Galaxy(35000)
galaxy.reset(13000, 4000, 0.0004, 0.90, 0.90, 0.5, 200, 300)
t0, t1 = 1000.0, 10000.0
n = 256
dt =  (t1-t0)/n
colors = np.zeros((n,3), dtype=np.float32)
for i in range(n):
    temperature = t0 + i*dt
    x,y,z = spectrum_to_xyz(bb_spectrum, temperature)
    r,g,b = xyz_to_rgb(SMPTEsystem, x, y, z)
    r = min((max(r,0),1))
    g = min((max(g,0),1))
    b = min((max(b,0),1))
    colors[i] = norm_rgb(r, g, b)

program = gloo.Program(vertex, fragment, count=len(galaxy))

view = np.eye(4, dtype=np.float32)
model = np.eye(4, dtype=np.float32)
projection = np.eye(4, dtype=np.float32)
glm.translate(view, 0, 0, -5)
program['u_model'] = model
program['u_view'] = view
program['u_colormap'] = colors
program['u_texture'] = data.get("particle.png")
program['u_texture'].interpolation = gl.GL_LINEAR

program['a_temperature'] = (galaxy['temperature'] - t0) / (t1-t0)
program['a_brightness'] = galaxy['brightness']
program['a_size'] = galaxy['size']
program['a_type'] = galaxy['type']

gl.glClearColor(0.0, 0.0, 0.03, 1.0)
gl.glDisable(gl.GL_DEPTH_TEST)
gl.glEnable(gl.GL_BLEND)
gl.glBlendFunc(gl.GL_SRC_ALPHA, gl.GL_ONE)

app.run(framerate=60)


text:
"""
Implementation of the class `CuIBMSimulation`.
"""

import os
import struct

import numpy

from ..barbaGroupSimulation import BarbaGroupSimulation
from ..field import Field
from ..force import Force

class CuIBMSimulation(BarbaGroupSimulation):
  """
  Contains info about a cuIBM simulation.
  Inherits from the class BarbaGroupSimulation.
  """

  def __init__(self, description=None, directory=os.getcwd(), **kwargs):
    """
    Initializes by calling the parent constructor.

    Parameters
    ----------
    description: string, optional
      Description of the simulation;
      default: None.
    directory: string, optional
      Directory of the simulation;
      default: <current working directory>.
    **kwargs: dictionary
      Extra keyword-arguments.
    """
    super(CuIBMSimulation, self).__init__(software='cuibm',
                                          description=description,
                                          directory=directory,
                                          **kwargs)

  def read_grid(self, file_path=None):
    """
    Reads the computational grid from file.

    Parameters
    ----------
    file_path: string, optional
      Path of the file containing grid stations along each direction;
      default: None.
    """
    print('[info] reading grid ...')
    if not file_path:
      file_path = os.path.join(self.directory, 'grid')
    # test if file written in binary format
    textchars = (bytearray({7, 8, 9, 10, 12, 13, 27}
                           | set(range(0x20, 0x100)) - {0x7f}))
    is_binary_string = lambda bytes: bool(bytes.translate(None, textchars))
    infile = open(file_path, 'rb')
    binary_format = is_binary_string(infile.read(1024))
    infile.close()
    if binary_format:
      with open(file_path, 'rb') as infile:
        # x-direction
        nx = struct.unpack('i', infile.read(4))[0]
        x = numpy.array(struct.unpack('d' * (nx + 1),
                                      infile.read(8 * (nx + 1))))
        # y-direction
        ny = struct.unpack('i', infile.read(4))[0]
        y = numpy.array(struct.unpack('d' * (ny + 1),
                                      infile.read(8 * (ny + 1))))
    else:
      with open(file_path, 'r') as infile:
        data = numpy.loadtxt(infile, dtype=numpy.float64)
        # x-direction
        nx = int(data[0])
        x, data = data[1:nx + 2], data[nx + 2:]
        # y-direction
        ny = int(data[0])
        y = data[1:]
    self.grid = x, y
    print('\tgrid-size: {}x{}'.format(x.size - 1, y.size - 1))

  def read_forces(self, file_path=None, labels=None, usecols=(0, 1, 2)):
    """
    Reads forces from files.

    Parameters
    ----------
    file_path: string, optional
      Path of the file containing the forces data;
      default: None.
    labels: list of strings, optional
      Label to give to each force that will be read from file;
      default: None
    usecols: tuple of integers, optional
      Index of each column to read in the forces file
      (including the time column);
      default: (0, 1, 2)
    """
    if not file_path:
      file_path = os.path.join(self.directory, 'forces')
    print('[info] reading forces ...')
    with open(file_path, 'r') as infile:
      data = numpy.loadtxt(infile,
                           dtype=numpy.float64,
                           usecols=usecols,
                           unpack=True)
    times = data[0]
    if not labels:
      labels = ['f_x', 'f_y']  # default labels
    self.forces = []  # reset forces if already present
    for index, values in enumerate(data[1:]):
      self.forces.append(Force(times, values, label=labels[index]))

  def read_fluxes(self, time_step, directory=None, **kwargs):
    """
    Reads the flux fields from file at a given time-step.

    Parameters
    ----------
    time_step: integer
      Time-step at which to read the fluxes.
    directory: string, optional
      Directory containing the saved time-step folders;
      default: None (will be <simulation-directory>/<time-step>).
    **kwargs: dictionary
      Extra keyword-arguments.

    Returns
    -------
    qx, qy: Field objects
      Fluxes in the x- and y-directions.
    """
    print('[time-step {}] reading fluxes from file ...'.format(time_step))
    # get grid-stations and number of cells along each direction
    x, y = self.grid
    nx, ny = x.size - 1, y.size - 1
    # read fluxes from file
    if not directory:
      directory = os.path.join(self.directory, '{:0>7}'.format(time_step))
    file_path = os.path.join(directory, 'q')
    # test if file written in binary format
    textchars = bytearray({7, 8, 9, 10, 12, 13, 27}
                          | set(range(0x20, 0x100)) - {0x7f})
    is_binary_string = lambda bytes: bool(bytes.translate(None, textchars))
    infile = open(file_path, 'rb')
    binary_format = is_binary_string(infile.read(1024))
    infile.close()
    if binary_format:
      with open(file_path, 'rb') as infile:
        nq = struct.unpack('i', infile.read(4))[0]
        q = numpy.array(struct.unpack('d' * nq, infile.read(8 * nq)))
    else:
      with open(file_path, 'r') as infile:
        nq = int(infile.readline())
        q = numpy.loadtxt(infile, dtype=numpy.float64)
    # set flux Field objects
    qx = Field(label='x-flux',
               time_step=time_step,
               x=x[1:-1],
               y=0.5 * (y[:-1] + y[1:]),
               values=q[:(nx - 1) * ny].reshape(ny, nx - 1))
    offset = qx.values.size
    qy = Field(label='y-flux',
               time_step=time_step,
               x=0.5 * (x[:-1] + x[1:]),
               y=y[1:-1],
               values=q[offset:].reshape(ny - 1, nx))
    return qx, qy

  def read_pressure(self, time_step, directory=None, **kwargs):
    """
    Reads pressure field from solution file at given time-step.

    Parameters
    ----------
    time_step: integer
      Time-step at which to read the pressure field.
    directory: string, optional
      Directory containing the saved time-step folders;
      default: None (will be <simulation-directory>/<time-step>).
    **kwargs: dictionary
      Extra keyword-arguments.

    Returns
    -------
    p: Field object
      The pressure field.
    """
    print('[time-step {}] reading pressure from file ...'.format(time_step))
    # get info about mesh-grid
    x, y = self.grid
    nx, ny = x.size - 1, y.size - 1
    # read pressure from file
    if not directory:
      directory = os.path.join(self.directory, '{:0>7}'.format(time_step))
    file_path = os.path.join(directory, 'lambda')
    # test if file written in binary format
    textchars = bytearray({7, 8, 9, 10, 12, 13, 27}
                          | set(range(0x20, 0x100)) - {0x7f})
    is_binary_string = lambda bytes: bool(bytes.translate(None, textchars))
    infile = open(file_path, 'rb')
    binary_format = is_binary_string(infile.read(1024))
    infile.close()
    if binary_format:
      with open(file_path, 'rb') as infile:
        nlambda = struct.unpack('i', infile.read(4))[0]
        p = numpy.array(struct.unpack('d' * nlambda,
                                      infile.read(8 * nlambda)))[:nx * ny]
    else:
      with open(file_path, 'r') as infile:
        nlambda = int(infile.readline())
        p = numpy.loadtxt(infile, dtype=numpy.float64)[:nx * ny]
    # set pressure Field object
    p = Field(label='pressure',
              time_step=time_step,
              x=0.5 * (x[:-1] + x[1:]),
              y=0.5 * (y[:-1] + y[1:]),
              values=p.reshape(ny, nx))
    return p

max_stars_repo_name: mesnardo/cuIBM

max_issues_repo_name: mesnardo/cuIBM

max_forks_repo_name: mesnardo/cuIBM


text:
# Copyright (c) 2022 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 os
import paddle
import numpy as np
from paddlenlp.utils.env import PPNLP_HOME
from paddlenlp.utils.log import logger
from paddlenlp.taskflow.utils import dygraph_mode_guard
from modeling import ErnieDocForSequenceClassification
from paddlenlp.transformers import ErnieDocTokenizer, ErnieDocBPETokenizer
from paddlenlp.datasets import load_dataset
from data import ClassifierIterator, ImdbTextPreprocessor, HYPTextPreprocessor, to_json_file
import paddle.nn as nn
from train import init_memory
from functools import partial
import argparse

# yapf: disable
parser = argparse.ArgumentParser()
parser.add_argument("--batch_size", default=16, type=int,
                    help="Batch size per GPU/CPU for predicting (In static mode, it should be the same as in model training process.)")
parser.add_argument("--model_name_or_path", type=str, default="ernie-doc-base-zh",
                    help="Pretraining or finetuned model name or path")
parser.add_argument("--max_seq_length", type=int, default=512,
                    help="The maximum total input sequence length after SentencePiece tokenization.")
parser.add_argument("--memory_length", type=int, default=128, help="Length of the retained previous heads.")
parser.add_argument("--device", type=str, default="gpu", choices=["cpu", "gpu"],
                    help="Select cpu, gpu devices to train model.")
parser.add_argument("--test_results_file", default="./test_restuls.json", type=str,
                    help="The file path you would like to save the model ouputs on test dataset.")
parser.add_argument("--static_mode", default=False, type=bool,
                    help="Whether you would like to perform predicting by static model or dynamic model.")
parser.add_argument("--dataset", default="iflytek", choices=["imdb", "iflytek", "thucnews", "hyp"], type=str,
                    help="The training dataset")
parser.add_argument("--static_path", default=None, type=str,
                    help="The path which your static model is at or where you want to save after converting.")

args = parser.parse_args()
# yapf: enable

DATASET_INFO = {
    "imdb": (ErnieDocBPETokenizer, "test", ImdbTextPreprocessor()),
    "hyp": (ErnieDocBPETokenizer, "test", HYPTextPreprocessor()),
    "iflytek": (ErnieDocTokenizer, "test", None),
    "thucnews": (ErnieDocTokenizer, "test", None)
}

def predict(model,
            test_dataloader,
            file_path,
            memories,
            label_list,
            static_mode,
            input_handles=None,
            output_handles=None):
    label_dict = dict()
    if not static_mode:
        model.eval()
        for _, batch in enumerate(test_dataloader, start=1):
            input_ids, position_ids, token_type_ids, attn_mask, _, qids, \
            gather_idxs, need_cal_loss = batch
            logits, memories = model(input_ids, memories, token_type_ids,
                                     position_ids, attn_mask)
            logits, qids = list(
                map(lambda x: paddle.gather(x, gather_idxs), [logits, qids]))
            probs = nn.functional.softmax(logits, axis=1)
            idx = paddle.argmax(probs, axis=1).numpy()
            idx = idx.tolist()
            labels = [label_list[i] for i in idx]
            for i, qid in enumerate(qids.numpy().flatten()):
                label_dict[str(qid)] = labels[i]
    else:
        for _, batch in enumerate(test_dataloader, start=1):
            input_ids, position_ids, token_type_ids, attn_mask, _, qids, \
            gather_idxs, need_cal_loss = batch
            input_handles[0].copy_from_cpu(input_ids.numpy())
            input_handles[1].copy_from_cpu(memories)
            input_handles[2].copy_from_cpu(token_type_ids.numpy())
            input_handles[3].copy_from_cpu(position_ids.numpy())
            input_handles[4].copy_from_cpu(attn_mask.numpy())
            model.run()
            logits = paddle.to_tensor(output_handles[0].copy_to_cpu())
            memories = paddle.to_tensor(output_handles[1].copy_to_cpu())
            logits, qids = list(
                map(lambda x: paddle.gather(x, gather_idxs), [logits, qids]))
            probs = nn.functional.softmax(logits, axis=1)
            idx = paddle.argmax(probs, axis=1).numpy()
            idx = idx.tolist()
            labels = [label_list[i] for i in idx]
            for i, qid in enumerate(qids.numpy().flatten()):
                label_dict[str(qid)] = labels[i]
    to_json_file("iflytek", label_dict, file_path)

class LongDocClassifier:
    def __init__(self,
                 model_name_or_path,
                 trainer_num=1,
                 rank=0,
                 batch_size=16,
                 max_seq_length=512,
                 memory_len=128,
                 static_mode=False,
                 dataset="iflytek",
                 static_path=None):
        self.model_name_or_path = model_name_or_path
        self.batch_size = batch_size
        self.trainer_num = trainer_num
        self.rank = rank
        self.max_seq_length = max_seq_length
        self.memory_len = memory_len
        self.static_mode = static_mode
        self.static_path = static_path if static_path else PPNLP_HOME

        tokenizer_class, test_name, preprocess_text_fn = DATASET_INFO[dataset]
        self._construct_tokenizer(tokenizer_class)
        self._input_preparation(args.dataset, test_name, preprocess_text_fn)
        self._construct_model()
        if static_mode:
            logger.info("Loading the static model from {}".format(
                self.static_path))
            self._load_static_model()

    def _input_preparation(self,
                           dataset="iflytek",
                           test_name="test",
                           preprocess_text_fn=None):
        test_ds = load_dataset("clue", name=dataset, splits=[test_name])
        self.label_list = test_ds.label_list
        self.num_classes = len(test_ds.label_list)
        self.test_ds_iter = ClassifierIterator(
            test_ds,
            self.batch_size,
            self._tokenizer,
            self.trainer_num,
            trainer_id=self.rank,
            memory_len=self.memory_len,
            max_seq_length=self.max_seq_length,
            mode="eval",
            preprocess_text_fn=preprocess_text_fn)
        self.test_dataloader = paddle.io.DataLoader.from_generator(
            capacity=70, return_list=True)
        self.test_dataloader.set_batch_generator(self.test_ds_iter,
                                                 paddle.get_device())

    def _construct_tokenizer(self, tokenizer_class):
        """
        Construct the tokenizer for the predictor.
        :return:
        """
        tokenizer_instance = tokenizer_class.from_pretrained(
            self.model_name_or_path)
        self._tokenizer = tokenizer_instance

    def _construct_model(self):
        """
        Construct the inference model for the predictor
        :param model_name_or_path: str
        :return: model instance
        """
        model_instance = ErnieDocForSequenceClassification.from_pretrained(
            self.model_name_or_path, num_classes=self.num_classes)
        self.model_config = model_instance.ernie_doc.config
        self._model = model_instance

    def _load_static_model(self, params_path=None):
        """Load static model"""
        inference_model_path = os.path.join(self.static_path, "static",
                                            "inference")
        with dygraph_mode_guard():
            self._construct_model()
            if params_path:
                state_dict = paddle.load(params_path)
                self._model.set_dict(state_dict)
            self._construct_input_spec()
            self._convert_dygraph_to_static()

        model_file = inference_model_path + ".pdmodel"
        params_file = inference_model_path + ".pdiparams"
        self._config = paddle.inference.Config(model_file, params_file)

    def _prepare_static_mode(self):
        """
        Construct the input data and predictor in the PaddlePaddele static mode.
        """
        place = paddle.get_device()
        if place == 'cpu':
            self._config.disable_gpu()
        else:
            self._config.enable_use_gpu(100)
        self._config.switch_use_feed_fetch_ops(False)
        self._config.disable_glog_info()
        self.predictor = paddle.inference.create_predictor(self._config)
        self.input_handles = [
            self.predictor.get_input_handle(name)
            for name in self.predictor.get_input_names()
        ]
        self.output_handle = [
            self.predictor.get_output_handle(name)
            for name in self.predictor.get_output_names()
        ]

    def _construct_input_spec(self):
        """
        Construct the input spec for the convert dygraph model to static model.
        """
        B, T, H, M, N = self.batch_size, self.max_seq_length, self.model_config["hidden_size"], self.memory_len, \
                        self.model_config["num_hidden_layers"]
        self._input_spec = [
            paddle.static.InputSpec(
                shape=[B, T, 1], dtype="int64", name="input_ids"),  # input_ids
            paddle.static.InputSpec(
                shape=[N, B, M, H], dtype="float32",
                name="memories"),  # memories
            paddle.static.InputSpec(
                shape=[B, T, 1], dtype="int64",
                name="token_type_ids"),  # token_type_ids
            paddle.static.InputSpec(
                shape=[B, 2 * T + M, 1], dtype="int64",
                name="position_ids"),  # position_ids
            paddle.static.InputSpec(
                shape=[B, T, 1], dtype="float32",
                name="attn_mask"),  # attn_mask
        ]

    def _convert_dygraph_to_static(self):
        """
        Convert the dygraph model to static model.
        """
        assert self._model is not None, 'The dygraph model must be created before converting the dygraph model to static model.'
        assert self._input_spec is not None, 'The input spec must be created before converting the dygraph model to static model.'
        logger.info("Converting to the inference model cost a little time.")
        static_model = paddle.jit.to_static(
            self._model, input_spec=self._input_spec)
        save_path = os.path.join(self.static_path, "static", "inference")
        paddle.jit.save(static_model, save_path)
        logger.info("The inference model save in the path:{}".format(save_path))

    def run_model(self, saved_path):
        if not self.static_mode:
            create_memory = partial(init_memory, self.batch_size,
                                    self.memory_len,
                                    self.model_config["hidden_size"],
                                    self.model_config["num_hidden_layers"])
            # Copy the memory
            memories = create_memory()
        else:
            memories = np.zeros(
                [
                    self.model_config["num_hidden_layers"], self.batch_size,
                    self.memory_len, self.model_config["hidden_size"]
                ],
                dtype="float32")
        file_path = saved_path
        if not self.static_mode:
            self.input_handles, self.output_handle, self.predictor = None, None, self._model
        else:
            self._prepare_static_mode()
        predict(self.predictor, self.test_dataloader, file_path, memories,
                self.label_list, self.static_mode, self.input_handles,
                self.output_handle)

def do_predict(args):
    # Initialize model
    paddle.set_device(args.device)
    trainer_num = paddle.distributed.get_world_size()
    if trainer_num > 1:
        paddle.distributed.init_parallel_env()
    rank = paddle.distributed.get_rank()
    if rank == 0:
        if os.path.exists(args.model_name_or_path):
            logger.info("init checkpoint from %s" % args.model_name_or_path)

    predictor = LongDocClassifier(
        model_name_or_path=args.model_name_or_path,
        rank=rank,
        trainer_num=trainer_num,
        batch_size=args.batch_size,
        max_seq_length=args.max_seq_length,
        memory_len=args.memory_length,
        static_mode=args.static_mode,
        static_path=args.static_path)
    predictor.run_model(saved_path=args.test_results_file)

if __name__ == "__main__":
    do_predict(args)

max_stars_repo_name: SunYanCN/PaddleNLP

max_issues_repo_name: SunYanCN/PaddleNLP

max_forks_repo_name: SunYanCN/PaddleNLP


text:
# -*- coding: utf-8 -*-
#
#  test_fft.py
#  aopy
#
#  Created by Jaberwocky on 2013-09-10.
#  Copyright 2013 Jaberwocky. All rights reserved.
#

from __future__ import (absolute_import, unicode_literals, division,
                        print_function)

import nose.tools as nt
from nose.plugins.skip import SkipTest
from nose.plugins.attrib import attr
import numpy as np

import os, os.path

import matplotlib.pyplot as plt

from .util import npeq_, PIDLYTests, npeq_or_save

@attr('fft')
class test_ffts(PIDLYTests):
    """Compare FFTs"""

    def setup(self):
        """Set up the FFTs"""
        super(test_ffts, self).setup()

        self.arrays = {}

        self.arrays["ones"] = np.ones((40))
        self.arrays["inds"] = np.arange(40)
        self.arrays["2d"] = np.arange((100)).reshape((10,10))

    def test_fft_arange(self):
        """FFT on 1D data"""

        import scipy.fftpack

        py_fft = scipy.fftpack.fft(self.arrays["inds"]) / self.arrays['inds'].shape[0]

        self.IDL.data = self.arrays["inds"]
        self.IDL('idl_fft = fft(data)')
        self.IDL('idl_fft_r = real_part(idl_fft)')
        self.IDL('idl_fft_i = imaginary(idl_fft)')
        idl_fft = self.IDL.idl_fft_r + 1j * self.IDL.idl_fft_i

        if not np.allclose(idl_fft, py_fft):
            plt.subplot(2,1,1)
            plt.plot(np.real(py_fft), label="Python")
            plt.plot(np.real(idl_fft), label="IDL")
            plt.legend()
            plt.subplot(2,1,2)
            plt.plot(np.imag(py_fft), label="Python")
            plt.plot(np.imag(idl_fft), label="IDL")
            plt.legend()
            plt.show()

        assert np.allclose(idl_fft, py_fft)

    def test_fft_2d(self):
        """FFT on 2D data"""
        import scipy.fftpack

        py_fft = scipy.fftpack.fftn(self.arrays["2d"]) / self.arrays['2d'].size

        self.IDL.data = self.arrays["2d"]
        self.IDL('idl_fft = fft(data)')
        self.IDL('idl_fft_r = real_part(idl_fft)')
        self.IDL('idl_fft_i = imaginary(idl_fft)')
        idl_fft = self.IDL.idl_fft_r + 1j * self.IDL.idl_fft_i

        if not np.allclose(idl_fft, py_fft, atol=1e-6):
            plt.subplot(2,1,1)
            plt.plot(np.real(py_fft[3,:]), label="Python")
            plt.plot(np.real(idl_fft[3,:]), label="IDL")
            plt.legend()
            plt.subplot(2,1,2)
            plt.plot(np.imag(py_fft[3,:]), label="Python")
            plt.plot(np.imag(idl_fft[3,:]), label="IDL")
            plt.legend()
            plt.figure()
            plt.subplot(1,2,1)
            plt.title("Real")
            plt.imshow(np.real(py_fft - idl_fft))
            plt.colorbar()
            plt.subplot(1,2,2)
            plt.title("Imaginary")
            plt.imshow(np.imag(py_fft - idl_fft))
            plt.colorbar()
            plt.show()

        idl_fft_c, py_fft_c = idl_fft[0,:], py_fft[0,:]
        assert np.allclose(idl_fft_c, py_fft_c)
        assert np.allclose(idl_fft, py_fft, atol=1e-6)

max_stars_repo_name: alexrudy/aopy

max_issues_repo_name: alexrudy/aopy

max_forks_repo_name: alexrudy/aopy

reason: import numpy,import scipy