Memory leak fixes

This commit is contained in:
Roman Krček
2025-07-30 16:36:06 +02:00
parent 1aa6cd53fa
commit 923300e49b
7 changed files with 213 additions and 86 deletions

View File

@@ -7,6 +7,8 @@
import PhotoCard from './subcomponents/PhotoCard.svelte';
import * as tf from '@tensorflow/tfjs';
import * as blazeface from '@tensorflow-models/blazeface';
import PQueue from 'p-queue';
import { set, clear } from 'idb-keyval';
let photos = $state<PhotoInfo[]>([]);
let isProcessing = $state(false);
@@ -14,6 +16,8 @@
let totalCount = $state(0);
let detector: blazeface.BlazeFaceModel | undefined;
let detectorPromise: Promise<void> | undefined;
let downloadQueue: PQueue;
let faceDetectionQueue: PQueue;
interface PhotoInfo {
name: string;
@@ -38,13 +42,59 @@
return detectorPromise;
}
// Force memory cleanup
async function forceMemoryCleanup() {
await tf.nextFrame(); // Wait for any pending GPU operations
// Log memory state without aggressive cleanup
const memInfo = tf.memory();
console.log('Memory status:', {
tensors: memInfo.numTensors,
dataBuffers: memInfo.numDataBuffers,
bytes: memInfo.numBytes
});
// Only run garbage collection if available, don't dispose variables
if (typeof window !== 'undefined' && 'gc' in window) {
(window as any).gc();
}
}
async function processPhotosInParallel() {
if (isProcessing) return;
console.log('Starting processPhotos in parallel...');
console.log('Starting processPhotos with queues...');
isProcessing = true;
processedCount = 0;
try {
// Clear previous session's images from IndexedDB
await clear();
console.log('Cleared IndexedDB.');
} catch (e) {
console.error('Could not clear IndexedDB:', e);
}
// Initialize queues with more conservative concurrency
downloadQueue = new PQueue({ concurrency: 3 }); // Reduced from 5
faceDetectionQueue = new PQueue({ concurrency: 1 }); // Keep at 1 for memory safety
// When both queues are idle, we're done
downloadQueue.on('idle', async () => {
if (faceDetectionQueue.size === 0 && faceDetectionQueue.pending === 0) {
await forceMemoryCleanup(); // Clean up memory when processing is complete
isProcessing = false;
console.log('All queues are idle. Processing complete.');
}
});
faceDetectionQueue.on('idle', async () => {
if (downloadQueue.size === 0 && downloadQueue.pending === 0) {
await forceMemoryCleanup(); // Clean up memory when processing is complete
isProcessing = false;
console.log('All queues are idle. Processing complete.');
}
});
const validRows = $filteredSheetData.filter((row) => row._isValid);
const photoUrls = new Set<string>();
const photoMap = new Map<string, any[]>();
@@ -62,7 +112,7 @@
});
totalCount = photoUrls.size;
console.log(`Found ${totalCount} unique photo URLs`);
console.log(`Found ${totalCount} unique photo URLs to process.`);
photos = Array.from(photoUrls).map((url) => ({
name: photoMap.get(url)![0].name + ' ' + photoMap.get(url)![0].surname,
@@ -72,26 +122,10 @@
faceDetectionStatus: 'pending' as const
}));
const concurrencyLimit = 5;
const promises = [];
// Add all photos to the download queue
for (let i = 0; i < photos.length; i++) {
const promise = (async () => {
await loadPhoto(i);
processedCount++;
})();
promises.push(promise);
if (promises.length >= concurrencyLimit) {
await Promise.all(promises);
promises.length = 0;
}
downloadQueue.add(() => loadPhoto(i));
}
await Promise.all(promises);
isProcessing = false;
console.log('All photos processed.');
}
// Initialize detector and process photos
@@ -148,6 +182,8 @@
} catch (error) {
console.error(`Failed to load photo for ${photo.name}:`, error);
photo.status = 'error';
// Since this step failed, we still count it as "processed" to not stall the progress bar
processedCount++;
}
}
@@ -175,6 +211,8 @@
} catch (e) {
console.error(`Failed to convert HEIC image for ${photo.name}:`, e);
photo.status = 'error';
// Since this step failed, we still count it as "processed" to not stall the progress bar
processedCount++;
}
}
@@ -198,12 +236,14 @@
photo.status = 'success';
console.log(`Photo loaded successfully: ${photo.name}`);
// Save to pictures store
// Save blob to IndexedDB instead of the store
await set(photo.url, blob);
// Save to pictures store, but without the blob to save memory
pictures.update((pics) => ({
...pics,
[photo.url]: {
id: photo.url,
blob: blob,
url: objectUrl,
downloaded: true,
faceDetected: false,
@@ -211,32 +251,47 @@
}
}));
// Automatically run face detection to generate crop
await detectFaceForPhoto(index);
// Add face detection to its queue
faceDetectionQueue.add(() => detectFaceForPhoto(index));
} catch (error) {
console.error(`Failed to process blob for ${photo.name}:`, error);
photo.status = 'error';
// Since this step failed, we still count it as "processed" to not stall the progress bar
processedCount++;
}
}
async function detectFaceForPhoto(index: number) {
const photo = photos[index];
let imageTensor;
try {
await initializeDetector(); // Ensure detector is loaded
if (!detector) {
photos[index].faceDetectionStatus = 'failed';
photo.faceDetectionStatus = 'failed';
console.error('Face detector not available.');
return;
}
photos[index].faceDetectionStatus = 'processing';
photo.faceDetectionStatus = 'processing';
const img = new Image();
img.crossOrigin = 'anonymous';
img.src = photos[index].objectUrl!;
img.src = photo.objectUrl!;
await new Promise((r, e) => {
img.onload = r;
img.onerror = e;
});
const predictions = await detector.estimateFaces(img, false);
// Create tensor and manually dispose it after use
imageTensor = tf.browser.fromPixels(img);
const predictions = await detector.estimateFaces(imageTensor, false);
// Log memory usage for debugging
const memInfo = tf.memory();
console.log(`TensorFlow.js memory after face detection for ${photo.name}:`, {
numTensors: memInfo.numTensors,
numDataBuffers: memInfo.numDataBuffers,
numBytes: memInfo.numBytes
});
if (predictions.length > 0) {
const getProbability = (p: number | tf.Tensor) =>
@@ -245,26 +300,27 @@
const face = predictions.sort(
(a, b) => getProbability(b.probability!) - getProbability(a.probability!)
)[0];
// Coordinates in displayed image space
let [x1, y1] = face.topLeft as [number, number];
let [x2, y2] = face.bottomRight as [number, number];
// Scale to natural image size
const topLeft = face.topLeft as [number, number];
const bottomRight = face.bottomRight as [number, number];
let [x1, y1] = topLeft;
let [x2, y2] = bottomRight;
const scaleX = img.naturalWidth / img.width;
const scaleY = img.naturalHeight / img.height;
const faceWidth = (x2 - x1) * scaleX;
const faceHeight = (y2 - y1) * scaleY;
const faceCenterX = (x1 + (x2 - x1) / 2) * scaleX;
const faceCenterY = (y1 + (y2 - y1) / 2) * scaleY;
// Load crop config from env
const cropRatio = parseFloat(env.PUBLIC_CROP_RATIO || '1.0');
const offsetX = parseFloat(env.PUBLIC_FACE_OFFSET_X || '0.0');
const offsetY = parseFloat(env.PUBLIC_FACE_OFFSET_Y || '0.0');
const cropScale = parseFloat(env.PUBLIC_CROP_SCALE || '2.5');
// Compute crop size and center
let cropWidth = faceWidth * cropScale;
let cropHeight = cropWidth / cropRatio;
// If crop is larger than image, scale it down while maintaining aspect ratio
if (cropWidth > img.naturalWidth || cropHeight > img.naturalHeight) {
const widthRatio = img.naturalWidth / cropWidth;
const heightRatio = img.naturalHeight / cropHeight;
@@ -276,9 +332,11 @@
let centerX = faceCenterX + cropWidth * offsetX;
let centerY = faceCenterY + cropHeight * offsetY;
// Clamp center to ensure crop fits
centerX = Math.max(cropWidth / 2, Math.min(centerX, img.naturalWidth - cropWidth / 2));
centerY = Math.max(cropHeight / 2, Math.min(centerY, img.naturalHeight - cropHeight / 2));
centerY = Math.max(
cropHeight / 2,
Math.min(centerY, img.naturalHeight - cropHeight / 2)
);
const cropX = centerX - cropWidth / 2;
const cropY = centerY - cropHeight / 2;
@@ -289,32 +347,40 @@
width: Math.round(cropWidth),
height: Math.round(cropHeight)
};
photos[index].cropData = crop;
photos[index].faceDetectionStatus = 'completed';
photo.cropData = crop;
photo.faceDetectionStatus = 'completed';
// Save crop data to store
cropRects.update((crops) => ({
...crops,
[photos[index].url]: crop
[photo.url]: crop
}));
// Update pictures store with face detection info
pictures.update((pics) => ({
...pics,
[photos[index].url]: {
...pics[photos[index].url],
[photo.url]: {
...pics[photo.url],
faceDetected: true,
faceCount: predictions.length
}
}));
} else {
photos[index].faceDetectionStatus = 'failed';
photo.faceDetectionStatus = 'failed';
}
} catch (error) {
console.error(`Face detection failed for ${photos[index].name}:`, error);
photos[index].faceDetectionStatus = 'failed';
console.error(`Face detection failed for ${photo.name}:`, error);
photo.faceDetectionStatus = 'failed';
} finally {
// Manually dispose of the input tensor to prevent memory leaks
if (imageTensor) {
imageTensor.dispose();
}
// Add a small delay to allow GPU memory to be freed before next operation
await new Promise(resolve => setTimeout(resolve, 100));
// This is the final step for a photo, so we increment the processed count here.
processedCount++;
}
// No need to reassign photos array with $state reactivity
}
async function retryPhoto(index: number) {
@@ -325,7 +391,8 @@
}
photo.retryCount++;
await loadPhoto(index, true);
// Add the retry attempt back to the download queue
downloadQueue.add(() => loadPhoto(index, true));
}
function handleCropUpdate(
@@ -364,6 +431,13 @@
// Cleanup on unmount using $effect
$effect(() => {
return () => {
// Clear queues on component unmount to stop any ongoing processing
if (downloadQueue) {
downloadQueue.clear();
}
if (faceDetectionQueue) {
faceDetectionQueue.clear();
}
cleanupObjectUrls();
};
});