-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathcreate-mask-matrix.py
More file actions
55 lines (41 loc) · 1.29 KB
/
create-mask-matrix.py
File metadata and controls
55 lines (41 loc) · 1.29 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
import argparse
import os
from pathlib import Path
import numpy as np
from utilities import load_yaml
SCRIPT_LABEL = f"\033[35m[{os.path.basename(__file__)}]\033[0m "
def zero_diagonal_band(M, thickness=1):
rows, cols = M.shape
for offset in range(-thickness, thickness + 1):
diag_indices = np.diag_indices(min(rows, cols))
row_idx = diag_indices[0]
col_idx = diag_indices[1] + offset
# Filter out indices that go out of bounds
mask = (col_idx >= 0) & (col_idx < cols)
M[row_idx[mask], col_idx[mask]] = 0
return M
def main():
parser = argparse.ArgumentParser(
description="Generate mask matrix from experiment config."
)
parser.add_argument(
"--exp_yaml",
type=str,
default="exp.yaml",
help="Path to experiment YAML config file",
)
args = parser.parse_args()
# Inputs
print(f"{SCRIPT_LABEL}: {args.exp_yaml}")
exp_yaml = load_yaml(args.exp_yaml)
# Experiment paths
GT_path = Path(exp_yaml["log_dir"]) / "GT.npy"
M_path = Path(exp_yaml["log_dir"]) / "M.npy"
# Create Mask
GT = np.load(GT_path)
M = np.ones_like(GT, dtype=np.uint8)
M = zero_diagonal_band(M, thickness=40)
# Save matrix
np.save(M_path, M)
if __name__ == "__main__":
main()