feat: call ML service for image classification
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 4.7
parent
9a642a8581
commit
1689f3f44e
@@ -0,0 +1,81 @@
|
||||
import { afterEach, describe, expect, test } from 'bun:test';
|
||||
import { classifyImage } from './image-model';
|
||||
|
||||
const originalFetch = globalThis.fetch;
|
||||
|
||||
afterEach(() => {
|
||||
globalThis.fetch = originalFetch;
|
||||
});
|
||||
|
||||
function makeImageFile(type = 'image/jpeg') {
|
||||
return new File([new Uint8Array([1, 2, 3])], 'leaf.jpg', { type });
|
||||
}
|
||||
|
||||
function mockFetch(handler: (input: Parameters<typeof fetch>[0], init?: Parameters<typeof fetch>[1]) => Promise<Response>) {
|
||||
globalThis.fetch = Object.assign(handler, { preconnect: originalFetch.preconnect });
|
||||
}
|
||||
|
||||
describe('classifyImage', () => {
|
||||
test('maps ML service prediction response to API classification result', async () => {
|
||||
mockFetch(async (input, init) => {
|
||||
expect(String(input)).toBe('http://127.0.0.1:8001/predict');
|
||||
expect(init?.method).toBe('POST');
|
||||
expect(init?.body).toBeInstanceOf(FormData);
|
||||
|
||||
return new Response(
|
||||
JSON.stringify({
|
||||
label: 'Daun Sehat',
|
||||
confidence: 0.92,
|
||||
probabilities: {
|
||||
'Bercak Daun': 0.02,
|
||||
'Daun Sehat': 0.92,
|
||||
'Karat Daun': 0.03,
|
||||
'Hawar Daun': 0.03,
|
||||
},
|
||||
}),
|
||||
{ status: 200, headers: { 'content-type': 'application/json' } },
|
||||
);
|
||||
});
|
||||
|
||||
const result = await classifyImage(makeImageFile());
|
||||
|
||||
expect(result.predictedDiseaseSlug).toBe('daun-sehat');
|
||||
expect(result.confidence).toBe(0.92);
|
||||
expect(result.probabilities).toEqual([
|
||||
{ diseaseSlug: 'daun-sehat', label: 'Daun Sehat', confidence: 0.92 },
|
||||
{ diseaseSlug: 'karat-daun', label: 'Karat Daun', confidence: 0.03 },
|
||||
{ diseaseSlug: 'hawar-daun', label: 'Hawar Daun', confidence: 0.03 },
|
||||
{ diseaseSlug: 'bercak-daun', label: 'Bercak Daun', confidence: 0.02 },
|
||||
]);
|
||||
});
|
||||
|
||||
test('rejects unsupported file types before calling ML service', async () => {
|
||||
let called = false;
|
||||
mockFetch(async () => {
|
||||
called = true;
|
||||
return new Response('{}');
|
||||
});
|
||||
|
||||
await expect(classifyImage(makeImageFile('image/webp'))).rejects.toThrow('File must be JPEG or PNG');
|
||||
expect(called).toBe(false);
|
||||
});
|
||||
|
||||
test('throws when ML service returns a non-success response', async () => {
|
||||
mockFetch(async () => new Response(JSON.stringify({ detail: 'Model is not loaded' }), { status: 503 }));
|
||||
|
||||
await expect(classifyImage(makeImageFile())).rejects.toThrow('ML service returned 503: Model is not loaded');
|
||||
});
|
||||
|
||||
test('throws when ML service returns an unknown label', async () => {
|
||||
mockFetch(async () => new Response(
|
||||
JSON.stringify({
|
||||
label: 'Unknown Disease',
|
||||
confidence: 0.7,
|
||||
probabilities: { 'Unknown Disease': 0.7 },
|
||||
}),
|
||||
{ status: 200, headers: { 'content-type': 'application/json' } },
|
||||
));
|
||||
|
||||
await expect(classifyImage(makeImageFile())).rejects.toThrow('Unknown ML service label: Unknown Disease');
|
||||
});
|
||||
});
|
||||
+83
-108
@@ -1,9 +1,5 @@
|
||||
import * as tf from '@tensorflow/tfjs';
|
||||
import * as jpeg from 'jpeg-js';
|
||||
import { PNG } from 'pngjs';
|
||||
import { existsSync } from 'fs';
|
||||
import { resolve } from 'path';
|
||||
import type { DiseaseSlug, DiseaseLabel, PredictionProbability } from '@zeavis/shared';
|
||||
import { env } from '../config/env';
|
||||
|
||||
const DISEASE_CLASSES: Array<{ slug: DiseaseSlug; label: DiseaseLabel }> = [
|
||||
{ slug: 'bercak-daun', label: 'Bercak Daun' },
|
||||
@@ -12,38 +8,15 @@ const DISEASE_CLASSES: Array<{ slug: DiseaseSlug; label: DiseaseLabel }> = [
|
||||
{ slug: 'hawar-daun', label: 'Hawar Daun' },
|
||||
];
|
||||
|
||||
function resolveModelPath() {
|
||||
const candidates = [
|
||||
resolve(process.cwd(), 'Machine_Learning/model/tfjs_model/model.json'),
|
||||
resolve(process.cwd(), '../../Machine_Learning/model/tfjs_model/model.json'),
|
||||
];
|
||||
const DISEASE_BY_LABEL = new Map<DiseaseLabel, { slug: DiseaseSlug; label: DiseaseLabel }>(
|
||||
DISEASE_CLASSES.map((disease) => [disease.label, disease]),
|
||||
);
|
||||
|
||||
const modelPath = candidates.find((candidate) => existsSync(candidate));
|
||||
if (!modelPath) {
|
||||
throw new Error('TFJS model file was not found');
|
||||
}
|
||||
|
||||
return modelPath;
|
||||
}
|
||||
|
||||
let modelPromise: Promise<tf.GraphModel> | null = null;
|
||||
|
||||
async function loadModel(): Promise<tf.GraphModel> {
|
||||
if (modelPromise) {
|
||||
return modelPromise;
|
||||
}
|
||||
|
||||
modelPromise = (async () => {
|
||||
try {
|
||||
const fileUrl = `file://${resolveModelPath()}`;
|
||||
return await tf.loadGraphModel(fileUrl);
|
||||
} catch (error) {
|
||||
throw new Error(`Failed to load TFJS model: ${error instanceof Error ? error.message : String(error)}`);
|
||||
}
|
||||
})();
|
||||
|
||||
return modelPromise;
|
||||
}
|
||||
type MlPredictionResponse = {
|
||||
label: unknown;
|
||||
confidence: unknown;
|
||||
probabilities: unknown;
|
||||
};
|
||||
|
||||
export type ClassificationResult = {
|
||||
predictedDiseaseSlug: DiseaseSlug;
|
||||
@@ -51,84 +24,86 @@ export type ClassificationResult = {
|
||||
probabilities: PredictionProbability[];
|
||||
};
|
||||
|
||||
function predictUrl() {
|
||||
return `${env.mlServiceUrl.replace(/\/+$/, '')}/predict`;
|
||||
}
|
||||
|
||||
function assertKnownLabel(label: unknown): DiseaseLabel {
|
||||
if (typeof label !== 'string' || !DISEASE_BY_LABEL.has(label as DiseaseLabel)) {
|
||||
throw new Error(`Unknown ML service label: ${String(label)}`);
|
||||
}
|
||||
|
||||
return label as DiseaseLabel;
|
||||
}
|
||||
|
||||
function assertConfidence(value: unknown, label: string): number {
|
||||
if (typeof value !== 'number' || !Number.isFinite(value)) {
|
||||
throw new Error(`Invalid ML service confidence for ${label}`);
|
||||
}
|
||||
|
||||
return Math.max(0, Math.min(1, value));
|
||||
}
|
||||
|
||||
function mapProbabilities(probabilities: unknown): PredictionProbability[] {
|
||||
if (!probabilities || typeof probabilities !== 'object' || Array.isArray(probabilities)) {
|
||||
throw new Error('Invalid ML service probabilities');
|
||||
}
|
||||
|
||||
return Object.entries(probabilities).map(([label, confidence]) => {
|
||||
const knownLabel = assertKnownLabel(label);
|
||||
const disease = DISEASE_BY_LABEL.get(knownLabel)!;
|
||||
|
||||
return {
|
||||
diseaseSlug: disease.slug,
|
||||
label: disease.label,
|
||||
confidence: assertConfidence(confidence, disease.label),
|
||||
};
|
||||
}).sort((a, b) => b.confidence - a.confidence);
|
||||
}
|
||||
|
||||
async function parseErrorResponse(response: Response): Promise<string> {
|
||||
try {
|
||||
const body = await response.json();
|
||||
if (body && typeof body === 'object' && 'detail' in body) {
|
||||
return String(body.detail);
|
||||
}
|
||||
} catch {
|
||||
return response.statusText || 'Unknown error';
|
||||
}
|
||||
|
||||
return response.statusText || 'Unknown error';
|
||||
}
|
||||
|
||||
export async function classifyImage(file: File): Promise<ClassificationResult> {
|
||||
if (file.type !== 'image/jpeg' && file.type !== 'image/png') {
|
||||
throw new Error('File must be JPEG or PNG');
|
||||
}
|
||||
|
||||
const buffer = await file.arrayBuffer();
|
||||
const uint8Array = new Uint8Array(buffer);
|
||||
|
||||
let imageData: { data: Uint8Array; width: number; height: number };
|
||||
|
||||
if (file.type === 'image/jpeg') {
|
||||
const decoded = jpeg.decode(uint8Array, { useTArray: true });
|
||||
imageData = {
|
||||
data: decoded.data,
|
||||
width: decoded.width,
|
||||
height: decoded.height,
|
||||
};
|
||||
} else {
|
||||
const png = new PNG();
|
||||
await new Promise<void>((resolve, reject) => {
|
||||
png.parse(Buffer.from(uint8Array), (err: Error | null) => {
|
||||
if (err) reject(err);
|
||||
else resolve();
|
||||
});
|
||||
});
|
||||
imageData = {
|
||||
data: png.data,
|
||||
width: png.width,
|
||||
height: png.height,
|
||||
};
|
||||
}
|
||||
|
||||
const imageTensor = tf.tidy(() => {
|
||||
const rgb = new Uint8Array(imageData.width * imageData.height * 3);
|
||||
for (let source = 0, target = 0; source < imageData.data.length; source += 4, target += 3) {
|
||||
rgb[target] = imageData.data[source];
|
||||
rgb[target + 1] = imageData.data[source + 1];
|
||||
rgb[target + 2] = imageData.data[source + 2];
|
||||
}
|
||||
|
||||
return tf
|
||||
.tensor3d(rgb, [imageData.height, imageData.width, 3], 'int32')
|
||||
.resizeBilinear([224, 224])
|
||||
.toFloat()
|
||||
.expandDims(0);
|
||||
});
|
||||
const formData = new FormData();
|
||||
formData.append('file', file, file.name || 'leaf-image');
|
||||
|
||||
let response: Response;
|
||||
try {
|
||||
const model = await loadModel();
|
||||
const predictions = model.predict(imageTensor) as tf.Tensor;
|
||||
|
||||
try {
|
||||
const scoresArray = await predictions.data();
|
||||
|
||||
let maxScore = -Infinity;
|
||||
let maxIndex = 0;
|
||||
for (let i = 0; i < scoresArray.length; i++) {
|
||||
if (scoresArray[i] > maxScore) {
|
||||
maxScore = scoresArray[i];
|
||||
maxIndex = i;
|
||||
}
|
||||
}
|
||||
|
||||
const probabilities: PredictionProbability[] = DISEASE_CLASSES.map((disease, index) => ({
|
||||
diseaseSlug: disease.slug,
|
||||
label: disease.label,
|
||||
confidence: Math.max(0, Math.min(1, scoresArray[index])),
|
||||
})).sort((a, b) => b.confidence - a.confidence);
|
||||
|
||||
return {
|
||||
predictedDiseaseSlug: DISEASE_CLASSES[maxIndex].slug,
|
||||
confidence: Math.max(0, Math.min(1, maxScore)),
|
||||
probabilities,
|
||||
};
|
||||
} finally {
|
||||
predictions.dispose();
|
||||
}
|
||||
} finally {
|
||||
imageTensor.dispose();
|
||||
response = await fetch(predictUrl(), {
|
||||
method: 'POST',
|
||||
body: formData,
|
||||
});
|
||||
} catch (error) {
|
||||
throw new Error(`ML service request failed: ${error instanceof Error ? error.message : String(error)}`);
|
||||
}
|
||||
|
||||
if (!response.ok) {
|
||||
const message = await parseErrorResponse(response);
|
||||
throw new Error(`ML service returned ${response.status}: ${message}`);
|
||||
}
|
||||
|
||||
const prediction = await response.json() as MlPredictionResponse;
|
||||
const predictedLabel = assertKnownLabel(prediction.label);
|
||||
const predictedDisease = DISEASE_BY_LABEL.get(predictedLabel)!;
|
||||
|
||||
return {
|
||||
predictedDiseaseSlug: predictedDisease.slug,
|
||||
confidence: assertConfidence(prediction.confidence, predictedDisease.label),
|
||||
probabilities: mapProbabilities(prediction.probabilities),
|
||||
};
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user