Add helper functions

Author: Abdul Dakkak

vision/coco.go

@@ -2,10 +2,13 @@ package vision
 
 import (
 	context "context"
+	"fmt"
 	"path"
 	"path/filepath"
 	"strings"
 
+	"github.com/rai-project/dldataset/vision/support/object_detection"
+
 	"github.com/Unknwon/com"
 	"github.com/pkg/errors"
 	"github.com/rai-project/config"
@@ -35,11 +38,13 @@ type CocoLabeledImage struct {
 // CocoValidationTFRecord ...
 type CocoValidationTFRecord struct {
 	base
-	name           string
-	baseURL        string
-	recordFileName string
-	md5sum         string
-	recordReader   *reader.TFRecordReader
+	name             string
+	baseURL          string
+	recordFileName   string
+	md5sum           string
+	labelMap         object_detection.StringIntLabelMap
+	completeLabelMap object_detection.StringIntLabelMap
+	recordReader     *reader.TFRecordReader
 }
 
 var (
@@ -210,27 +215,41 @@ func init() {
 
 		const baseURLPrefix = "https://s3.amazonaws.com/store.carml.org/datasets"
 
+		labelMap, err := object_detection.Get("mscoco_label_map.pbtxt")
+		if err != nil {
+			panic(fmt.Sprintf("failed to get mscoco_label_map.pbtxt due to %v", err))
+		}
+
+		completeLabelMap, err := object_detection.Get("mscoco_complete_label_map.pbtxt")
+		if err != nil {
+			panic(fmt.Sprintf("failed to get mscoco_complete_label_map.pbtxt due to %v", err))
+		}
+
 		baseWorkingDir := filepath.Join(dldataset.Config.WorkingDirectory, "dldataset")
 		coco2014ValidationTFRecord = &CocoValidationTFRecord{
 			base: base{
 				ctx:            context.Background(),
 				baseWorkingDir: baseWorkingDir,
 			},
-			name:           "coco2014",
-			baseURL:        baseURLPrefix + "/coco2014",
-			recordFileName: "coco_val.record-00000-of-00001",
-			md5sum:         "b1f63512f72d3c84792a1f53ec40062a",
+			name:             "coco2014",
+			baseURL:          baseURLPrefix + "/coco2014",
+			labelMap:         labelMap,
+			completeLabelMap: completeLabelMap,
+			recordFileName:   "coco_val.record-00000-of-00001",
+			md5sum:           "b1f63512f72d3c84792a1f53ec40062a",
 		}
 
 		coco2017ValidationTFRecord = &CocoValidationTFRecord{
 			base: base{
 				ctx:            context.Background(),
 				baseWorkingDir: baseWorkingDir,
 			},
-			name:           "coco2017",
-			baseURL:        baseURLPrefix + "/coco2017",
-			recordFileName: "coco_val.record-00000-of-00001",
-			md5sum:         "b8a0cfed5ad569d4572b4ad8645acb5b",
+			name:             "coco2017",
+			baseURL:          baseURLPrefix + "/coco2017",
+			labelMap:         labelMap,
+			completeLabelMap: completeLabelMap,
+			recordFileName:   "coco_val.record-00000-of-00001",
+			md5sum:           "b8a0cfed5ad569d4572b4ad8645acb5b",
 		}
 
 		dldataset.Register(coco2014ValidationTFRecord)

vision/pascal.go

@@ -2,6 +2,7 @@ package vision
 
 import (
 	context "context"
+	"fmt"
 	"path"
 	"path/filepath"
 	"strings"
@@ -12,6 +13,7 @@ import (
 	"github.com/rai-project/dldataset"
 	"github.com/rai-project/dldataset/reader"
 	"github.com/rai-project/dldataset/reader/tfrecord"
+	"github.com/rai-project/dldataset/vision/support/object_detection"
 	"github.com/rai-project/dlframework"
 	"github.com/rai-project/dlframework/framework/feature"
 	"github.com/rai-project/downloadmanager"
@@ -40,6 +42,7 @@ type PascalValidationTFRecord struct {
 	baseURL        string
 	recordFileName string
 	md5sum         string
+	labelMap       object_detection.StringIntLabelMap
 	recordReader   *reader.TFRecordReader
 }
 
@@ -216,13 +219,19 @@ func init() {
 
 		const baseURLPrefix = "https://s3.amazonaws.com/store.carml.org/datasets"
 
+		labelMap, err := object_detection.Get("pascal_label_map.pbtxt")
+		if err != nil {
+			panic(fmt.Sprintf("failed to get pascal_label_map.pbtxt due to %v", err))
+		}
+
 		baseWorkingDir := filepath.Join(dldataset.Config.WorkingDirectory, "dldataset")
 		Pascal2007ValidationTFRecord = &PascalValidationTFRecord{
 			base: base{
 				ctx:            context.Background(),
 				baseWorkingDir: baseWorkingDir,
 			},
 			name:           "Pascal2007",
+			labelMap:       labelMap,
 			baseURL:        baseURLPrefix + "/pascal2007",
 			recordFileName: "validation.tfrecord",
 			md5sum:         "e646ecf0bf838fa39d34e58d87c3e914",
@@ -234,6 +243,7 @@ func init() {
 				baseWorkingDir: baseWorkingDir,
 			},
 			name:           "Pascal2012",
+			labelMap:       labelMap,
 			baseURL:        baseURLPrefix + "/pascal2012",
 			recordFileName: "validation.tfrecord",
 			md5sum:         "9a59d26492103b8635ba0c916d68535a",

vision/support/object_detection/label_map.go

@@ -0,0 +1,27 @@
+package object_detection
+
+func (s *StringIntLabelMap) Get(idx int) *StringIntLabelMapItem {
+	items := s.GetItem()
+	if idx < 0 {
+		return nil
+	}
+	if idx >= len(items) {
+		return nil
+	}
+	return items[idx]
+}
+
+func (s *StringIntLabelMap) GetName(idx int) string {
+	item := s.Get(idx)
+	return item.GetName()
+}
+
+func (s *StringIntLabelMap) GetId(idx int) int32 {
+	item := s.Get(idx)
+	return item.GetId()
+}
+
+func (s *StringIntLabelMap) GetDisplayName(idx int) string {
+	item := s.Get(idx)
+	return item.GetDisplayName()
+}