Image-test/src/get_hyperparameters.py at main · LEGO/Image-test · GitHub

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
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
import mlflow
import optuna
from mlflow.tracking import MlflowClient
from json import load, dumps
from sklearn.model_selection import train_test_split

# Custom Python Files
from all_parts import process_files_vtk
import helpers
from tesseract_train_pipeline import evaluate_ocr

# Process images
with open("JSON_files/wsl_vtk_paths_remote.json", "r") as f:
    data = load(f)

mlflow.set_experiment("OCR_VTK_Tuning")
client = MlflowClient()
experiment = client.get_experiment_by_name("OCR_VTK_Tuning")
# Fixed ranges
bounding_box = (-5, -2.5, -33, -28, 1.5, 3.0)
# x_center = -3.6
# threshold_range = (0.06, 0.09)
# delta_range = [0.00, 0.05, 0.1, 0.15, 0.2]
# stroke = 300

def run_ocr_pipeline(params):
    # Extract images
    process_files_vtk(params["stl_files_folder"], params["output_folder_filled"], params["bounding_box"],
                      params["delta_range"], params["threshold_range"], params["x_center"],
                      params["min_stroke_thickness"])
    helpers.write_ground_truth(params["output_folder_filled"])
    # Evaluate both sets
    image_gt_pairs = helpers.collect_image_gt_pairs(params["output_folder_filled"])
    # Split into train and validation sets
    train_data, val_data = train_test_split(image_gt_pairs, test_size=0.2, random_state=42)
    train_accuracy, _, train_summary = evaluate_ocr(train_data, params['lang'])
    val_accuracy, _, val_summary = evaluate_ocr(val_data,  params['lang'])
    return val_accuracy, val_summary

def r2c(rows):
    all_keys = set()
    for r in rows: all_keys.update(r.keys())
    all_keys = list(all_keys)

    cols = {k:[] for k in all_keys}
    for r in rows:
        for k in all_keys:
            v = r.get(k, None)
            if isinstance(v,(list, dict)):
                v = dumps(v, ensure_ascii=False)
            cols[k].append(v)
    return cols

all_runs = []

def objective(trial):
    # To get variable range values
    threshold = trial.suggest_categorical("threshold", [0.06, 0.075])
    x_center = trial.suggest_uniform("x_center", -3.8, -3.5)
    stroke = trial.suggest_int("min_stroke_thickness", 250, 350)
    delta_range = trial.suggest_categorical("delta_range", [
            [0.00, 0.05],
            [0.00, 0.05, 0.1],
            [0.05, 0.1, 0.15],
            [0.00, 0.05, 0.1, 0.15, 0.2]
        ])
    lang = trial.suggest_categorical("lang", ["eng", "hattingam_legacy_single", "Hattingam_legacy_indiechars", "Hattingam_legacy_realdata",
                                              "Hattingam_legacy_text2img", "hattingam_lstm", "indie-hattings", "hattingam_lstmtest",
                                              "hattingam_legacy_psm7", "hattingam_lstm1"])

    params = {
        "bounding_box": bounding_box,
        "delta_range": delta_range,
        "threshold_range": threshold,
        "x_center": x_center,
        "min_stroke_thickness": stroke,
        "lang": lang,
        "stl_files_folder": data["stl_files_folder"],
        "output_folder_filled": data["output_folder_filled"]
    }

    accuracy, summary = run_ocr_pipeline(params)
    per_char = helpers.char_accuracy_dataframe(summary)
    confusion = helpers.confusion_dataframe(summary)
    plot = helpers.plot_heatmap(per_char, path="plots/cer_heatmap.png")

    # Build evaluation table for this trial
    table_dict = {
        "trial_id": [trial.number],
        "threshold": threshold,
        "min_stroke_thickness": [stroke],
        "lang": [lang],
        "delta_range": [delta_range],
        "accuracy": [accuracy]
    }

    all_runs.append(table_dict)
    column_run = r2c(all_runs)

    with mlflow.start_run():
        for k, v in params.items():
            mlflow.log_param(k, v)
        mlflow.log_metric("accuracy", accuracy)

        mlflow.log_table(data=column_run, artifact_file=f"all_results.json")
        mlflow.log_table(data=per_char, artifact_file="metrics/invidiual_character.json")
        mlflow.log_table(data=confusion, artifact_file="metrics/confusion_df.json")
        mlflow.log_artifact(plot, artifact_path="plots")

    return accuracy

# Running study for optuna
study = optuna.create_study(direction="maximize")
study.optimize(objective, n_trials=20)

print("✅ Optuna tuning complete. Visualizations generated and logged to MLflow.")