From 2480bf49eadff7a4cb023e1ca1a709cca8bed5a4 Mon Sep 17 00:00:00 2001
From: yuyuanshifu <747342561@qq.com>
Date: Tue, 16 Mar 2021 20:28:05 +0800
Subject: [PATCH 1/2] add benchmark category

---
 models/cloudbrain.go               | 9 +++++++++
 modules/setting/setting.go         | 2 ++
 routers/repo/cloudbrain.go         | 5 ++++-
 templates/repo/cloudbrain/new.tmpl | 9 +++++++++
 4 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/models/cloudbrain.go b/models/cloudbrain.go
index 9784a58b5..f053552db 100755
--- a/models/cloudbrain.go
+++ b/models/cloudbrain.go
@@ -265,6 +265,15 @@ type ImageInfo struct {
 	PlaceView   string
 }
 
+type Categories struct {
+	Category []*Category	`json:"category"`
+}
+
+type Category struct {
+	Id int	`json:"id"`
+	Value string `json:"value"`
+}
+
 type CommitImageParams struct {
 	Ip               string `json:"ip"`
 	TaskContainerId  string `json:"taskContainerId"`
diff --git a/modules/setting/setting.go b/modules/setting/setting.go
index 65fddc7a0..9926c154e 100755
--- a/modules/setting/setting.go
+++ b/modules/setting/setting.go
@@ -443,6 +443,7 @@ var (
 	IsBenchmarkEnabled  bool
 	BenchmarkCode       string
 	BenchmarkServerHost string
+	BenchmarkCategory string
 
 	//snn4imagenet config
 	IsSnn4imagenetEnabled  bool
@@ -1150,6 +1151,7 @@ func NewContext() {
 	IsBenchmarkEnabled = sec.Key("ENABLED").MustBool(false)
 	BenchmarkCode = sec.Key("BENCHMARKCODE").MustString("https://yangzhx:justfortest123@git.openi.org.cn/yangzhx/detection_benchmark_script.git")
 	BenchmarkServerHost = sec.Key("HOST").MustString("http://192.168.202.90:3366/")
+	BenchmarkCategory = sec.Key("CATEGORY").MustString("")
 
 	sec = Cfg.Section("snn4imagenet")
 	IsSnn4imagenetEnabled = sec.Key("ENABLED").MustBool(false)
diff --git a/routers/repo/cloudbrain.go b/routers/repo/cloudbrain.go
index 30542e26e..c710f8c5b 100755
--- a/routers/repo/cloudbrain.go
+++ b/routers/repo/cloudbrain.go
@@ -108,7 +108,6 @@ func CloudBrainNew(ctx *context.Context) {
 	}
 
 	for i, payload := range resultPublic.Payload.ImageInfo {
-		log.Info(resultPublic.Payload.ImageInfo[i].Place)
 		if strings.HasPrefix(resultPublic.Payload.ImageInfo[i].Place, "192.168") {
 			resultPublic.Payload.ImageInfo[i].PlaceView = payload.Place[strings.Index(payload.Place, "/"):len(payload.Place)]
 		} else {
@@ -131,6 +130,10 @@ func CloudBrainNew(ctx *context.Context) {
 	ctx.Data["model_path"] = cloudbrain.ModelMountPath
 	ctx.Data["benchmark_path"] = cloudbrain.BenchMarkMountPath
 	ctx.Data["is_benchmark_enabled"] = setting.IsBenchmarkEnabled
+
+	var categories *models.Categories
+	json.Unmarshal([]byte(setting.BenchmarkCategory), &categories)
+	ctx.Data["benchmark_categories"] = categories.Category
 	ctx.Data["snn4imagenet_path"] = cloudbrain.Snn4imagenetMountPath
 	ctx.Data["is_snn4imagenet_enabled"] = setting.IsSnn4imagenetEnabled
 	ctx.HTML(200, tplCloudBrainNew)
diff --git a/templates/repo/cloudbrain/new.tmpl b/templates/repo/cloudbrain/new.tmpl
index ddbe087d3..28e33cec2 100755
--- a/templates/repo/cloudbrain/new.tmpl
+++ b/templates/repo/cloudbrain/new.tmpl
@@ -125,6 +125,15 @@
                         </select>
                     </div>
 
+                    <div class="inline required field cloudbrain_benchmark">
+                        <label>数据集类别</label>
+                        <select class="ui search dropdown" id="cloudbrain_benchmark_category" placeholder="选择数据集类别" style='width:385px' name="benchmark_category">
+                            {{range .benchmark_categories}}
+                                <option name="category" value="{{.Value}}">{{.Value}}</option>
+                            {{end}}
+                        </select>
+                    </div>
+
                     <div class="inline required field">
                         <label>镜像</label>
                         <select class="ui search dropdown" id="cloudbrain_image" placeholder="选择镜像" style='width:385px' name="image">

From 9f1592f4d4734a678a9e477960ad7f96f8d047ae Mon Sep 17 00:00:00 2001
From: yuyuanshifu <747342561@qq.com>
Date: Wed, 17 Mar 2021 10:52:29 +0800
Subject: [PATCH 2/2] benchmark_category

---
 models/cloudbrain.go               | 1 +
 modules/auth/cloudbrain.go         | 1 +
 routers/repo/cloudbrain.go         | 7 ++++---
 templates/repo/cloudbrain/new.tmpl | 2 +-
 4 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/models/cloudbrain.go b/models/cloudbrain.go
index f053552db..d9db9ef23 100755
--- a/models/cloudbrain.go
+++ b/models/cloudbrain.go
@@ -165,6 +165,7 @@ type TaskInfo struct {
 	Username string `json:"username"`
 	TaskName string `json:"task_name"`
 	CodeName string `json:"code_name"`
+	BenchmarkCategory string `json:"selected_category"`
 }
 
 func ConvertToTaskPod(input map[string]interface{}) (TaskPod, error) {
diff --git a/modules/auth/cloudbrain.go b/modules/auth/cloudbrain.go
index 2470c2ad6..e7268052d 100755
--- a/modules/auth/cloudbrain.go
+++ b/modules/auth/cloudbrain.go
@@ -11,6 +11,7 @@ type CreateCloudBrainForm struct {
 	Command    string `form:"command" binding:"Required"`
 	Attachment string `form:"attachment" binding:"Required"`
 	JobType    string `form:"job_type" binding:"Required"`
+	BenchmarkCategory    string `form:"benchmark_category"`
 }
 
 type CommitImageCloudBrainForm struct {
diff --git a/routers/repo/cloudbrain.go b/routers/repo/cloudbrain.go
index c710f8c5b..8b072e1e6 100755
--- a/routers/repo/cloudbrain.go
+++ b/routers/repo/cloudbrain.go
@@ -165,12 +165,12 @@ func CloudBrainCreate(ctx *context.Context, form auth.CreateCloudBrainForm) {
 
 	benchmarkPath := setting.JobPath + jobName + cloudbrain.BenchMarkMountPath
 	if setting.IsBenchmarkEnabled && jobType == string(models.JobTypeBenchmark) {
-		downloadRateCode(repo, jobName, setting.BenchmarkCode, benchmarkPath)
+		downloadRateCode(repo, jobName, setting.BenchmarkCode, benchmarkPath, form.BenchmarkCategory)
 	}
 
 	snn4imagenetPath := setting.JobPath + jobName + cloudbrain.Snn4imagenetMountPath
 	if setting.IsSnn4imagenetEnabled && jobType == string(models.JobTypeSnn4imagenet) {
-		downloadRateCode(repo, jobName, setting.Snn4imagenetCode, snn4imagenetPath)
+		downloadRateCode(repo, jobName, setting.Snn4imagenetCode, snn4imagenetPath, "")
 	}
 
 	err = cloudbrain.GenerateTask(ctx, jobName, image, command, uuid, codePath, modelPath, benchmarkPath, snn4imagenetPath, jobType)
@@ -340,7 +340,7 @@ func downloadCode(repo *models.Repository, codePath string) error {
 	return nil
 }
 
-func downloadRateCode(repo *models.Repository, taskName, gitPath, codePath string) error {
+func downloadRateCode(repo *models.Repository, taskName, gitPath, codePath, benchmarkCategory string) error {
 	err := os.MkdirAll(codePath, os.ModePerm)
 	if err != nil {
 		log.Error("mkdir codePath failed", err.Error())
@@ -369,6 +369,7 @@ func downloadRateCode(repo *models.Repository, taskName, gitPath, codePath strin
 		Username: repo.Owner.Name,
 		TaskName: taskName,
 		CodeName: repo.Name,
+		BenchmarkCategory: benchmarkCategory,
 	})
 	if err != nil {
 		log.Error("json.Marshal failed", err.Error())
diff --git a/templates/repo/cloudbrain/new.tmpl b/templates/repo/cloudbrain/new.tmpl
index 28e33cec2..2c4832736 100755
--- a/templates/repo/cloudbrain/new.tmpl
+++ b/templates/repo/cloudbrain/new.tmpl
@@ -129,7 +129,7 @@
                         <label>数据集类别</label>
                         <select class="ui search dropdown" id="cloudbrain_benchmark_category" placeholder="选择数据集类别" style='width:385px' name="benchmark_category">
                             {{range .benchmark_categories}}
-                                <option name="category" value="{{.Value}}">{{.Value}}</option>
+                                <option name="benchmark_category" value="{{.Value}}">{{.Value}}</option>
                             {{end}}
                         </select>
                     </div>