beenull/ente
1
0
Fork 0
Code Issues Pull requests Actions Packages Projects Releases Wiki Activity

configure @tensorflow-models/coco-ssd for testing

Browse source
This commit is contained in:
Abhinav 2022-02-23 14:40:46 +05:30
parent fa8a0ce620
commit 8b3a5d96f8
1 changed files with 6 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
src/services/machineLearning/ssdMobileNetV2Service.ts
View file

@ -12,7 +12,7 @@ import {
// SSDMobileNetV2Model, // SSDMobileNetV2Model,
// } from './modelWrapper/SSDMobileNetV2'; // } from './modelWrapper/SSDMobileNetV2';
import * as SSDMobileNet from 'ssd-mobilenet'; import * as SSDMobileNet from '@tensorflow-models/coco-ssd';
class SSDMobileNetV2 implements ObjectDetectionService { class SSDMobileNetV2 implements ObjectDetectionService {
private ssdMobileNetV2Model: SSDMobileNet.ObjectDetection; private ssdMobileNetV2Model: SSDMobileNet.ObjectDetection;
@ -27,9 +27,9 @@ class SSDMobileNetV2 implements ObjectDetectionService {
private async init() { private async init() {
this.ssdMobileNetV2Model = await SSDMobileNet.load({ this.ssdMobileNetV2Model = await SSDMobileNet.load({
base: 'mobilenet_v2', // base: 'mobilenet_v2',
dataset: 'open_images', // dataset: 'open_images',
modelUrl: '/models/open-images-ssd-mobilenet-v2/model.json', // modelUrl: '/models/open-images-ssd-mobilenet-v2/model.json',
}); });
console.log( console.log(
'loaded ssdMobileNetV2Model', 'loaded ssdMobileNetV2Model',
@ -78,7 +78,8 @@ class SSDMobileNetV2 implements ObjectDetectionService {
public async detectObjectUsingModel(imageBitmap: ImageBitmap) { public async detectObjectUsingModel(imageBitmap: ImageBitmap) {
const ssdMobileNetV2Model = await this.getSSDMobileNetV2Model(); const ssdMobileNetV2Model = await this.getSSDMobileNetV2Model();
const predictions = await ssdMobileNetV2Model.detect(imageBitmap); const tfImage = tf.browser.fromPixels(imageBitmap);
const predictions = await ssdMobileNetV2Model.detect(tfImage);
return predictions; return predictions;
} }