diff --git a/imageOperations.go b/imageOperations.go
index 1344c25..687cd3e 100644
--- a/imageOperations.go
+++ b/imageOperations.go
@@ -19,9 +19,7 @@ func dataToImage(data []byte, imageExtension string) (image.Image, error) {
 	switch imageExtension {
 	case "png":
 		img, err = png.Decode(reader)
-	case "jpg":
-		img, err = jpeg.Decode(reader)
-	case "jpeg":
+	case "jpg", "jpeg":
 		img, err = jpeg.Decode(reader)
 	default:
 		img = nil
@@ -38,9 +36,7 @@ func imageToData(img image.Image, imageExtension string) ([]byte, error) {
 	switch imageExtension {
 	case "png":
 		err = png.Encode(buf, img)
-	case "jpg":
-		err = jpeg.Encode(buf, img, nil)
-	case "jpeg":
+	case "jpg", "jpeg":
 		err = jpeg.Encode(buf, img, nil)
 	default:
 		img = nil
diff --git a/knn.go b/knn.go
index 7789835..1177070 100644
--- a/knn.go
+++ b/knn.go
@@ -5,6 +5,11 @@ import (
 	"sort"
 )
 
+type Neighbour struct {
+	Dist  float64
+	Label string
+}
+
 func euclideanDist(img1, img2 [][]float64) float64 {
 	var dist float64
 	for i := 0; i < len(img1); i++ {
@@ -16,11 +21,6 @@ func euclideanDist(img1, img2 [][]float64) float64 {
 	return dist
 }
 
-type Neighbour struct {
-	Dist  float64
-	Label string
-}
-
 func isNeighbour(neighbours []Neighbour, dist float64, label string) []Neighbour {
 	var temp []Neighbour
 
@@ -71,30 +71,39 @@ func averageLabel(neighbours []Neighbour) string {
 	//send the most appeared neighbour in k
 	return a[0].Label
 }
-func knn(dataset Dataset, input [][]float64) string {
-	k := 10
-	var neighbours []Neighbour
-	label := getMapKey(dataset)
-	for i := 0; i < k; i++ {
-		/*neighbours[i].Dist = euclideanDist(dataset["leopard"][0], input)
-		neighbours[i].Label = "leopard"*/
-		neighbours = append(neighbours, Neighbour{euclideanDist(dataset[label][0], input), label})
-	}
+
+func distNeighboursFromDataset(dataset Dataset, neighbours []Neighbour, input [][]float64) []Neighbour {
+	//check the complete dataset, checking if each entry is a k nearest neighbour
 	for l, v := range dataset {
 		for i := 0; i < len(v); i++ {
 			dNew := euclideanDist(v[i], input)
-			/*if dNew < d {
-				d = dNew
-				label = l
-			}*/
 			neighbours = isNeighbour(neighbours, dNew, l)
 		}
 	}
+	return neighbours
+}
+func knn(dataset Dataset, input [][]float64) string {
+	k := 6
+	var neighbours []Neighbour
+	var neighboursED []Neighbour
+
+	//get a key from map dataset, the key is a label
+	label := getMapKey(dataset)
+	//fill the first k neighbours
+	for i := 0; i < k; i++ {
+		neighbours = append(neighbours, Neighbour{euclideanDist(dataset[label][0], input), label})
+		neighboursED = append(neighbours, Neighbour{euclideanDist(dataset[label][0], input), label})
+	}
+
+	neighbours = distNeighboursFromDataset(dataset, neighbours, input)
+	neighboursED = distNeighboursFromDataset(datasetED, neighbours, input)
+	neighbours = append(neighbours, neighboursED...)
+
 	for i := 0; i < len(neighbours); i++ {
 		fmt.Print(neighbours[i].Label + " - ")
 		fmt.Println(neighbours[i].Dist)
 	}
-
+	//from the k nearest neighbours, get the more frequent neighbour
 	r := averageLabel(neighbours)
 	return r
 }
diff --git a/main.go b/main.go
index cfd7476..8e1946e 100644
--- a/main.go
+++ b/main.go
@@ -10,13 +10,14 @@ import (
 
 //dataset := make(Dataset)
 var dataset Dataset
+var datasetED Dataset
 
 func main() {
 	readConfig("./config.json")
 
 	c.Cyan("reading images datasets")
 	tStart := time.Now()
-	dataset = readDataset("./dataset")
+	dataset, datasetED = readDataset("./dataset")
 	fmt.Print("time spend reading images: ")
 	fmt.Println(time.Since(tStart))
 	fmt.Println("total folders scanned: " + strconv.Itoa(len(dataset)))
diff --git a/readDataset.go b/readDataset.go
index f66545e..ec4f6fb 100644
--- a/readDataset.go
+++ b/readDataset.go
@@ -42,33 +42,35 @@ func readImage(path string) image.Image {
 		return histogram
 	*/
 }
-func readDataset(path string) map[string]ImgDataset {
+func readDataset(path string) (Dataset, Dataset) {
 	dataset := make(Dataset)
+	datasetED := make(Dataset)
 
 	folders, _ := ioutil.ReadDir(path)
 	for _, folder := range folders {
 		fmt.Println(folder.Name())
 
 		var imgDataset ImgDataset
+		var imgDatasetED ImgDataset
 
 		folderFiles, _ := ioutil.ReadDir(path + "/" + folder.Name())
 		for _, file := range folderFiles {
 			//get the image as original
 			image := readImage(path + "/" + folder.Name() + "/" + file.Name())
 			histogram := imageToHistogram(image)
+			imgDataset = append(imgDataset, histogram)
 
 			//get the image with EdgeDetection filter
 			imageED := EdgeDetection(image)
 			histogramED := imageToHistogram(imageED)
-
-			imgDataset = append(imgDataset, histogram)
-			imgDataset = append(imgDataset, histogramED)
+			imgDatasetED = append(imgDatasetED, histogramED)
 
 		}
 
 		//add the foldername to the Dataset map
 		dataset[folder.Name()] = imgDataset
+		datasetED[folder.Name()] = imgDatasetED
 	}
 
-	return dataset
+	return dataset, datasetED
 }