diff --git a/WPF/03-TreeViewsSimpleViewModel/WpfTreeView/Directory/ViewModels/DirectoryItemViewModel.cs b/WPF/03-TreeViewsSimpleViewModel/WpfTreeView/Directory/ViewModels/DirectoryItemViewModel.cs
index ee777ef..1949be2 100644
--- a/WPF/03-TreeViewsSimpleViewModel/WpfTreeView/Directory/ViewModels/DirectoryItemViewModel.cs
+++ b/WPF/03-TreeViewsSimpleViewModel/WpfTreeView/Directory/ViewModels/DirectoryItemViewModel.cs
@@ -16,6 +16,8 @@ namespace WpfTreeView
         /// </summary>
         public DirectoryItemType Type { get; set; }
 
+        public string ImageName => Type == DirectoryItemType.Drive ? "drive" : (Type == DirectoryItemType.File ? "file" : (IsExpanded ? "folder-open" : "folder-closed"));
+
         /// <summary>
         /// The full path to the item
         /// </summary>
diff --git a/WPF/03-TreeViewsSimpleViewModel/WpfTreeView/HeaderToImageConverter.cs b/WPF/03-TreeViewsSimpleViewModel/WpfTreeView/HeaderToImageConverter.cs
index a2553b9..be02107 100644
--- a/WPF/03-TreeViewsSimpleViewModel/WpfTreeView/HeaderToImageConverter.cs
+++ b/WPF/03-TreeViewsSimpleViewModel/WpfTreeView/HeaderToImageConverter.cs
@@ -16,20 +16,7 @@ namespace WpfTreeView
 
         public object Convert(object value, Type targetType, object parameter, CultureInfo culture)
         {
-            // By default, we presume an image
-            var image = "Images/file.png";
-
-            switch ((DirectoryItemType)value)
-            {
-                case DirectoryItemType.Drive:
-                    image = "Images/drive.png";
-                    break;
-                case DirectoryItemType.Folder:
-                    image = "Images/folder-closed.png";
-                    break;
-            }
-
-            return new BitmapImage(new Uri($"pack://application:,,,/{image}"));
+            return new BitmapImage(new Uri($"pack://application:,,,/Images/{value}.png"));
         }
 
         public object ConvertBack(object value, Type targetType, object parameter, CultureInfo culture)
diff --git a/WPF/03-TreeViewsSimpleViewModel/WpfTreeView/MainWindow.xaml b/WPF/03-TreeViewsSimpleViewModel/WpfTreeView/MainWindow.xaml
index 06bb258..90f288b 100644
--- a/WPF/03-TreeViewsSimpleViewModel/WpfTreeView/MainWindow.xaml
+++ b/WPF/03-TreeViewsSimpleViewModel/WpfTreeView/MainWindow.xaml
@@ -20,7 +20,7 @@
                 <HierarchicalDataTemplate ItemsSource="{Binding Children}">
                     <StackPanel Orientation="Horizontal">
                         <Image Width="20" Margin="3" 
-                                    Source="{Binding Type,
+                                    Source="{Binding ImageName,
                                         Converter={x:Static local:HeaderToImageConverter.Instance}}" />
                         <TextBlock VerticalAlignment="Center" Text="{Binding Name}" />
                     </StackPanel>