vignesh-99/nfl

1 2from lightGBT import get_test_df3import lightgbm as lgb4import numpy as np5import pandas as pd6 7dx_model = lgb.Booster(model_file='lightGBT_dx_model.txt')8dy_model = lgb.Booster(model_file='lightGBT_dy_model.txt')9 10BASE_COLS = ['game_id', 'play_id', 'player_to_predict', 'nfl_id', 'frame_id',11       'play_direction', 'absolute_yardline_number', 'player_name',12       'player_height', 'player_weight', 'player_birth_date',13       'player_position', 'player_side', 'player_role', 'x', 'y', 's', 'a',14       'dir', 'o', 'num_frames_output', 'ball_land_x', 'ball_land_y']15 16MANDATORY_COLS=['game_id', 'play_id', 'nfl_id']17 18def predict(df):19    given_input_cols = set(df.columns)20    for c in MANDATORY_COLS:21        if c not in given_input_cols:22            raise Exception(f'{c} is missing in input')23        elif df[c].isna().any():24            raise Exception(f'{c} in input contains nan')25        26    27    for c in BASE_COLS:28        if c not in df:29            raise Exception(f'{c} is not there in input')30        if df[c].isna().all():31            raise Exception(f'{c} in input contains all nan')32    33   34    X_df = get_test_df(df)35 36    dx_features = dx_model.feature_name()37    dy_features = dy_model.feature_name()38 39    pred_dx = dx_model.predict(X_df[dx_features])40    pred_dy = dy_model.predict(X_df[dy_features])41    42    game_id = X_df['game_id'].values43    play_id = X_df['play_id'].values44    nfl_id =  X_df['nfl_id'].values45    frame_id = X_df['frame_id'].values46    player_position = X_df['player_position'].values47    player_role = X_df['player_role'].values48    x_last = X_df['x_last'].values49    y_last = X_df['y_last'].values50    51    play_direction = X_df['play_direction'].values52 53    pred_x = pred_dx + x_last54    pred_y = pred_dy + y_last55 56    mask = (play_direction == 'left')57 58    if np.any(mask):59        pred_x[mask] =  120 - pred_x[mask]60 61    pred_x = np.clip(pred_x, 0.0, 120.0)62    pred_y = np.clip(pred_y, 0.0, 53.3)63  64 65    preds = np.column_stack([game_id, play_id, nfl_id, frame_id, player_position, 66                    player_role, play_direction, x_last, y_last, pred_x, pred_y])67    68        69    df = pd.DataFrame(preds, columns=['game_id', 'play_id', 'nfl_id','frame_id', 'player_position', 70                                    'player_role','play_direction', 71                                    'x_last', 'y_last', 'pred_x', 'pred_y'])72    73    return df.sort_values(by=['game_id', 'play_id', 'nfl_id', 'frame_id'])74