def single_frame_loss():
JOB = "20210124114528"
JOB_DIR = f"{EXPERIMENTS_DIR}/{JOB}"
device = torch.device("cuda:0")
# Load model.
opts = yaml.safe_load(open(f"{JOB_DIR}/{JOB}.yaml"))
(
train_dataset,
train_loader,
valid_dataset,
valid_loader,
test_dataset,
test_loader,
) = init_datasets(opts)
model = init_model(opts, train_dataset)
model.load_state_dict(torch.load(f"{JOB_DIR}/best_params.pth"))
model = model.to(device)
model.eval()
criterion = nn.CrossEntropyLoss()
test_loss_one = 0.0
test_loss_all = 0.0
n_test = 0
with torch.no_grad():
for test_tensors in test_loader:
# Skip bad sequences.
if len(test_tensors["player_idxs"]) < model.seq_len:
continue
player_trajs = test_tensors["player_trajs"].flatten()
n_player_trajs = len(player_trajs)
labels = player_trajs.to(device)
preds = model(test_tensors)["player"][:n_player_trajs]
test_loss_one += criterion(preds[:10], labels[:10]).item()
test_loss_all += criterion(preds, labels).item()
n_test += 1
print(test_loss_one / n_test)
print(test_loss_all / n_test)
def plot_attn_through_time():
JOB = "20210124114950"
JOB_DIR = f"{EXPERIMENTS_DIR}/{JOB}"
# Load model.
opts = yaml.safe_load(open(f"{JOB_DIR}/{JOB}.yaml"))
(train_dataset, _, _, _, test_dataset, _) = init_datasets(opts)
model = init_model(opts, train_dataset)
model.load_state_dict(torch.load(f"{JOB_DIR}/best_params.pth"))
model.eval()
court = plt.imread("court.png")
width = 5
height = width * court.shape[0] / float(court.shape[1])
pt_scale = 1.4
home_dir = os.path.expanduser("~")
os.makedirs(f"{home_dir}/test", exist_ok=True)
(fig, ax) = plt.subplots(figsize=(width, height))
save_output = SaveOutput()
hook_handles = []
for layer in model.transformer.layers:
hook_handles.append(layer.self_attn.register_forward_hook(save_output))
test_idx = 851
tensors = test_dataset[test_idx]
with torch.no_grad():
_ = model(tensors)
colors = []
markers = []
for p_idx in range(10):
if tensors["player_hoop_sides"][0, p_idx]:
colors.append("white")
markers.append("s")
else:
colors.append("gray")
markers.append("^")
ball_start = 10 * len(tensors["player_idxs"])
for time_step in range(19):
print(time_step, flush=True)
ball_pos = ball_start + time_step
ax.imshow(court, zorder=0, extent=[X_MIN, X_MAX - DIFF, Y_MAX, Y_MIN])
ax.axis("off")
ax.grid(False)
for p_idx in range(10):
ax.scatter(
[tensors["player_xs"][time_step, p_idx]],
[tensors["player_ys"][time_step, p_idx]],
s=(pt_scale * matplotlib.rcParams["lines.markersize"])**2,
c=colors[p_idx],
marker=markers[p_idx],
edgecolors="black",
)
ax.scatter(
[tensors["ball_xs"][time_step]],
[tensors["ball_ys"][time_step]],
s=(pt_scale * matplotlib.rcParams["lines.markersize"])**2,
c="#ff8c00",
edgecolors="black",
)
plt.subplots_adjust(0, 0, 1, 1)
fig.canvas.draw()
court_img = Image.frombytes("RGB", fig.canvas.get_width_height(),
fig.canvas.tostring_rgb())
plt.cla()
time_step = str(time_step).zfill(2)
court_img.save(f"{home_dir}/test/court_{time_step}.png")
for attn_layer_idx in range(len(save_output.outputs)):
attn_scores = (save_output.outputs[attn_layer_idx][1]
[0].detach().cpu().numpy())
for attn_head_idx in range(len(attn_scores)):
attn_head_scores = attn_scores[attn_head_idx, ball_pos]
player_scores = (attn_head_scores[:ball_start].reshape(
-1, 10).sum(axis=0))
vmax = player_scores.max()
player_scores_normed = player_scores / vmax
player_scores_normed = np.nan_to_num(player_scores_normed, 0)
player_colors = [
str(player_score) for player_score in player_scores_normed
]
ax.imshow(court,
zorder=0,
extent=[X_MIN, X_MAX - DIFF, Y_MAX, Y_MIN])
ax.axis("off")
ax.grid(False)
for p_idx in range(10):
ax.scatter(
[tensors["player_xs"][time_step, p_idx]],
[tensors["player_ys"][time_step, p_idx]],
s=(pt_scale *
matplotlib.rcParams["lines.markersize"])**2,
c=player_colors[p_idx],
marker=markers[p_idx],
cmap="gray",
vmin=0,
vmax=1,
edgecolors="black",
)
ball_score = attn_head_scores[ball_start:].sum()
ax.scatter(
[tensors["ball_xs"][time_step]],
[tensors["ball_ys"][time_step]],
s=(pt_scale * matplotlib.rcParams["lines.markersize"])**2,
c="#ff8c00",
edgecolors="black",
)
plt.subplots_adjust(0, 0, 1, 1)
fig.canvas.draw()
court_img = Image.frombytes("RGB",
fig.canvas.get_width_height(),
fig.canvas.tostring_rgb())
plt.cla()
ball_score = str(ball_score).split(".")[1][:2]
img_id = f"{str(attn_layer_idx).zfill(2)}_{str(attn_head_idx).zfill(2)}"
time_step = str(time_step).zfill(2)
court_img.save(
f"{home_dir}/test/{img_id}_{time_step}_{ball_score}.png")
shutil.make_archive(f"{home_dir}/test", "zip", f"{home_dir}/test")
shutil.rmtree(f"{home_dir}/test")
def naive_player_test_set_perplexity():
JOB = "20210124114528"
JOB_DIR = f"{EXPERIMENTS_DIR}/{JOB}"
device = torch.device("cuda:0")
# Load model.
opts = yaml.safe_load(open(f"{JOB_DIR}/{JOB}.yaml"))
(
train_dataset,
train_loader,
valid_dataset,
valid_loader,
test_dataset,
test_loader,
) = init_datasets(opts)
default_hz = 25
hz = 5
skip = default_hz // hz
skip_secs = 1 / skip
all_player_trajs = []
for gameid in train_dataset.gameids:
X = np.load(f"{GAMES_DIR}/{gameid}_X.npy")
seq_data = X[::skip]
player_xs = seq_data[:, 20:30]
player_ys = seq_data[:, 30:40]
player_x_diffs = np.diff(player_xs, axis=0)
player_y_diffs = np.diff(player_ys, axis=0)
wall_clock_diffs = np.diff(seq_data[:, -1]) / 1000
keep_diffs = wall_clock_diffs <= 1.5 * skip_secs
player_x_diffs = player_x_diffs[keep_diffs].flatten()
player_y_diffs = player_y_diffs[keep_diffs].flatten()
player_traj_rows = np.digitize(player_y_diffs,
train_dataset.player_traj_bins)
player_traj_cols = np.digitize(player_x_diffs,
train_dataset.player_traj_bins)
player_trajs = player_traj_rows * train_dataset.player_traj_n + player_traj_cols
all_player_trajs.append(player_trajs)
all_player_trajs = np.concatenate(all_player_trajs)
(unique, counts) = np.unique(all_player_trajs, return_counts=True)
d_counts = dict(zip(unique, counts))
probs = np.zeros(train_dataset.ball_traj_n**3)
for (traj, count) in d_counts.items():
probs[traj] = count
probs = probs / probs.sum()
model = init_model(opts, train_dataset)
model.load_state_dict(torch.load(f"{JOB_DIR}/best_params.pth"))
model = model.to(device)
model.eval()
test_loss = 0.0
n_test = 0
with torch.no_grad():
for test_tensors in test_loader:
# Skip bad sequences.
if len(test_tensors["player_idxs"]) < model.seq_len:
continue
labels = test_tensors["player_trajs"].flatten()
test_loss -= np.log(probs[labels]).mean()
n_test += 1
print(test_loss)
print(np.exp(test_loss / n_test))
def plot_generated_trajectories():
JOB = "20210124114528"
JOB_DIR = f"{EXPERIMENTS_DIR}/{JOB}"
# Load model.
opts = yaml.safe_load(open(f"{JOB_DIR}/{JOB}.yaml"))
(train_dataset, _, _, _, test_dataset, _) = init_datasets(opts)
model = init_model(opts, train_dataset)
model.load_state_dict(torch.load(f"{JOB_DIR}/best_params.pth"))
model.eval()
cand_test_idxs = []
for test_idx in range(len(test_dataset.gameids)):
tensors = test_dataset[test_idx]
if len(tensors["player_idxs"]) == model.seq_len:
cand_test_idxs.append(test_idx)
court = plt.imread("court.png")
width = 5
height = width * court.shape[0] / float(court.shape[1])
pt_scale = 1.4
player_traj_bins = np.array(list(test_dataset.player_traj_bins) + [5.5])
home_dir = os.path.expanduser("~")
os.makedirs(f"{home_dir}/test", exist_ok=True)
(fig, ax) = plt.subplots(figsize=(width, height))
saved = 0
np.random.seed(2010)
while saved < 20:
test_idx = np.random.choice(cand_test_idxs)
cand_test_idxs.remove(test_idx)
print(test_idx, flush=True)
for start_time_step in [0, 5, 9, 19]:
tensors = test_dataset[test_idx]
(player_idxs, p_idxs) = tensors["player_idxs"][0].sort()
player_traj_n = test_dataset.player_traj_n
seq_len = len(tensors["player_idxs"])
grid_gap = np.diff(test_dataset.player_traj_bins)[-1] / 2
with torch.no_grad():
for step in range(seq_len - start_time_step - 1):
preds_start = 10 * (start_time_step + step)
preds_stop = preds_start + 10
preds = model(tensors)["player"][preds_start:preds_stop]
probs = torch.softmax(preds, dim=1)
samp_trajs = torch.multinomial(probs, 1)
samp_rows = samp_trajs // player_traj_n
samp_cols = samp_trajs % player_traj_n
samp_xs = (player_traj_bins[samp_cols] - grid_gap +
np.random.uniform(-grid_gap, grid_gap))
samp_ys = (player_traj_bins[samp_rows] - grid_gap +
np.random.uniform(-grid_gap, grid_gap))
tensors["player_xs"][start_time_step + step + 1] = (
tensors["player_xs"][start_time_step + step] +
samp_xs.flatten())
tensors["player_ys"][start_time_step + step + 1] = (
tensors["player_ys"][start_time_step + step] +
samp_ys.flatten())
markers = []
for p_idx in range(10):
if tensors["player_hoop_sides"][0, p_idx]:
markers.append("s")
else:
markers.append("^")
ax.imshow(court,
zorder=0,
extent=[X_MIN, X_MAX - DIFF, Y_MAX, Y_MIN])
ax.axis("off")
ax.grid(False)
for (idx, p_idx) in enumerate(p_idxs):
ax.scatter(
[tensors["player_xs"][0, p_idx]],
[tensors["player_ys"][0, p_idx]],
s=(pt_scale * matplotlib.rcParams["lines.markersize"])**2,
c=colors[idx],
marker=markers[p_idx],
edgecolors="black",
)
plt.xlim(auto=False)
plt.ylim(auto=False)
ax.plot(
tensors["player_xs"][:start_time_step + 1, p_idx],
tensors["player_ys"][:start_time_step + 1, p_idx],
c=colors[idx],
)
ax.plot(
tensors["player_xs"][start_time_step:, p_idx],
tensors["player_ys"][start_time_step:, p_idx],
c=colors[idx],
linestyle=":",
)
ax.plot(
tensors["ball_xs"],
tensors["ball_ys"],
c="#ff8c00",
)
ax.scatter(
[tensors["ball_xs"][0]],
[tensors["ball_ys"][0]],
s=(pt_scale * matplotlib.rcParams["lines.markersize"])**2,
c="#ff8c00",
edgecolors="black",
)
plt.subplots_adjust(0, 0, 1, 1)
fig.canvas.draw()
court_img = Image.frombytes("RGB", fig.canvas.get_width_height(),
fig.canvas.tostring_rgb())
plt.cla()
court_img.save(
f"{home_dir}/test/trajectories_{test_idx}_{start_time_step}.png"
)
saved += 1
shutil.make_archive(f"{home_dir}/test", "zip", f"{home_dir}/test")
shutil.rmtree(f"{home_dir}/test")
def plot_traj_preds():
JOB = "20210124114528"
JOB_DIR = f"{EXPERIMENTS_DIR}/{JOB}"
# Load model.
opts = yaml.safe_load(open(f"{JOB_DIR}/{JOB}.yaml"))
(train_dataset, _, _, _, test_dataset, _) = init_datasets(opts)
model = init_model(opts, train_dataset)
model.load_state_dict(torch.load(f"{JOB_DIR}/best_params.pth"))
model.eval()
cand_test_idxs = []
for test_idx in range(len(test_dataset.gameids)):
tensors = test_dataset[test_idx]
if len(tensors["player_idxs"]) == model.seq_len:
cand_test_idxs.append(test_idx)
court = plt.imread("court.png")
width = 5
height = width * court.shape[0] / float(court.shape[1])
vert_buffer = 10
pt_scale = 1.4
scale = 8
pred_player_idx = 5
home_dir = os.path.expanduser("~")
os.makedirs(f"{home_dir}/test", exist_ok=True)
(fig, ax) = plt.subplots(figsize=(width, height))
saved = 0
np.random.seed(2010)
while saved < 20:
test_idx = np.random.choice(cand_test_idxs)
cand_test_idxs.remove(test_idx)
print(test_idx, flush=True)
tensors = test_dataset[test_idx]
player_traj_n = test_dataset.player_traj_n
with torch.no_grad():
seq_len = len(tensors["player_idxs"])
labels = tensors["player_trajs"].flatten()
preds = model(tensors)["player"][:len(labels)]
probs = torch.softmax(preds, dim=1)
pred_player_probs = probs.view(seq_len, -1,
player_traj_n**2)[:,
pred_player_idx]
colors = []
markers = []
for p_idx in range(10):
if p_idx == pred_player_idx:
colors.append("r")
elif tensors["player_hoop_sides"][0, p_idx]:
colors.append("white")
else:
colors.append("gray")
if tensors["player_hoop_sides"][0, p_idx]:
markers.append("s")
else:
markers.append("^")
for time_step in range(len(tensors["player_idxs"])):
ax.imshow(court,
zorder=0,
extent=[X_MIN, X_MAX - DIFF, Y_MAX, Y_MIN])
ax.axis("off")
ax.grid(False)
for p_idx in range(10):
ax.scatter(
[tensors["player_xs"][time_step, p_idx]],
[tensors["player_ys"][time_step, p_idx]],
s=(pt_scale * matplotlib.rcParams["lines.markersize"])**2,
c=colors[p_idx],
marker=markers[p_idx],
edgecolors="black",
)
plt.xlim(auto=False)
plt.ylim(auto=False)
ax.plot(
tensors["player_xs"][:time_step + 1, pred_player_idx],
tensors["player_ys"][:time_step + 1, pred_player_idx],
c="red",
)
ax.plot(
tensors["player_xs"][time_step:, pred_player_idx],
tensors["player_ys"][time_step:, pred_player_idx],
linestyle=":",
c="red",
)
ax.scatter(
[tensors["ball_xs"][time_step]],
[tensors["ball_ys"][time_step]],
s=(pt_scale * matplotlib.rcParams["lines.markersize"])**2,
c="#ff8c00",
edgecolors="black",
)
plt.subplots_adjust(0, 0, 1, 1)
fig.canvas.draw()
court_img = Image.frombytes("RGB", fig.canvas.get_width_height(),
fig.canvas.tostring_rgb())
plt.cla()
traj_data = np.full((player_traj_n, player_traj_n, 3),
255,
dtype=np.uint8)
player_traj = tensors["player_trajs"][time_step, pred_player_idx]
traj_data[player_traj // player_traj_n,
player_traj % player_traj_n] = 0
traj_img = Image.fromarray(traj_data).resize(
(scale * player_traj_n, scale * player_traj_n), resample=0)
add_grid(traj_img, player_traj_n)
traj_pred = np.zeros((player_traj_n, player_traj_n, 3))
traj_pred[:, :] = pred_player_probs[time_step].view(
player_traj_n, player_traj_n, 1)
pred_img = Image.fromarray(np.uint8(255 - 255 * traj_pred)).resize(
(scale * player_traj_n, scale * player_traj_n), resample=0)
add_grid(pred_img, player_traj_n)
final_img = Image.new(
"RGBA",
(
court_img.size[0],
court_img.size[1] + scale * player_traj_n + vert_buffer,
),
)
final_img.paste(court_img, (0, 0))
full_width = 3 * traj_img.size[0]
horiz_buffer = court_img.size[0] // 2 - full_width // 2
final_img.paste(traj_img,
(horiz_buffer, court_img.size[1] + vert_buffer))
final_img.paste(
pred_img,
(2 * traj_img.size[0] + horiz_buffer,
court_img.size[1] + vert_buffer),
)
time_step = str(time_step).zfill(2)
final_img.save(f"{home_dir}/test/{test_idx}_{time_step}.png")
saved += 1
shutil.make_archive(f"{home_dir}/test", "zip", f"{home_dir}/test")
shutil.rmtree(f"{home_dir}/test")
def get_nearest_neighbors():
JOB = "20210124114950"
JOB_DIR = f"{EXPERIMENTS_DIR}/{JOB}"
# Load model.
opts = yaml.safe_load(open(f"{JOB_DIR}/{JOB}.yaml"))
train_dataset = init_datasets(opts)[0]
model = init_model(opts, train_dataset)
model.load_state_dict(torch.load(f"{JOB_DIR}/best_params.pth"))
model.eval()
player_embeddings = model.player_embedding.weight
with torch.no_grad():
if opts["model"]["sigmoid"] == "logistic":
player_embeddings = torch.sigmoid(player_embeddings)
elif opts["model"]["sigmoid"] == "tanh":
player_embeddings = torch.tanh(player_embeddings)
player_embeddings = player_embeddings.detach().numpy()
all_player_dists = {
"cosine": pairwise_distances(player_embeddings, metric="cosine"),
"euclidean": pairwise_distances(player_embeddings, metric="euclidean"),
}
players = [
"LeBron James",
"Stephen Curry",
"Russell Westbrook",
"James Harden",
"Pau Gasol",
"Anthony Davis",
"Serge Ibaka",
"Giannis Antetokounmpo",
"Kawhi Leonard",
"Jimmy Butler",
]
k = 20
min_playing_time = 60000 # 0/13314/39917/60000 --> 100%/75%/50%/starters.
baller2vec_config = pickle.load(
open(f"{DATA_DIR}/baller2vec_config.pydict", "rb"))
player_idx2props = baller2vec_config["player_idx2props"]
for player in players:
print(f"\n{player}")
player_idx = get_player_idx(player, player_idx2props)
for metric in ["cosine", "euclidean"]:
print(f"\n{metric}\n")
player_dists = all_player_dists[metric]
sim_player_idxs = np.argsort(player_dists[player_idx])
neighbors = 0
for sim_player_idx in sim_player_idxs[1:]:
if sim_player_idx in player_idx2props:
if "playing_time" not in player_idx2props[sim_player_idx]:
continue
if (player_idx2props[sim_player_idx]["playing_time"] <
min_playing_time):
continue
sim_player_score = player_dists[player_idx, sim_player_idx]
name = player_idx2props.get(sim_player_idx,
{"name": "Baller"})["name"]
print(f"{name}: {sim_player_score:.4f}")
neighbors += 1
if neighbors == k:
break
keep_embeddings = []
player_names = []
for (player_idx, player_embedding) in enumerate(player_embeddings):
if player_idx != len(player_embeddings) - 1:
if "playing_time" not in player_idx2props[player_idx]:
continue
if player_idx2props[player_idx]["playing_time"] < min_playing_time:
continue
keep_embeddings.append(player_embedding)
if player_idx == len(player_embeddings) - 1:
player_names.append("Baller")
else:
player_names.append(player_idx2props[player_idx]["name"])
keep_embeddings = np.vstack(keep_embeddings)
home_dir = os.path.expanduser("~")
os.makedirs(f"{home_dir}/test", exist_ok=True)
np.savetxt(f"{home_dir}/test/embeddings.tsv",
keep_embeddings,
delimiter="\t")
with open(f"{home_dir}/test/names.tsv", "w") as f:
print("\n".join(player_names), file=f)
shutil.make_archive(f"{home_dir}/test", "zip", f"{home_dir}/test")
shutil.rmtree(f"{home_dir}/test")
def plot_ball_traj_preds_specific_game():
JOB = "20210124114950"
JOB_DIR = f"{EXPERIMENTS_DIR}/{JOB}"
# Load model.
opts = yaml.safe_load(open(f"{JOB_DIR}/{JOB}.yaml"))
(train_dataset, _, _, _, _, _) = init_datasets(opts)
model = init_model(opts, train_dataset)
model.load_state_dict(torch.load(f"{JOB_DIR}/best_params.pth"))
model.eval()
court = plt.imread("court.png")
width = 5
height = width * court.shape[0] / float(court.shape[1])
vert_buffer = 10
pt_scale = 1.4
scale = 4
home_dir = os.path.expanduser("~")
os.makedirs(f"{home_dir}/test", exist_ok=True)
(fig, ax) = plt.subplots(figsize=(width, height))
save_output = SaveOutput()
hook_handles = []
for layer in model.transformer.layers:
hook_handles.append(layer.self_attn.register_forward_hook(save_output))
gameid = "0021500622"
X = np.load(f"{GAMES_DIR}/{gameid}_X.npy")
y = np.load(f"{GAMES_DIR}/{gameid}_y.npy")
period = 3
period_time = "1:55"
start = get_start_idx_from_time(X, period, period_time)
tensors = train_dataset.get_sample(X, y, start)
sub_rubio = False
if sub_rubio:
baller2vec_config = pickle.load(
open(f"{DATA_DIR}/baller2vec_config.pydict", "rb"))
player_idx2props = baller2vec_config["player_idx2props"]
lebron_idx = get_player_idx("LeBron James", player_idx2props)
rubio_idx = get_player_idx("Ricky Rubio", player_idx2props)
player_idxs = tensors["player_idxs"]
player_idxs[player_idxs == lebron_idx] = rubio_idx
with torch.no_grad():
ball_trajs = model(tensors)["ball"][-len(tensors["player_idxs"]):]
ball_probs = torch.softmax(ball_trajs, dim=1)
ball_traj_n = train_dataset.ball_traj_n
ball_probs = ball_probs.reshape(len(ball_probs), ball_traj_n,
ball_traj_n, ball_traj_n)
ball_xy_probs = ball_probs.sum(dim=1).permute(0, 2, 1)
colors = []
markers = []
for p_idx in range(10):
if tensors["player_hoop_sides"][0, p_idx]:
colors.append("white")
markers.append("s")
else:
colors.append("gray")
markers.append("^")
ball_start = 10 * len(tensors["player_idxs"])
for time_step in range(19):
print(time_step, flush=True)
ax.imshow(court, zorder=0, extent=[X_MIN, X_MAX - DIFF, Y_MAX, Y_MIN])
ax.axis("off")
ax.grid(False)
for p_idx in range(10):
ax.scatter(
[tensors["player_xs"][time_step, p_idx]],
[tensors["player_ys"][time_step, p_idx]],
s=(pt_scale * matplotlib.rcParams["lines.markersize"])**2,
c=colors[p_idx],
marker=markers[p_idx],
edgecolors="black",
)
ax.scatter(
[tensors["ball_xs"][time_step]],
[tensors["ball_ys"][time_step]],
s=(pt_scale * matplotlib.rcParams["lines.markersize"])**2,
c="#ff8c00",
edgecolors="black",
)
plt.subplots_adjust(0, 0, 1, 1)
fig.canvas.draw()
court_img = Image.frombytes("RGB", fig.canvas.get_width_height(),
fig.canvas.tostring_rgb())
plt.cla()
traj_data = np.full((ball_traj_n, ball_traj_n, 3), 255, dtype=np.uint8)
ball_traj = tensors["ball_trajs"][time_step]
(ball_row, ball_col_dep) = (
ball_traj // ball_traj_n**2,
ball_traj % ball_traj_n**2,
)
(ball_col, ball_dep) = (ball_col_dep // ball_traj_n,
ball_col_dep % ball_traj_n)
traj_data[ball_dep, ball_col] = 0
traj_img = Image.fromarray(traj_data).resize(
(scale * ball_traj_n, scale * ball_traj_n), resample=0)
add_grid(traj_img, ball_traj_n)
traj_pred = np.zeros((ball_traj_n, ball_traj_n, 3))
traj_pred[:, :] = ball_xy_probs[time_step].view(
ball_traj_n, ball_traj_n, 1)
pred_img = Image.fromarray(np.uint8(255 - 255 * traj_pred)).resize(
(scale * ball_traj_n, scale * ball_traj_n), resample=0)
add_grid(pred_img, ball_traj_n)
final_img = Image.new(
"RGBA",
(
court_img.size[0],
court_img.size[1] + scale * ball_traj_n + vert_buffer,
),
)
final_img.paste(court_img, (0, 0))
full_width = 3 * traj_img.size[0]
horiz_buffer = court_img.size[0] // 2 - full_width // 2
final_img.paste(traj_img,
(horiz_buffer, court_img.size[1] + vert_buffer))
final_img.paste(
pred_img,
(2 * traj_img.size[0] + horiz_buffer,
court_img.size[1] + vert_buffer),
)
true_prob = traj_pred[ball_dep, ball_col, 0]
print(true_prob)
true_prob = f"{true_prob:.2f}".split(".")[1]
time_step = str(time_step).zfill(2)
final_img.save(f"{home_dir}/test/court_{time_step}_{true_prob}.png")
shutil.make_archive(f"{home_dir}/test", "zip", f"{home_dir}/test")
shutil.rmtree(f"{home_dir}/test")
if __name__ == "__main__":
JOB = sys.argv[1]
JOB_DIR = f"{EXPERIMENTS_DIR}/{JOB}"
try:
os.environ["CUDA_VISIBLE_DEVICES"] = sys.argv[2]
except IndexError:
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
opts = yaml.safe_load(open(f"{JOB_DIR}/{JOB}.yaml"))
patience = opts["train"]["patience"]
# Initialize datasets.
(
train_dataset,
train_loader,
valid_dataset,
valid_loader,
test_dataset,
test_loader,
) = init_datasets(opts)
# Initialize model.
device = torch.device("cuda:0")
model = init_model(opts, train_dataset).to(device)
print(model)
n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"Parameters: {n_params}")
train_model()
