fixed vtk containerbox handler

src/Vtk/HEP/Detectors/testing/vtkDetectorChamberTest.cpp

@@ -31,28 +31,31 @@
 #define BOOST_TEST_MODULE vtkDetectorChamberTest
 #include <boost/test/unit_test.hpp>
 
+using namespace uLib;
+
 BOOST_AUTO_TEST_CASE(vtkDetectorChamberTest) {
-  uLib::DetectorChamber d1, d2;
-  d1.SetSize(uLib::Vector3f(1, 1, 1));
-  d1.SetPosition(uLib::Vector3f(0, 0, 0));
-  d1.Scale(uLib::Vector3f(5, 10, 2));
-  d1.Translate(uLib::Vector3f(0, 0, 0));
+  DetectorChamber d1, d2;
+  d1.SetSize(Vector3f(1, 1, 1));
+  d1.SetPosition(Vector3f(0, 0, 0));
+  d1.Scale(Vector3f(5, 10, 2));
+  d1.Translate(Vector3f(0, 0, 0));
 
-  d2.SetSize(uLib::Vector3f(1, 1, 1));
-  d2.SetPosition(uLib::Vector3f(0, 0, 0));
-  d2.Scale(uLib::Vector3f(5, 10, 2));
-  d2.Translate(uLib::Vector3f(0, 0, 10));
+  d2.SetSize(Vector3f(1, 1, 1));
+  d2.SetPosition(Vector3f(0, 0, 0));
+  d2.Scale(Vector3f(5, 10, 2));
+  d2.Translate(Vector3f(0, 0, 10));
 
 
-  uLib::Vtk::vtkDetectorChamber vtkDetectorChamber(&d1);
-  uLib::Vtk::vtkDetectorChamber vtkDetectorChamber2(&d2);
+  Vtk::vtkDetectorChamber vtkDetectorChamber(&d1);
+  Vtk::vtkDetectorChamber vtkDetectorChamber2(&d2);
 
   if (!vtkDetectorChamber.GetProp()) {
     BOOST_FAIL("vtkDetectorChamber::GetProp() returned NULL");
   }
 
   if (std::getenv("CTEST_PROJECT_NAME") == nullptr) {
-    uLib::Vtk::Viewer viewer;
+    Vtk::Viewer viewer;
+    viewer.SetGridAxis(Vtk::Viewport::Y);
     viewer.AddPuppet(vtkDetectorChamber);
     viewer.AddPuppet(vtkDetectorChamber2);
     viewer.Start();

src/Vtk/HEP/Detectors/vtkDetectorChamber.cxx

@@ -46,15 +46,7 @@ namespace uLib {
 namespace Vtk {
 
 vtkDetectorChamber::vtkDetectorChamber(DetectorChamber *content)
-    : vtkContainerBox(content), m_Actor(vtkActor::New()),
-      m_Widget(vtkBoxWidget::New()) {
-  m_Callback = vtkWidgetCallback::New();
-  m_PickerCallback = vtkSelectionCallback::New();
-
-  m_Callback->SetChamber(this);
-  m_PickerCallback->SetChamber(this);
-  m_Widget->AddObserver(vtkCommand::InteractionEvent, m_Callback);
-
+    : vtkContainerBox(content), m_Actor(vtkActor::New()) {
   m_InitialTransform = vtkSmartPointer<vtkTransform>::New();
   m_RelativeTransform = vtkSmartPointer<vtkTransform>::New();
   m_TotalTransform = vtkSmartPointer<vtkTransform>::New();
@@ -64,9 +56,6 @@ vtkDetectorChamber::vtkDetectorChamber(DetectorChamber *content)
 
 vtkDetectorChamber::~vtkDetectorChamber() {
   m_Actor->Delete();
-  m_Widget->Delete();
-  m_Callback->Delete();
-  m_PickerCallback->Delete();
 }
 
 DetectorChamber *vtkDetectorChamber::GetContent() {
@@ -77,20 +66,8 @@ void vtkDetectorChamber::PrintSelf(std::ostream &o) const {
   vtkContainerBox::PrintSelf(o);
 }
 
-/**
- * Connect the interactor to the widget
- */
-void vtkDetectorChamber::ConnectInteractor(
-    vtkRenderWindowInteractor *interactor) {
-  if (!interactor)
-    return;
-  m_Widget->SetInteractor(interactor);
-  m_Widget->SetProp3D(m_Actor);
-  interactor->AddObserver(vtkCommand::LeftButtonPressEvent, m_PickerCallback);
-}
-
 void vtkDetectorChamber::SetTransform(vtkTransform *t) {
-  
+  if (!t) return;
   
   m_RelativeTransform->SetMatrix(t->GetMatrix());
   m_RelativeTransform->Update();
@@ -107,11 +84,22 @@ void vtkDetectorChamber::SetTransform(vtkTransform *t) {
   this->Update();
 }
 
-vtkBoxWidget *vtkDetectorChamber::GetWidget() { return m_Widget; }
-
 void vtkDetectorChamber::Update() {
   if (m_Actor->GetMapper())
     m_Actor->GetMapper()->Update();
+
+  // If the actor has a UserMatrix, we update the content.
+  if (m_Actor->GetUserMatrix()) {
+      vtkMatrix4x4* vmat = m_Actor->GetUserMatrix();
+      Matrix4f transform;
+      for (int i = 0; i < 4; ++i)
+          for (int j = 0; j < 4; ++j)
+              transform(i, j) = vmat->GetElement(i, j);
+      
+      this->GetContent()->SetMatrix(transform);
+      this->GetContent()->Updated();
+  }
+
   BaseClass::Update();
 }
 
@@ -144,56 +132,14 @@ void vtkDetectorChamber::InstallPipe() {
   m_Actor->GetProperty()->SetOpacity(0.4);
   m_Actor->GetProperty()->SetAmbient(0.7);
 
-  // Temporarily disable UserTransform to place widget on local base
-  m_Widget->SetProp3D(m_Actor);
-
   m_TotalTransform->SetInput(m_RelativeTransform);
   m_TotalTransform->Concatenate(m_InitialTransform);
   m_Actor->SetUserTransform(m_TotalTransform);
   m_TotalTransform->Update();
 
-  m_Widget->PlaceWidget();
-  m_Widget->SetPlaceFactor(2);
-
   this->SetProp(m_Actor);
   this->Update();
 }
 
-void vtkDetectorChamber::vtkWidgetCallback::Execute(vtkObject *caller,
-                                                    unsigned long, void *) {
-  vtkBoxWidget *widget = reinterpret_cast<vtkBoxWidget *>(caller);
-
-  // Get the Relative transform from the widget //
-  vtkSmartPointer<vtkTransform> t = vtkSmartPointer<vtkTransform>::New();
-  widget->GetTransform(t);
-  chamber->SetTransform(t);
-
-  // Apply to both the content and the actor state //
-  chamber->Update();
-}
-
-void vtkDetectorChamber::vtkSelectionCallback::Execute(vtkObject *caller,
-                                                       unsigned long, void *) {
-  vtkRenderWindowInteractor *interactor =
-      reinterpret_cast<vtkRenderWindowInteractor *>(caller);
-  vtkSmartPointer<vtkPropPicker> picker = vtkSmartPointer<vtkPropPicker>::New();
-  int *pos = interactor->GetEventPosition();
-  picker->Pick(
-      pos[0], pos[1], 0,
-      interactor->GetRenderWindow()->GetRenderers()->GetFirstRenderer());
-
-  vtkProp *picked = picker->GetViewProp();
-  if (picked == chamber->m_Actor) {
-    if (!chamber->m_Widget->GetEnabled()) {
-      chamber->m_Widget->SetInteractor(interactor);
-      chamber->m_Widget->On();
-    }
-  } else {
-    if (chamber->m_Widget->GetEnabled()) {
-      chamber->m_Widget->Off();
-    }
-  }
-}
-
 } // namespace Vtk
 } // namespace uLib

src/Vtk/HEP/Detectors/vtkDetectorChamber.h

@@ -56,11 +56,7 @@ public:
 
   void SetTransform(class vtkTransform *t);
 
-  class vtkBoxWidget *GetWidget();
-
-  void Update();
-
-  void ConnectInteractor(vtkRenderWindowInteractor *interactor) override;
+  void Update() override;
 
   void PrintSelf(std::ostream &o) const;
 
@@ -68,30 +64,7 @@ protected:
   void InstallPipe() override;
 
 private:
-  class vtkWidgetCallback : public vtkCommand {
-  public:
-    static vtkWidgetCallback *New() { return new vtkWidgetCallback; }
-    void SetChamber(uLib::Vtk::vtkDetectorChamber *ch) { this->chamber = ch; }
-    virtual void Execute(vtkObject *caller, unsigned long, void *) override;
-
-  private:
-    uLib::Vtk::vtkDetectorChamber *chamber;
-  };
-
-  class vtkSelectionCallback : public vtkCommand {
-  public:
-    static vtkSelectionCallback *New() { return new vtkSelectionCallback; }
-    void SetChamber(uLib::Vtk::vtkDetectorChamber *ch) { this->chamber = ch; }
-    virtual void Execute(vtkObject *caller, unsigned long, void *) override;
-
-  private:
-    uLib::Vtk::vtkDetectorChamber *chamber;
-  };
-
   vtkActor *m_Actor;
-  vtkBoxWidget *m_Widget;
-  vtkWidgetCallback *m_Callback;
-  vtkSelectionCallback *m_PickerCallback;
 
   vtkSmartPointer<vtkTransform> m_InitialTransform;
   vtkSmartPointer<vtkTransform> m_RelativeTransform;

src/Vtk/HEP/MuonTomography/testing/vtkVoxRaytracerTest.cpp

@@ -110,21 +110,13 @@ BOOST_AUTO_TEST_CASE(vtkVoxRaytracerRepresentationTest) {
     // renderer //
     Vtk::Viewer viewer;
 
-    // widget //
-    vtkBoxWidget *widget = v_grid.GetWidget();
-
-    vtkWidgetCallback *cbk = vtkWidgetCallback::New();
-    cbk->SetTracer(&v_rt);
-    cbk->SetMuon(&muon);
-    cbk->SetAnnotation(viewer.GetAnnotation());
-    widget->AddObserver(vtkCommand::InteractionEvent, cbk);
-    widget->SetInteractor(viewer.GetInteractor());
-    widget->PlaceWidget();
-    widget->On();
-
     viewer.AddPuppet(v_grid);
     viewer.AddPuppet(v_rt);
     viewer.AddPuppet(v_muon);
+    
+    // Select grid to show handler widget
+    viewer.SelectPuppet(&v_grid);
+    
     viewer.Start();
   }
 

src/Vtk/HEP/MuonTomography/vtkVoxRaytracerRepresentation.cpp

@@ -42,35 +42,52 @@ namespace Vtk {
 ////////////////////////////////////////////////////////////////////////////////
 
 vtkVoxRaytracerRepresentation::vtkVoxRaytracerRepresentation(Content &content)
-    : m_Content(&content), m_Assembly(vtkAssembly::New()),
+    : m_Content(&content),
       m_Sphere1(vtkSphereSource::New()), m_Sphere2(vtkSphereSource::New()),
       m_Line1(vtkLineSource::New()), m_Line2(vtkLineSource::New()),
       m_Line3(vtkLineSource::New()), m_RayLine(vtkAppendPolyData::New()),
       m_RayLineActor(vtkActor::New()),
       m_RayRepresentation(vtkAppendPolyData::New()),
       m_RayRepresentationActor(vtkActor::New()),
-      m_Transform(vtkTransform::New()) {
+      m_Transform(vtkTransform::New()),
+      m_HasMuon(false), m_HasPoca(false) {
   default_radius = content.GetImage()->GetSpacing()(0) / 4;
   m_Sphere1->SetRadius(default_radius);
   m_Sphere2->SetRadius(default_radius);
   m_SelectedElement = m_RayLine;
 
+  this->SetSelectable(false);
   InstallPipe();
+  
+  if (m_Content && m_Content->GetImage()) {
+      Object::connect(m_Content->GetImage(), &StructuredGrid::Updated, this, &vtkVoxRaytracerRepresentation::imageUpdate);
+  }
 }
 
 vtkVoxRaytracerRepresentation::~vtkVoxRaytracerRepresentation() {
-  m_Assembly->Delete();
   m_RayLine->Delete();
   m_RayLineActor->Delete();
   m_RayRepresentationActor->Delete();
   m_Transform->Delete();
 }
 
-VoxRaytracer *vtkVoxRaytracerRepresentation::GetRaytracerAlgorithm() {
+uLib::VoxRaytracer *vtkVoxRaytracerRepresentation::GetRaytracerAlgorithm() {
   return m_Content;
 }
 
-vtkProp *vtkVoxRaytracerRepresentation::GetProp() { return m_Assembly; }
+void vtkVoxRaytracerRepresentation::Update() {
+   this->imageUpdate();
+}
+
+void vtkVoxRaytracerRepresentation::imageUpdate() {
+   if (m_HasMuon) {
+       if (m_HasPoca) {
+           this->SetMuon(m_Muon, m_Poca);
+       } else {
+           this->SetMuon(m_Muon);
+       }
+   }
+}
 
 vtkPolyData *vtkVoxRaytracerRepresentation::GetPolyData() const {
   std::cout << "get Raytracer polydata\n";
@@ -94,6 +111,10 @@ void vtkVoxRaytracerRepresentation::SetRepresentationElements(
 }
 
 void vtkVoxRaytracerRepresentation::SetMuon(MuonScatter &muon) {
+  m_Muon = muon;
+  m_HasMuon = true;
+  m_HasPoca = false;
+
   HPoint3f pt1, pt2, src;
   src = muon.LineIn().origin;
   m_Content->GetEntryPoint(muon.LineIn(), pt1);
@@ -152,6 +173,11 @@ void vtkVoxRaytracerRepresentation::SetMuon(MuonScatter &muon) {
 }
 
 void vtkVoxRaytracerRepresentation::SetMuon(MuonScatter &muon, HPoint3f poca) {
+  m_Muon = muon;
+  m_Poca = poca;
+  m_HasMuon = true;
+  m_HasPoca = true;
+
   HPoint3f pt1, pt2, src;
   src = muon.LineIn().origin;
   m_Content->GetEntryPoint(muon.LineIn(), pt1);

src/Vtk/HEP/MuonTomography/vtkVoxRaytracerRepresentation.h

@@ -58,7 +58,7 @@ class vtkActor;
 namespace uLib {
 namespace Vtk {
 
-class vtkVoxRaytracerRepresentation : public Puppet {
+class vtkVoxRaytracerRepresentation : public Puppet, public Object {
   typedef VoxRaytracer Content;
 
 public:
@@ -67,7 +67,9 @@ public:
 
   uLib::VoxRaytracer *GetRaytracerAlgorithm();
 
-  vtkProp *GetProp();
+  virtual void Update() override;
+
+  void imageUpdate();
 
   vtkPolyData *GetPolyData() const;
 
@@ -99,9 +101,12 @@ private:
   void SetColor(vtkActor *actor, Vector4f rgba);
 
   VoxRaytracer *m_Content;
+  MuonScatter m_Muon;
+  HPoint3f m_Poca;
+  bool m_HasMuon;
+  bool m_HasPoca;
 
   Scalarf default_radius;
-  vtkAssembly *m_Assembly;
   vtkAppendPolyData *m_RayLine;
   vtkActor *m_RayLineActor;
   vtkActor *m_RayRepresentationActor;

src/Vtk/Math/testing/vtkStructuredGridTest.cpp

@@ -44,5 +44,5 @@ BOOST_AUTO_TEST_CASE(vtkStructuredGridTest) {
     viewer.Start();
   }
 
-  BOOST_CHECK(grid_viewer.GetWidget() != nullptr);
+  BOOST_CHECK(grid_viewer.GetProp() != nullptr);
 }

src/Vtk/Math/vtkDense.h

@@ -0,0 +1,54 @@
+#ifndef U_VTK_DENSE_H
+#define U_VTK_DENSE_H
+
+#include "Math/Dense.h"
+#include <vtkMatrix4x4.h>
+
+namespace uLib {
+namespace Vtk {
+
+/**
+ * @brief Converts a uLib::Matrix4f to an existing vtkMatrix4x4.
+ * @param src The source Eigen matrix.
+ * @param dst The destination vtkMatrix4x4.
+ */
+inline void Matrix4fToVtk(const Matrix4f& src, vtkMatrix4x4* dst) {
+  if (!dst)
+    return;
+  for (int i = 0; i < 4; ++i)
+    for (int j = 0; j < 4; ++j)
+      dst->SetElement(i, j, src(i, j));
+}
+
+/**
+ * @brief Converts a uLib::Matrix4f to a new vtkMatrix4x4.
+ * @param src The source Eigen matrix.
+ * @return A new vtkMatrix4x4 (caller is responsible for management/deletion).
+ */
+inline vtkMatrix4x4* Matrix4fToVtk(const Matrix4f& src) {
+  vtkMatrix4x4* dst = vtkMatrix4x4::New();
+  Matrix4fToVtk(src, dst);
+  return dst;
+}
+
+/**
+ * @brief Converts a vtkMatrix4x4 to a uLib::Matrix4f.
+ * @param src The source vtkMatrix4x4.
+ * @return The converted uLib::Matrix4f matrix.
+ */
+inline Matrix4f VtkToMatrix4f(vtkMatrix4x4* src) {
+  Matrix4f dst;
+  if (!src) {
+    dst.setIdentity();
+    return dst;
+  }
+  for (int i = 0; i < 4; ++i)
+    for (int j = 0; j < 4; ++j)
+      dst(i, j) = src->GetElement(i, j);
+  return dst;
+}
+
+} // namespace Vtk
+} // namespace uLib
+
+#endif // U_VTK_DENSE_H

src/Vtk/Math/vtkStructuredGrid.cpp

@@ -30,6 +30,8 @@
 #include "Math/StructuredGrid.h"
 #include "Vtk/Math/vtkStructuredGrid.h"
 
+#include "Vtk/Math/vtkDense.h"
+
 namespace uLib {
 namespace Vtk {
 
@@ -39,73 +41,64 @@ namespace Vtk {
 
 vtkStructuredGrid::vtkStructuredGrid(Content &content)
     : m_Content(&content), m_Actor(vtkActor::New()),
-      m_Widget(vtkBoxWidget::New()), m_Transform(vtkTransform::New()) {
-  vtkSmartPointer<vtkWidgetCallback> callback =
-      vtkSmartPointer<vtkWidgetCallback>::New();
-  callback->SetGrid(this);
-  m_Widget->AddObserver(vtkCommand::InteractionEvent, callback);
+      m_Transform(vtkTransform::New()) {
 
   this->InstallPipe();
 }
 
 vtkStructuredGrid::~vtkStructuredGrid() {
   m_Actor->Delete();
-  m_Widget->Delete();
   m_Transform->Delete();
 }
 
 void vtkStructuredGrid::SetTransform(vtkTransform *t) {
   vtkMatrix4x4 *vmat = t->GetMatrix();
-  Matrix4f mat;
-  for (int i = 0; i < 4; ++i)
-    for (int j = 0; j < 4; ++j)
-      mat(i, j) = vmat->GetElement(i, j);
+  Matrix4f mat = VtkToMatrix4f(vmat);
   m_Content->SetMatrix(mat);
 
   vtkSmartPointer<vtkMatrix4x4> vmat2 = vtkSmartPointer<vtkMatrix4x4>::New();
   mat = m_Content->GetWorldMatrix();
-  for (int i = 0; i < 4; ++i)
-    for (int j = 0; j < 4; ++j)
-      vmat2->SetElement(i, j, mat(i, j));
+  Matrix4fToVtk(mat, vmat2);
   m_Transform->SetMatrix(vmat2);
   m_Transform->Update();
   this->Update();
 }
 
-vtkBoxWidget *vtkStructuredGrid::GetWidget() { return m_Widget; }
+void vtkStructuredGrid::Update() {
+  if (!m_Content) return;
 
-void vtkStructuredGrid::Update() { m_Actor->GetMapper()->Update(); }
+  vtkProp3D* actor = vtkProp3D::SafeDownCast(this->GetProp());
+  if (!actor) return;
+
+  vtkMatrix4x4* vmat = actor->GetUserMatrix();
+  if (!vmat) return;
+  
+  Matrix4f transform = VtkToMatrix4f(vmat);
+  m_Content->SetMatrix(transform);
+
+  m_Content->Updated(); // Notify others (like raytracer)
+}
 
 void vtkStructuredGrid::InstallPipe() {
   vtkSmartPointer<vtkCubeSource> cube = vtkSmartPointer<vtkCubeSource>::New();
-  vtkSmartPointer<vtkTransformPolyDataFilter> filter =
-      vtkSmartPointer<vtkTransformPolyDataFilter>::New();
-
-  vtkSmartPointer<vtkMatrix4x4> vmat = vtkSmartPointer<vtkMatrix4x4>::New();
-  Matrix4f mat = m_Content->GetWorldMatrix();
-  for (int i = 0; i < 4; ++i)
-    for (int j = 0; j < 4; ++j)
-      vmat->SetElement(i, j, mat(i, j));
-  m_Transform->SetMatrix(vmat);
-  filter->SetTransform(m_Transform);
-  filter->SetInputConnection(cube->GetOutputPort());
-
+ 
   Vector3i dims = m_Content->GetDims();
   cube->SetBounds(0, dims(0), 0, dims(1), 0, dims(2));
   cube->Update();
-  filter->Update();
 
   vtkSmartPointer<vtkPolyDataMapper> mapper =
       vtkSmartPointer<vtkPolyDataMapper>::New();
-  mapper->SetInputConnection(filter->GetOutputPort());
+  mapper->SetInputConnection(cube->GetOutputPort());
 
   m_Actor->SetMapper(mapper);
   m_Actor->GetProperty()->SetRepresentationToSurface();
   m_Actor->GetProperty()->SetEdgeVisibility(true);
   m_Actor->GetProperty()->SetOpacity(0.4);
   m_Actor->GetProperty()->SetAmbient(0.7);
-  this->Update();
-  m_Widget->SetProp3D(m_Actor);
+
+  vtkNew<vtkMatrix4x4> vmat;
+  Matrix4fToVtk(m_Content->GetWorldMatrix(), vmat);
+  m_Actor->SetUserMatrix(vmat);
 
   this->SetProp(m_Actor);
 }

src/Vtk/Math/vtkStructuredGrid.h

@@ -56,32 +56,12 @@ public:
 
   void SetTransform(class vtkTransform *t);
 
-  class vtkBoxWidget *GetWidget();
-
-  void Update();
+  virtual void Update() override;
 
 private:
   void InstallPipe();
 
-  class vtkWidgetCallback : public vtkCommand {
-  public:
-    static vtkWidgetCallback *New() { return new vtkWidgetCallback; }
-
-    void SetGrid(uLib::Vtk::vtkStructuredGrid *grid) { this->grid = grid; }
-
-    virtual void Execute(vtkObject *caller, unsigned long, void *) {
-      vtkSmartPointer<vtkTransform> t = vtkSmartPointer<vtkTransform>::New();
-      vtkBoxWidget *widget = reinterpret_cast<vtkBoxWidget *>(caller);
-      widget->GetTransform(t);
-      grid->SetTransform(t);
-    }
-
-  private:
-    uLib::Vtk::vtkStructuredGrid *grid;
-  };
-
   vtkActor *m_Actor;
-  vtkBoxWidget *m_Widget;
   StructuredGrid *m_Content;
   vtkTransform *m_Transform;
 };

src/Vtk/testing/vtkContainerBoxTest.cpp

@@ -36,9 +36,16 @@ int main() {
   BEGIN_TESTING(vtk ContainerBox Test);
 
   ContainerBox box;
-  box.SetSize(Vector3f(2, 3, 4));
-  box.SetPosition(Vector3f(1, 2, 3));
+  box.Scale(Vector3f(1,5,1));
+  box.SetPosition(Vector3f(0,1,0));
   Vtk::vtkContainerBox v_box(&box);
+  v_box.SetRepresentation(Vtk::Puppet::Surface);
+  v_box.SetOpacity(0.5);
+  v_box.SetSelectable(true);
+  
+  box.findOrAddSignal(&ContainerBox::Updated)->connect([&box](){
+    std::cout << "box updated: " << box.GetWorldPoint(HPoint3f(1,1,1)) << std::endl;
+  });
 
   if (std::getenv("CTEST_PROJECT_NAME") == nullptr) {
     Vtk::Viewer v_viewer;

src/Vtk/uLibVtkInterface.cxx

@@ -45,8 +45,9 @@
 
 #include <vtkActor.h>
 #include <vtkPropCollection.h>
+#include <vtkProp3DCollection.h>
 #include <vtkRendererCollection.h>
-#include <vtkPropAssembly.h>
+#include <vtkAssembly.h>
 #include <vtkOutlineSource.h>
 #include <vtkPolyDataMapper.h>
 #include <vtkCubeAxesActor.h>
@@ -82,7 +83,7 @@ class PuppetData {
 public:
     PuppetData() :
         m_Renderers(vtkSmartPointer<vtkRendererCollection>::New()),
-        m_Assembly(vtkSmartPointer<vtkPropAssembly>::New()),
+        m_Assembly(vtkSmartPointer<vtkAssembly>::New()),
         m_ShowBoundingBox(false),
         m_ShowScaleMeasures(false),
         m_Representation(-1),
@@ -99,7 +100,7 @@ public:
 
     // members //
     vtkSmartPointer<vtkRendererCollection> m_Renderers;
-    vtkSmartPointer<vtkPropAssembly>       m_Assembly;
+    vtkSmartPointer<vtkAssembly>       m_Assembly;
 
     vtkSmartPointer<vtkOutlineSource>      m_OutlineSource;
     vtkSmartPointer<vtkActor>              m_OutlineActor;
@@ -160,19 +161,16 @@ public:
                 m_HighlightActor->GetProperty()->SetLighting(0);
             }
 
-            // Update highlight data and transform from first actor
-            vtkPropCollection *parts = m_Assembly->GetParts();
-            parts->InitTraversal();
-            for (int i = 0; i < parts->GetNumberOfItems(); ++i) {
-                vtkActor *actor = vtkActor::SafeDownCast(parts->GetNextProp());
-                if (actor) {
-                    // Sync transform
-                    m_HighlightActor->SetUserTransform(actor->GetUserTransform());
-                    m_HighlightActor->SetPosition(actor->GetPosition());
-                    m_HighlightActor->SetOrientation(actor->GetOrientation());
-                    m_HighlightActor->SetScale(actor->GetScale());
-                    break;
-                }
+            // Update highlight matrix from the root prop
+            vtkProp3D* root = nullptr;
+            if (m_Assembly->GetParts()->GetNumberOfItems() == 1) {
+                root = vtkProp3D::SafeDownCast(m_Assembly->GetParts()->GetLastProp());
+            } else {
+                root = m_Assembly;
+            }
+
+            if (root) {
+                m_HighlightActor->SetUserMatrix(root->GetMatrix());
             }
 
             m_Renderers->InitTraversal();
@@ -209,18 +207,17 @@ Puppet::~Puppet()
 vtkProp *Puppet::GetProp()
 {
     if (d->m_Assembly->GetParts()->GetNumberOfItems() == 1)
-    {
-        //        d->m_Assembly->GetParts()->InitTraversal();
-        //        return d->m_Assembly->GetParts()->GetNextProp();
         return d->m_Assembly->GetParts()->GetLastProp();
-    }
-    else return d->m_Assembly;
+    else 
+        return d->m_Assembly;
 }
 
 void Puppet::SetProp(vtkProp *prop)
 {
     if(prop) {
-        d->m_Assembly->AddPart(prop);
+        if (auto* p3d = vtkProp3D::SafeDownCast(prop)) {
+            d->m_Assembly->AddPart(p3d);
+        }
         d->ApplyAppearance(prop);
     }
 }
@@ -240,14 +237,8 @@ void Puppet::ConnectRenderer(vtkRenderer *renderer)
 {
     if(renderer) {
         this->GetRenderers()->AddItem(renderer);
-        if(d->m_Assembly->GetParts()->GetNumberOfItems() == 1)
-            renderer->AddActor(this->GetProp());
-        else if(d->m_Assembly->GetParts()->GetNumberOfItems() >0)
-        {
-            vtkPropCollection *props = d->m_Assembly->GetParts();
-            props->InitTraversal();
-            for (int i=0; i<props->GetNumberOfItems(); ++i)
-                renderer->AddActor(props->GetNextProp());
+        if(vtkProp* prop = this->GetProp()) {
+            renderer->AddViewProp(prop);
         }
 
         if (d->m_ShowBoundingBox && d->m_OutlineActor) renderer->AddActor(d->m_OutlineActor);
@@ -265,15 +256,8 @@ void Puppet::ConnectRenderer(vtkRenderer *renderer)
 void Puppet::DisconnectRenderer(vtkRenderer *renderer)
 {
     if(renderer) {
-        if(this->GetProp())
-            renderer->RemoveViewProp(this->GetProp());
-        else if(d->m_Assembly->GetParts()->GetNumberOfItems() >0)
-        {
-            vtkPropCollection *props = d->m_Assembly->GetParts();
-            props->InitTraversal();
-            for (int i=0; i<props->GetNumberOfItems(); ++i)
-                renderer->RemoveViewProp(props->GetNextProp());
-        }
+        if(vtkProp* prop = this->GetProp())
+            renderer->RemoveViewProp(prop);
 
         if (d->m_ShowBoundingBox && d->m_OutlineActor) renderer->RemoveActor(d->m_OutlineActor);
         if (d->m_ShowScaleMeasures && d->m_CubeAxesActor) renderer->RemoveActor(d->m_CubeAxesActor);
@@ -382,10 +366,10 @@ void Puppet::SetRepresentation(Representation mode)
     }
     d->m_Representation = rep;
 
-    vtkPropCollection *props = d->m_Assembly->GetParts();
+    vtkProp3DCollection *props = d->m_Assembly->GetParts();
     props->InitTraversal();
     for (int i = 0; i < props->GetNumberOfItems(); ++i) {
-        d->ApplyAppearance(props->GetNextProp());
+        d->ApplyAppearance(props->GetNextProp3D());
     }
 }
 
@@ -403,10 +387,10 @@ void Puppet::SetColor(double r, double g, double b)
     d->m_Color[1] = g;
     d->m_Color[2] = b;
 
-    vtkPropCollection *props = d->m_Assembly->GetParts();
+    vtkProp3DCollection *props = d->m_Assembly->GetParts();
     props->InitTraversal();
     for (int i = 0; i < props->GetNumberOfItems(); ++i) {
-        d->ApplyAppearance(props->GetNextProp());
+        d->ApplyAppearance(props->GetNextProp3D());
     }
 }
 
@@ -414,10 +398,10 @@ void Puppet::SetOpacity(double alpha)
 {
     d->m_Opacity = alpha;
 
-    vtkPropCollection *props = d->m_Assembly->GetParts();
+    vtkProp3DCollection *props = d->m_Assembly->GetParts();
     props->InitTraversal();
     for (int i = 0; i < props->GetNumberOfItems(); ++i) {
-        d->ApplyAppearance(props->GetNextProp());
+        d->ApplyAppearance(props->GetNextProp3D());
     }
 }
 
@@ -450,6 +434,24 @@ bool Puppet::IsSelected() const
     return d->m_Selected;
 }
 
+void Puppet::Update()
+{
+    if (d->m_Selected) {
+        d->UpdateHighlight();
+    }
+    
+    if (d->m_ShowBoundingBox) {
+        double* bounds = d->m_Assembly->GetBounds();
+        d->m_OutlineSource->SetBounds(bounds);
+        d->m_OutlineSource->Update();
+    }
+    
+    if (d->m_ShowScaleMeasures) {
+        double* bounds = d->m_Assembly->GetBounds();
+        d->m_CubeAxesActor->SetBounds(bounds);
+    }
+}
+
 void Puppet::ConnectInteractor(vtkRenderWindowInteractor *interactor)
 {
 }

src/Vtk/uLibVtkInterface.h

@@ -86,6 +86,8 @@ public:
   void SetSelected(bool selected = true);
   bool IsSelected() const;
 
+  virtual void Update();
+
   virtual void ConnectInteractor(class vtkRenderWindowInteractor *interactor);
 
 protected:

src/Vtk/uLibVtkViewer.cpp

@@ -38,6 +38,13 @@
 #include <vtkObjectFactory.h>
 #include <vtkTextProperty.h>
 #include <vtkRenderWindow.h>
+#include <vtkButtonWidget.h>
+#include <vtkTexturedButtonRepresentation2D.h>
+#include <vtkImageData.h>
+#include <vtkImageCanvasSource2D.h>
+#include <vtkCallbackCommand.h>
+#include <vtkCoordinate.h>
+#include <vtkRenderer.h>
 
 #include "uLibVtkViewer.h"
 
@@ -82,6 +89,9 @@ void Viewer::InstallPipe() {
 
   // Common setup
   Viewport::SetupPipeline(renderWindowInteractor);
+  
+  // Setup native grid button
+  SetupGridButton();
 
   // BUT we want to override the style with our custom NoSpin version
   vtkSmartPointer<vtkInteractorStyleNoSpin> style =
@@ -114,6 +124,89 @@ Viewer::MakeCameraOrientationWidget(vtkRenderWindowInteractor *interactor,
   return widget;
 }
 
+void Viewer::SetupGridButton() {
+  if (!m_RenderWindow || !m_RenderWindow->GetInteractor()) return;
+
+  // Create procedural textures for the button using canvas
+  vtkNew<vtkImageCanvasSource2D> canvas;
+  canvas->SetScalarTypeToUnsignedChar();
+  canvas->SetNumberOfScalarComponents(4);
+  canvas->SetExtent(0, 63, 0, 63, 0, 0);
+
+  // State 0: OFF (Gray circle, transparent background)
+  canvas->SetDrawColor(0, 0, 0, 0);
+  canvas->FillBox(0, 63, 0, 63);
+  canvas->SetDrawColor(120, 120, 120, 255);
+  canvas->DrawCircle(32, 32, 25);
+  canvas->Update();
+
+  vtkNew<vtkImageData> imgOff;
+  imgOff->DeepCopy(canvas->GetOutput());
+
+  // State 1: ON (White circle, transparent background)
+  canvas->SetDrawColor(0, 0, 0, 0);
+  canvas->FillBox(0, 63, 0, 63);
+  canvas->SetDrawColor(255, 255, 255, 255);
+  canvas->DrawCircle(32, 32, 25);
+  canvas->Update();
+
+  vtkNew<vtkImageData> imgOn;
+  imgOn->DeepCopy(canvas->GetOutput());
+
+  vtkNew<vtkTexturedButtonRepresentation2D> rep;
+  rep->SetNumberOfStates(2);
+  rep->SetButtonTexture(0, imgOff);
+  rep->SetButtonTexture(1, imgOn);
+
+  m_GridButton = vtkSmartPointer<vtkButtonWidget>::New();
+  m_GridButton->SetInteractor(m_RenderWindow->GetInteractor());
+  m_GridButton->SetRepresentation(rep);
+
+  // Position it initially
+  UpdateGridButtonPosition();
+
+  // Callback for resize (ModifiedEvent on RenderWindow)
+  vtkNew<vtkCallbackCommand> resizeCallback;
+  resizeCallback->SetClientData(this);
+  resizeCallback->SetCallback([](vtkObject*, unsigned long, void* clientdata, void*){
+      auto* v = static_cast<Viewer*>(clientdata);
+      v->UpdateGridButtonPosition();
+  });
+  m_RenderWindow->AddObserver(vtkCommand::ModifiedEvent, resizeCallback);
+
+  // Callback for state change
+  vtkNew<vtkCallbackCommand> stateCallback;
+  stateCallback->SetClientData(this);
+  stateCallback->SetCallback([](vtkObject* caller, unsigned long, void* clientdata, void*){
+      auto* btn = vtkButtonWidget::SafeDownCast(caller);
+      auto* v = static_cast<Viewer*>(clientdata);
+      auto* r = vtkTexturedButtonRepresentation2D::SafeDownCast(btn->GetRepresentation());
+      v->SetGridVisible(r->GetState() == 1);
+  });
+  
+  m_GridButton->AddObserver(vtkCommand::StateChangedEvent, stateCallback);
+  m_GridButton->On();
+  
+  // Set initial state
+  rep->SetState(GetGridVisible() ? 1 : 0);
+}
+
+void Viewer::UpdateGridButtonPosition() {
+  if (!m_GridButton || !m_RenderWindow) return;
+  auto* rep = vtkTexturedButtonRepresentation2D::SafeDownCast(m_GridButton->GetRepresentation());
+  if (!rep) return;
+
+  int *sz = m_RenderWindow->GetSize();
+  if (sz[0] == 0 || sz[1] == 0) return; // Window not yet sized or hidden
+
+  int margin_rigth = 23;
+  int margin_top = 170;
+  int btnSz = 100; // Button size in display coordinates
+  double bds[6] = { (double)sz[0] - btnSz - margin_rigth, (double)sz[0] - margin_rigth,
+                    (double)sz[1] - margin_top  - btnSz/2.0, (double)sz[1] - margin_top + btnSz/2.0, 0, 0 };
+  rep->PlaceWidget(bds);
+}
+
 void Viewer::Start() { m_RenderWindow->GetInteractor()->Start(); }
 
 vtkRenderWindow *Viewer::GetRenderWindow() { return m_RenderWindow; }

src/Vtk/uLibVtkViewer.h

@@ -62,7 +62,11 @@ private:
   void InstallPipe();
   void UninstallPipe();
 
+  void SetupGridButton();
+  void UpdateGridButtonPosition();
+
   vtkRenderWindow *m_RenderWindow;
+  vtkSmartPointer<class vtkButtonWidget> m_GridButton;
 };
 
 // template <> class Tie<Viewer> {

src/Vtk/vtkContainerBox.cpp

@@ -41,21 +41,24 @@
 #include <vtkSmartPointer.h>
 #include <vtkTransform.h>
 
+#include "Math/vtkDense.h"
+
 
 namespace uLib {
 namespace Vtk {
 
 vtkContainerBox::vtkContainerBox(vtkContainerBox::Content *content)
     : m_Cube(vtkActor::New()), m_Axes(vtkActor::New()),
-      m_Pivot(vtkActor::New()), 
+      // m_Pivot(vtkActor::New()), 
       m_Content(content) {
   this->InstallPipe();
+  Object::connect(m_Content, &Content::Updated, this, &vtkContainerBox::contentUpdate);
 }
 
 vtkContainerBox::~vtkContainerBox() {
   m_Cube->Delete();
   m_Axes->Delete();
-  m_Pivot->Delete();
+//   m_Pivot->Delete();
 }
 
 vtkPolyData *vtkContainerBox::GetPolyData() const {
@@ -63,21 +66,52 @@ vtkPolyData *vtkContainerBox::GetPolyData() const {
   return NULL;
 }
 
-void vtkContainerBox::Update() {
+
+void vtkContainerBox::contentUpdate() {
   if (!m_Content)
     return;
-  vtkSmartPointer<vtkMatrix4x4> vmat = vtkSmartPointer<vtkMatrix4x4>::New();
+
+  vtkProp3D* root = vtkProp3D::SafeDownCast(this->GetProp());
+  if (!root) return;
+
+  vtkMatrix4x4* vmat = root->GetUserMatrix();
+  if (!vmat) {
+    // Should have been set in InstallPipe, but let's be safe
+    vtkNew<vtkMatrix4x4> mat;
+    root->SetUserMatrix(mat);
+    vmat = mat;
+  }
+
   Matrix4f transform = m_Content->GetMatrix();
   for (int i = 0; i < 4; ++i)
-    for (int j = 0; j < 4; ++j)
+    for (int j = 0; j < 4; ++j) {
       vmat->SetElement(i, j, transform(i, j));
+    }
 
-  // std::cout << "transform: " << transform << std::endl;
-  // m_RelativeTransform->SetMatrix(vmat);
-  // m_RelativeTransform->Update();
+  root->Modified();
+  Puppet::Update();
+}
 
-  m_Cube->SetUserMatrix(vmat);
-  m_Axes->SetUserMatrix(vmat);
+
+void vtkContainerBox::Update() {
+  if (!m_Content) return;
+
+  vtkProp3D* assembly = vtkProp3D::SafeDownCast(this->GetProp());
+  if (!assembly) return;
+
+  vtkMatrix4x4* vmat = assembly->GetUserMatrix();
+  if (!vmat) return;
+  Matrix4f transform = VtkToMatrix4f(vmat);
+
+  // Update uLib model's affine transform
+  if (m_Content->GetParent()) {
+    Matrix4f localT = m_Content->GetParent()->GetWorldMatrix().inverse() * transform;
+    m_Content->SetMatrix(localT);
+  } else {
+    m_Content->SetMatrix(transform);
+  }
+
+  m_Content->Updated(); // Notify change
 }
 
 
@@ -91,24 +125,7 @@ void vtkContainerBox::InstallPipe() {
 
   // CUBE
   vtkSmartPointer<vtkCubeSource> cube = vtkSmartPointer<vtkCubeSource>::New();
-
-  Vector3f p = c->GetPosition();
-  // cube->SetCenter(p(0), p(1), p(2));
-  // Vector4f p1 = c->GetWorldPoint(HPoint3f(0, 0, 0));
-  // Vector4f p2 = c->GetWorldPoint(HPoint3f(1, 1, 1));
-  // vtkSmartPointer<vtkLineSource> line =
-  // vtkSmartPointer<vtkLineSource>::New(); line->SetPoint1(p1(0), p1(1),
-  // p1(2)); line->SetPoint2(p2(0), p2(1), p2(2)); line->Update();
-  // cube->SetBounds(line->GetOutput()->GetBounds());
-
-  vtkSmartPointer<vtkMatrix4x4> vmat = vtkSmartPointer<vtkMatrix4x4>::New();
-  Matrix4f transform = c->GetMatrix();
-  for (int i = 0; i < 4; ++i)
-    for (int j = 0; j < 4; ++j)
-      vmat->SetElement(i, j, transform(i, j));
-
-  m_Cube->SetUserMatrix(vmat);
-
+  
   vtkSmartPointer<vtkPolyDataMapper> mapper =
       vtkSmartPointer<vtkPolyDataMapper>::New();
   cube->SetBounds(0, 1, 0, 1, 0, 1);
@@ -125,9 +142,6 @@ void vtkContainerBox::InstallPipe() {
   mapper->SetInputConnection(axes->GetOutputPort());
   mapper->Update();
   m_Axes->SetMapper(mapper);
-  m_Axes->SetUserMatrix(vmat);
-  Vector3f s = c->GetSize();
-  //    m_Axes->SetScale(s(0),s(1),s(2));
   m_Axes->GetProperty()->SetLineWidth(3);
   m_Axes->GetProperty()->SetAmbient(0.4);
   m_Axes->GetProperty()->SetSpecular(0);
@@ -138,24 +152,16 @@ void vtkContainerBox::InstallPipe() {
   mapper = vtkSmartPointer<vtkPolyDataMapper>::New();
   mapper->SetInputConnection(axes->GetOutputPort());
   mapper->Update();
-  m_Pivot->SetUserMatrix(vmat);
-  m_Pivot->SetMapper(mapper);
-
-  Matrix4f pivotTransform = c->AffineTransform::GetWorldMatrix();
-  vtkSmartPointer<vtkMatrix4x4> pmat = vtkSmartPointer<vtkMatrix4x4>::New();
-  for (int i = 0; i < 4; ++i)
-    for (int j = 0; j < 4; ++j)
-      pmat->SetElement(i, j, pivotTransform(i, j));
-  m_Pivot->SetUserMatrix(pmat);
-  s = c->GetScale();
-  //    m_Pivot->SetScale(s(0),s(1),s(2));
-  m_Pivot->GetProperty()->SetLineWidth(3);
-  m_Pivot->GetProperty()->SetAmbient(0.4);
-  m_Pivot->GetProperty()->SetSpecular(0);
-
+  
   this->SetProp(m_Cube);
   this->SetProp(m_Axes);
-  this->SetProp(m_Pivot);
+  
+  vtkProp3D* root = vtkProp3D::SafeDownCast(this->GetProp());
+  if (root) {
+      vtkNew<vtkMatrix4x4> vmat;
+      Matrix4fToVtk(c->GetMatrix(), vmat);
+      root->SetUserMatrix(vmat);
+  }
 }
 
 } // namespace Vtk

src/Vtk/vtkContainerBox.h

@@ -43,15 +43,16 @@ public:
 
   virtual class vtkPolyData *GetPolyData() const;
 
-  virtual void Update();
+  virtual void contentUpdate();
 
+  virtual void Update();
 
 protected:
   virtual void InstallPipe();
 
   vtkActor *m_Cube;
   vtkActor *m_Axes;
-  vtkActor *m_Pivot;
+  // vtkActor *m_Pivot;
 
   Content *m_Content;
 };

src/Vtk/vtkHandlerWidget.cpp

@@ -545,12 +545,37 @@ void vtkHandlerWidget::UpdateGizmoPosition() {
   if (!this->Prop3D)
     return;
 
+  // Calculate scaling factor: min(object_bbox, 1/5 of viewport)
+  double bboxSize = 1.0;
+  double bounds[6];
+  this->Prop3D->GetBounds(bounds);
+  if (vtkMath::AreBoundsInitialized(bounds)) {
+    bboxSize = std::max({bounds[1] - bounds[0], bounds[3] - bounds[2], bounds[5] - bounds[4]});
+  }
+  if (bboxSize < 1e-6) bboxSize = 1.0;
+
+  double screenLimit = bboxSize * 2.0; // Default if no renderer
+  if (this->CurrentRenderer) {
+    int *sz = this->CurrentRenderer->GetSize();
+    if (sz[1] > 0) {
+      double pixelSize = std::min(sz[0], sz[1]) / 5.0;
+      vtkCamera *cam = this->CurrentRenderer->GetActiveCamera();
+      if (cam->GetParallelProjection()) {
+        screenLimit = (pixelSize / (double)sz[1]) * 2.0 * cam->GetParallelScale();
+      } else {
+        double dist = cam->GetDistance();
+        double angleRad = vtkMath::Pi() * cam->GetViewAngle() / 180.0;
+        double viewHeightAtDist = 2.0 * dist * tan(angleRad / 2.0);
+        screenLimit = (pixelSize / (double)sz[1]) * viewHeightAtDist;
+      }
+    }
+  }
+  double scaleFactor = std::min(bboxSize, screenLimit);
+
   vtkNew<vtkMatrix4x4> mat_gizmo;
   mat_gizmo->Identity();
 
   double center[3];
-  double bounds[6];
-  this->Prop3D->GetBounds(bounds);
   center[0] = (bounds[0] + bounds[1]) / 2.0;
   center[1] = (bounds[2] + bounds[3]) / 2.0;
   center[2] = (bounds[4] + bounds[5]) / 2.0;
@@ -591,10 +616,11 @@ void vtkHandlerWidget::UpdateGizmoPosition() {
     mat_gizmo->SetElement(1, 3, center[1]);
     mat_gizmo->SetElement(2, 3, center[2]);
   } else if (m_Frame == GLOBAL) {
+    ::vtkMatrix4x4 *mat = this->Prop3D->GetMatrix();
     mat_gizmo->Identity();
-    mat_gizmo->SetElement(0, 3, pos[0]);
-    mat_gizmo->SetElement(1, 3, pos[1]);
-    mat_gizmo->SetElement(2, 3, pos[2]);
+    mat_gizmo->SetElement(0, 3, mat->GetElement(0, 3));
+    mat_gizmo->SetElement(1, 3, mat->GetElement(1, 3));
+    mat_gizmo->SetElement(2, 3, mat->GetElement(2, 3));
   } else if (m_Frame == CENTER) {
     mat_gizmo->Identity();
     mat_gizmo->SetElement(0, 3, center[0]);
@@ -602,6 +628,13 @@ void vtkHandlerWidget::UpdateGizmoPosition() {
     mat_gizmo->SetElement(2, 3, center[2]);
   }
 
+  // Apply scaleFactor to the gizmo matrix (only top-left 3x3)
+  for (int j = 0; j < 3; ++j) {
+    for (int i = 0; i < 3; ++i) {
+      mat_gizmo->SetElement(i, j, mat_gizmo->GetElement(i, j) * scaleFactor);
+    }
+  }
+
   m_AxesX->SetUserMatrix(mat_gizmo);
   m_AxesY->SetUserMatrix(mat_gizmo);
   m_AxesZ->SetUserMatrix(mat_gizmo);
@@ -635,6 +668,7 @@ void vtkHandlerWidget::UpdateGizmoPosition() {
     // Orient RotCam actor to face 'dir'
     vtkNew<vtkTransform> tcam;
     tcam->PostMultiply();
+    tcam->Scale(scaleFactor, scaleFactor, scaleFactor);
     // Default circle is in XY plane (Normal Z: 0,0,1)
     double z[3] = {0, 0, 1};
     double cross[3];

src/Vtk/vtkQViewport.cpp

@@ -1,6 +1,7 @@
 #include "vtkQViewport.h"
 
 #include <QVBoxLayout>
+#include <QResizeEvent>
 
 #include <vtkAxesActor.h>
 #include <vtkCamera.h>
@@ -17,6 +18,7 @@ QViewport::QViewport(QWidget* parent)
     : QWidget(parent)
     , Viewport()
     , m_VtkWidget(nullptr)
+    , m_GridButton(nullptr)
 {
     // Build the layout – zero margins so VTK fills the entire widget
     auto* layout = new QVBoxLayout(this);
@@ -26,6 +28,36 @@ QViewport::QViewport(QWidget* parent)
     m_VtkWidget = new QVTKOpenGLNativeWidget(this);
     layout->addWidget(m_VtkWidget);
 
+    // Grid Toggle Button
+    m_GridButton = new QPushButton(m_VtkWidget);
+    m_GridButton->setText("#");
+    m_GridButton->setFixedSize(40, 40);
+    m_GridButton->setToolTip("Toggle Grid");
+    m_GridButton->setStyleSheet(
+        "QPushButton {"
+        "  border-radius: 20px;"              // Perfectly circular
+        "  background-color: rgba(40, 40, 40, 180);"
+        "  color: white;"
+        "  font-size: 22px;"
+        "  border: 1.5px solid rgba(255, 255, 255, 60);"
+        "}"
+        "QPushButton:hover {"
+        "  background-color: rgba(70, 70, 70, 200);"
+        "  border: 1.5px solid rgba(255, 255, 255, 100);"
+        "}"
+        "QPushButton:checked {"
+        "  background-color: rgba(0, 120, 215, 200);" // Nice "active" blue
+        "  color: white;"
+        "  border: 1.5px solid rgba(255, 255, 255, 120);"
+        "}"
+        "QPushButton:pressed {"
+        "  background-color: rgba(0, 90, 160, 220);"
+        "}"
+    );
+    m_GridButton->setCheckable(true);
+    m_GridButton->setChecked(true); // Grid is on by default
+    connect(m_GridButton, &QPushButton::clicked, this, &QViewport::onGridButtonClicked);
+
     // After the Qt widget exists but before the first paint,
     // attach the renderer and configure the pipeline.
     SetupPipeline();
@@ -64,5 +96,22 @@ vtkRenderWindowInteractor* QViewport::GetInteractor()
     return m_VtkWidget->renderWindow()->GetInteractor();
 }
 
+void QViewport::onGridButtonClicked()
+{
+    SetGridVisible(m_GridButton->isChecked());
+}
+
+void QViewport::resizeEvent(QResizeEvent* event)
+{
+    QWidget::resizeEvent(event);
+    if (m_GridButton) {
+        // Position under the gizmo (top-right corner).
+        // Standard CameraOrientationWidget is usually 150-180px.
+        int x = width() - m_GridButton->width() - 10;
+        int y = 160; 
+        m_GridButton->move(x, y);
+    }
+}
+
 } // namespace Vtk
 } // namespace uLib

src/Vtk/vtkQViewport.h

@@ -3,6 +3,7 @@
 
 #include <QWidget>
 #include <QVTKOpenGLNativeWidget.h>
+#include <QPushButton>
 
 #include <vtkCornerAnnotation.h>
 #include <vtkOrientationMarkerWidget.h>
@@ -41,10 +42,17 @@ public:
     virtual vtkRenderWindowInteractor*  GetInteractor() override;
     QVTKOpenGLNativeWidget*     GetWidget()      { return m_VtkWidget; }
 
+protected:
+    virtual void resizeEvent(QResizeEvent* event) override;
+
+private slots:
+    void onGridButtonClicked();
+
 private:
     void SetupPipeline();
 
     QVTKOpenGLNativeWidget*             m_VtkWidget;
+    QPushButton*                        m_GridButton;
 };
 
 } // namespace Vtk

src/Vtk/vtkViewport.cpp

@@ -1,7 +1,9 @@
 #include "vtkViewport.h"
 #include <vtkPropAssembly.h>
+#include <vtkAssembly.h>
 #include <vtkPropCollection.h>
 #include <vtkProp3D.h>
+#include <vtkProp3DCollection.h>
 #include <vtkCamera.h>
 #include <algorithm>
 #include <vtkInteractorStyleTrackballCamera.h>
@@ -29,6 +31,7 @@ Viewport::Viewport()
     , m_Marker(vtkSmartPointer<vtkOrientationMarkerWidget>::New())
     , m_CameraWidget(nullptr)
     , m_Colors(vtkSmartPointer<vtkNamedColors>::New())
+    , m_GridAxis(Y)
 {
 }
 
@@ -125,6 +128,19 @@ void Viewport::SetupPipeline(vtkRenderWindowInteractor* iren)
         m_HandlerWidget->GetOverlayRenderer()->SetLayer(1);
     }
 
+    // Observe InteractionEvent to update the selected puppet when the widget moves it
+    vtkNew<vtkCallbackCommand> widgetInteractionCallback;
+    widgetInteractionCallback->SetClientData(this);
+    widgetInteractionCallback->SetCallback([](vtkObject*, unsigned long, void* clientdata, void*){
+        auto* self = static_cast<Viewport*>(clientdata);
+        for (auto* p : self->m_Puppets) {
+            if (p->IsSelected()) {
+                p->Update();
+            }
+        }
+    });
+    m_HandlerWidget->AddObserver(vtkCommand::InteractionEvent, widgetInteractionCallback);
+
     // Picking for selection
     m_Picker = vtkSmartPointer<vtkCellPicker>::New();
     vtkNew<vtkCallbackCommand> clickCallback;
@@ -144,10 +160,14 @@ void Viewport::SetupPipeline(vtkRenderWindowInteractor* iren)
                     target = p;
                     break;
                 }
-                auto* assembly = vtkPropAssembly::SafeDownCast(p->GetProp());
-                if (assembly) {
+                auto* propAssembly = vtkPropAssembly::SafeDownCast(p->GetProp());
+                auto* actorAssembly = vtkAssembly::SafeDownCast(p->GetProp());
+                vtkPropCollection* parts = nullptr;
+                if (propAssembly) parts = propAssembly->GetParts();
+                else if (actorAssembly) parts = actorAssembly->GetParts();
+
+                if (parts) {
                     bool found = false;
-                    auto* parts = assembly->GetParts();
                     parts->InitTraversal();
                     for (int i=0; i<parts->GetNumberOfItems(); ++i) {
                         if (parts->GetNextProp() == picked) {
@@ -251,6 +271,29 @@ void Viewport::SelectPuppet(Puppet* prop)
     Render();
 }
 
+void Viewport::SetGridVisible(bool visible)
+{
+    if (m_GridActor) {
+        m_GridActor->SetVisibility(visible);
+        Render();
+    }
+}
+
+bool Viewport::GetGridVisible() const
+{
+    if (m_GridActor) {
+        return m_GridActor->GetVisibility() != 0;
+    }
+    return false;
+}
+
+void Viewport::SetGridAxis(Axis axis)
+{
+    m_GridAxis = axis;
+    UpdateGrid();
+    Render();
+}
+
 void Viewport::addProp(vtkProp* prop)
 {
     if (m_Renderer) {
@@ -270,6 +313,7 @@ void Viewport::RemoveProp(vtkProp* prop)
 void Viewport::UpdateGrid()
 {
     if (!m_Renderer || !m_GridSource) return;
+    if (m_GridActor && !m_GridActor->GetVisibility()) return;
 
     vtkCamera* camera = m_Renderer->GetActiveCamera();
     if (!camera) return;
@@ -296,24 +340,35 @@ void Viewport::UpdateGrid()
     double focalPoint[3];
     camera->GetFocalPoint(focalPoint);
 
-    // Align center to spacing
-    double centerX = std::round(focalPoint[0] / spacing) * spacing;
-    double centerY = std::round(focalPoint[1] / spacing) * spacing;
+    // Indices for the two dimensions of the grid plane
+    int idxH, idxV, idxN;
+    if (m_GridAxis == X) { idxH = 1; idxV = 2; idxN = 0; }
+    else if (m_GridAxis == Y) { idxH = 0; idxV = 2; idxN = 1; }
+    else { idxH = 0; idxV = 1; idxN = 2; }
 
-    // Number of lines: enough to cover the screen even if aspect ratio is wide
-    // or if we rotate. 20x20 is usually plenty.
+    // Align center to spacing
+    double centerH = std::round(focalPoint[idxH] / spacing) * spacing;
+    double centerV = std::round(focalPoint[idxV] / spacing) * spacing;
+    double centerN = 0.0; // Grid plane typically passes through the origin
+
+    // Number of lines
     int numLines = 20;
     double halfSize = (numLines / 2.0) * spacing;
 
-    double xmin = centerX - halfSize;
-    double xmax = centerX + halfSize;
-    double ymin = centerY - halfSize;
-    double ymax = centerY + halfSize;
+    double minH = centerH - halfSize;
+    double maxH = centerH + halfSize;
+    double minV = centerV - halfSize;
+    double maxV = centerV + halfSize;
 
-    // Update Plane Source
-    m_GridSource->SetOrigin(xmin, ymin, 0.0);
-    m_GridSource->SetPoint1(xmax, ymin, 0.0);
-    m_GridSource->SetPoint2(xmin, ymax, 0.0);
+    // Update Plane Source mapping axes to origin/point1/point2
+    double origin[3] = {0,0,0}, p1[3] = {0,0,0}, p2[3] = {0,0,0};
+    origin[idxH] = minH; origin[idxV] = minV; origin[idxN] = centerN;
+    p1[idxH]     = maxH; p1[idxV]     = minV; p1[idxN]     = centerN;
+    p2[idxH]     = minH; p2[idxV]     = maxV; p2[idxN]     = centerN;
+
+    m_GridSource->SetOrigin(origin);
+    m_GridSource->SetPoint1(p1);
+    m_GridSource->SetPoint2(p2);
     m_GridSource->SetXResolution(numLines);
     m_GridSource->SetYResolution(numLines);
     m_GridSource->Update();

src/Vtk/vtkViewport.h

@@ -61,6 +61,14 @@ public:
     vtkCornerAnnotation*        GetAnnotation()  { return m_Annotation; }
     vtkCameraOrientationWidget* GetCameraWidget(){ return m_CameraWidget; }
 
+    // Grid control
+    void SetGridVisible(bool visible);
+    bool GetGridVisible() const;
+
+    enum Axis { X=0, Y, Z };
+    void SetGridAxis(Axis axis);
+    Axis GetGridAxis() const { return m_GridAxis; }
+
 protected:
     void SetupPipeline(vtkRenderWindowInteractor* iren);
 
@@ -78,6 +86,7 @@ protected:
 
     vtkSmartPointer<vtkNamedColors>              m_Colors;
     
+    Axis                                         m_GridAxis;
     vtkSmartPointer<vtkHandlerWidget>            m_HandlerWidget;
     std::vector<Puppet*>                         m_Puppets;
     vtkSmartPointer<vtkCellPicker>               m_Picker;