AI_Photographer/Modes/AI/camera_module.py
2026-04-12 18:52:37 +04:00

334 lines
11 KiB
Python

"""
Camera capture module for Unitree G1 robot
Handles OpenCV camera interface and image processing for AI analysis
"""
import cv2
import numpy as np
import base64
import io
import os
import time
from pathlib import Path
from PIL import Image
from typing import Optional, Union, List
import logging
from Core.settings import CAMERA_INDEX, FRAME_WIDTH, FRAME_HEIGHT, FPS
class CameraCapture:
"""
Camera capture class for Unitree G1 robot
Provides methods to capture frames and prepare them for AI analysis
"""
def __init__(self, camera_index: Optional[Union[int, str]] = None):
"""
Initialize camera capture
Args:
camera_index: int index, /dev/videoX path, "auto", or None
"""
self.camera_source = self._resolve_camera_source(camera_index)
# Keep legacy field name for backward compatibility with older modules/tests
self.camera_index = self.camera_source
self.cap = None
self.is_initialized = False
self.logger = logging.getLogger(__name__)
self._last_read_warn_ts = 0.0
def _resolve_camera_source(self, camera_source: Optional[Union[int, str]]) -> Union[int, str]:
"""
Resolve requested camera source from argument/env/defaults.
Priority:
1) explicit argument
2) CAMERA_DEVICE env (supports /dev/videoX or "auto")
3) CAMERA_INDEX env/config int
"""
if camera_source is None:
env_device = os.environ.get("CAMERA_DEVICE", "").strip()
if env_device:
camera_source = env_device
else:
camera_source = os.environ.get("CAMERA_INDEX", str(CAMERA_INDEX))
if isinstance(camera_source, str):
s = camera_source.strip()
if not s:
return int(CAMERA_INDEX)
if s.lower() == "auto":
return "auto"
try:
return int(s)
except ValueError:
return s
return int(camera_source)
def _candidate_sources(self) -> List[Union[int, str]]:
"""
Build candidate capture sources. When CAMERA_FALLBACK_SCAN=1 (default),
we also try available /dev/video* nodes if the primary source fails.
"""
src = self.camera_source
candidates: List[Union[int, str]] = [src]
# If primary source is /dev/videoX, also try numeric index X.
if isinstance(src, str) and src.startswith("/dev/video"):
try:
idx = int(src.replace("/dev/video", ""))
candidates.append(idx)
except Exception:
pass
do_scan = os.environ.get("CAMERA_FALLBACK_SCAN", "1").strip().lower() not in (
"0", "false", "no", "off"
)
if src == "auto" or do_scan:
for p in sorted(Path("/dev").glob("video*")):
sp = str(p)
if sp not in candidates:
candidates.append(sp)
try:
idx = int(sp.replace("/dev/video", ""))
if idx not in candidates:
candidates.append(idx)
except Exception:
pass
return candidates
def _open_capture(self, source: Union[int, str]):
if isinstance(source, str) and source.startswith("/dev/video"):
# V4L2 backend is more reliable for /dev/videoX on Linux.
return cv2.VideoCapture(source, cv2.CAP_V4L2)
return cv2.VideoCapture(source)
def initialize(self) -> bool:
"""
Initialize camera connection
Returns:
bool: True if successful, False otherwise
"""
try:
self.release()
tried: List[str] = []
for source in self._candidate_sources():
tried.append(str(source))
cap = None
try:
cap = self._open_capture(source)
except Exception:
cap = None
if cap is None or not cap.isOpened():
try:
if cap is not None:
cap.release()
except Exception:
pass
continue
# Set camera properties
cap.set(cv2.CAP_PROP_FRAME_WIDTH, FRAME_WIDTH)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, FRAME_HEIGHT)
cap.set(cv2.CAP_PROP_FPS, FPS)
# Test capture
ret, _frame = cap.read()
if not ret:
cap.release()
continue
self.cap = cap
self.camera_source = source
self.camera_index = source
self.is_initialized = True
self.logger.info(f"Camera initialized successfully: {source}")
return True
self.is_initialized = False
self.cap = None
self.logger.error(f"Failed to open camera. Tried sources: {', '.join(tried)}")
return False
except Exception as e:
self.logger.error(f"Camera initialization failed: {e}")
return False
def capture_frame(self) -> Optional[np.ndarray]:
"""
Capture a single frame from camera
Returns:
numpy.ndarray: Captured frame or None if failed
"""
if not self.is_initialized or self.cap is None:
return None
try:
ret, frame = self.cap.read()
if not ret:
now = time.time()
# Throttle noisy warnings when camera stream is unstable/unavailable.
if now - self._last_read_warn_ts >= 5.0:
self.logger.warning("Failed to capture frame from camera")
self._last_read_warn_ts = now
return None
return frame
except Exception as e:
self.logger.error(f"Frame capture failed: {e}")
return None
def frame_to_base64(self, frame: np.ndarray, format: str = 'JPEG') -> Optional[str]:
"""
Convert OpenCV frame to base64 string for API transmission
Args:
frame: OpenCV frame (BGR format)
format: Image format ('JPEG' or 'PNG')
Returns:
str: Base64 encoded image or None if failed
"""
try:
# Convert BGR to RGB
rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
# Convert to PIL Image
pil_image = Image.fromarray(rgb_frame)
# Convert to bytes
buffer = io.BytesIO()
pil_image.save(buffer, format=format, quality=85)
buffer.seek(0)
# Encode to base64
img_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
return img_base64
except Exception as e:
self.logger.error(f"Base64 conversion failed: {e}")
return None
def frame_to_pil(self, frame: np.ndarray) -> Optional[Image.Image]:
"""
Convert OpenCV frame to PIL Image for Gemini API
Args:
frame: OpenCV frame (BGR format)
Returns:
PIL.Image: Converted image or None if failed
"""
try:
# Convert BGR to RGB
rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
# Convert to PIL Image
pil_image = Image.fromarray(rgb_frame)
return pil_image
except Exception as e:
self.logger.error(f"PIL conversion failed: {e}")
return None
def get_frame_info(self) -> dict:
"""
Get current camera frame information
Returns:
dict: Frame information including resolution, FPS
"""
if not self.is_initialized or self.cap is None:
return {}
try:
info = {
'width': int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
'height': int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT)),
'fps': self.cap.get(cv2.CAP_PROP_FPS),
'format': int(self.cap.get(cv2.CAP_PROP_FORMAT)),
'brightness': self.cap.get(cv2.CAP_PROP_BRIGHTNESS),
'contrast': self.cap.get(cv2.CAP_PROP_CONTRAST),
'saturation': self.cap.get(cv2.CAP_PROP_SATURATION)
}
return info
except Exception as e:
self.logger.error(f"Failed to get frame info: {e}")
return {}
def release(self):
"""
Release camera resources
"""
try:
if self.cap is not None:
self.cap.release()
self.cap = None
self.logger.info("Camera released successfully")
self.is_initialized = False
except Exception as e:
self.logger.error(f"Camera release failed: {e}")
finally:
# Headless systems often ship OpenCV without HighGUI window backends.
# Do not fail camera lifecycle because window teardown is unavailable.
try:
import cv2
cv2.destroyAllWindows()
except Exception:
pass
def __enter__(self):
"""Context manager entry"""
self.initialize()
return self
def __exit__(self, exc_type, exc_val, exc_tb):
"""Context manager exit"""
self.release()
return False # Don't suppress exceptions
# Utility functions for camera testing
def test_camera_capture():
"""
Test camera capture functionality
"""
logging.basicConfig(level=logging.INFO)
with CameraCapture() as camera:
if not camera.is_initialized:
print("Failed to initialize camera")
return
print("Camera info:", camera.get_frame_info())
# Capture test frame
frame = camera.capture_frame()
if frame is not None:
print(f"Captured frame shape: {frame.shape}")
# Test base64 conversion
base64_str = camera.frame_to_base64(frame)
if base64_str:
print(f"Base64 conversion successful, length: {len(base64_str)}")
# Test PIL conversion
pil_img = camera.frame_to_pil(frame)
if pil_img:
print(f"PIL conversion successful, size: {pil_img.size}")
# Save test image
cv2.imwrite('test_capture.jpg', frame)
print("Test image saved as test_capture.jpg")
else:
print("Failed to capture frame")