diff --git a/config-streaming-runpod.yaml b/config-streaming-runpod.yaml index 768ed4d..413e54b 100644 --- a/config-streaming-runpod.yaml +++ b/config-streaming-runpod.yaml @@ -29,7 +29,7 @@ matting: memory_offload: true # Critical for streaming - offload to CPU when needed fp16: false # Disable FP16 to avoid type mismatch with compiled models for memory efficiency continuous_correction: true # Periodically refine tracking - correction_interval: 300 # Correct every 5 seconds at 60fps + correction_interval: 30 # Correct every 0.5 seconds at 60fps (for testing) stereo: mode: "master_slave" # Left eye detects, right eye follows diff --git a/vr180_streaming/streaming_processor.py b/vr180_streaming/streaming_processor.py index e371122..92d7f45 100644 --- a/vr180_streaming/streaming_processor.py +++ b/vr180_streaming/streaming_processor.py @@ -183,20 +183,36 @@ class VR180StreamingProcessor: # Split into eyes left_eye, right_eye = self.stereo_manager.split_frame(frame) - # Process frames with simple approach (no detections in regular frames) - left_masks = self.sam2_processor.add_frame_and_detections(left_eye, [], frame_idx) - right_masks = self.sam2_processor.add_frame_and_detections(right_eye, [], frame_idx) + # Check if we need to run detection for continuous correction + detections = [] + if (self.config.matting.continuous_correction and + frame_idx % self.config.matting.correction_interval == 0): + print(f"\nšŸ”„ Running YOLO detection for correction at frame {frame_idx}") + master_eye = left_eye if self.stereo_manager.master_eye == 'left' else right_eye + detections = self.detector.detect_persons(master_eye) + if detections: + print(f" Detected {len(detections)} person(s) for correction") + else: + print(f" No persons detected for correction") + + # Process frames (with detections if this is a correction frame) + left_masks = self.sam2_processor.add_frame_and_detections(left_eye, detections, frame_idx) + + # For right eye, transfer detections if we have them + if detections: + transferred_detections = self.stereo_manager.transfer_detections( + detections, + 'left_to_right' if self.stereo_manager.master_eye == 'left' else 'right_to_left' + ) + right_masks = self.sam2_processor.add_frame_and_detections(right_eye, transferred_detections, frame_idx) + else: + right_masks = self.sam2_processor.add_frame_and_detections(right_eye, [], frame_idx) # Validate stereo consistency right_masks = self.stereo_manager.validate_masks( left_masks, right_masks, frame_idx ) - # Apply continuous correction if enabled - if (self.config.matting.continuous_correction and - frame_idx % self.config.matting.correction_interval == 0): - self._apply_continuous_correction(left_eye, right_eye, frame_idx) - # Apply masks and write frame processed_frame = self._apply_masks_to_frame(frame, left_masks, right_masks) self.frame_writer.write_frame(processed_frame)