tomrikert commited on
Commit
c9c7adc
·
1 Parent(s): ab3858a

Add local face tracking (YOLO) and vision processing (SmolVLM2)

Browse files

Face Tracking:
- Initialize YOLO HeadTracker for local face detection
- Face tracking now works out of the box with --no-face-tracking to disable
- CameraWorker properly integrates with head tracker

Local Vision (optional):
- Add VisionProcessor using SmolVLM2 for on-device image understanding
- VisionManager for periodic scene description
- Enable with --local-vision flag or ENABLE_LOCAL_VISION=true
- Falls back to OpenClaw for vision if local not available

Configuration:
- ENABLE_FACE_TRACKING (default: true)
- ENABLE_LOCAL_VISION (default: false)
- LOCAL_VISION_MODEL (default: SmolVLM2-256M-Video-Instruct)
- VISION_DEVICE (auto, cuda, mps, cpu)
- HF_HOME for model cache directory

Based on pollen-robotics/reachy_mini_conversation_app implementation.

src/reachy_mini_openclaw/config.py CHANGED
@@ -42,7 +42,13 @@ class Config:
42
 
43
  # Face Tracking Configuration
44
  # Options: "yolo", "mediapipe", or None for auto-detect
45
- HEAD_TRACKER_TYPE: Optional[str] = field(default_factory=lambda: os.getenv("HEAD_TRACKER_TYPE"))
 
 
 
 
 
 
46
 
47
  # Custom Profile (for personality customization)
48
  CUSTOM_PROFILE: Optional[str] = field(default_factory=lambda: os.getenv("REACHY_MINI_CUSTOM_PROFILE"))
@@ -64,3 +70,15 @@ def set_custom_profile(profile: Optional[str]) -> None:
64
  global config
65
  config.CUSTOM_PROFILE = profile
66
  os.environ["REACHY_MINI_CUSTOM_PROFILE"] = profile or ""
 
 
 
 
 
 
 
 
 
 
 
 
 
42
 
43
  # Face Tracking Configuration
44
  # Options: "yolo", "mediapipe", or None for auto-detect
45
+ HEAD_TRACKER_TYPE: Optional[str] = field(default_factory=lambda: os.getenv("HEAD_TRACKER_TYPE", "yolo"))
46
+
47
+ # Local Vision Processing
48
+ ENABLE_LOCAL_VISION: bool = field(default_factory=lambda: os.getenv("ENABLE_LOCAL_VISION", "false").lower() == "true")
49
+ LOCAL_VISION_MODEL: str = field(default_factory=lambda: os.getenv("LOCAL_VISION_MODEL", "HuggingFaceTB/SmolVLM2-256M-Video-Instruct"))
50
+ VISION_DEVICE: str = field(default_factory=lambda: os.getenv("VISION_DEVICE", "auto")) # "auto", "cuda", "mps", "cpu"
51
+ HF_HOME: str = field(default_factory=lambda: os.getenv("HF_HOME", os.path.expanduser("~/.cache/huggingface")))
52
 
53
  # Custom Profile (for personality customization)
54
  CUSTOM_PROFILE: Optional[str] = field(default_factory=lambda: os.getenv("REACHY_MINI_CUSTOM_PROFILE"))
 
70
  global config
71
  config.CUSTOM_PROFILE = profile
72
  os.environ["REACHY_MINI_CUSTOM_PROFILE"] = profile or ""
73
+
74
+
75
+ def set_face_tracking_enabled(enabled: bool) -> None:
76
+ """Enable or disable face tracking at runtime."""
77
+ global config
78
+ config.ENABLE_FACE_TRACKING = enabled
79
+
80
+
81
+ def set_local_vision_enabled(enabled: bool) -> None:
82
+ """Enable or disable local vision processing at runtime."""
83
+ global config
84
+ config.ENABLE_LOCAL_VISION = enabled
src/reachy_mini_openclaw/main.py CHANGED
@@ -111,6 +111,16 @@ Examples:
111
  action="store_true",
112
  help="Disable OpenClaw integration"
113
  )
 
 
 
 
 
 
 
 
 
 
