diff --git a/images.go b/images.go
index 31f6ae3..c2893e6 100644
--- a/images.go
+++ b/images.go
@@ -85,36 +85,10 @@ func (x *ImagesCommand) Execute(args []string) error {
 	if imagesCommand.Tree || imagesCommand.Dot {
 		var startImage *Image
 		if len(args) > 0 {
+			startImage, err = findStartImage(args[0], images)
 
-			// attempt to find the start image, which can be specified as an
-			// image ID or a repository name
-			startImageArg := args[0]
-			startImageRepo := args[0]
-
-			// if tag is not defined, find by :latest tag
-			if strings.Index(startImageRepo, ":") == -1 {
-				startImageRepo = fmt.Sprintf("%s:latest", startImageRepo)
-			}
-
-		IMAGES:
-			for _, image := range *images {
-				// find by image id
-				if strings.Index(image.Id, startImageArg) == 0 {
-					startImage = &image
-					break IMAGES
-				}
-
-				// find by image name (name and tag)
-				for _, repotag := range image.RepoTags {
-					if repotag == startImageRepo {
-						startImage = &image
-						break IMAGES
-					}
-				}
-			}
-
-			if startImage == nil {
-				return fmt.Errorf("Unable to find image %s = %s.", startImageArg, startImageRepo)
+			if err != nil {
+				return err
 			}
 		}
 
@@ -128,30 +102,20 @@ func (x *ImagesCommand) Execute(args []string) error {
 		}
 
 		// build helper map (image -> children)
-		var imagesByParent = make(map[string][]Image)
-		imagesByParent = collectChildren(images)
-
-		// image ids truncate
-		// initialize image informations
+		imagesByParent := collectChildren(images)
 
 		// filter images
 		if imagesCommand.OnlyLabelled {
 			*images, imagesByParent = filterImages(images, &imagesByParent)
 		}
 
-		var buffer bytes.Buffer
-
 		if imagesCommand.Tree {
-			jsonToText(&buffer, imagesCommand.NoTruncate, roots, imagesByParent, "")
+			fmt.Print(jsonToTree(imagesCommand.NoTruncate, roots, imagesByParent))
 		}
 		if imagesCommand.Dot {
-			buffer.WriteString("digraph docker {\n")
-			imagesToDot(&buffer, roots, imagesByParent)
-			buffer.WriteString(" base [style=invisible]\n}\n")
-			buffer.String()
+			fmt.Print(jsonToDot(roots, imagesByParent))
 		}
 
-		fmt.Print(buffer.String())
 	} else if imagesCommand.Short {
 		fmt.Printf(jsonToShort(images))
 	} else {
@@ -161,6 +125,62 @@ func (x *ImagesCommand) Execute(args []string) error {
 	return nil
 }
 
+func findStartImage(name string, images *[]Image) (*Image, error) {
+
+	var startImage *Image
+
+	// attempt to find the start image, which can be specified as an
+	// image ID or a repository name
+	startImageArg := name
+	startImageRepo := name
+
+	// if tag is not defined, find by :latest tag
+	if strings.Index(startImageRepo, ":") == -1 {
+		startImageRepo = fmt.Sprintf("%s:latest", startImageRepo)
+	}
+
+IMAGES:
+	for _, image := range *images {
+		// find by image id
+		if strings.Index(image.Id, startImageArg) == 0 {
+			startImage = &image
+			break IMAGES
+		}
+
+		// find by image name (name and tag)
+		for _, repotag := range image.RepoTags {
+			if repotag == startImageRepo {
+				startImage = &image
+				break IMAGES
+			}
+		}
+	}
+
+	if startImage == nil {
+		return nil, fmt.Errorf("Unable to find image %s = %s.", startImageArg, startImageRepo)
+	}
+
+	return startImage, nil
+}
+
+func jsonToTree(noTrunc bool, images []Image, byParent map[string][]Image) string {
+	var buffer bytes.Buffer
+
+	jsonToText(&buffer, noTrunc, images, byParent, "")
+
+	return buffer.String()
+}
+
+func jsonToDot(roots []Image, byParent map[string][]Image) string {
+	var buffer bytes.Buffer
+
+	buffer.WriteString("digraph docker {\n")
+	imagesToDot(&buffer, roots, byParent)
+	buffer.WriteString(" base [style=invisible]\n}\n")
+
+	return buffer.String()
+}
+
 func collectChildren(images *[]Image) map[string][]Image {
 	var imagesByParent = make(map[string][]Image)
 	for _, image := range *images {
diff --git a/images_test.go b/images_test.go
index 09b32c3..b88f154 100644
--- a/images_test.go
+++ b/images_test.go
@@ -56,7 +56,12 @@ func Test_Dot(t *testing.T) {
 
 	for _, dotTest := range dotTests {
 		im, _ := parseImagesJSON([]byte(dotTest.json))
-		result := jsonToDot(im)
+		byParent := collectChildren(im)
+		roots := collectRoots(im)
+
+		// TODO: test start image limiting
+
+		result := jsonToDot(roots, byParent)
 
 		for _, regexp := range allRegex {
 			if !regexp.MatchString(result) {
@@ -109,7 +114,16 @@ func Test_Tree(t *testing.T) {
 
 	for _, treeTest := range treeTests {
 		im, _ := parseImagesJSON([]byte(treeTest.json))
-		result := jsonToTree(im, treeTest.startImage, treeTest.noTrunc)
+		byParent := collectChildren(im)
+		var roots []Image
+		if len(treeTest.startImage) > 0 {
+			startImage, _ := findStartImage(treeTest.startImage, im)
+			startImage.ParentId = ""
+			roots = []Image{*startImage}
+		} else {
+			roots = collectRoots(im)
+		}
+		result := jsonToTree(treeTest.noTrunc, roots, byParent)
 
 		for _, regexp := range compileRegexps(t, treeTest.regexps) {
 			if !regexp.MatchString(result) {