diff --git a/qiskit_research/protein_folding/protein_folding_result.py b/qiskit_research/protein_folding/protein_folding_result.py
index e96a37c5..c8fcd439 100644
--- a/qiskit_research/protein_folding/protein_folding_result.py
+++ b/qiskit_research/protein_folding/protein_folding_result.py
@@ -21,7 +21,7 @@
 
 if _optionals.HAS_MATPLOTLIB:
     # pylint: disable=import-error,unused-import
-    from matplotlib.pyplot import figure
+    from matplotlib.figure import Figure
 
 
 class ProteinFoldingResult:
@@ -168,7 +168,7 @@ def save_xyz_file(
     @_optionals.HAS_MATPLOTLIB.require_in_call
     def get_figure(
         self, title: str = "Protein Structure", ticks: bool = False, grid: bool = False
-    ) -> "figure":
+    ) -> Figure:
         """
         Generates a figure of the molecule in 3D.
 
diff --git a/qiskit_research/protein_folding/utils/protein_plotter.py b/qiskit_research/protein_folding/utils/protein_plotter.py
index 3504423b..c1a7bc12 100644
--- a/qiskit_research/protein_folding/utils/protein_plotter.py
+++ b/qiskit_research/protein_folding/utils/protein_plotter.py
@@ -11,18 +11,20 @@
 """An auxiliary class that plots aminoacids of a molecule
  in a ProteinFoldingResult."""
 from typing import Optional
+
 import numpy as np
 from qiskit.utils import optionals as _optionals
 
-
 from qiskit_research.protein_folding.utils.protein_shape_file_gen import (
     ProteinShapeFileGen,
 )
 
 if _optionals.HAS_MATPLOTLIB:
     # pylint: disable=import-error,unused-import
-    from matplotlib.pyplot import figure
     from matplotlib.axes import Axes
+    from matplotlib.figure import Figure
+    from mpl_toolkits.mplot3d.axes3d import Axes3D
+
 # pylint: disable=too-few-public-methods
 
 
@@ -53,7 +55,7 @@ def __init__(self, shape_gen: ProteinShapeFileGen) -> None:
         )
 
         self._fig = plt.figure()
-        self._ax_graph = self._fig.add_subplot(projection="3d")
+        self._ax_graph: Axes3D = self._fig.add_subplot(projection="3d")
 
     def _draw_main_chain(self):
         """
@@ -146,7 +148,7 @@ def _format_graph(
 
     def get_figure(
         self, title: str = "Protein Structure", ticks: bool = False, grid: bool = False
-    ) -> "figure":
+    ) -> Figure:
         """
         Plots the molecule in 3D.