114
  parser.add_argument(
115
  "--profile",
116
  type=str,
@@ -208,16 +218,29 @@ class ClawBodyCore:
208
 
209
  # Camera worker for video streaming and frame capture
210
  self.camera_worker = None
 
 
 
211
  if enable_camera:
212
  logger.info("Initializing camera worker...")
213
  from reachy_mini_openclaw.camera_worker import CameraWorker
214
- # Initialize without face tracking (face tracking is optional)
 
 
 
 
 
215
  self.camera_worker = CameraWorker(
216
  reachy_mini=self.robot,
217
- head_tracker=None, # No face tracking by default
218
  )
219
- # Disable head tracking since we don't have a tracker
220
- self.camera_worker.set_head_tracking_enabled(False)
 
 
 
 
 
221
 
222
  # Create tool dependencies
223
  self.deps = ToolDependencies(
@@ -226,6 +249,7 @@ class ClawBodyCore:
226
  robot=self.robot,
227
  camera_worker=self.camera_worker,
228
  openclaw_bridge=self.openclaw_bridge,
 
229
  )
230
 
231
  # Initialize OpenAI Realtime handler with OpenClaw bridge
@@ -238,6 +262,88 @@ class ClawBodyCore:
238
  self._stop_event = asyncio.Event()
239
  self._tasks: list[asyncio.Task] = []
240
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
241
  def _should_stop(self) -> bool:
242
  """Check if we should stop."""
243
  if self._stop_event.is_set():
@@ -305,6 +411,11 @@ class ClawBodyCore:
305
  logger.info("Starting camera worker...")
306
  self.camera_worker.start()
307
 
 
 
 
 
 
308
  # Start audio
309
  logger.info("Starting audio...")
310
  self.robot.media.start_recording()
@@ -342,6 +453,10 @@ class ClawBodyCore:
342
  self.head_wobbler.stop()
343
  self.movement_manager.stop()
344
 
 
 
 
 
345
  # Stop camera worker
346
  if self.camera_worker is not None:
347
  self.camera_worker.stop()
@@ -404,6 +519,16 @@ def main() -> None:
404
  from reachy_mini_openclaw.config import set_custom_profile
405
  set_custom_profile(args.profile)
406
 
 
 
 
 
 
 
 
 
 
 
407
  if args.gradio:
408
  # Launch Gradio UI
409
  logger.info("Starting Gradio UI...")
 
111
  action="store_true",
112
  help="Disable OpenClaw integration"
113
  )
