Create An API for Image Object Detection

Guvanch Nurgeldiyev · Guvanch Nurgeldiyev · commit 3635816dff60 · 2020-05-24T17:07:46.000+03:00
diff --git a/.env.sample b/.env.sample
@@ -0,0 +1,3 @@
+NODE_ENV=production
+HOST=localhost
+PORT=3000
diff --git a/.gitignore b/.gitignore
@@ -0,0 +1,3 @@
+coverage/
+node_modules/
+*.env
diff --git a/LICENSE b/LICENSE
@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2020 Guvanch Nurgeldiyev <guvanch@nurgeldiyev.com>
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
diff --git a/README.md b/README.md
@@ -0,0 +1,86 @@
+# Building a Object Detection API with Tensorflow.js & COCO SSD pretrained model
+
+Basic API for posting images and getting the result of detection.
+Uses COCO SSD pretrained model to detect objects in 300x300 px images. List of object [names](https://github.com/tensorflow/tfjs-models/blob/master/coco-ssd/src/classes.ts)
+
+## Getting Started
+
+These instructions will get you a copy of the project up and running on your local machine for development and testing purposes.
+
+### Prerequisites
+
+- Installed local [Node.js](https://nodejs.org/) environment
+- Package manager to install packages. [Yarn](https://yarnpkg.com/) or [npm](https://www.npmjs.com/)
+
+Clone the repo in your environment
+
+```bash
+git clone https://github.com/gnurgeldiyev/object-detection-api.git
+```
+
+### Installing
+
+```bash
+# move into project folder
+cd object-detection-api
+
+# install the dependencies
+yarn install
+
+# rename the .env.sample
+mv .env.sample .env
+
+# add your variables
+nano .env
+```
+
+#### Running the server
+
+```bash
+yarn start
+```
+
+## Example
+
+```bash
+# POST - /detection
+curl -H "Content-Type: application/octet-stream" --data-binary "@image/bicycle 300x300.jpg" "http://localhost:3000/detection"
+```
+
+##### Response
+
+```json
+{
+  "duration": 171,
+  "result": [
+    {
+      "bbox": [
+        5.610904097557068,
+        85.95118224620819,
+        295.6936866044998,
+        171.45759165287018
+      ],
+      "class": "bicycle",
+      "score": 0.9325430393218994
+    }
+  ]
+}
+```
+
+## Running the tests
+
+```bash
+yarn test
+```
+
+## Built With
+
+* [Express](https://github.com/expressjs/express/) - The web framework used
+* [Tensorflow.js](https://github.com/tensorflow/tfjs) - For using ML model
+* [COCO SSD](https://github.com/tensorflow/tfjs-models/tree/master/coco-ssd) - Pretrained Model for Object Detection
+* [Jest](https://github.com/facebook/jest) - For testing
+* [SuperTest](https://github.com/visionmedia/supertest) - For testing HTTP API
+
+## License
+
+This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details
diff --git a/config/index.js b/config/index.js
@@ -0,0 +1,12 @@
+const dotenv = require('dotenv')
+
+const result = dotenv.config()
+if (result.error) {
+  throw result.error
+}
+
+module.exports = {
+  NODE_ENV: process.env.NODE_ENV,
+  HOST: process.env.HOST,
+  PORT: process.env.PORT
+}
diff --git a/controller/index.js b/controller/index.js
@@ -0,0 +1,25 @@
+const Model = require('../util/model')
+const model = new Model()
+
+/**
+ * @returns prediction result & duration
+ */
+async function detectObjects(req, res) {
+  const start = new Date()
+  const rawImage = req.body
+  let prediction = []
+  try {
+    prediction = await model.detect(rawImage)
+  } catch(err) {
+    console.log(err)
+    return res.status(500).json({ message: 'Internal server error' })
+  }
+  return res.status(200).json({
+    duration: new Date() - start,
+    result: prediction
+  })
+}
+
+module.exports = {
+  detectObjects
+}
diff --git a/image/bicycle 300x300.jpg b/image/bicycle 300x300.jpg
diff --git a/image/door 300x300.png b/image/door 300x300.png
diff --git a/image/plant 300x300.jpeg b/image/plant 300x300.jpeg
diff --git a/index.js b/index.js
@@ -0,0 +1,7 @@
+const app = require('./server')
+const { HOST, PORT } = require('./config')
+
+ // run the server
+app.listen(PORT, HOST, () => {
+  console.log(`Server listening on http://${HOST}:${PORT}`)
+})
diff --git a/jest.config.js b/jest.config.js
@@ -0,0 +1,5 @@
+module.exports = {
+  testEnvironment: 'node',
+  collectCoverage: true,
+  testTimeout: 30000
+}
diff --git a/middleware/index.js b/middleware/index.js
@@ -0,0 +1,21 @@
+function validateNewDetection(req, res, next) {
+  const imageBuffer = req.body
+  // validate input data length
+  if (!imageBuffer.length) {
+    return res.status(400).json({ message: 'Bad request' })
+  }
+  // validate image format
+  const head = imageBuffer.slice(0, 2)
+  const tail = imageBuffer.slice(-2)
+  if (
+    !head.equals(Buffer.from('ffd8', 'hex'))
+    || !tail.equals(Buffer.from('ffd9', 'hex'))
+  ) {
+    return res.status(400).json({ message: 'Accepts only jpg/jpeg format' })
+  }
+  return next()
+}
+
+module.exports = {
+  validateNewDetection
+}
diff --git a/package.json b/package.json
@@ -0,0 +1,28 @@
+{
+  "name": "tensorflow-node",
+  "version": "1.0.0",
+  "main": "index.js",
+  "license": "MIT",
+  "scripts": {
+    "dev": "nodemon index.js",
+    "start": "node index.js",
+    "test": "jest"
+  },
+  "dependencies": {
+    "@tensorflow-models/coco-ssd": "^2.0.3",
+    "@tensorflow/tfjs": "^1.7.4",
+    "@tensorflow/tfjs-node": "^1.7.4",
+    "canvas": "^2.6.1",
+    "body-parser": "^1.19.0",
+    "compression": "^1.7.4",
+    "cors": "^2.8.5",
+    "dotenv": "^8.2.0",
+    "express": "^4.17.1",
+    "helmet": "^3.22.0"
+  },
+  "devDependencies": {
+    "jest": "^26.0.1",
+    "nodemon": "^2.0.4",
+    "supertest": "^4.0.2"
+  }
+}
diff --git a/route/index.js b/route/index.js
@@ -0,0 +1,7 @@
+const router = require('express').Router()
+const { detectObjects } = require('../controller')
+const { validateNewDetection } = require('../middleware')
+
+router.post('/detection', validateNewDetection, detectObjects)
+
+module.exports = router
diff --git a/server/index.js b/server/index.js
@@ -0,0 +1,37 @@
+const express = require('express')
+const cors = require('cors')
+const bodyParser = require('body-parser')
+const compression = require('compression')
+const helmet = require('helmet')
+
+// import API routes
+const routes = require('../route')
+
+const app = express()
+
+// CORS config
+app.use(
+  cors({
+    origin: '*',
+    optionsSuccessStatus: 200,
+    method: ['POST'],
+    allowedHeaders: ['Content-Type']
+  })
+)
+
+// security headers
+app.use(helmet())
+
+// request body parser
+app.use(bodyParser.json())
+app.use(bodyParser.raw({
+  limit: '2mb'
+}))
+
+// compress with gzip
+app.use(compression())
+
+// attach routes
+app.use('/', routes)
+
+module.exports = app
diff --git a/test/api.test.js b/test/api.test.js
@@ -0,0 +1,86 @@
+const request = require('supertest')
+const server = require('../server')
+const fs = require('fs')
+const { join } = require('path')
+
+const waitForModelLoad = (done) => {
+  setTimeout(() => {
+    done()
+  }, 20000)
+}
+
+describe('Test Object Detection API', () => {
+  let api
+  beforeAll(async (done) => {
+    api = await request(server)
+    await waitForModelLoad(done)
+  })
+
+  afterAll((done) => {
+    done()
+  })
+
+  test('Post empty body - shoudl fail with 400', async () => {
+    const expectedRes = {
+      statusCode: 400,
+      body: {
+        message: 'Bad request'
+      }
+    }
+    await api
+      .post('/detection')
+      .set('Content-Type', 'application/octet-stream')
+      .set('Accept', 'application/json')
+      .send()
+      .expect('Content-Type', /json/)
+      .expect(expectedRes.statusCode)
+      .then(res => {
+        expect(res.body).toEqual(expectedRes.body)
+      })
+  })
+
+  test('Post an unsupperted image format - should fail with 400', async () => {
+    const testImagePath = join(__dirname, '../image/door 300x300.png')
+    const expectedRes = {
+      statusCode: 400,
+      body: {
+        message: 'Accepts only jpg/jpeg format'
+      }
+    }
+    
+    await api
+      .post('/detection')
+      .set('Content-Type', 'application/octet-stream')
+      .set('Accept', 'application/json')
+      .send(fs.readFileSync(testImagePath))
+      .expect('Content-Type', /json/)
+      .expect(expectedRes.statusCode)
+      .then(res => {
+        expect(res.body).toEqual(expectedRes.body)
+      })
+  })
+
+  test('Post an image - should succeed with 200', async () => {
+    const testImagePath = join(__dirname, '../image/bicycle 300x300.jpg')
+    const expectedRes = {
+      statusCode: 200,
+      body: {
+        duration: 10,
+        resultClass: 'bicycle'
+      }
+    }
+    
+    await api
+      .post('/detection')
+      .set('Content-Type', 'application/octet-stream')
+      .set('Accept', 'application/json')
+      .send(fs.readFileSync(testImagePath))
+      .expect('Content-Type', /json/)
+      .expect(expectedRes.statusCode)
+      .then(res => {
+        expect(res.body.duration).toBeGreaterThan(expectedRes.body.duration)
+        expect(res.body.result).not.toHaveLength(0)
+        expect(res.body.result[0].class).toEqual(expectedRes.body.resultClass)
+      })
+  })
+})
diff --git a/util/model.js b/util/model.js
@@ -0,0 +1,25 @@
+const tf = require('@tensorflow/tfjs-node')
+const cocoSsd = require('@tensorflow-models/coco-ssd')
+
+class Model {
+  constructor() {
+    this.model = this.loadDetectionModel()
+  }
+
+  async loadDetectionModel() {
+    const start = new Date()
+    console.log('Loading model...')
+    this.model = await cocoSsd.load()
+    console.log(`🤖 Model is loaded ${(new Date() - start)/1000}s`)
+  }
+
+  async detect(imageBuffer) {
+    const arrByte = Uint8Array.from(imageBuffer)
+    const decodedImage = await tf.node.decodeJpeg(arrByte)
+    // make a prediction
+    const prediction = await this.model.detect(decodedImage)
+    return prediction
+  }
+}
+
+module.exports = Model
diff --git a/yarn.lock b/yarn.lock

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+NODE_ENV=production`
	`2`	`+HOST=localhost`
	`3`	`+PORT=3000`
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+coverage/`
	`2`	`+node_modules/`
	`3`	`+*.env`