From 091d2074ca6350a47c33bcaf69e89b5fee67fd12 Mon Sep 17 00:00:00 2001
From: Laurent FAINSIN <laurent.fainsin@safrangroup.com>
Date: Mon, 17 Apr 2023 10:35:46 +0200
Subject: [PATCH] feat: modify test_generation to translate back sampled blade
 from norminal deformations

---
 compare_samples.py           | 6 +++---
 dataset/test_rotor37_data.py | 2 +-
 test_generation.py           | 6 +++++-
 3 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/compare_samples.py b/compare_samples.py
index e98c22a..9f186ab 100644
--- a/compare_samples.py
+++ b/compare_samples.py
@@ -9,7 +9,7 @@ VTKFILE_NOMINAL = Path("~/data/stage-laurent-f/datasets/Rotor37/processed/nomina
 nominal = pv.read(VTKFILE_NOMINAL)
 
 # for each generated/sampled blade
-gen_files = Path.cwd().glob("gen*.txt")
+gen_files = Path("./output").glob("gen*.txt")
 for gen_file in gen_files:
     # load numpy txt
     blade = np.loadtxt(gen_file)
@@ -28,10 +28,10 @@ for gen_file in gen_files:
     blade -= center
 
     # save to txt
-    np.savetxt(f"test_{gen_file.stem}.txt", blade)
+    np.savetxt(f"output/test_{gen_file.stem}.txt", blade)
 
     # swap nominal points to blade points
     nominal.points = blade
 
     # save altered blade to vtk
-    nominal.save(f"test_{gen_file.stem}.vtk")
+    nominal.save(f"output/test_{gen_file.stem}.vtk")
diff --git a/dataset/test_rotor37_data.py b/dataset/test_rotor37_data.py
index b8c340d..b54716b 100644
--- a/dataset/test_rotor37_data.py
+++ b/dataset/test_rotor37_data.py
@@ -18,7 +18,7 @@ for idx, blade in enumerate(test_ds):
     pc = pc * STD + MEAN
 
     print(f"Saving point cloud {idx}...")
-    np.savetxt(f"pc_{idx}.txt", pc)
+    np.savetxt(f"output/pc_{idx}.txt", pc)
 
     if idx >= 10:
         break
diff --git a/test_generation.py b/test_generation.py
index 548dfa4..f7d3a83 100644
--- a/test_generation.py
+++ b/test_generation.py
@@ -2,6 +2,7 @@ import argparse
 from pprint import pprint
 
 import datasets
+import pyvista as pv
 import torch
 import torch.nn as nn
 import torch.utils.data
@@ -501,6 +502,8 @@ def generate(model, opt):
     test_dataloader = torch.utils.data.DataLoader(
         test_dataset, batch_size=opt.batch_size, shuffle=False, num_workers=int(opt.workers), drop_last=False
     )
+    VTKFILE_NOMINAL = Path("~/data/stage-laurent-f/datasets/Rotor37/processed/nominal_blade_rotated.vtk")
+    nominal = pv.read(VTKFILE_NOMINAL)
 
     with torch.no_grad():
         samples = []
@@ -514,6 +517,7 @@ def generate(model, opt):
 
             gen = gen.transpose(1, 2).contiguous()
             x = x.transpose(1, 2).contiguous()
+            x = x + nominal.points
 
             # gen = gen * s + m
             # x = x * s + m
@@ -528,7 +532,7 @@ def generate(model, opt):
                 pc = pc * STD + MEAN
 
                 print(f"Saving point cloud {idx}...")
-                np.savetxt(f"gen_{i}_{idx}.txt", pc)
+                np.savetxt(f"output/gen_{i}_{idx}.txt", pc)
 
                 if idx >= 10:
                     break