114
+ parser.add_argument(
115
+ "--no-face-tracking",
116
+ action="store_true",
117
+ help="Disable face tracking"
118
+ )
119
+ parser.add_argument(
120
+ "--local-vision",
121
+ action="store_true",
122
+ help="Enable local vision processing with SmolVLM2"
123
+ )
124
  parser.add_argument(
125
  "--profile",
126
  type=str,
 
218
 
219
  # Camera worker for video streaming and frame capture
220
  self.camera_worker = None
221
+ self.head_tracker = None
222
+ self.vision_manager = None
223
+
224
  if enable_camera:
225
  logger.info("Initializing camera worker...")
226
  from reachy_mini_openclaw.camera_worker import CameraWorker
227
+
228
+ # Initialize head tracker for local face tracking
229
+ if config.ENABLE_FACE_TRACKING:
230
+ self.head_tracker = self._initialize_head_tracker(config.HEAD_TRACKER_TYPE)
231
+
232
+ # Initialize camera worker with head tracker
233
  self.camera_worker = CameraWorker(
234
  reachy_mini=self.robot,
235
+ head_tracker=self.head_tracker,
236
  )
237
+
238
+ # Enable/disable head tracking based on whether we have a tracker
239
+ self.camera_worker.set_head_tracking_enabled(self.head_tracker is not None)
240
+
241
+ # Initialize local vision processor if enabled
242
+ if config.ENABLE_LOCAL_VISION:
243
+ self.vision_manager = self._initialize_vision_manager()
244
 
245
  # Create tool dependencies
246
  self.deps = ToolDependencies(
 
249
  robot=self.robot,
250
  camera_worker=self.camera_worker,
251
  openclaw_bridge=self.openclaw_bridge,
252
+ vision_manager=self.vision_manager,
253
  )
254
 
255
  # Initialize OpenAI Realtime handler with OpenClaw bridge
 
262
  self._stop_event = asyncio.Event()
263
  self._tasks: list[asyncio.Task] = []
264
 
265
+ def _initialize_vision_manager(self) -> Optional[Any]:
266
+ """Initialize local vision processor (SmolVLM2).
267
+
268
+ Returns:
269
+ VisionManager instance or None if initialization fails
270
+ """
271
+ if self.camera_worker is None:
272
+ logger.warning("Cannot initialize vision manager without camera worker")
273
+ return None
274
+
275
+ try:
276
+ from reachy_mini_openclaw.vision.processors import (
277
+ VisionConfig,
278
+ initialize_vision_manager,
279
+ )
280
+ from reachy_mini_openclaw.config import config
281
+
282
+ vision_config = VisionConfig(
283
+ model_path=config.LOCAL_VISION_MODEL,
284
+ device_preference=config.VISION_DEVICE,
285
+ hf_home=config.HF_HOME,
286
+ )
287
+
288
+ logger.info("Initializing local vision processor (SmolVLM2)...")
289
+ vision_manager = initialize_vision_manager(self.camera_worker, vision_config)
290
+
291
+ if vision_manager is not None:
292
+ logger.info("Local vision processor initialized")
293
+ else:
294
+ logger.warning("Local vision processor failed to initialize")
295
+
296
+ return vision_manager
297
+
298
+ except ImportError as e:
299
+ logger.warning(f"Local vision not available: {e}")
300
+ logger.warning("Install with: pip install torch transformers")
301
+ return None
302
+ except Exception as e:
303
+ logger.error(f"Failed to initialize vision manager: {e}")
304
+ return None
305
+
306
+ def _initialize_head_tracker(self, tracker_type: Optional[str] = None) -> Optional[Any]:
307
+ """Initialize head tracker for local face tracking.
308
+
309
+ Args:
310
+ tracker_type: Type of tracker ("yolo", "mediapipe", or None for auto)
311
+
312
+ Returns:
313
+ Initialized head tracker or None if initialization fails
314
+ """
315
+ # Default to YOLO if not specified
316
+ if tracker_type is None:
317
+ tracker_type = "yolo"
318
+
319
+ if tracker_type == "yolo":
320
+ try:
321
+ from reachy_mini_openclaw.vision.yolo_head_tracker import HeadTracker
322
+ logger.info("Initializing YOLO face tracker...")
323
+ tracker = HeadTracker(device="cpu") # CPU is fast enough for face detection
324
+ logger.info("YOLO face tracker initialized")
325
+ return tracker
326
+ except ImportError as e:
327
+ logger.warning(f"YOLO tracker not available: {e}")
328
+ logger.warning("Install with: pip install ultralytics supervision")
329
+ except Exception as e:
330
+ logger.error(f"Failed to initialize YOLO tracker: {e}")
331
+
332
+ elif tracker_type == "mediapipe":
333
+ try:
334
+ from reachy_mini_openclaw.vision.mediapipe_tracker import HeadTracker
335
+ logger.info("Initializing MediaPipe face tracker...")
336
+ tracker = HeadTracker()
337
+ logger.info("MediaPipe face tracker initialized")
338
+ return tracker
339
+ except ImportError as e:
340
+ logger.warning(f"MediaPipe tracker not available: {e}")
341
+ except Exception as e:
342
+ logger.error(f"Failed to initialize MediaPipe tracker: {e}")
343
+
344
+ logger.warning("No face tracker available - face tracking disabled")
345
+ return None
346
+
347
  def _should_stop(self) -> bool:
348
  """Check if we should stop."""
349
  if self._stop_event.is_set():
 
411
  logger.info("Starting camera worker...")
412
  self.camera_worker.start()
413
 
414
+ # Start local vision processor if available
415
+ if self.vision_manager is not None:
416
+ logger.info("Starting local vision processor...")
417
+ self.vision_manager.start()
418
+
419
  # Start audio
420
  logger.info("Starting audio...")
421
  self.robot.media.start_recording()
 
453
  self.head_wobbler.stop()
454
  self.movement_manager.stop()
455
 
456
+ # Stop vision manager
457
+ if self.vision_manager is not None:
458
+ self.vision_manager.stop()
459
+
460
  # Stop camera worker
461
  if self.camera_worker is not None:
462
  self.camera_worker.stop()
 
519
  from reachy_mini_openclaw.config import set_custom_profile
520
  set_custom_profile(args.profile)
521
 
522
+ # Configure face tracking and local vision from args
523
+ from reachy_mini_openclaw.config import (
524
+ set_face_tracking_enabled,
525
+ set_local_vision_enabled,
526
+ )
527
+ if args.no_face_tracking:
528
+ set_face_tracking_enabled(False)
529
+ if args.local_vision:
530
+ set_local_vision_enabled(True)
531
+
532
  if args.gradio:
533
  # Launch Gradio UI
534
  logger.info("Starting Gradio UI...")
src/reachy_mini_openclaw/tools/core_tools.py CHANGED
@@ -36,6 +36,7 @@ class ToolDependencies:
36
  robot: Any # ReachyMini instance
37
  camera_worker: Optional[Any] = None
38
  openclaw_bridge: Optional["OpenClawBridge"] = None
 
39
 
40
 
41
  # Tool specifications in OpenAI format
@@ -211,12 +212,12 @@ async def _handle_look(args: dict, deps: ToolDependencies) -> dict:
211
 
212
 
213
  async def _handle_camera(args: dict, deps: ToolDependencies) -> dict:
214
- """Handle the camera tool - capture image and get description from OpenClaw.
215
 
216
- Since OpenAI Realtime doesn't have vision, we send the image to OpenClaw
217
- for analysis and return the description.
218
  """
219
- logger.info("Camera tool called, camera_worker=%s", deps.camera_worker is not None)
 
220
 
221
  if deps.camera_worker is None:
222
  logger.warning("Camera worker is None")
@@ -239,16 +240,31 @@ async def _handle_camera(args: dict, deps: ToolDependencies) -> dict:
239
  if frame is None:
240
  return {"error": "No frame available from camera"}
241
 
242
- # Encode frame as JPEG base64
243
- import cv2
244
- logger.info("Encoding frame, shape=%s", frame.shape)
245
- _, buffer = cv2.imencode('.jpg', frame, [cv2.IMWRITE_JPEG_QUALITY, 85])
246
- b64_image = base64.b64encode(buffer).decode('utf-8')
247
- logger.info("Frame encoded, size=%d bytes", len(b64_image))
248
 
249
- # Send to OpenClaw for vision analysis (OpenAI Realtime doesn't have vision)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
250
  if deps.openclaw_bridge is not None and deps.openclaw_bridge.is_connected:
251
- logger.info("Sending image to OpenClaw for vision analysis...")
 
 
 
 
252
  response = await deps.openclaw_bridge.chat(
253
  "Describe what you see in this image. Be specific about people, objects, and the environment. Keep it concise (2-3 sentences).",
254
  image_b64=b64_image,
@@ -259,14 +275,15 @@ async def _handle_camera(args: dict, deps: ToolDependencies) -> dict:
259
  return {
260
  "status": "success",
261
  "description": response.content,
 
262
  }
263
  else:
264
  logger.warning("OpenClaw vision failed: %s", response.error)
265
 
266
- # Fallback if OpenClaw not available
267
  return {
268
- "status": "success",
269
- "description": "I captured an image but couldn't analyze it. My vision system needs OpenClaw to process what I see."
270
  }
271
  except Exception as e:
272
  logger.error("Camera tool error: %s", e, exc_info=True)
@@ -281,11 +298,12 @@ async def _handle_face_tracking(args: dict, deps: ToolDependencies) -> dict:
281
  return {"error": "Camera not available for face tracking"}
282
 
283
  try:
284
- if hasattr(deps.camera_worker, 'set_face_tracking'):
285
- deps.camera_worker.set_face_tracking(enabled)
286
- return {"status": "success", "face_tracking": enabled}
287
- else:
288
- return {"error": "Face tracking not supported"}
 
289
  except Exception as e:
290
  return {"error": str(e)}
291
 
 
36
  robot: Any # ReachyMini instance
37
  camera_worker: Optional[Any] = None
38
  openclaw_bridge: Optional["OpenClawBridge"] = None
39
+ vision_manager: Optional[Any] = None # Local vision processor (SmolVLM2)
40
 
41
 
42
  # Tool specifications in OpenAI format
 
212
 
213
 
214
  async def _handle_camera(args: dict, deps: ToolDependencies) -> dict:
215
+ """Handle the camera tool - capture image and get description.
216
 
217
+ Uses local vision (SmolVLM2) if available, otherwise falls back to OpenClaw.
 
218
  """
219
+ logger.info("Camera tool called, camera_worker=%s, vision_manager=%s",
220
+ deps.camera_worker is not None, deps.vision_manager is not None)
221
 
222
  if deps.camera_worker is None:
223
  logger.warning("Camera worker is None")
 
240
  if frame is None:
241
  return {"error": "No frame available from camera"}
242
 
243
+ logger.info("Got frame, shape=%s", frame.shape)
 
 
 
 
 
244
 
245
+ # Option 1: Use local vision processor (SmolVLM2) if available
246
+ if deps.vision_manager is not None:
247
+ logger.info("Using local vision processor (SmolVLM2)...")
248
+ description = deps.vision_manager.process_now(
249
+ "Describe what you see in this image. Be specific about people, objects, and the environment. Keep it concise (2-3 sentences)."
250
+ )
251
+ if description and not description.startswith(("Vision", "Failed", "Error", "GPU", "No camera")):
252
+ logger.info("Local vision response: %s", description[:100])
253
+ return {
254
+ "status": "success",
255
+ "description": description,
256
+ "source": "local_vision"
257
+ }
258
+ else:
259
+ logger.warning("Local vision failed: %s", description)
260
+
261
+ # Option 2: Fall back to OpenClaw for vision analysis
262
  if deps.openclaw_bridge is not None and deps.openclaw_bridge.is_connected:
263
+ logger.info("Using OpenClaw for vision analysis...")
264
+ import cv2
265
+ _, buffer = cv2.imencode('.jpg', frame, [cv2.IMWRITE_JPEG_QUALITY, 85])
266
+ b64_image = base64.b64encode(buffer).decode('utf-8')
267
+
268
  response = await deps.openclaw_bridge.chat(
269
  "Describe what you see in this image. Be specific about people, objects, and the environment. Keep it concise (2-3 sentences).",
270
  image_b64=b64_image,
 
275
  return {
276
  "status": "success",
277
  "description": response.content,
278
+ "source": "openclaw"
279
  }
280
  else:
281
  logger.warning("OpenClaw vision failed: %s", response.error)
282
 
283
+ # Fallback if neither is available
284
  return {
285
+ "status": "partial",
286
+ "description": "I captured an image but couldn't analyze it. No vision processing available."
287
  }
288
  except Exception as e:
289
  logger.error("Camera tool error: %s", e, exc_info=True)
 
298
  return {"error": "Camera not available for face tracking"}
299
 
300
  try:
301
+ # Check if head tracker is available
302
+ if deps.camera_worker.head_tracker is None:
303
+ return {"error": "Face tracking not available - no head tracker initialized"}
304
+
305
+ deps.camera_worker.set_head_tracking_enabled(enabled)
306
+ return {"status": "success", "face_tracking": enabled}
307
  except Exception as e:
308
  return {"error": str(e)}
309
 
src/reachy_mini_openclaw/vision/__init__.py CHANGED
@@ -1,5 +1,18 @@
1
- """Vision modules for face tracking and detection."""
2
 
3
  from reachy_mini_openclaw.vision.head_tracker import get_head_tracker
4
 
5
- __all__ = ["get_head_tracker"]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Vision modules for face tracking, detection, and image understanding."""
2
 
3
  from reachy_mini_openclaw.vision.head_tracker import get_head_tracker
4
 
5
+ __all__ = [
6
+ "get_head_tracker",
7
+ ]
8
+
9
+ # Lazy imports for optional heavy dependencies
10
+ def get_vision_processor():
11
+ """Get the VisionProcessor class (requires torch, transformers)."""
12
+ from reachy_mini_openclaw.vision.processors import VisionProcessor
13
+ return VisionProcessor
14
+
15
+ def get_vision_manager():
16
+ """Get the VisionManager class (requires torch, transformers)."""
17
+ from reachy_mini_openclaw.vision.processors import VisionManager
18
+ return VisionManager
src/reachy_mini_openclaw/vision/processors.py ADDED
@@ -0,0 +1,419 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Local vision processing with SmolVLM2.
2
+
3
+ Provides on-device image understanding using the SmolVLM2 model
4
+ for scene description and visual analysis.
5
+
6
+ Based on pollen-robotics/reachy_mini_conversation_app vision processors.
7
+ """
8
+
9
+ import os
10
+ import time
11
+ import base64
12
+ import logging
13
+ import threading
14
+ from typing import Any, Dict, Optional
15
+ from dataclasses import dataclass, field
16
+
17
+ import cv2
18
+ import numpy as np
19
+ from numpy.typing import NDArray
20
+
21
+ try:
22
+ import torch
23
+ from transformers import AutoProcessor, AutoModelForImageTextToText
24
+ from huggingface_hub import snapshot_download
25
+ VISION_AVAILABLE = True
26
+ except ImportError:
27
+ VISION_AVAILABLE = False
28
+
29
+ logger = logging.getLogger(__name__)
30
+
31
+
32
+ @dataclass
33
+ class VisionConfig:
34
+ """Configuration for vision processing."""
35
+
36
+ model_path: str = "HuggingFaceTB/SmolVLM2-256M-Video-Instruct"
37
+ vision_interval: float = 5.0
38
+ max_new_tokens: int = 64
39
+ jpeg_quality: int = 85
40
+ max_retries: int = 3
41
+ retry_delay: float = 1.0
42
+ device_preference: str = "auto" # "auto", "cuda", "mps", "cpu"
43
+ hf_home: str = field(default_factory=lambda: os.path.expanduser("~/.cache/huggingface"))
44
+
45
+
46
+ class VisionProcessor:
47
+ """Handles SmolVLM2 model loading and inference for local vision."""
48
+
49
+ def __init__(self, vision_config: Optional[VisionConfig] = None):
50
+ """Initialize the vision processor.
51
+
52
+ Args:
53
+ vision_config: Vision configuration settings
54
+ """
55
+ if not VISION_AVAILABLE:
56
+ raise ImportError(
57
+ "Vision processing requires: pip install torch transformers huggingface-hub"
58
+ )
59
+
60
+ self.vision_config = vision_config or VisionConfig()
61
+ self.model_path = self.vision_config.model_path
62
+ self.device = self._determine_device()
63
+ self.processor = None
64
+ self.model = None
65
+ self._initialized = False
66
+
67
+ def _determine_device(self) -> str:
68
+ """Determine the best device for inference."""
69
+ pref = self.vision_config.device_preference
70
+
71
+ if pref == "cpu":
72
+ return "cpu"
73
+ if pref == "cuda":
74
+ return "cuda" if torch.cuda.is_available() else "cpu"
75
+ if pref == "mps":
76
+ return "mps" if torch.backends.mps.is_available() else "cpu"
77
+
78
+ # auto: prefer mps on Apple, then cuda, else cpu
79
+ if torch.backends.mps.is_available():
80
+ return "mps"
81
+ return "cuda" if torch.cuda.is_available() else "cpu"
82
+
83
+ def initialize(self) -> bool:
84
+ """Load model and processor onto the selected device.
85
+
86
+ Returns:
87
+ True if initialization successful, False otherwise
88
+ """
89
+ try:
90
+ cache_dir = self.vision_config.hf_home
91
+ os.makedirs(cache_dir, exist_ok=True)
92
+ os.environ["HF_HOME"] = cache_dir
93
+
94
+ logger.info(f"Loading SmolVLM2 model on {self.device} (HF_HOME={cache_dir})")
95
+
96
+ # Download model to cache first
97
+ logger.info(f"Downloading vision model {self.model_path}...")
98
+ snapshot_download(
99
+ repo_id=self.model_path,
100
+ repo_type="model",
101
+ cache_dir=cache_dir,
102
+ )
103
+
104
+ self.processor = AutoProcessor.from_pretrained(self.model_path)
105
+
106
+ # Select dtype depending on device
107
+ if self.device == "cuda":
108
+ dtype = torch.bfloat16
109
+ elif self.device == "mps":
110
+ dtype = torch.float32 # best for MPS
111
+ else:
112
+ dtype = torch.float32
113
+
114
+ model_kwargs: Dict[str, Any] = {"torch_dtype": dtype}
115
+
116
+ # flash_attention_2 is CUDA-only; skip on MPS/CPU
117
+ if self.device == "cuda":
118
+ model_kwargs["_attn_implementation"] = "flash_attention_2"
119
+
120
+ # Load model weights
121
+ self.model = AutoModelForImageTextToText.from_pretrained(
122
+ self.model_path, **model_kwargs
123
+ ).to(self.device)
124
+
125
+ if self.model is not None:
126
+ self.model.eval()
127
+ self._initialized = True
128
+ logger.info(f"Vision model loaded successfully on {self.device}")
129
+ return True
130
+
131
+ except Exception as e:
132
+ logger.error(f"Failed to initialize vision model: {e}")
133
+ return False
134
+
135
+ return False
136
+
137
+ def process_image(
138
+ self,
139
+ cv2_image: NDArray[np.uint8],
140
+ prompt: str = "Briefly describe what you see in one sentence.",
141
+ ) -> str:
142
+ """Process CV2 image and return description with retry logic.
143
+
144
+ Args:
145
+ cv2_image: OpenCV image (BGR format)
146
+ prompt: Question/prompt to ask about the image
147
+
148
+ Returns:
149
+ Text description of the image
150
+ """
151
+ if not self._initialized or self.processor is None or self.model is None:
152
+ return "Vision model not initialized"
153
+
154
+ for attempt in range(self.vision_config.max_retries):
155
+ try:
156
+ # Convert to JPEG bytes
157
+ success, jpeg_buffer = cv2.imencode(
158
+ ".jpg",
159
+ cv2_image,
160
+ [cv2.IMWRITE_JPEG_QUALITY, self.vision_config.jpeg_quality],
161
+ )
162
+ if not success:
163
+ return "Failed to encode image"
164
+
165
+ # Convert to base64
166
+ image_base64 = base64.b64encode(jpeg_buffer.tobytes()).decode("utf-8")
167
+
168
+ messages = [
169
+ {
170
+ "role": "user",
171
+ "content": [
172
+ {
173
+ "type": "image",
174
+ "url": f"data:image/jpeg;base64,{image_base64}",
175
+ },
176
+ {"type": "text", "text": prompt},
177
+ ],
178
+ },
179
+ ]
180
+
181
+ inputs = self.processor.apply_chat_template(
182
+ messages,
183
+ add_generation_prompt=True,
184
+ tokenize=True,
185
+ return_dict=True,
186
+ return_tensors="pt",
187
+ )
188
+
189
+ # Move tensors to device WITHOUT forcing dtype (keeps input_ids as torch.long)
190
+ inputs = {
191
+ k: (v.to(self.device) if hasattr(v, "to") else v)
192
+ for k, v in inputs.items()
193
+ }
194
+
195
+ with torch.no_grad():
196
+ generated_ids = self.model.generate(
197
+ **inputs,
198
+ do_sample=False,
199
+ max_new_tokens=self.vision_config.max_new_tokens,
200
+ pad_token_id=self.processor.tokenizer.eos_token_id,
201
+ )
202
+
203
+ generated_texts = self.processor.batch_decode(
204
+ generated_ids,
205
+ skip_special_tokens=True,
206
+ )
207
+
208
+ # Extract just the response part
209
+ full_text = generated_texts[0]
210
+ response = self._extract_response(full_text)
211
+
212
+ # Clean up GPU memory if using CUDA
213
+ if self.device == "cuda":
214
+ torch.cuda.empty_cache()
215
+ elif self.device == "mps":
216
+ torch.mps.empty_cache()
217
+
218
+ return response.replace(chr(10), " ").strip()
219
+
220
+ except Exception as e:
221
+ if "OutOfMemory" in str(type(e).__name__):
222
+ logger.error(f"GPU OOM on attempt {attempt + 1}: {e}")
223
+ if self.device == "cuda":
224
+ torch.cuda.empty_cache()
225
+ if attempt < self.vision_config.max_retries - 1:
226
+ time.sleep(self.vision_config.retry_delay * (attempt + 1))
227
+ else:
228
+ return "GPU out of memory - vision processing failed"
229
+ else:
230
+ logger.error(f"Vision processing failed (attempt {attempt + 1}): {e}")
231
+ if attempt < self.vision_config.max_retries - 1:
232
+ time.sleep(self.vision_config.retry_delay)
233
+ else:
234
+ return f"Vision processing error after {self.vision_config.max_retries} attempts"
235
+
236
+ return "Vision processing failed"
237
+
238
+ def _extract_response(self, full_text: str) -> str:
239
+ """Extract the assistant's response from the full generated text."""
240
+ # Handle different response formats
241
+ markers = ["assistant\n", "Assistant:", "Response:", "\n\n"]
242
+
243
+ for marker in markers:
244
+ if marker in full_text:
245
+ response = full_text.split(marker)[-1].strip()
246
+ if response: # Ensure we got a meaningful response
247
+ return response
248
+
249
+ # Fallback: return the full text cleaned up
250
+ return full_text.strip()
251
+
252
+ def get_model_info(self) -> Dict[str, Any]:
253
+ """Get information about the loaded model."""
254
+ info = {
255
+ "initialized": self._initialized,
256
+ "device": self.device,
257
+ "model_path": self.model_path,
258
+ "cuda_available": torch.cuda.is_available() if VISION_AVAILABLE else False,
259
+ }
260
+
261
+ if VISION_AVAILABLE and torch.cuda.is_available():
262
+ info["gpu_memory_gb"] = torch.cuda.get_device_properties(0).total_memory // (1024**3)
263
+ else:
264
+ info["gpu_memory_gb"] = "N/A"
265
+
266
+ return info
267
+
268
+
269
+ class VisionManager:
270
+ """Manages periodic vision processing and scene understanding.
271
+
272
+ This runs in the background, periodically capturing frames and
273
+ generating scene descriptions that can be queried.
274
+ """
275
+
276
+ def __init__(
277
+ self,
278
+ camera_worker: Any,
279
+ vision_config: Optional[VisionConfig] = None,
280
+ ):
281
+ """Initialize vision manager.
282
+
283
+ Args:
284
+ camera_worker: CameraWorker instance for frame capture
285
+ vision_config: Vision configuration settings
286
+ """
287
+ self.camera_worker = camera_worker
288
+ self.vision_config = vision_config or VisionConfig()
289
+ self.vision_interval = self.vision_config.vision_interval
290
+ self.processor = VisionProcessor(self.vision_config)
291
+
292
+ self._last_processed_time = 0.0
293
+ self._last_description = ""
294
+ self._description_lock = threading.Lock()
295
+ self._stop_event = threading.Event()
296
+ self._thread: Optional[threading.Thread] = None
297
+
298
+ # Initialize processor
299
+ if not self.processor.initialize():
300
+ logger.error("Failed to initialize vision processor")
301
+ raise RuntimeError("Vision processor initialization failed")
302
+
303
+ def start(self) -> None:
304
+ """Start the vision processing loop in a background thread."""
305
+ self._stop_event.clear()
306
+ self._thread = threading.Thread(target=self._working_loop, daemon=True)
307
+ self._thread.start()
308
+ logger.info("Local vision processing started")
309
+
310
+ def stop(self) -> None:
311
+ """Stop the vision processing loop."""
312
+ self._stop_event.set()
313
+ if self._thread is not None:
314
+ self._thread.join(timeout=5.0)
315
+ logger.info("Local vision processing stopped")
316
+
317
+ def get_latest_description(self) -> str:
318
+ """Get the most recent scene description.
319
+
320
+ Returns:
321
+ Latest scene description or empty string if none available
322
+ """
323
+ with self._description_lock:
324
+ return self._last_description
325
+
326
+ def process_now(self, prompt: str = "Briefly describe what you see in one sentence.") -> str:
327
+ """Process the current frame immediately with a custom prompt.
328
+
329
+ Args:
330
+ prompt: Question/prompt to ask about the image
331
+
332
+ Returns:
333
+ Description of what the camera sees
334
+ """
335
+ frame = self.camera_worker.get_latest_frame()
336
+ if frame is None:
337
+ return "No camera frame available"
338
+
339
+ return self.processor.process_image(frame, prompt)
340
+
341
+ def _working_loop(self) -> None:
342
+ """Vision processing loop (runs in separate thread)."""
343
+ while not self._stop_event.is_set():
344
+ try:
345
+ current_time = time.time()
346
+
347
+ if current_time - self._last_processed_time >= self.vision_interval:
348
+ frame = self.camera_worker.get_latest_frame()
349
+ if frame is not None:
350
+ description = self.processor.process_image(
351
+ frame,
352
+ "Briefly describe what you see in one sentence.",
353
+ )
354
+
355
+ # Only update if we got a valid response
356
+ if description and not description.startswith(
357
+ ("Vision", "Failed", "Error", "GPU")
358
+ ):
359
+ with self._description_lock:
360
+ self._last_description = description
361
+ self._last_processed_time = current_time
362
+ logger.debug(f"Vision update: {description}")
363
+ else:
364
+ logger.warning(f"Invalid vision response: {description}")
365
+
366
+ time.sleep(1.0) # Check every second
367
+
368
+ except Exception:
369
+ logger.exception("Vision processing loop error")
370
+ time.sleep(5.0) # Longer sleep on error
371
+
372
+ logger.info("Vision loop finished")
373
+
374
+ def get_status(self) -> Dict[str, Any]:
375
+ """Get comprehensive status information."""
376
+ return {
377
+ "last_processed": self._last_processed_time,
378
+ "last_description": self.get_latest_description(),
379
+ "processor_info": self.processor.get_model_info(),
380
+ "config": {
381
+ "interval": self.vision_interval,
382
+ },
383
+ }
384
+
385
+
386
+ def initialize_vision_manager(
387
+ camera_worker: Any,
388
+ config: Optional[VisionConfig] = None,
389
+ ) -> Optional[VisionManager]:
390
+ """Initialize vision manager with model download and configuration.
391
+
392
+ Args:
393
+ camera_worker: CameraWorker instance for frame capture
394
+ config: Optional vision configuration
395
+
396
+ Returns:
397
+ VisionManager instance or None if initialization fails
398
+ """
399
+ if not VISION_AVAILABLE:
400
+ logger.warning("Vision dependencies not available. Install: pip install torch transformers")
401
+ return None
402
+
403
+ try:
404
+ vision_config = config or VisionConfig()
405
+
406
+ # Initialize vision manager
407
+ vision_manager = VisionManager(camera_worker, vision_config)
408
+
409
+ # Log device info
410
+ device_info = vision_manager.processor.get_model_info()
411
+ logger.info(
412
+ f"Local vision enabled: {device_info.get('model_path')} on {device_info.get('device')}"
413
+ )
414
+
415
+ return vision_manager
416
+
417
+ except Exception as e:
418
+ logger.error(f"Failed to initialize vision manager: {e}")
419
+ return None