preprocess_and_augment.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "import os\n",
    "import sys\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from datetime import datetime\n",
    "import laspy as lp\n",
    "import subprocess\n",
    "import glob\n",
    "import matplotlib.pyplot as plt\n",
    "plt.rcParams['font.size'] = 14\n",
    "import shutil\n",
    "from collections import Counter\n",
    "from sklearn.model_selection import StratifiedKFold\n",
    "from scipy.interpolate import LinearNDInterpolator\n",
    "from scipy.spatial import cKDTree\n",
    "from jakteristics import las_utils, compute_features, FEATURE_NAMES\n",
    "from tqdm import tqdm\n",
    "import re\n",
    "import tempfile"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "input_ALS_folder = '/home/davidhersh/Dropbox/Uni/ThesisHersh/ALS_data'\n",
    "output_folder = '/media/davidhersh/T7 Shield/pre-processing/tmp_Dec26'\n",
    "dataset_dir = '/media/davidhersh/T7 Shield/Datasets_Dec26'\n",
    "\n",
    "# Running each \n",
    "max_epoch = 50\n",
    "first_kpconv_subsampling_dl = 0.4\n",
    "\n",
    "timestamp = datetime.now().strftime(\"%Y_%m_%d_%H_%M\")\n",
    "saving_path = f'/media/davidhersh/T7 Shield/results_{timestamp}'\n"
   ],
   "id": "fdf22083b0d25cfe",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "point_threshold = 500\n",
    "\n",
    "# k-fold\n",
    "n_splits = 2\n",
    "\n",
    "# Augmentation values\n",
    "min_subsample_distance = 0.4\n",
    "rotations = [-15, 15]\n",
    "normals_search_radius = 2\n"
   ],
   "id": "1a53d5a26594bca3",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "def copy_folder(input_folder, output_folder=output_folder, point_threshold=point_threshold, plot=True, redo=True):\n",
    "    \"\"\"\n",
    "    Copy single tree ALS files (ending with _g meaning ground classified) for a specific list of trees if they have a minimum point count\n",
    "    \"\"\"\n",
    "    \n",
    "    if redo:\n",
    "    \n",
    "        if not os.path.exists(output_folder):\n",
    "            os.makedirs(output_folder)\n",
    "            \n",
    "        # Save figs\n",
    "        figdir = output_folder.replace('/tmp', '/processing_figures')\n",
    "        if not os.path.exists(figdir):\n",
    "            os.makedirs(figdir)\n",
    "        \n",
    "        # Only copy certain species\n",
    "        species_to_copy = [\n",
    "            \"PseMen\",\n",
    "            \"FagSyl\",\n",
    "            \"PinSyl\",\n",
    "            \"QueRub\",\n",
    "            \"PicAbi\",\n",
    "            \"QuePet\"\n",
    "        ]\n",
    "    \n",
    "        files = []\n",
    "        species_counter = Counter()\n",
    "        \n",
    "        for root, dirs, filenames in os.walk(input_folder):\n",
    "            for filename in filenames:\n",
    "                try:\n",
    "                    species_name = filename.split('_')[0]\n",
    "        \n",
    "                    # Study area BR06 has an issue with ground. Do not copy for further processing\n",
    "                    # File format: FagSyl_BR02_04_2019-07-05_q2_ALS-on_g.laz\n",
    "                    study_area = filename.split('_')[1] # BR02\n",
    "        \n",
    "                    if filename.endswith('ALS-on_g.laz') and species_name in species_to_copy and study_area != \"BR06\":\n",
    "                            las = lp.read(os.path.join(root, filename))\n",
    "                            number_of_nonground_points = len(las.points[las.classification !=2])\n",
    "                            if number_of_nonground_points <= point_threshold:\n",
    "                                species_counter[species_name] += 1\n",
    "                                files.append(os.path.join(root, filename))  # Store full file paths\n",
    "                                shutil.copyfile(os.path.join(root, filename), os.path.join(output_folder, filename))\n",
    "                except: \n",
    "                    pass\n",
    "        print(f'Copying {len(files)} files')\n",
    "        print(f'Unique species: {species_counter}')\n",
    "        \n",
    "        if plot:\n",
    "            fig, ax = plt.subplots(figsize=(12,8))\n",
    "            plt.bar(species_counter.keys(), species_counter.values())\n",
    "            plt.ylabel('Number of trees')\n",
    "            plt.xlabel('Species')\n",
    "            plt.title(f\"Number of trees with more than {point_threshold} non-ground points\")\n",
    "            plt.savefig(os.path.join(figdir, f'species_count_greaterthan_{point_threshold}_points.png'), bbox_inches='tight')\n",
    "            plt.show()\n",
    "        \n",
    "        return output_folder, species_counter\n",
    "    \n",
    "    else:\n",
    "        print(f'Skipping copying...')\n"
   ],
   "id": "2c2aaeabc71c8b32",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "copied_als_folder, species_counter = copy_folder(input_ALS_folder)",
   "id": "7861e14769e968f7",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "def calculate_hag(las_file):\n",
    "    las = lp.read(las_file)\n",
    "    x = las.x\n",
    "    y = las.y\n",
    "    z = las.z\n",
    "    \n",
    "    classification = las.classification\n",
    "    \n",
    "    ground_indices = classification == 2\n",
    "    x_ground = x[ground_indices]\n",
    "    y_ground = y[ground_indices]\n",
    "    z_ground = z[ground_indices]\n",
    "    \n",
    "    X, Y = np.meshgrid(x_ground, y_ground)\n",
    "    ground_interpolated = LinearNDInterpolator(list(zip(x_ground, y_ground)), z_ground)\n",
    "    Z = ground_interpolated(X, Y)\n",
    "    z_ground_at_points = ground_interpolated(x, y)\n",
    "    min_ground = np.nanmin(z_ground)\n",
    "    z_ground_at_points = np.where(np.isnan(z_ground_at_points), min_ground, z_ground_at_points)\n",
    "    hag = z - z_ground_at_points\n",
    "    las.z = hag\n",
    "    las = las[las.classification != 2]\n",
    "    \n",
    "    las.write(las_file)"
   ],
   "id": "4dca802491ce3d23",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "# Can do HAG calculation before split\n",
    "\n",
    "def convert_hag_to_z(input_folder=copied_als_folder, redo=True):\n",
    "    \n",
    "    if redo:\n",
    "        files = os.listdir(copied_als_folder)\n",
    "        for file in tqdm(files, desc='Processing files', unit='file'):\n",
    "            full_path = os.path.join(copied_als_folder, file)\n",
    "            calculate_hag(full_path)\n",
    "    else:\n",
    "        print(\"Skipping HAG step...\")\n",
    "        \n",
    "convert_hag_to_z(input_folder=copied_als_folder)"
   ],
   "id": "66dad78927f5f4a0",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "def stratified_k_fold_split(input_folder, output_folder, n_splits=5):\n",
    "    input_folder = copied_als_folder\n",
    "    output_folder = input_folder.replace('/tmp', '/kfolders')\n",
    "\n",
    "    files = []\n",
    "    labels = []\n",
    "    \n",
    "    for file in os.listdir(input_folder):\n",
    "        class_name = file.split('_')[0] #Fagsyl etc...\n",
    "        files.append(os.path.join(input_folder, file))\n",
    "        labels.append(class_name)\n",
    "    class_counts = Counter(labels)\n",
    "    print(f'Class counts: {class_counts}')\n",
    "    \n",
    "    skf = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=32)\n",
    "    \n",
    "    train_folders = []\n",
    "    test_folders = []\n",
    "    \n",
    "    for fold, (train_index, test_index) in enumerate(skf.split(files, labels)):\n",
    "        fold_train_dir = os.path.join(output_folder, f'fold_{fold+1}_train')\n",
    "        train_folders.append(fold_train_dir)\n",
    "        fold_test_dir = os.path.join(output_folder, f'fold_{fold+1}_val')\n",
    "        test_folders.append(fold_test_dir)\n",
    "        os.makedirs(fold_train_dir, exist_ok=True)\n",
    "        os.makedirs(fold_test_dir, exist_ok=True)\n",
    "        \n",
    "        for train_idx in train_index:\n",
    "            shutil.copy(files[train_idx], fold_train_dir)\n",
    "        for test_idx in test_index:\n",
    "            shutil.copy(files[test_idx], fold_test_dir)\n",
    "            \n",
    "    return train_folders, test_folders\n",
    "    \n",
    "train_folders, test_folders = stratified_k_fold_split(copied_als_folder, output_folder, n_splits=n_splits)"
   ],
   "id": "828540b56e86d123",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "def poisson_subsample(las_file, min_distance=min_subsample_distance):\n",
    "    las = lp.read(las_file)\n",
    "    points = np.vstack((las.x, las.y, las.z)).transpose()\n",
    "\n",
    "    kdtree = cKDTree(points)\n",
    "\n",
    "    selected = np.zeros(len(points), dtype=bool)\n",
    "    selected_indices = []\n",
    "    for i, point in enumerate(points):\n",
    "        if selected[i]:\n",
    "            continue\n",
    "\n",
    "        selected_indices.append(i)\n",
    "        selected[i] = True\n",
    "\n",
    "        indices = kdtree.query_ball_point(point, min_distance)\n",
    "        selected[indices] = True\n",
    "\n",
    "    header = lp.LasHeader(point_format=las.header.point_format, version=las.header.version)\n",
    "    header.offsets = las.header.offsets\n",
    "    header.scales = las.header.scales\n",
    "\n",
    "    subsampled_las = lp.LasData(header)\n",
    "    subsampled_las.points = las.points[selected_indices]\n",
    "\n",
    "    subsampled_las.write(las_file)"
   ],
   "id": "d58dd06dcaa3137f",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "def normalize_xy(las_file):\n",
    "        las = lp.read(las_file)\n",
    "        original_offset_x = las.header.offsets[0]\n",
    "        original_offset_y = las.header.offsets[1]\n",
    "        scale_x = las.header.scales[0]\n",
    "        scale_y = las.header.scales[1]\n",
    "\n",
    "        x = las.x\n",
    "        y = las.y\n",
    "\n",
    "        mean_x = np.mean(x)\n",
    "        mean_y = np.mean(y)\n",
    "        \n",
    "        # Normalize coordinates\n",
    "        normalized_x = x - mean_x\n",
    "        normalized_y = y - mean_y\n",
    "\n",
    "        new_header = lp.LasHeader(point_format=las.header.point_format, version=las.header.version)\n",
    "        new_header.scales = las.header.scales\n",
    "\n",
    "        new_header.offsets = [\n",
    "            original_offset_x - mean_x,\n",
    "            original_offset_y - mean_y,\n",
    "            las.header.offsets[2]  # Same z\n",
    "        ]\n",
    "\n",
    "        new_las = lp.LasData(new_header)\n",
    "        new_las.x = normalized_x\n",
    "        new_las.y = normalized_y\n",
    "        new_las.z = las.z\n",
    "\n",
    "        for dim_name in las.point_format.dimension_names:\n",
    "            if dim_name not in [\"X\", \"Y\", \"Z\"]:\n",
    "                setattr(new_las, dim_name, getattr(las, dim_name))\n",
    "\n",
    "        new_las.write(las_file)"
   ],
   "id": "f9127ad9978190d7",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "def rotate_z(pointcloud, degrees):\n",
    "    theta = np.deg2rad(degrees)\n",
    "\n",
    "    rotation_matrix = np.array([\n",
    "        [np.cos(theta), -np.sin(theta), 0],\n",
    "        [np.sin(theta), np.cos(theta), 0],\n",
    "        [0, 0, 1]\n",
    "    ])\n",
    "\n",
    "    return np.dot(pointcloud, rotation_matrix.T)\n",
    "\n",
    "\n",
    "def rotate_las(las_file, rotations):\n",
    "        las = lp.read(las_file)\n",
    "\n",
    "        point_data = np.vstack((las.x, las.y, las.z)).T\n",
    "\n",
    "        for i, rotation in enumerate(rotations):\n",
    "\n",
    "            rotated_points = rotate_z(point_data, rotation)\n",
    "\n",
    "            header = lp.LasHeader(point_format=las.header.point_format, version=las.header.version)\n",
    "            header.offsets = las.header.offsets\n",
    "            header.scales = las.header.scales\n",
    "            # header.crs = lp.crs.CRS.from_epsg(25832)\n",
    "\n",
    "            rotated_las = lp.LasData(header)\n",
    "            rotated_las.x = rotated_points[:, 0]\n",
    "            rotated_las.y = rotated_points[:, 1]\n",
    "            rotated_las.z = rotated_points[:, 2]\n",
    "\n",
    "            for dim_name in las.point_format.dimension_names:\n",
    "                if dim_name not in [\"X\", \"Y\", \"Z\"]:\n",
    "                    setattr(rotated_las, dim_name, getattr(las, dim_name))\n",
    "\n",
    "            unique_filename = las_file.replace('.laz', f'_rot_{rotation}.laz')\n",
    "            output_filename = os.path.join(output_folder, unique_filename)\n",
    "\n",
    "            os.makedirs(output_folder, exist_ok=True)\n",
    "            rotated_las.write(output_filename)"
   ],
   "id": "44f5634f32094f43",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "def calculate_normals(las_file):\n",
    "    las = lp.read(las_file)\n",
    "    FEATURE_NAMES = ['nx', 'ny', 'nz']\n",
    "    \n",
    "    # Remove nx, ny, nz if existing\n",
    "    for feature_name in FEATURE_NAMES:\n",
    "        if feature_name in las.point_format.dimension_names:\n",
    "            las.point_format.remove_extra_dimension(feature_name)\n",
    "    \n",
    "    xyz = las_utils.read_las_xyz(las_file)\n",
    "    \n",
    "    features = compute_features(xyz, search_radius=normals_search_radius, feature_names=FEATURE_NAMES)\n",
    "    \n",
    "    output_file = las_file.replace('.laz', '_normals.laz')\n",
    "    \n",
    "    if not os.path.exists(output_file):\n",
    "        las_utils.write_with_extra_dims(las_file, output_file, features, FEATURE_NAMES)\n",
    "        os.remove(las_file)"
   ],
   "id": "6c2a091fefd20749",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "# Augment each folder\n",
    "\n",
    "def augmentation(train_folders=train_folders, test_folders=test_folders, \n",
    "                 redo_xy_normalization=True,\n",
    "                 redo_subsample=True,\n",
    "                 redo_rotation=True):\n",
    "\n",
    "    all_folders = train_folders + test_folders\n",
    "    for folder in all_folders:\n",
    "        for las_file in os.listdir(folder):\n",
    "            \n",
    "            # Normalize x,y by means\n",
    "            if redo_xy_normalization:\n",
    "                normalize_xy(las_file=os.path.join(folder, las_file))\n",
    "            else:\n",
    "                pass\n",
    "            \n",
    "            # Downsample using minimum distance\n",
    "            if redo_subsample:\n",
    "                poisson_subsample(las_file=os.path.join(folder, las_file),\n",
    "                                  min_distance=min_subsample_distance)\n",
    "            else:\n",
    "                pass\n",
    "            \n",
    "            # Rotate on z-axis\n",
    "            if redo_rotation:\n",
    "                rotate_las(os.path.join(folder, las_file), rotations=rotations)\n",
    "            else:\n",
    "                pass\n",
    "            \n",
    "    # Need to relist the folders with the new rotated files            \n",
    "    for folder in all_folders:\n",
    "        for las_file in os.listdir(folder):\n",
    "            \n",
    "            # Calculate normals\n",
    "            calculate_normals(os.path.join(folder, las_file))\n",
    "    \n",
    "augmentation()"
   ],
   "id": "b15af38c5b60886c",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "# Need to deal with nan normal values\n",
    "def convert_to_txt(train_folders=train_folders, test_folders=test_folders):\n",
    "    all_folders = train_folders + test_folders\n",
    "\n",
    "    for folder in all_folders:\n",
    "        for las_file in os.listdir(folder):\n",
    "            if not las_file.endswith(('laz', 'las')):\n",
    "                continue\n",
    "            las = lp.read(os.path.join(folder, las_file))\n",
    "            output_file = las_file.replace('.laz', '.txt')\n",
    "            \n",
    "            with open(os.path.join(folder, output_file), 'w') as f:\n",
    "                for x, y, z, nx, ny, nz in zip(las.x, las.y, las.z, las.nx, las.ny, las.nz):\n",
    "                    f.write(f'{x:.6f}, {y:.6f}, {z:.6f}, {nx}, {ny}, {nz}\\n')\n",
    "            # os.remove(os.path.join(folder, las_file))\n",
    "convert_to_txt()\n"
   ],
   "id": "3dc93712e00d6a49",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "# Convert to KPConv repository dataset format\n",
    "\n",
    "def copy_to_datasets(train_folders=train_folders, test_folders=test_folders, n_splits=n_splits):\n",
    "    all_folders = train_folders + test_folders\n",
    "    base_folder = train_folders[0]\n",
    "    base_folder_test = test_folders[0]\n",
    "    print(f'base test folder: {base_folder_test}')\n",
    "    print(base_folder)\n",
    "    \n",
    "    # Get the folders of the species\n",
    "    species_set = set()\n",
    "    for folder in all_folders:\n",
    "        for las_file in os.listdir(folder):\n",
    "            species = las_file.split('_')[0]\n",
    "            species_set.add(species)\n",
    "    print(f'species_set: {species_set}')\n",
    "    \n",
    "    all_data_dirs = []\n",
    "    for i in range(1, n_splits+1):\n",
    "        \n",
    "        # Make a folder for each dataset\n",
    "        datadir = os.path.join(dataset_dir, f'data_{i}_subsample_{min_subsample_distance}')\n",
    "        all_data_dirs.append(datadir)\n",
    "        os.makedirs(datadir, exist_ok=True)\n",
    "        \n",
    "        # Make a subfolder for each species\n",
    "        species_dirs = {species: os.path.join(datadir, species) for species in species_set}\n",
    "        for species_dir in species_dirs.values():\n",
    "            os.makedirs(species_dir, exist_ok=True)\n",
    "        \n",
    "        # Create filelist.txt, train.txt and test.txt files in the folder\n",
    "        filelist_txt_file = os.path.join(datadir, 'filelist.txt')\n",
    "        train_txt_file = os.path.join(datadir, 'train.txt')\n",
    "        test_txt_file = os.path.join(datadir, 'test.txt')\n",
    "        \n",
    "        # Process training files\n",
    "        train_folder = re.sub(r\"fold_\\d+_train\", f'fold_{i}_train', base_folder)\n",
    "        with open (train_txt_file, 'w') as f:\n",
    "            for file in os.listdir(train_folder):\n",
    "                if file.endswith('.txt'):\n",
    "                    f.write(file[:-4] + '\\n')\n",
    "                else:\n",
    "                    pass\n",
    "            # print(f'Processed: {i} train files')\n",
    "        \n",
    "        for file in os.listdir(train_folder):\n",
    "            species = file.split('_')[0]\n",
    "            if file.endswith('.txt'):\n",
    "                shutil.copy(os.path.join(train_folder, file), species_dirs[species])\n",
    "            else:\n",
    "                pass\n",
    "                \n",
    "        # Process test files\n",
    "        test_folder = re.sub(r\"fold_\\d+_val\", f'fold_{i}_val', base_folder_test)\n",
    "        test_files = [file for file in os.listdir(test_folder) if file.endswith('.txt')]\n",
    "        with open(test_txt_file, 'w') as f:\n",
    "            for file in os.listdir(test_folder):\n",
    "                if file.endswith('.txt'):\n",
    "                    f.write(file[:-4] + '\\n')\n",
    "                else:\n",
    "                    pass\n",
    "                \n",
    "        for file in os.listdir(test_folder):\n",
    "            species = file.split('_')[0]\n",
    "            if file.endswith('.txt'):\n",
    "                shutil.copy(os.path.join(test_folder, file), species_dirs[species])\n",
    "            else:\n",
    "                pass\n",
    "                \n",
    "        with open (filelist_txt_file, 'w') as f:\n",
    "            for species_dir in species_dirs.values():\n",
    "                for file in os.listdir(species_dir):\n",
    "                    relative_path = os.path.join(os.path.basename(species_dir), file)\n",
    "                    f.write(relative_path + '\\n')\n",
    "        \n",
    "    return all_data_dirs\n",
    "        \n",
    "all_data_dirs = copy_to_datasets()\n",
    "# Loop over each dataset and run trainNeuesPalaisTrees.py\n",
    "# Need to pass config parameters here and update class_w"
   ],
   "id": "aa91d6ed998cfc69",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "\n",
    "# all_data_dirs\n",
    "# \n",
    "# base_dir = '/media/davidhersh/T7 Shield/Datasets/'\n",
    "# all_data_dirs = os.listdir('/media/davidhersh/T7 Shield/Datasets/')\n",
    "# all_data_dirs = [os.path.join(base_dir, d) for d in all_data_dirs]\n",
    "\n",
    "print(all_data_dirs)"
   ],
   "id": "7ffbc490c3fc116b",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "def calculate_number_of_train_and_test_files(data_folder):\n",
    "    with open(os.path.join(data_folder, 'train.txt')) as f:\n",
    "        num_train_files = len(f.readlines())\n",
    "    \n",
    "    with open(os.path.join(data_folder, 'test.txt')) as f:\n",
    "        num_test_files = len(f.readlines())\n",
    "    return num_train_files, num_test_files"
   ],
   "id": "42f59db89abf2f33",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "# KPConv parameters\n",
    "# parser.add_argument('--do_subsample', action='store_false', help='Enable subsampling or not')\n",
    "# parser.add_argument('--first_subsampling_dl', type=float, default=None, help='The KPConv subsampling value')\n",
    "# parser.add_argument('--max_epoch', type=int, default=None, help='Override max epoch')\n",
    "# parser.add_argument('--data_path', type=str, default=None, help='Data path')\n",
    "# parser.add_argument('--saving_path', type=str, default=None, help='Saving path')\n",
    "\n",
    "# Need to add the len of files\n",
    "# parser.add_argument('--num_train_models', type=int, help='# of trees')\n",
    "# parser.add_argument('--num_test_models', type=int,  help='# of trees')\n",
    "# \n",
    "# if args.num_train_models and self.mode == 'train':\n",
    "#     self.num_models = args.num_train_models\n",
    "# if args.num_test_models and self.mode == 'test':\n",
    "#     self.num_models = args.num_test_models\n",
    "\n",
    "if not os.path.exists(saving_path):\n",
    "    os.makedirs(saving_path)\n",
    "\n",
    "# Save all parameters in results\n",
    "parameters = os.path.join(saving_path, 'parameters.txt')\n",
    "with open(parameters, 'w') as f:\n",
    "    f.write('Augmentation parameters:\\n')\n",
    "    f.write('-----------------------\\n')\n",
    "    f.write(f'Minimum distance between points: {min_subsample_distance}\\n')\n",
    "    f.write(f'Rotations: {rotations}\\n')\n",
    "    f.write(f'Normals search radius: {normals_search_radius}\\n')\n",
    "    f.write('\\n')\n",
    "    \n",
    "    # k-fold\n",
    "    f.write(f'k-folds: {n_splits}\\n')\n",
    "    f.write('\\n')\n",
    "    \n",
    "    #KPConv values\n",
    "    f.write('KPConv Parameters:\\n')\n",
    "    f.write('-----------------------\\n')\n",
    "    f.write(f'max_epoch: {max_epoch}\\n')\n",
    "    f.write(f'first_kpconv_subsampling_dl: {first_kpconv_subsampling_dl}\\n')\n",
    "        \n",
    "    \n",
    "\n",
    "for folder in all_data_dirs:\n",
    "    # Save each fold within the results\n",
    "    kfold_folder_name = folder.split('/')[-1]\n",
    "    kfold_folder_path = os.path.join(saving_path, kfold_folder_name) # Save path\n",
    "    if not os.path.exists(kfold_folder_path):\n",
    "        os.makedirs(kfold_folder_path, exist_ok=True)\n",
    "    \n",
    "    # Set num train,test files\n",
    "    num_train_files, num_test_files = calculate_number_of_train_and_test_files(folder)\n",
    "    \n",
    "    args = ['--max_epoch',str(max_epoch),\n",
    "            '--first_subsampling_dl',str(first_kpconv_subsampling_dl),\n",
    "            '--data_path', folder,\n",
    "            '--saving_path', kfold_folder_path,\n",
    "            '--num_train_models', str(num_train_files),\n",
    "            '--num_test_models', str(num_test_files)]\n",
    "    \n",
    "    results = subprocess.run(\n",
    "        ['python', 'train_NeuesPalaisTrees.py'] + args,\n",
    "        capture_output=True,\n",
    "        text=True\n",
    "    )\n",
    "    print(f'Running: {results}')\n",
    "    print(f'Results folder: {kfold_folder_path}')\n",
    "    print(results.stdout)\n"
   ],
   "id": "9298a85596ddb0e1",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "def plot_train_and_val_accuracy_for_all_folds(saving_path=saving_path, num_classes=6):\n",
    "    folders = [os.path.join(saving_path, f) for f in os.listdir(saving_path) if os.path.isdir(os.path.join(saving_path, f))]\n",
    "\n",
    "    fold=0\n",
    "    alpha = .2\n",
    "    fig, ax = plt.subplots(figsize=(14, 8))\n",
    "    for folder in folders:\n",
    "        training_file = os.path.join(folder, 'training.txt')\n",
    "        validation_file = os.path.join(folder, 'val_confs.txt')\n",
    "        df = pd.read_csv(training_file, sep='\\s+')\n",
    "        df = df.groupby('epochs').mean().reset_index()\n",
    "        ax.plot(df['epochs'], df['train_accuracy'] * 100, 'k', label=f'Train: {fold+1}', alpha=alpha)\n",
    "\n",
    "        # Val plot\n",
    "        with open(validation_file, 'r') as f:\n",
    "            val_confs = f.readlines()\n",
    "        val_accuracies = []\n",
    "        for val in val_confs:\n",
    "            matrix = np.array(list(map(int, val.split()))).reshape(num_classes, num_classes)\n",
    "            accuracy = np.sum(np.diag(matrix)) / np.sum(matrix)\n",
    "            val_accuracies.append(accuracy * 100)\n",
    "        ax.plot(df['epochs'], val_accuracies, label=f'Val: {fold+1}', color='b', linestyle='--', alpha=alpha)\n",
    "\n",
    "        fold += 1\n",
    "        alpha +=.2\n",
    "\n",
    "\n",
    "    plt.legend(title='Metric and fold')\n",
    "    plt.xlabel('Epoch')\n",
    "    plt.ylabel('Accuracy [%]')\n",
    "    plt.ylim([0, 100])\n",
    "    plt.title(f'Training and validation accuracy for {n_splits} folds')\n",
    "    plt.grid()\n",
    "    plt.savefig(os.path.join(saving_path, 'train_accuracy.png'), bbox_inches='tight', dpi=400)\n",
    "    plt.show();\n",
    "\n",
    "    return folders\n",
    "\n",
    "results_folders = plot_train_and_val_accuracy_for_all_folds(saving_path=saving_path);"
   ],
   "id": "6277fd04a37117e3",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "# parser.add_argument('--results_path', type=str, default=None, help='Results path')\n",
    "\n",
    "# Pass each of the folders after training to test_models.py\n",
    "# Within each folder, a new folder /test will be saved with test confusion matrices\n",
    "# Should add an arg for the tester.py num_votes value. Current is 10 votes\n",
    "\n",
    "def test_models(data_folders=all_data_dirs):\n",
    "    for data_dir in data_folders:\n",
    "        print(results_folders)\n",
    "        args = ['--data_path', data_dir]\n",
    "        \n",
    "    results = subprocess.run(\n",
    "        ['python', 'test_models.py'] + args,\n",
    "        capture_output=True,\n",
    "        text=True\n",
    "    )\n",
    "    print(f'Path to data: {data_dir}')\n",
    "    print(results.stdout)\n",
    "    print(results.stderr)\n",
    "\n",
    "    test_results_folders = [os.path.join(saving_path, folder) for folder in saving_path]\n",
    "\n",
    "    return test_results_folders\n",
    "\n",
    "test_results_folders = test_models(data_folders=all_data_dirs)"
   ],
   "id": "620fc9ad0163a13a",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "",
   "id": "2bed92ea919f5aa1",
   "outputs": [],
   "execution_count": null
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}