diff --git a/.idea/misc.xml b/.idea/misc.xml
index e1c3f7e..ad08689 100644
--- a/.idea/misc.xml
+++ b/.idea/misc.xml
@@ -4,7 +4,7 @@
   <component name="FrameworkDetectionExcludesConfiguration">
     <file type="web" url="file://$PROJECT_DIR$" />
   </component>
-  <component name="ProjectRootManager" version="2" languageLevel="JDK_17" project-jdk-name="corretto-21" project-jdk-type="JavaSDK">
+  <component name="ProjectRootManager" version="2" languageLevel="JDK_17" project-jdk-name="17" project-jdk-type="JavaSDK">
     <output url="file://$PROJECT_DIR$/out" />
   </component>
 </project>
\ No newline at end of file
diff --git a/src/main/java/module-info.java b/src/main/java/module-info.java
index abf62b6..b2baeae 100644
--- a/src/main/java/module-info.java
+++ b/src/main/java/module-info.java
@@ -1,5 +1,6 @@
 module com.nuculabs.dev.imagetagger.ui {
     requires javafx.controls;
+    requires javafx.graphics;
     requires javafx.fxml;
     requires kotlin.stdlib;
 
@@ -11,6 +12,6 @@ module com.nuculabs.dev.imagetagger.ui {
     requires java.logging;
     requires java.desktop;
 
-    opens com.nuculabs.dev.imagetagger.ui to javafx.fxml;
+    opens com.nuculabs.dev.imagetagger.ui to javafx.fxml, javafx.graphics;
     exports com.nuculabs.dev.imagetagger.ui;
 }
\ No newline at end of file
diff --git a/src/main/kotlin/com/nuculabs/dev/imagetagger/ui/controls/ImageTagsEntry.kt b/src/main/kotlin/com/nuculabs/dev/imagetagger/ui/ImageTagsEntryControl.kt
similarity index 57%
rename from src/main/kotlin/com/nuculabs/dev/imagetagger/ui/controls/ImageTagsEntry.kt
rename to src/main/kotlin/com/nuculabs/dev/imagetagger/ui/ImageTagsEntryControl.kt
index ccd308e..6990550 100644
--- a/src/main/kotlin/com/nuculabs/dev/imagetagger/ui/controls/ImageTagsEntry.kt
+++ b/src/main/kotlin/com/nuculabs/dev/imagetagger/ui/ImageTagsEntryControl.kt
@@ -1,16 +1,22 @@
-package com.nuculabs.dev.imagetagger.ui.controls
+package com.nuculabs.dev.imagetagger.ui
 import javafx.fxml.FXML
 import javafx.fxml.FXMLLoader
 import javafx.scene.control.TextArea
 import javafx.scene.image.Image
 import javafx.scene.image.ImageView
 import javafx.scene.layout.HBox
+import java.io.File
 import java.io.IOException
+import java.util.logging.Logger
+
 
 /**
  * This class is used to create a custom control for the image prediction entry.
  */
-class ImageTagsEntry(imagePath: String, predictions: List<String>) : HBox() {
+class ImageTagsEntryControl
+    (imagePath: String, predictions: List<String>) : HBox() {
+    private val logger: Logger = Logger.getLogger("ImageTagsEntryControl")
+
     /**
      * The image view for the image prediction entry.
      */
@@ -23,23 +29,17 @@ class ImageTagsEntry(imagePath: String, predictions: List<String>) : HBox() {
     @FXML
     private lateinit var predictedImageTags: TextArea
 
-    /**
-     * Constructor for the image prediction entry.
-     *
-     * @param imagePath   The image path.
-     * @param predictions The prediction list.
-     */
     init {
-        val fxmlLoader = FXMLLoader(ImageTagsEntry::class.java.getResource("/controls/image-tags-entry.fxml"))
+        val resource = ImageTagsEntryControl::class.java.getResource("controls/image-tags-entry.fxml")
+        logger.info("Using resource URL: $resource")
+        val fxmlLoader = FXMLLoader(resource)
         fxmlLoader.setRoot(this)
         fxmlLoader.setController(this)
-
         try {
-            fxmlLoader.load<Any>()
+            fxmlLoader.load()
         } catch (exception: IOException) {
             throw RuntimeException(exception)
         }
-
         setImage(imagePath)
         setText(predictions)
     }
@@ -56,12 +56,14 @@ class ImageTagsEntry(imagePath: String, predictions: List<String>) : HBox() {
     /**
      * Setter for setting the image.
      */
-    fun setImage(imagePath: String?) {
-        imageView.image = Image(imagePath)
-        imageView.resize(244.0, 244.0)
-        imageView.fitHeight = 244.0
-        imageView.fitWidth = 244.0
-        imageView.isSmooth = true
-        imageView.isCache = true
+    fun setImage(imagePath: String) {
+        File(imagePath).inputStream().use {
+            imageView.image = Image(it)
+            imageView.resize(244.0, 244.0)
+            imageView.fitHeight = 244.0
+            imageView.fitWidth = 244.0
+            imageView.isSmooth = true
+            imageView.isCache = true
+        }
     }
 }
\ No newline at end of file
diff --git a/src/main/kotlin/com/nuculabs/dev/imagetagger/ui/MainPageController.kt b/src/main/kotlin/com/nuculabs/dev/imagetagger/ui/MainPageController.kt
index 78683f3..5b1f788 100644
--- a/src/main/kotlin/com/nuculabs/dev/imagetagger/ui/MainPageController.kt
+++ b/src/main/kotlin/com/nuculabs/dev/imagetagger/ui/MainPageController.kt
@@ -1,7 +1,6 @@
 package com.nuculabs.dev.imagetagger.ui
 
 import com.nuculabs.dev.imagetagger.tag_prediction.ImageTagsPrediction
-import com.nuculabs.dev.imagetagger.ui.controls.ImageTagsEntry
 import javafx.application.Platform
 import javafx.fxml.FXML
 import javafx.scene.control.Label
@@ -14,7 +13,7 @@ import javax.imageio.ImageIO
 
 
 class MainPageController {
-    private val logger: Logger = Logger.getLogger("InfoLogging")
+    private val logger: Logger = Logger.getLogger("MainPageController")
 
     @FXML
     private lateinit var verticalBox: VBox
@@ -44,14 +43,12 @@ class MainPageController {
                     val tags: List<String> = imageTagsPrediction.predictTags(imageFile)
                     Platform.runLater {
                         // Add image and prediction to the view.
-                        verticalBox.children.add(
-                            Label("${filePath.name} - $tags")
-                        )
+                        verticalBox.children.add(ImageTagsEntryControl(filePath.absolutePath, tags))
                         verticalBox.children.add(Separator())
                     }
                 }
                 Platform.runLater {
-                    statusLabel.setVisible(false)
+                    statusLabel.isVisible = false
                 }
             } catch (e: Exception) {
                 logger.warning("Error while predicting images $e")