diff --git a/src/main/scala/report/donut/transformers/cucumber/CucumberModel.scala b/src/main/scala/report/donut/transformers/cucumber/CucumberModel.scala
index 24c5f86..9aef848 100644
--- a/src/main/scala/report/donut/transformers/cucumber/CucumberModel.scala
+++ b/src/main/scala/report/donut/transformers/cucumber/CucumberModel.scala
@@ -10,7 +10,7 @@ case class Match(location: String = "", arguments: Option[List[Argument]])
 
 case class BeforeHook(result: Result, output: List[String], `match`: Match)
 
-case class AfterHook(result: Result, output: List[String], `match`: Match)
+case class AfterHook(result: Result, output: List[String], `match`: Match,embeddings: List[Embedding])
 
 case class Row(cells: List[String])
 
@@ -58,4 +58,4 @@ case class Feature(keyword: String,
   val scenariosExcludeBackground = {
     elements.filterNot(e => e.keyword == "Background")
   }
-}
\ No newline at end of file
+}
diff --git a/src/main/scala/report/donut/transformers/cucumber/CucumberTransformer.scala b/src/main/scala/report/donut/transformers/cucumber/CucumberTransformer.scala
index 66e34ea..cbf8fb8 100644
--- a/src/main/scala/report/donut/transformers/cucumber/CucumberTransformer.scala
+++ b/src/main/scala/report/donut/transformers/cucumber/CucumberTransformer.scala
@@ -159,7 +159,7 @@ object CucumberTransformer extends Log {
   }
 
   private[cucumber] def donutScenarioScreenshots(e: Element) = {
-    val elementScreenshots: List[Embedding] = e.steps.flatMap(s => s.embeddings)
+        val elementScreenshots: List[Embedding] = e.after.flatMap(s => if (s.embeddings != Nil) s.embeddings else e.steps.flatMap(t => t.embeddings))
     val screenshotsSize = elementScreenshots.size
     val screenshotStyle = if (elementScreenshots.nonEmpty) "" else "display:none;"
     val screenshots = elementScreenshots.map(e => DonutEmbedding(e.mime_type, e.data, e.id))