{ "cells": [ { "cell_type": "code", "execution_count": null, "id": "69f89461", "metadata": {}, "outputs": [], "source": [ "import glob\n", "import os\n", "\n", "\n", "import matplotlib\n", "matplotlib.use(\"Agg\")\n", "\n", "\n", "import matplotlib.pyplot as plt\n", "import numpy as np\n", "import pandas as pd\n", "from PIL import Image\n", "from scipy.signal import spectrogram\n", "\n", "\n", "# ============================================================\n", "# CONFIGURATION\n", "# ============================================================\n", "\n", "\n", "BASE_DIR = os.getcwd()\n", "\n", "\n", "INPUT_BASE = os.path.join(BASE_DIR, \"Reorganized-Data\")\n", "\n", "\n", "SPECTROGRAM_OUTPUT_BASE = os.path.join(\n", " BASE_DIR,\n", " \"CNN-Images\",\n", " \"Spectrogram\",\n", ")\n", "\n", "\n", "# Lowest sampling rate\n", "FS = 12800\n", "FS_FOLDER = \"12,8kHz\"\n", "\n", "\n", "# Higher-load points only\n", "LOADS = [\"323Nm_2000RPM\"]\n", "\n", "\n", "# Folder names created by your reorganization script\n", "STATES = [\"Healthy\", \"IR\", \"OR\", \"Corrosion\"]\n", "\n", "\n", "# Spectrogram settings\n", "SPECTROGRAM_SIZE = (128, 128) # width, height in pixels\n", "SPECTROGRAM_DPI = 150\n", "NPERSEG = 1024\n", "NOVERLAP = int(NPERSEG * 0.75)\n", "WINDOW = \"hann\"\n", "\n", "\n", "# ============================================================\n", "# HELPER FUNCTIONS\n", "# ============================================================\n", "\n", "\n", "def load_vibration_csv(path):\n", " \"\"\"\n", " Load vibration data from a CSV file.\n", "\n", "\n", " Tries common vibration-column names first. If none are found,\n", " it uses the first numeric column.\n", " \"\"\"\n", " df = pd.read_csv(path)\n", "\n", "\n", " for col in [\n", " \"Vibration\",\n", " \"vibration\",\n", " \"Acceleration\",\n", " \"acceleration\",\n", " \"Signal\",\n", " \"signal\",\n", " ]:\n", " if col in df.columns:\n", " return df[col].to_numpy(dtype=float)\n", "\n", "\n", " numeric_cols = df.select_dtypes(include=[np.number]).columns\n", "\n", "\n", " if len(numeric_cols) == 0:\n", " raise ValueError(f\"No numeric columns found in: {path}\")\n", "\n", "\n", " return df[numeric_cols[0]].to_numpy(dtype=float)\n", "\n", "\n", "\n", "def save_spectrogram_image(signal, fs, out_path):\n", " \"\"\"\n", " Compute and save a grayscale spectrogram image without axes,\n", " titles, colorbars, or white borders.\n", "\n", "\n", " The final image is resized to SPECTROGRAM_SIZE.\n", " \"\"\"\n", " f, t, sxx = spectrogram(\n", " signal,\n", " fs=fs,\n", " window=WINDOW,\n", " nperseg=NPERSEG,\n", " noverlap=NOVERLAP,\n", " detrend=\"constant\",\n", " scaling=\"density\",\n", " mode=\"magnitude\",\n", " )\n", "\n", "\n", " # Convert magnitude to dB\n", " sxx_db = 20 * np.log10(sxx + 1e-12)\n", "\n", "\n", " # Normalize each spectrogram independently\n", " sxx_db = (sxx_db - sxx_db.min()) / (\n", " sxx_db.max() - sxx_db.min() + 1e-12\n", " )\n", "\n", "\n", " # Create borderless figure\n", " fig = plt.figure(figsize=(6, 4), frameon=False)\n", " ax = plt.Axes(fig, [0.0, 0.0, 1.0, 1.0])\n", " ax.set_axis_off()\n", " fig.add_axes(ax)\n", "\n", "\n", " ax.imshow(\n", " sxx_db,\n", " aspect=\"auto\",\n", " origin=\"lower\",\n", " cmap=\"gray\",\n", " interpolation=\"bilinear\",\n", " )\n", "\n", "\n", " temp_path = out_path + \".tmp.png\"\n", "\n", "\n", " fig.savefig(\n", " temp_path,\n", " dpi=SPECTROGRAM_DPI,\n", " bbox_inches=None,\n", " pad_inches=0,\n", " )\n", " plt.close(fig)\n", "\n", "\n", " # Resize to the exact CNN input size\n", " with Image.open(temp_path) as img:\n", " img = img.convert(\"L\") # force grayscale\n", " img = img.resize(\n", " SPECTROGRAM_SIZE,\n", " Image.Resampling.LANCZOS,\n", " )\n", " img.save(out_path)\n", "\n", "\n", " os.remove(temp_path)\n", "\n", "\n", "\n", "# ============================================================\n", "# MAIN PROCESSING\n", "# ============================================================\n", "\n", "\n", "def main():\n", " os.makedirs(SPECTROGRAM_OUTPUT_BASE, exist_ok=True)\n", "\n", "\n", " total_csv_files = 0\n", " total_spectrograms = 0\n", "\n", "\n", " for load in LOADS:\n", " for state in STATES:\n", " input_folder = os.path.join(\n", " INPUT_BASE,\n", " FS_FOLDER,\n", " load,\n", " state,\n", " )\n", "\n", "\n", " spectrogram_folder = os.path.join(\n", " SPECTROGRAM_OUTPUT_BASE,\n", " FS_FOLDER,\n", " load,\n", " state,\n", " )\n", "\n", "\n", " if not os.path.isdir(input_folder):\n", " print(f\"Skipping missing input folder: {input_folder}\")\n", " continue\n", "\n", "\n", " os.makedirs(spectrogram_folder, exist_ok=True)\n", "\n", "\n", " csv_files = sorted(\n", " glob.glob(os.path.join(input_folder, \"*.csv\"))\n", " )\n", "\n", "\n", " if not csv_files:\n", " print(f\"No CSV files found in: {input_folder}\")\n", " continue\n", "\n", "\n", " print(\"=\" * 70)\n", " print(f\"Processing: {input_folder}\")\n", " print(f\"Files found: {len(csv_files)}\")\n", "\n", "\n", " total_csv_files += len(csv_files)\n", "\n", "\n", " for csv_path in csv_files:\n", " base_name = os.path.splitext(\n", " os.path.basename(csv_path)\n", " )[0]\n", "\n", "\n", " spectrogram_path = os.path.join(\n", " spectrogram_folder,\n", " f\"{base_name}_Spectrogram.png\",\n", " )\n", "\n", "\n", " try:\n", " signal = load_vibration_csv(csv_path)\n", "\n", "\n", " save_spectrogram_image(\n", " signal,\n", " FS,\n", " spectrogram_path,\n", " )\n", "\n", "\n", " total_spectrograms += 1\n", "\n", "\n", " print(\n", " f\" Saved spectrogram: \"\n", " f\"{os.path.basename(spectrogram_path)}\"\n", " )\n", "\n", "\n", " except Exception as e:\n", " print(\n", " f\" ERROR processing {csv_path}: {e}\"\n", " )\n", "\n", "\n", " print(\"=\" * 70)\n", " print(\"Processing complete!\")\n", " print(f\"Total CSV files processed: {total_csv_files}\")\n", " print(f\"Total spectrogram images: {total_spectrograms}\")\n", " print(f\"Spectrogram output: {SPECTROGRAM_OUTPUT_BASE}\")\n", "\n", "\n", "\n", "if __name__ == \"__main__\":\n", " main()" ] }, { "cell_type": "code", "execution_count": null, "id": "807a3cb2", "metadata": {}, "outputs": [], "source": [ "import json\n", "from pathlib import Path\n", "\n", "\n", "import matplotlib\n", "matplotlib.use(\"Agg\")\n", "\n", "\n", "import matplotlib.pyplot as plt\n", "import numpy as np\n", "import pandas as pd\n", "import tensorflow as tf\n", "\n", "\n", "from sklearn.metrics import (\n", " accuracy_score,\n", " classification_report,\n", " ConfusionMatrixDisplay,\n", " confusion_matrix,\n", " precision_recall_fscore_support,\n", ")\n", "from sklearn.model_selection import train_test_split\n", "\n", "\n", "# ============================================================\n", "# CONFIGURATION\n", "# ============================================================\n", "\n", "\n", "BASE_DIR = Path.cwd()\n", "\n", "\n", "DATA_DIR = (\n", " BASE_DIR / \"CNN-Images\" / \"Spectrogram\" / \"12,8kHz\"\n", ")\n", "\n", "\n", "OUTPUT_DIR = (\n", " BASE_DIR / \"CNN-Results\" / \"Spectrogram_2D_CNN_70_15_15\"\n", ")\n", "\n", "\n", "LOADS = [\"323Nm_2000RPM\"]\n", "\n", "\n", "CLASS_NAMES = [\"Healthy\", \"IR\", \"OR\", \"Corrosion\"]\n", "\n", "\n", "CLASS_FOLDERS = {\n", " \"Healthy\": [\"Healthy\"],\n", " \"IR\": [\"IR\"],\n", " \"OR\": [\"OR\"],\n", " \"Corrosion\": [\"Corrosion\", \"C\"],\n", "}\n", "\n", "\n", "IMAGE_SIZE = (128, 128)\n", "INPUT_SHAPE = IMAGE_SIZE + (1,)\n", "\n", "\n", "TRAIN_FRACTION = 0.70\n", "VALIDATION_FRACTION = 0.15\n", "TEST_FRACTION = 0.15\n", "\n", "\n", "BATCH_SIZE = 32\n", "EPOCHS = 50\n", "LEARNING_RATE = 1e-3\n", "SEED = 42\n", "\n", "\n", "LABEL_IDS = np.arange(len(CLASS_NAMES))\n", "AUTOTUNE = tf.data.AUTOTUNE\n", "\n", "\n", "# Architecture source:\n", "# https://keras.io/examples/vision/mnist_convnet/\n", "\n", "\n", "\n", "# ============================================================\n", "# DATA DISCOVERY\n", "# ============================================================\n", "\n", "\n", "def collect_images():\n", " paths = []\n", " labels = []\n", " operating_points = []\n", "\n", "\n", " for load in LOADS:\n", " for class_id, class_name in enumerate(CLASS_NAMES):\n", " count = 0\n", "\n", "\n", " for folder_name in CLASS_FOLDERS[class_name]:\n", " folder = DATA_DIR / load / folder_name\n", "\n", "\n", " if not folder.is_dir():\n", " continue\n", "\n", "\n", " files = sorted(folder.glob(\"*.png\"))\n", "\n", "\n", " for path in files:\n", " paths.append(str(path.resolve()))\n", " labels.append(class_id)\n", " operating_points.append(load)\n", "\n", "\n", " count += len(files)\n", "\n", "\n", " print(f\"{load:20s} | {class_name:10s} | {count:4d} images\")\n", "\n", "\n", " if not paths:\n", " raise FileNotFoundError(\n", " f\"No PNG spectrograms found under:\\n{DATA_DIR}\\n\"\n", " \"Expected: ///*.png\"\n", " )\n", "\n", "\n", " paths = np.asarray(paths, dtype=str)\n", " labels = np.asarray(labels, dtype=np.int32)\n", " operating_points = np.asarray(operating_points, dtype=str)\n", "\n", "\n", " counts = np.bincount(labels, minlength=len(CLASS_NAMES))\n", "\n", "\n", " print(\"\\nTotal samples by class:\")\n", " for name, count in zip(CLASS_NAMES, counts):\n", " print(f\" {name:10s}: {count}\")\n", "\n", "\n", " return paths, labels, operating_points\n", "\n", "\n", "\n", "# ============================================================\n", "# TENSORFLOW DATA PIPELINE\n", "# ============================================================\n", "\n", "\n", "def load_image(path, label):\n", " image_bytes = tf.io.read_file(path)\n", "\n", "\n", " image = tf.io.decode_png(image_bytes, channels=1)\n", " image = tf.image.resize(image, IMAGE_SIZE, method=\"bilinear\")\n", "\n", "\n", " image = tf.cast(image, tf.float32) / 255.0\n", " image = tf.ensure_shape(image, INPUT_SHAPE)\n", "\n", "\n", " return image, label\n", "\n", "\n", "\n", "def make_dataset(paths, labels, training=False, seed=SEED):\n", " dataset = tf.data.Dataset.from_tensor_slices((paths, labels))\n", "\n", "\n", " dataset = dataset.map(\n", " load_image,\n", " num_parallel_calls=AUTOTUNE,\n", " )\n", "\n", "\n", " if training:\n", " dataset = dataset.shuffle(\n", " buffer_size=len(paths),\n", " seed=seed,\n", " reshuffle_each_iteration=True,\n", " )\n", "\n", "\n", " return dataset.batch(BATCH_SIZE).prefetch(AUTOTUNE)\n", "\n", "\n", "\n", "# ============================================================\n", "# MODEL\n", "# ============================================================\n", "\n", "\n", "def build_model():\n", " \"\"\"\n", " Architecture from the official Keras \"Simple MNIST convnet\".\n", "\n", "\n", " Adaptations:\n", " - Input: 128 x 128 x 1 spectrograms\n", " - Output: four bearing conditions\n", " - Normalization is performed in the data pipeline\n", " \"\"\"\n", " model = tf.keras.Sequential(\n", " [\n", " tf.keras.Input(shape=INPUT_SHAPE),\n", "\n", "\n", " tf.keras.layers.Conv2D(\n", " 32,\n", " kernel_size=(3, 3),\n", " activation=\"relu\",\n", " ),\n", " tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),\n", "\n", "\n", " tf.keras.layers.Conv2D(\n", " 64,\n", " kernel_size=(3, 3),\n", " activation=\"relu\",\n", " ),\n", " tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),\n", "\n", "\n", " tf.keras.layers.Flatten(),\n", " tf.keras.layers.Dropout(0.5),\n", "\n", "\n", " tf.keras.layers.Dense(\n", " len(CLASS_NAMES),\n", " activation=\"softmax\",\n", " ),\n", " ],\n", " name=\"keras_example_spectrogram_cnn\",\n", " )\n", "\n", "\n", " model.compile(\n", " optimizer=tf.keras.optimizers.Adam(\n", " learning_rate=LEARNING_RATE\n", " ),\n", " loss=\"sparse_categorical_crossentropy\",\n", " metrics=[\"accuracy\"],\n", " )\n", "\n", "\n", " return model\n", "\n", "\n", "\n", "# ============================================================\n", "# PLOTS AND METRICS\n", "# ============================================================\n", "\n", "\n", "def save_history_plot(history, path):\n", " fig, axes = plt.subplots(1, 2, figsize=(11, 4))\n", "\n", "\n", " for ax, metric in zip(axes, [\"accuracy\", \"loss\"]):\n", " ax.plot(history.history[metric], label=\"Training\")\n", " ax.plot(\n", " history.history[f\"val_{metric}\"],\n", " label=\"Validation\",\n", " )\n", " ax.set_xlabel(\"Epoch (zero-based)\")\n", " ax.set_ylabel(metric.capitalize())\n", " ax.legend()\n", " ax.grid(alpha=0.3)\n", "\n", "\n", " fig.tight_layout()\n", " fig.savefig(path, dpi=160)\n", " plt.close(fig)\n", "\n", "\n", "\n", "def save_confusion_plot(y_true, y_pred, path, title):\n", " cm = confusion_matrix(\n", " y_true,\n", " y_pred,\n", " labels=LABEL_IDS,\n", " )\n", "\n", "\n", " fig, ax = plt.subplots(figsize=(7, 6))\n", "\n", "\n", " ConfusionMatrixDisplay(\n", " confusion_matrix=cm,\n", " display_labels=CLASS_NAMES,\n", " ).plot(\n", " ax=ax,\n", " cmap=\"Blues\",\n", " values_format=\"d\",\n", " colorbar=False,\n", " )\n", "\n", "\n", " ax.set_title(title)\n", " fig.tight_layout()\n", " fig.savefig(path, dpi=160)\n", " plt.close(fig)\n", "\n", "\n", "\n", "def calculate_metrics(y_true, y_pred):\n", " metrics = {\n", " \"accuracy\": float(accuracy_score(y_true, y_pred))\n", " }\n", "\n", "\n", " for average in [\"macro\", \"weighted\"]:\n", " precision, recall, f1, _ = precision_recall_fscore_support(\n", " y_true,\n", " y_pred,\n", " labels=LABEL_IDS,\n", " average=average,\n", " zero_division=0,\n", " )\n", "\n", "\n", " metrics[f\"precision_{average}\"] = float(precision)\n", " metrics[f\"recall_{average}\"] = float(recall)\n", " metrics[f\"f1_{average}\"] = float(f1)\n", "\n", "\n", " return metrics\n", "\n", "\n", "\n", "# ============================================================\n", "# MAIN\n", "# ============================================================\n", "\n", "\n", "def main():\n", " OUTPUT_DIR.mkdir(parents=True, exist_ok=True)\n", "\n", "\n", " paths, labels, loads = collect_images()\n", "\n", "\n", " config = {\n", " \"model_source\": (\n", " \"https://keras.io/examples/vision/mnist_convnet/\"\n", " ),\n", " \"class_names\": CLASS_NAMES,\n", " \"image_size\": list(IMAGE_SIZE),\n", " \"loads\": LOADS,\n", " \"train_fraction\": TRAIN_FRACTION,\n", " \"validation_fraction\": VALIDATION_FRACTION,\n", " \"test_fraction\": TEST_FRACTION,\n", " \"batch_size\": BATCH_SIZE,\n", " \"epochs\": EPOCHS,\n", " \"learning_rate\": LEARNING_RATE,\n", " \"seed\": SEED,\n", " \"pixel_normalization\": \"float32 / 255.0\",\n", " }\n", "\n", "\n", " with open(OUTPUT_DIR / \"config.json\", \"w\") as f:\n", " json.dump(config, f, indent=2)\n", "\n", "\n", " # ========================================================\n", " # SINGLE STRATIFIED 70 / 15 / 15 SPLIT\n", " # ========================================================\n", "\n", "\n", " all_indices = np.arange(len(paths))\n", "\n", "\n", " # First split:\n", " # 70% training, 30% temporary validation + test\n", " train_idx, temp_idx = train_test_split(\n", " all_indices,\n", " test_size=1.0 - TRAIN_FRACTION,\n", " stratify=labels,\n", " random_state=SEED,\n", " )\n", "\n", "\n", " # Second split:\n", " # Split the temporary 30% evenly:\n", " # 15% validation and 15% test\n", " val_idx, test_idx = train_test_split(\n", " temp_idx,\n", " test_size=0.50,\n", " stratify=labels[temp_idx],\n", " random_state=SEED,\n", " )\n", "\n", "\n", " print(\"\\n\" + \"=\" * 70)\n", " print(\"SINGLE STRATIFIED SPLIT\")\n", " print(\"=\" * 70)\n", " print(\n", " f\"Train: {len(train_idx)} \"\n", " f\"({len(train_idx) / len(paths):.1%})\"\n", " )\n", " print(\n", " f\"Validation: {len(val_idx)} \"\n", " f\"({len(val_idx) / len(paths):.1%})\"\n", " )\n", " print(\n", " f\"Test: {len(test_idx)} \"\n", " f\"({len(test_idx) / len(paths):.1%})\"\n", " )\n", "\n", "\n", " # Save the exact split for reproducibility.\n", " split_rows = []\n", "\n", "\n", " for split_name, indices in [\n", " (\"train\", train_idx),\n", " (\"validation\", val_idx),\n", " (\"test\", test_idx),\n", " ]:\n", " for idx in indices:\n", " split_rows.append(\n", " {\n", " \"path\": paths[idx],\n", " \"label\": int(labels[idx]),\n", " \"class\": CLASS_NAMES[labels[idx]],\n", " \"load\": loads[idx],\n", " \"split\": split_name,\n", " }\n", " )\n", "\n", "\n", " pd.DataFrame(split_rows).to_csv(\n", " OUTPUT_DIR / \"split_manifest.csv\",\n", " index=False,\n", " )\n", "\n", "\n", " tf.keras.backend.clear_session()\n", " tf.keras.utils.set_random_seed(SEED)\n", "\n", "\n", " train_ds = make_dataset(\n", " paths[train_idx],\n", " labels[train_idx],\n", " training=True,\n", " seed=SEED,\n", " )\n", "\n", "\n", " val_ds = make_dataset(\n", " paths[val_idx],\n", " labels[val_idx],\n", " )\n", "\n", "\n", " test_ds = make_dataset(\n", " paths[test_idx],\n", " labels[test_idx],\n", " )\n", "\n", "\n", " model = build_model()\n", " model.summary()\n", "\n", "\n", " best_model_path = OUTPUT_DIR / \"best_model.keras\"\n", "\n", "\n", " callbacks = [\n", " tf.keras.callbacks.ModelCheckpoint(\n", " str(best_model_path),\n", " monitor=\"val_loss\",\n", " mode=\"min\",\n", " save_best_only=True,\n", " verbose=0,\n", " ),\n", " tf.keras.callbacks.EarlyStopping(\n", " monitor=\"val_loss\",\n", " mode=\"min\",\n", " patience=10,\n", " restore_best_weights=True,\n", " verbose=1,\n", " ),\n", " tf.keras.callbacks.ReduceLROnPlateau(\n", " monitor=\"val_loss\",\n", " mode=\"min\",\n", " factor=0.5,\n", " patience=5,\n", " min_lr=1e-6,\n", " verbose=1,\n", " ),\n", " ]\n", "\n", "\n", " history = model.fit(\n", " train_ds,\n", " validation_data=val_ds,\n", " epochs=EPOCHS,\n", " callbacks=callbacks,\n", " verbose=2,\n", " )\n", "\n", "\n", " pd.DataFrame(history.history).to_csv(\n", " OUTPUT_DIR / \"history.csv\",\n", " index=False,\n", " )\n", "\n", "\n", " save_history_plot(\n", " history,\n", " OUTPUT_DIR / \"training_curves.png\",\n", " )\n", "\n", "\n", " # ========================================================\n", " # TEST EVALUATION\n", " # ========================================================\n", "\n", "\n", " best_model = tf.keras.models.load_model(\n", " str(best_model_path)\n", " )\n", "\n", "\n", " evaluation = best_model.evaluate(\n", " test_ds,\n", " return_dict=True,\n", " verbose=0,\n", " )\n", "\n", "\n", " probabilities = best_model.predict(test_ds, verbose=0)\n", " predictions = probabilities.argmax(axis=1)\n", " y_true = labels[test_idx]\n", "\n", "\n", " metrics = calculate_metrics(y_true, predictions)\n", " metrics.update(\n", " {\n", " \"test_loss\": float(evaluation[\"loss\"]),\n", " \"n_train\": len(train_idx),\n", " \"n_validation\": len(val_idx),\n", " \"n_test\": len(test_idx),\n", " \"epochs_run\": len(history.history[\"loss\"]),\n", " \"best_epoch\": (\n", " int(np.argmin(history.history[\"val_loss\"])) + 1\n", " ),\n", " }\n", " )\n", "\n", "\n", " results_df = pd.DataFrame([metrics])\n", " results_df.to_csv(\n", " OUTPUT_DIR / \"metrics.csv\",\n", " index=False,\n", " )\n", "\n", "\n", " report = classification_report(\n", " y_true,\n", " predictions,\n", " labels=LABEL_IDS,\n", " target_names=CLASS_NAMES,\n", " digits=4,\n", " zero_division=0,\n", " )\n", "\n", "\n", " print(\"\\n\" + \"=\" * 70)\n", " print(\"TEST-SET RESULTS\")\n", " print(\"=\" * 70)\n", " print(f\"Test accuracy: {metrics['accuracy']:.4f}\")\n", " print(f\"Test loss: {metrics['test_loss']:.4f}\")\n", " print(report)\n", "\n", "\n", " (OUTPUT_DIR / \"classification_report.txt\").write_text(\n", " report,\n", " encoding=\"utf-8\",\n", " )\n", "\n", "\n", " save_confusion_plot(\n", " y_true,\n", " predictions,\n", " OUTPUT_DIR / \"confusion_matrix.png\",\n", " \"Held-out test predictions\",\n", " )\n", "\n", "\n", " # ========================================================\n", " # SAVE TEST PREDICTIONS\n", " # ========================================================\n", "\n", "\n", " predictions_df = pd.DataFrame(\n", " {\n", " \"path\": paths[test_idx],\n", " \"load\": loads[test_idx],\n", " \"true_label\": y_true,\n", " \"predicted_label\": predictions,\n", " \"true_class\": [CLASS_NAMES[i] for i in y_true],\n", " \"predicted_class\": [\n", " CLASS_NAMES[i] for i in predictions\n", " ],\n", " }\n", " )\n", "\n", "\n", " for class_id, name in enumerate(CLASS_NAMES):\n", " predictions_df[f\"probability_{name}\"] = (\n", " probabilities[:, class_id]\n", " )\n", "\n", "\n", " predictions_df.to_csv(\n", " OUTPUT_DIR / \"test_predictions.csv\",\n", " index=False,\n", " )\n", "\n", "\n", " # ========================================================\n", " # SAVE SUMMARY\n", " # ========================================================\n", "\n", "\n", " score_columns = [\n", " \"accuracy\",\n", " \"precision_macro\",\n", " \"recall_macro\",\n", " \"f1_macro\",\n", " \"precision_weighted\",\n", " \"recall_weighted\",\n", " \"f1_weighted\",\n", " \"test_loss\",\n", " ]\n", "\n", "\n", " summary = {\n", " metric: float(metrics[metric])\n", " for metric in score_columns\n", " }\n", "\n", "\n", " with open(OUTPUT_DIR / \"summary.json\", \"w\") as f:\n", " json.dump(summary, f, indent=2)\n", "\n", "\n", " print(f\"\\nResults saved to: {OUTPUT_DIR}\")\n", "\n", "\n", "\n", "if __name__ == \"__main__\":\n", " main()" ] } ], "metadata": { "kernelspec": { "display_name": "tf (3.12.3)", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.12.3" } }, "nbformat": 4, "nbformat_minor": 5 }