diff --git a/vr180_matting/video_processor.py b/vr180_matting/video_processor.py index 30f11e9..0b8d077 100644 --- a/vr180_matting/video_processor.py +++ b/vr180_matting/video_processor.py @@ -12,6 +12,7 @@ import subprocess import gc import psutil import os +import sys from .config import VR180Config from .detector import YOLODetector @@ -131,6 +132,15 @@ class VideoProcessor: except ImportError: pass + # Force Linux to release memory back to OS + if sys.platform == 'linux': + try: + import ctypes + libc = ctypes.CDLL("libc.so.6") + libc.malloc_trim(0) + except Exception as e: + print(f" Warning: Could not trim memory: {e}") + # Brief pause to allow cleanup time.sleep(0.1) @@ -565,58 +575,90 @@ class VideoProcessor: chunk_size, overlap_frames = self.calculate_optimal_chunking() # Process video in chunks - chunk_results = [] + chunk_files = [] # Store file paths instead of frame data + temp_chunk_dir = Path(tempfile.mkdtemp(prefix="vr180_chunks_")) - for start_frame in range(0, self.total_frames, chunk_size - overlap_frames): - end_frame = min(start_frame + chunk_size, self.total_frames) - frames_to_read = end_frame - start_frame + try: + for start_frame in range(0, self.total_frames, chunk_size - overlap_frames): + end_frame = min(start_frame + chunk_size, self.total_frames) + frames_to_read = end_frame - start_frame + + chunk_idx = len(chunk_files) + print(f"\nProcessing chunk {chunk_idx}: frames {start_frame}-{end_frame}") + + # Read chunk frames + frames = self.read_video_frames( + self.config.input.video_path, + start_frame=start_frame, + num_frames=frames_to_read, + scale_factor=self.config.processing.scale_factor + ) + + # Process chunk + matted_frames = self.process_chunk(frames, chunk_idx) + + # Save chunk to disk immediately to free memory + chunk_path = temp_chunk_dir / f"chunk_{chunk_idx:04d}.npz" + print(f"Saving chunk {chunk_idx} to disk...") + np.savez_compressed(str(chunk_path), frames=matted_frames) + chunk_files.append(chunk_path) + + # Free the frames from memory immediately + del matted_frames + del frames + + # Update statistics + self.processing_stats['chunks_processed'] += 1 + self.processing_stats['frames_processed'] += frames_to_read + + # Aggressive memory cleanup after each chunk + self._aggressive_memory_cleanup(f"chunk {chunk_idx} completion") + + # Also use memory manager cleanup + self.memory_manager.cleanup_memory() + + if self.memory_manager.should_emergency_cleanup(): + self.memory_manager.emergency_cleanup() - chunk_idx = len(chunk_results) - print(f"\nProcessing chunk {chunk_idx}: frames {start_frame}-{end_frame}") + # Load and merge chunks from disk + print("\nLoading and merging chunks...") + chunk_results = [] + for chunk_file in chunk_files: + print(f"Loading {chunk_file.name}...") + chunk_data = np.load(str(chunk_file)) + chunk_results.append(chunk_data['frames']) + chunk_data.close() # Close the file - # Read chunk frames - frames = self.read_video_frames( - self.config.input.video_path, - start_frame=start_frame, - num_frames=frames_to_read, - scale_factor=self.config.processing.scale_factor - ) + # Merge chunks + final_frames = self.merge_overlapping_chunks(chunk_results, overlap_frames) - # Process chunk - matted_frames = self.process_chunk(frames, chunk_idx) - chunk_results.append(matted_frames) + # Free chunk results after merging + del chunk_results + self._aggressive_memory_cleanup("after merging chunks") - # Update statistics - self.processing_stats['chunks_processed'] += 1 - self.processing_stats['frames_processed'] += len(frames) + # Save results + print(f"Saving {len(final_frames)} processed frames...") + self.save_video(final_frames, self.config.output.path) - # Memory cleanup - self.memory_manager.cleanup_memory() + # Calculate final statistics + self.processing_stats['end_time'] = time.time() + self.processing_stats['total_duration'] = self.processing_stats['end_time'] - self.processing_stats['start_time'] + if self.processing_stats['total_duration'] > 0: + self.processing_stats['processing_fps'] = self.processing_stats['frames_processed'] / self.processing_stats['total_duration'] - if self.memory_manager.should_emergency_cleanup(): - self.memory_manager.emergency_cleanup() - - # Merge chunks if multiple - print("\nMerging chunks...") - final_frames = self.merge_overlapping_chunks(chunk_results, overlap_frames) - - # Save results - print(f"Saving {len(final_frames)} processed frames...") - self.save_video(final_frames, self.config.output.path) - - # Calculate final statistics - self.processing_stats['end_time'] = time.time() - self.processing_stats['total_duration'] = self.processing_stats['end_time'] - self.processing_stats['start_time'] - if self.processing_stats['total_duration'] > 0: - self.processing_stats['processing_fps'] = self.processing_stats['frames_processed'] / self.processing_stats['total_duration'] - - # Print processing statistics - self._print_processing_statistics() - - # Print final memory report - self.memory_manager.print_memory_report() - - print("Video processing completed!") + # Print processing statistics + self._print_processing_statistics() + + # Print final memory report + self.memory_manager.print_memory_report() + + print("Video processing completed!") + + finally: + # Clean up temporary chunk files + if temp_chunk_dir.exists(): + print("Cleaning up temporary chunk files...") + shutil.rmtree(temp_chunk_dir) def _print_processing_statistics(self): """Print detailed processing statistics"""