mirror of
https://github.com/KevinMidboe/immich.git
synced 2025-12-08 20:29:05 +00:00
feat(machine-learning)!: move machine learning to Python based image (#1774)
BREAKING CHANGES * Users have to update the docker-compose file, machine-learning portion. * Temporary dropping machine-learning support for Arm64 and Armv7
This commit is contained in:
Alex
GitHub
@@ -1,10 +0,0 @@
|
||||
import { Module } from '@nestjs/common';
|
||||
import { ImageClassifierModule } from './image-classifier/image-classifier.module';
|
||||
import { ObjectDetectionModule } from './object-detection/object-detection.module';
|
||||
|
||||
@Module({
|
||||
imports: [ImageClassifierModule, ObjectDetectionModule],
|
||||
controllers: [],
|
||||
providers: [],
|
||||
})
|
||||
export class AppModule {}
|
||||
@@ -1,14 +0,0 @@
|
||||
import { Body, Controller, Post } from '@nestjs/common';
|
||||
import { ImageClassifierService } from './image-classifier.service';
|
||||
|
||||
@Controller('image-classifier')
|
||||
export class ImageClassifierController {
|
||||
constructor(
|
||||
private readonly imageClassifierService: ImageClassifierService,
|
||||
) { }
|
||||
|
||||
@Post('/tag-image')
|
||||
async tagImage(@Body('thumbnailPath') thumbnailPath: string) {
|
||||
return await this.imageClassifierService.tagImage(thumbnailPath);
|
||||
}
|
||||
}
|
||||
@@ -1,9 +0,0 @@
|
||||
import { Module } from '@nestjs/common';
|
||||
import { ImageClassifierService } from './image-classifier.service';
|
||||
import { ImageClassifierController } from './image-classifier.controller';
|
||||
|
||||
@Module({
|
||||
controllers: [ImageClassifierController],
|
||||
providers: [ImageClassifierService],
|
||||
})
|
||||
export class ImageClassifierModule {}
|
||||
@@ -1,49 +0,0 @@
|
||||
import { Injectable, Logger } from '@nestjs/common';
|
||||
import * as mobilenet from '@tensorflow-models/mobilenet';
|
||||
import * as cocoSsd from '@tensorflow-models/coco-ssd';
|
||||
import * as tf from '@tensorflow/tfjs-node';
|
||||
import * as fs from 'fs';
|
||||
|
||||
@Injectable()
|
||||
export class ImageClassifierService {
|
||||
private readonly MOBILENET_VERSION = 2;
|
||||
private readonly MOBILENET_ALPHA = 1.0;
|
||||
|
||||
private mobileNetModel: mobilenet.MobileNet;
|
||||
|
||||
constructor() {
|
||||
Logger.log(
|
||||
`Running Node TensorFlow Version : ${tf.version['tfjs']}`,
|
||||
'ImageClassifier',
|
||||
);
|
||||
mobilenet
|
||||
.load({
|
||||
version: this.MOBILENET_VERSION,
|
||||
alpha: this.MOBILENET_ALPHA,
|
||||
})
|
||||
.then((mobilenetModel) => (this.mobileNetModel = mobilenetModel));
|
||||
}
|
||||
|
||||
async tagImage(thumbnailPath: string) {
|
||||
try {
|
||||
const isExist = fs.existsSync(thumbnailPath);
|
||||
if (isExist) {
|
||||
const tags = [];
|
||||
const image = fs.readFileSync(thumbnailPath);
|
||||
const decodedImage = tf.node.decodeImage(image, 3) as tf.Tensor3D;
|
||||
const predictions = await this.mobileNetModel.classify(decodedImage);
|
||||
|
||||
for (const prediction of predictions) {
|
||||
if (prediction.probability >= 0.1) {
|
||||
tags.push(...prediction.className.split(',').map((e) => e.trim()));
|
||||
}
|
||||
}
|
||||
|
||||
tf.dispose(decodedImage);
|
||||
return tags;
|
||||
}
|
||||
} catch (e) {
|
||||
console.log('Error reading file ', e);
|
||||
}
|
||||
}
|
||||
}
|
||||
61
machine-learning/src/main.py
Normal file
61
machine-learning/src/main.py
Normal file
@@ -0,0 +1,61 @@
|
||||
import os
|
||||
from flask import Flask, request
|
||||
from transformers import pipeline
|
||||
|
||||
|
||||
server = Flask(__name__)
|
||||
|
||||
|
||||
classifier = pipeline(
|
||||
task="image-classification",
|
||||
model="microsoft/resnet-50"
|
||||
)
|
||||
|
||||
detector = pipeline(
|
||||
task="object-detection",
|
||||
model="hustvl/yolos-tiny"
|
||||
)
|
||||
|
||||
|
||||
# Environment resolver
|
||||
is_dev = os.getenv('NODE_ENV') == 'development'
|
||||
server_port = os.getenv('MACHINE_LEARNING_PORT') or 3003
|
||||
|
||||
|
||||
@server.route("/ping")
|
||||
def ping():
|
||||
return "pong"
|
||||
|
||||
|
||||
@server.route("/object-detection/detect-object", methods=['POST'])
|
||||
def object_detection():
|
||||
assetPath = request.json['thumbnailPath']
|
||||
return run_engine(detector, assetPath), 201
|
||||
|
||||
|
||||
@server.route("/image-classifier/tag-image", methods=['POST'])
|
||||
def image_classification():
|
||||
assetPath = request.json['thumbnailPath']
|
||||
return run_engine(classifier, assetPath), 201
|
||||
|
||||
|
||||
def run_engine(engine, path):
|
||||
result = []
|
||||
predictions = engine(path)
|
||||
|
||||
for index, pred in enumerate(predictions):
|
||||
tags = pred['label'].split(', ')
|
||||
if (index == 0):
|
||||
result = tags
|
||||
else:
|
||||
if (pred['score'] > 0.5):
|
||||
result = [*result, *tags]
|
||||
|
||||
if (len(result) > 1):
|
||||
result = list(set(result))
|
||||
|
||||
return result
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
server.run(debug=is_dev, host='0.0.0.0', port=server_port)
|
||||
@@ -1,27 +0,0 @@
|
||||
import { NestFactory } from '@nestjs/core';
|
||||
import { AppModule } from './app.module';
|
||||
import { Logger } from '@nestjs/common';
|
||||
|
||||
async function bootstrap() {
|
||||
const app = await NestFactory.create(AppModule);
|
||||
|
||||
const port = Number(process.env.MACHINE_LEARNING_PORT) || 3003;
|
||||
|
||||
await app.listen(port, () => {
|
||||
if (process.env.NODE_ENV == 'development') {
|
||||
Logger.log(
|
||||
'Running Immich Machine Learning in DEVELOPMENT environment',
|
||||
'IMMICH MICROSERVICES',
|
||||
);
|
||||
}
|
||||
|
||||
if (process.env.NODE_ENV == 'production') {
|
||||
Logger.log(
|
||||
'Running Immich Machine Learning in PRODUCTION environment',
|
||||
'IMMICH MICROSERVICES',
|
||||
);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
bootstrap();
|
||||
@@ -1,15 +0,0 @@
|
||||
import { Body, Controller, Post } from '@nestjs/common';
|
||||
import { ObjectDetectionService } from './object-detection.service';
|
||||
import { Logger } from '@nestjs/common';
|
||||
|
||||
@Controller('object-detection')
|
||||
export class ObjectDetectionController {
|
||||
constructor(
|
||||
private readonly objectDetectionService: ObjectDetectionService,
|
||||
) { }
|
||||
|
||||
@Post('/detect-object')
|
||||
async detectObject(@Body('thumbnailPath') thumbnailPath: string) {
|
||||
return await this.objectDetectionService.detectObject(thumbnailPath);
|
||||
}
|
||||
}
|
||||
@@ -1,9 +0,0 @@
|
||||
import { Module } from '@nestjs/common';
|
||||
import { ObjectDetectionService } from './object-detection.service';
|
||||
import { ObjectDetectionController } from './object-detection.controller';
|
||||
|
||||
@Module({
|
||||
controllers: [ObjectDetectionController],
|
||||
providers: [ObjectDetectionService],
|
||||
})
|
||||
export class ObjectDetectionModule {}
|
||||
@@ -1,39 +0,0 @@
|
||||
import { Injectable, Logger } from '@nestjs/common';
|
||||
import * as cocoSsd from '@tensorflow-models/coco-ssd';
|
||||
import * as tf from '@tensorflow/tfjs-node';
|
||||
import * as fs from 'fs';
|
||||
|
||||
@Injectable()
|
||||
export class ObjectDetectionService {
|
||||
private cocoSsdModel: cocoSsd.ObjectDetection;
|
||||
|
||||
constructor() {
|
||||
Logger.log(
|
||||
`Running Node TensorFlow Version : ${tf.version['tfjs']}`,
|
||||
'ObjectDetection',
|
||||
);
|
||||
cocoSsd.load().then((model) => (this.cocoSsdModel = model));
|
||||
}
|
||||
async detectObject(thumbnailPath: string) {
|
||||
try {
|
||||
const isExist = fs.existsSync(thumbnailPath);
|
||||
if (isExist) {
|
||||
const tags = new Set();
|
||||
const image = fs.readFileSync(thumbnailPath);
|
||||
const decodedImage = tf.node.decodeImage(image, 3) as tf.Tensor3D;
|
||||
const predictions = await this.cocoSsdModel.detect(decodedImage);
|
||||
|
||||
for (const result of predictions) {
|
||||
if (result.score > 0.5) {
|
||||
tags.add(result.class);
|
||||
}
|
||||
}
|
||||
|
||||
tf.dispose(decodedImage);
|
||||
return [...tags];
|
||||
}
|
||||
} catch (e) {
|
||||
console.log('Error reading file ', e);
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user