Skip to content

Commit 4c89ff0

Browse files
CorentinBCorentinB
committed
Add: ResNet-50 support through the -n / --network flag.
1 parent df16943 commit 4c89ff0

File tree

6 files changed

+98
-17
lines changed

6 files changed

+98
-17
lines changed

‎README.md‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ chmod +x deepsort
3939

4040
# Usage
4141

42-
Right now, DeepSort doesnt support a lot of different parameters, you're obliged to fill two of them:
42+
DeepSort support few different parameters, you're obliged to fill two of them:
4343
`--url` or `-u` that correspond to the URL of your DeepDetect server.
4444
`--input` or `-i` that correspond to your local folder full of images.
4545

@@ -49,7 +49,8 @@ For more informations, refeer to the helper:
4949
5050
[-u|--url] is required
5151
usage: deepsort [-h|--help] -u|--url "<value>" -i|--input "<value>"
52-
[-R|--recursive] [-j|--jobs <integer>] [-d|--dry-run]
52+
[-n|--network (resnet-50|googlenet)] [-R|--recursive]
53+
[-j|--jobs <integer>] [-d|--dry-run]
5354
5455
AI powered image tagger backed by DeepDetect
5556
@@ -58,6 +59,8 @@ Arguments:
5859
-h --help Print help information
5960
-u --url URL of your DeepDetect instance (i.e: http://localhost:8080)
6061
-i --input Your input folder.
62+
-n --network The pre-trained deep neural network you want to use, can be
63+
resnet-50 or googlenet. Default: resnet-50
6164
-R --recursive Process files recursively.
6265
-j --jobs Number of parallel jobs. Default: 1
6366
-d --dry-run Just classify images and return results, do not apply.

‎cmd/deepsort/arguments.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ func argumentParsing(args []string, argument *Arguments) {
1414
// Create flags
1515
URL := parser.String("u", "url", &argparse.Options{Required: true, Help: "URL of your DeepDetect instance (i.e: http://localhost:8080)"})
1616
input := parser.String("i", "input", &argparse.Options{Required: true, Help: "Your input folder."})
17+
network := parser.Selector("n", "network", []string{"resnet-50", "googlenet"}, &argparse.Options{Required: false, Help: "The pre-trained deep neural network you want to use, can be resnet-50 or googlenet", Default: "resnet-50"})
1718
recursive := parser.Flag("R", "recursive", &argparse.Options{Required: false, Help: "Process files recursively."})
1819
jobs := parser.Int("j", "jobs", &argparse.Options{Required: false, Help: "Number of parallel jobs", Default: 1})
1920
dryRun := parser.Flag("d", "dry-run", &argparse.Options{Required: false, Help: "Just classify images and return results, do not apply."})
@@ -28,6 +29,7 @@ func argumentParsing(args []string, argument *Arguments) {
2829
// Handle the input flag
2930
inputFolder, _ := filepath.Abs(*input)
3031
// Finally save the collected flags
32+
argument.Network = *network
3133
argument.Jobs = *jobs
3234
argument.DryRun = *dryRun
3335
argument.Recursive = *recursive

‎cmd/deepsort/classify.go‎

Lines changed: 26 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,32 @@ func googleNetClassification(path string, arguments *Arguments, wg *sync.WaitGro
1616
defer wg.Done()
1717
url := arguments.URL + "/predict"
1818
path, _ = filepath.Abs(path)
19-
var jsonStr = []byte(`{"service":"deepsort-googlenet","parameters":{"input":{"width":224,"height":224},"output":{"best":1},"mllib":{"gpu":false}},"data":["` + path + `"]}`)
19+
var jsonStr = []byte(`{"service":"deepsort-resnet","parameters":{"input":{"width":224,"height":224},"output":{"best":1},"mllib":{"gpu":false}},"data":["` + path + `"]}`)
20+
// DEBUG
21+
//fmt.Println("Request: " + string(jsonStr))
22+
req, err := http.NewRequest("POST", url, bytes.NewBuffer(jsonStr))
23+
req.Close = true
24+
client := &http.Client{}
25+
resp, err := client.Do(req)
26+
if err != nil {
27+
logging.Error("Unable to classify this file.", "["+filepath.Base(path)+"]")
28+
os.Exit(1)
29+
}
30+
defer resp.Body.Close()
31+
32+
body, _ := ioutil.ReadAll(resp.Body)
33+
parsedResponse := parseResponse(body)
34+
color.Println(color.Yellow("[") + color.Cyan(filepath.Base(path)) + color.Yellow("]") + color.Yellow(" Response: ") + color.Green(parsedResponse))
35+
if arguments.DryRun != true {
36+
renameFile(path, arguments, parsedResponse)
37+
}
38+
}
39+
40+
func resNet50Classification(path string, arguments *Arguments, wg *sync.WaitGroup) {
41+
defer wg.Done()
42+
url := arguments.URL + "/predict"
43+
path, _ = filepath.Abs(path)
44+
var jsonStr = []byte(`{"service":"deepsort-resnet-50","parameters":{"input":{"width":224,"height":224},"output":{"best":1},"mllib":{"gpu":false}},"data":["` + path + `"]}`)
2045
// DEBUG
2146
//fmt.Println("Request: " + string(jsonStr))
2247
req, err := http.NewRequest("POST", url, bytes.NewBuffer(jsonStr))

‎cmd/deepsort/deepdetect.go‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,3 +31,27 @@ func startGoogleNet(arguments *Arguments) {
3131
logging.Success("Successfully started the image classification service.", "[GoogleNet]")
3232
}
3333
}
34+
35+
func startResNet50(arguments *Arguments) {
36+
// Starting the image classification service
37+
logging.Success("Starting the classification service..", "[ResNet-50]")
38+
url := arguments.URL + "/services/deepsort-resnet-50"
39+
var jsonStr = []byte(`{"mllib":"caffe","description":"DeepSort-ResNet-50","type":"supervised","parameters":{"input":{"connector":"image"},"mllib":{"nclasses":1000}},"model":{"repository":"/opt/models/resnet_50/"}}`)
40+
req, err := http.NewRequest("PUT", url, bytes.NewBuffer(jsonStr))
41+
client := &http.Client{}
42+
resp, err := client.Do(req)
43+
if err != nil {
44+
logging.Error("Error while starting the classification service, please check if DeepDetect is running.", "[ResNet-50]")
45+
os.Exit(1)
46+
}
47+
defer resp.Body.Close()
48+
if resp.Status != "201 Created" && resp.Status != "500 Internal Server Error" {
49+
logging.Error("Error while starting the classification service, please check if DeepDetect is running.", "[ResNet-50]")
50+
os.Exit(1)
51+
}
52+
if resp.Status == "500 Internal Server Error" {
53+
logging.Success("Looks like you already have the deepsort-googlenet service started, no need to create a new one.", "[ResNet-50]")
54+
} else {
55+
logging.Success("Successfully started the image classification service.", "[ResNet-50]")
56+
}
57+
}

‎cmd/deepsort/main.go‎

Lines changed: 30 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -14,28 +14,48 @@ type Arguments struct {
1414
DryRun bool
1515
Recursive bool
1616
Jobs int
17+
Network string
1718
}
1819

1920
func main() {
2021
start := time.Now()
2122
arguments := new(Arguments)
2223
arguments.Jobs = 1
2324
argumentParsing(os.Args, arguments)
24-
startGoogleNet(arguments)
25-
if arguments.Recursive == true {
26-
if arguments.Jobs == 1 {
27-
logging.Success("Starting image classification recursively..", "[GoogleNet]")
25+
if arguments.Network == "resnet-50" {
26+
startResNet50(arguments)
27+
if arguments.Recursive == true {
28+
if arguments.Jobs == 1 {
29+
logging.Success("Starting image classification recursively..", "[ResNet-50]")
30+
} else {
31+
logging.Success("Starting image classification recursively with "+string(arguments.Jobs)+" parallel jobs..", "[ResNet-50]")
32+
}
33+
runRecursively(arguments)
2834
} else {
29-
logging.Success("Starting image classification recursively with "+string(arguments.Jobs)+" parallel jobs..", "[GoogleNet]")
35+
if arguments.Jobs == 1 {
36+
logging.Success("Starting image classification..", "[ResNet-50]")
37+
} else {
38+
logging.Success("Starting image classification with "+string(arguments.Jobs)+" parallel jobs..", "[ResNet-50]")
39+
}
40+
run(arguments)
3041
}
31-
runRecursively(arguments)
3242
} else {
33-
if arguments.Jobs == 1 {
34-
logging.Success("Starting image classification..", "[GoogleNet]")
43+
startGoogleNet(arguments)
44+
if arguments.Recursive == true {
45+
if arguments.Jobs == 1 {
46+
logging.Success("Starting image classification recursively..", "[GoogleNet]")
47+
} else {
48+
logging.Success("Starting image classification recursively with "+string(arguments.Jobs)+" parallel jobs..", "[GoogleNet]")
49+
}
50+
runRecursively(arguments)
3551
} else {
36-
logging.Success("Starting image classification with "+string(arguments.Jobs)+" parallel jobs..", "[GoogleNet]")
52+
if arguments.Jobs == 1 {
53+
logging.Success("Starting image classification..", "[GoogleNet]")
54+
} else {
55+
logging.Success("Starting image classification with "+string(arguments.Jobs)+" parallel jobs..", "[GoogleNet]")
56+
}
57+
run(arguments)
3758
}
38-
run(arguments)
3959
}
4060
color.Println(color.Cyan("Done in ") + color.Yellow(time.Since(start)) + color.Cyan("!"))
4161
}

‎cmd/deepsort/process.go‎

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,11 @@ func run(arguments *Arguments) {
3434
if filetype.IsImage(buf) {
3535
count++
3636
wg.Add(1)
37-
go googleNetClassification(path, arguments, &wg)
37+
if arguments.Network == "resnet-50" {
38+
go resNet50Classification(path, arguments, &wg)
39+
} else {
40+
go googleNetClassification(path, arguments, &wg)
41+
}
3842
if count == arguments.Jobs {
3943
wg.Wait()
4044
count = 0
@@ -53,24 +57,27 @@ func runRecursively(arguments *Arguments) ([]string, error) {
5357
fileList = append(fileList, path)
5458
return err
5559
})
56-
5760
if e != nil {
5861
logging.Error("Unable to process this directory.", "["+arguments.Input+"]")
5962
os.Exit(1)
6063
}
6164

65+
// Process files in the input folder
6266
for _, file := range fileList {
6367
buf, _ := ioutil.ReadFile(file)
6468
if filetype.IsImage(buf) {
6569
count++
6670
wg.Add(1)
67-
go googleNetClassification(file, arguments, &wg)
71+
if arguments.Network == "resnet-50" {
72+
go resNet50Classification(file, arguments, &wg)
73+
} else {
74+
go googleNetClassification(file, arguments, &wg)
75+
}
6876
}
6977
if count == arguments.Jobs {
7078
wg.Wait()
7179
count = 0
7280
}
7381
}
74-
7582
return fileList, nil
7683
}

0 commit comments

Comments
 (0)