Jake Kalstad преди 4 години
ревизия
874b05c24f
променени са 4 файла, в които са добавени 182 реда и са изтрити 0 реда
  1. 18 0
      birds.json
  2. 9 0
      config.json
  3. 128 0
      datalabel.go
  4. 27 0
      readme.md

+ 18 - 0
birds.json

@@ -0,0 +1,18 @@
+[
+    {
+        "name": "Parrot",
+        "images": [
+            "parrot_image_1.jpg",
+            "parrot_image_2.jpg",
+            "parrot_image_3.jpg"
+        ]
+    },
+    {
+        "name": "Cockatoo",
+        "images": [
+            "cockatoo_image_1.jpg",
+            "cockatoo_image_2.jpg",
+            "cockatoo_image_3.jpg"
+        ]
+    }
+]

+ 9 - 0
config.json

@@ -0,0 +1,9 @@
+{
+    "dataFile": "birds.json",
+    "imagePrefix": "images",
+    "defaultPrefix": "pr_set",
+    "secondaryPrefix": "pr_eggset",
+    "secondaryPredicate": "egg",
+    "ignorePredicate": "map",
+    "outfile": "classes"
+}

+ 128 - 0
datalabel.go

@@ -0,0 +1,128 @@
+package main
+
+import (
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"io/ioutil"
+	"os"
+	"strconv"
+	"strings"
+)
+
+type ImageSet struct {
+	Name   string
+	Images []string
+}
+
+func moveFile(source, dest string) {
+	sourceFile, err := os.Open(source)
+	if err != nil {
+		panic(err)
+	}
+	defer sourceFile.Close()
+
+	newFile, err := os.Create(dest)
+	if err != nil {
+		panic(err)
+	}
+	defer newFile.Close()
+
+	_, err = io.Copy(newFile, sourceFile)
+	if err != nil {
+		panic(err)
+	}
+}
+
+type Config struct {
+	DataFile           string
+	ImagePrefix        string
+	DefaultPrefix      string
+	SecondaryPrefix    string
+	SecondaryPredicate string
+	IgnorePredicate    string
+	Outfile            string
+}
+
+func main() {
+	combineString := "%s/%s"
+	configBytes, err := ioutil.ReadFile("config.json")
+	if err != nil {
+		panic(err)
+	}
+	config := Config{}
+	err = json.Unmarshal(configBytes, &config)
+	if err != nil {
+		panic(err)
+	}
+	if len(config.Outfile) == 0 {
+		panic(errors.New("No output file name defined `classes`"))
+	}
+	if len(config.ImagePrefix) == 0 {
+		panic(errors.New("No image directory defined `imagefolder`"))
+	}
+	if len(config.DataFile) == 0 {
+		panic(errors.New("No data file defined `myinput.json`"))
+	}
+	if len(config.DefaultPrefix) == 0 {
+		panic(errors.New("No default directory defined `my_data_set`"))
+	}
+	if len(config.SecondaryPredicate) > 0 && len(config.SecondaryPrefix) == 0 {
+		panic(errors.New("No secondary directory defined with a predicate defined:" + config.SecondaryPredicate))
+	}
+	if _, err := os.Stat(config.DefaultPrefix); os.IsNotExist(err) {
+		err = os.Mkdir(config.DefaultPrefix, os.FileMode(0777))
+		if err != nil {
+			panic(err)
+		}
+	}
+	if len(config.SecondaryPredicate) > 0 {
+		if _, err := os.Stat(config.SecondaryPrefix); os.IsNotExist(err) {
+			err = os.Mkdir(config.SecondaryPrefix, os.FileMode(0777))
+			if err != nil {
+				panic(err)
+			}
+		}
+	}
+
+	data, err := ioutil.ReadFile(config.DataFile)
+	if err != nil {
+		panic(err)
+	}
+	imageSets := []ImageSet{}
+	err = json.Unmarshal(data, &imageSets)
+	if err != nil {
+		panic(err)
+	}
+	classes := map[int]string{}
+
+	for idx, b := range imageSets {
+		classes[idx] = b.Name
+		for _, img := range b.Images {
+			if strings.Contains(img, config.IgnorePredicate) {
+				continue
+			}
+			dir := fmt.Sprintf(combineString, config.DefaultPrefix, strconv.Itoa(idx))
+			containsPredicate := strings.Contains(strings.ToLower(img), strings.ToLower(config.SecondaryPredicate))
+			if len(config.SecondaryPredicate) > 0 && containsPredicate {
+				dir = fmt.Sprintf(combineString, config.SecondaryPrefix, strconv.Itoa(idx))
+			}
+			if _, err := os.Stat(dir); os.IsNotExist(err) {
+				err = os.Mkdir(dir, os.FileMode(0777))
+				if err != nil {
+					panic(err)
+				}
+			}
+			moveFile(fmt.Sprintf(combineString, config.ImagePrefix, img), fmt.Sprintf(combineString, dir, img))
+		}
+	}
+	classBytes, err := json.Marshal(classes)
+	if err != nil {
+		panic(err)
+	}
+	err = ioutil.WriteFile(config.Outfile+".json", classBytes, os.FileMode(0777))
+	if err != nil {
+		panic(err)
+	}
+}

+ 27 - 0
readme.md

@@ -0,0 +1,27 @@
+# go-image-label
+
+Reads a JSON file and organizes images into subfolders to be easily used with Pytorch `ImageFolder`
+
+    {
+        "dataFile": "birds.json",
+        "imagePrefix": "images",
+        "defaultPrefix": "pr_set",
+        "secondaryPrefix": "pr_eggset",
+        "secondaryPredicate": "egg",
+        "ignorePredicate": "map",
+        "outfile": "classes"
+    }
+
+dataFile is the input file - see birds.json as example format
+
+imagePrefix is the directory that the entire image set is
+
+defaultPrefix is where we place the subfolders and images
+
+secondaryPrefix is where we place the subfolders and images of any data that matches our secondaryPredicate
+
+secondaryPredicate is a string that is checked if it's contained in the image name and if so it will use the secondaryPrefix instead of the default
+
+ignorePredicate is a string that is checked if it's contained in the image name and if so it will be discarded
+
+outfile is the file with the unique classes listed out as a json file