diff --git a/common_utility/capture.py b/common_utility/capture.py index 8b486aa..a5de2b7 100644 --- a/common_utility/capture.py +++ b/common_utility/capture.py @@ -31,6 +31,10 @@ INDEX_ENTRY_SIZE = struct.calcsize(INDEX_ENTRY_FORMAT) TRIGGER_FRAME_FLAG: int = 1 +RAW_IMAGE_FLAG: int = 2 + +RAW_IMAGE_HEADER_FORMAT = " None: @@ -66,6 +70,12 @@ def add_args(parser: argparse.ArgumentParser, **defaults: Any) -> None: metavar="0-9", help="PNG compression level for captured images (0=none, 9=max); default 1 favours throughput", ) + parser.add_argument( + "--blob-raw-capture", + action="store_true", + default=bool(defaults.get("blob_raw_capture", False)), + help="Store images as raw pixel arrays instead of PNG (faster capture, decoded to PNG on extraction)", + ) class BlobCompletionHandler(ABC): @@ -122,12 +132,14 @@ def __init__( max_image_bytes: int = 4 * 1024 * 1024, png_compression: int = 1, capture_folder: Optional[pathlib.Path] = None, + raw_capture: bool = False, ) -> None: self._blob_path = blob_path self._num_slots = num_slots self._max_image_bytes = max_image_bytes self._png_compression = png_compression self._capture_folder = capture_folder + self._raw_capture = raw_capture self._index_base = SUPERBLOCK_SIZE self._data_base = SUPERBLOCK_SIZE + num_slots * INDEX_ENTRY_SIZE self._file_size = self._data_base + num_slots * max_image_bytes @@ -179,26 +191,38 @@ def capture(self, image: NDArray[np.uint8], filename: str, flags: int = 0) -> No if len(filename_bytes) > MAX_FILENAME_BYTES: raise ValueError(f"Filename too long: {len(filename_bytes)} bytes, max {MAX_FILENAME_BYTES}") - ok, buf = cv2.imencode(".png", image, [cv2.IMWRITE_PNG_COMPRESSION, self._png_compression]) - if not ok: - raise ValueError("cv2.imencode failed to encode image as PNG") - png_bytes = buf.tobytes() + if self._raw_capture: + h, w = image.shape[:2] + c = image.shape[2] if image.ndim == 3 else 1 + header = struct.pack(RAW_IMAGE_HEADER_FORMAT, h, w, c) + image_size = RAW_IMAGE_HEADER_SIZE + image.nbytes + actual_flags = flags | RAW_IMAGE_FLAG + else: + ok, buf = cv2.imencode(".png", image, [cv2.IMWRITE_PNG_COMPRESSION, self._png_compression]) + if not ok: + raise ValueError("cv2.imencode failed to encode image as PNG") + image_size = len(buf) + actual_flags = flags - if len(png_bytes) > self._max_image_bytes: - raise ValueError(f"PNG too large: {len(png_bytes)} bytes, max {self._max_image_bytes}") + if image_size > self._max_image_bytes: + raise ValueError(f"Image data too large: {image_size} bytes, max {self._max_image_bytes}") slot = self._write_head data_offset = self._slot_data_offset(slot) idx_offset = self._slot_index_offset(slot) - self._mm[data_offset : data_offset + len(png_bytes)] = png_bytes + if self._raw_capture: + self._mm[data_offset : data_offset + RAW_IMAGE_HEADER_SIZE] = header + self._mm[data_offset + RAW_IMAGE_HEADER_SIZE : data_offset + image_size] = image + else: + self._mm[data_offset : data_offset + image_size] = buf filename_padded = filename_bytes.ljust(MAX_FILENAME_BYTES, b"\x00") entry = struct.pack( INDEX_ENTRY_FORMAT, data_offset, - len(png_bytes), - flags, + image_size, + actual_flags, len(filename_bytes), 0, filename_padded, @@ -223,7 +247,8 @@ def update_last_flags(self, flags: int) -> None: assert self._mm is not None last_slot = (self._write_head - 1) % self._num_slots idx_offset = self._slot_index_offset(last_slot) - struct.pack_into(" BlobFsCapture: @@ -277,13 +302,20 @@ def __init__( follow_up_count: int, completion_handler: BlobCompletionHandler, png_compression: int = 1, + raw_capture: bool = False, ) -> None: self._capture_folder = capture_folder self._num_slots = num_slots self._max_image_bytes = max_image_bytes self._png_compression = png_compression + self._raw_capture = raw_capture self._blob = BlobFsCapture( - blob_path, num_slots, max_image_bytes, png_compression, capture_folder=capture_folder + blob_path, + num_slots, + max_image_bytes, + png_compression, + capture_folder=capture_folder, + raw_capture=raw_capture, ) self._follow_up_count = follow_up_count self._completion_handler = completion_handler @@ -330,7 +362,9 @@ def rotate(self, event_id: str, immediate: bool = False) -> CompositeBlobCapture self._completion_handler(event_id, self._capture_folder / f"{event_id}.blob", []) else: post_path = self._capture_folder / f"{event_id}-post.blob" - post_blob = BlobFsCapture(post_path, self._num_slots, self._max_image_bytes, self._png_compression) + post_blob = BlobFsCapture( + post_path, self._num_slots, self._max_image_bytes, self._png_compression, raw_capture=self._raw_capture + ) self._post_blob = (post_blob, event_id, self._follow_up_count) return self @@ -456,38 +490,69 @@ def __enter__(self) -> BlobExtractor: def __exit__(self, *_: Any) -> None: self.close() - def _extract_blob(self, blob: _OpenBlob, dest_dir: pathlib.Path) -> list[pathlib.Path]: + def _extract_png(self, raw_data: bytes) -> Optional[bytes]: + arr = np.frombuffer(raw_data, dtype=np.uint8) + if cv2.imdecode(arr, cv2.IMREAD_COLOR) is None: + return None + return raw_data + + def _extract_raw(self, raw_data: bytes) -> Optional[bytes]: + if len(raw_data) < RAW_IMAGE_HEADER_SIZE: + return None + h, w, c = struct.unpack_from(RAW_IMAGE_HEADER_FORMAT, raw_data, 0) + pixel_data = raw_data[RAW_IMAGE_HEADER_SIZE:] + if len(pixel_data) != h * w * c: + return None + arr = np.frombuffer(pixel_data, dtype=np.uint8).reshape(h, w, c) + ok, buf = cv2.imencode(".png", arr) + if not ok: + return None + return buf.tobytes() + + def _extract_one(self, i: int, blob: _OpenBlob, dest_dir: pathlib.Path) -> Optional[pathlib.Path]: + idx_offset = SUPERBLOCK_SIZE + i * blob.index_entry_size assert blob.mm is not None - extracted: list[pathlib.Path] = [] - for i in range(blob.num_slots): - idx_offset = SUPERBLOCK_SIZE + i * blob.index_entry_size - if idx_offset + blob.index_entry_size > len(blob.mm): - break + if idx_offset + blob.index_entry_size > len(blob.mm): + return None - image_offset, image_size, _flags, filename_len, _reserved, filename_raw = struct.unpack_from( - INDEX_ENTRY_FORMAT, blob.mm, idx_offset - ) + image_offset, image_size, _flags, filename_len, _reserved, filename_raw = struct.unpack_from( + INDEX_ENTRY_FORMAT, blob.mm, idx_offset + ) - if image_size == 0 or image_offset + image_size > len(blob.mm): - continue + if image_size == 0 or image_offset + image_size > len(blob.mm): + return None - png_bytes = bytes(blob.mm[image_offset : image_offset + image_size]) - arr = np.frombuffer(png_bytes, dtype=np.uint8) - if cv2.imdecode(arr, cv2.IMREAD_COLOR) is None: - continue + raw_data = bytes(blob.mm[image_offset : image_offset + image_size]) + + if _flags & RAW_IMAGE_FLAG: + png_bytes = self._extract_raw(raw_data) + else: + png_bytes = self._extract_png(raw_data) - raw_fn = filename_raw[:filename_len] - try: - filename = raw_fn.decode("utf-8") - except UnicodeDecodeError: - filename = raw_fn.decode("latin-1") + if png_bytes is None: + return None - if not filename: - filename = f"slot_{i:04d}.png" + raw_fn = filename_raw[:filename_len] + try: + filename = raw_fn.decode("utf-8") + except UnicodeDecodeError: + filename = raw_fn.decode("latin-1") - dest_path = dest_dir / filename - dest_path.write_bytes(png_bytes) - extracted.append(dest_path) + if not filename: + filename = f"slot_{i:04d}.png" + + dest_path = dest_dir / str(filename) + dest_path.write_bytes(png_bytes) + return dest_path + + def _extract_blob(self, blob: _OpenBlob, dest_dir: pathlib.Path) -> list[pathlib.Path]: + assert blob.mm is not None + extracted: list[pathlib.Path] = [] + for i in range(blob.num_slots): + dst = self._extract_one(i, blob, dest_dir) + if dst is None: + continue + extracted.append(dst) return extracted @@ -616,7 +681,8 @@ def make_capture_backend( num_slots: int = args.blob_num_slots max_image_bytes: int = args.blob_max_image_bytes png_compression: int = args.blob_png_compression + raw_capture: bool = args.blob_raw_capture handler = completion_handler if completion_handler is not None else NoOpBlobCompletionHandler() return CompositeBlobCapture( - blob_path, capture_folder, num_slots, max_image_bytes, follow_up_count, handler, png_compression + blob_path, capture_folder, num_slots, max_image_bytes, follow_up_count, handler, png_compression, raw_capture ) diff --git a/tests/captureTest.py b/tests/captureTest.py index a07dd99..c5151b1 100644 --- a/tests/captureTest.py +++ b/tests/captureTest.py @@ -9,6 +9,8 @@ from unittest.mock import MagicMock, patch +import cv2 + from common_utility.capture import ( MAGIC, SUPERBLOCK_FORMAT, @@ -16,6 +18,9 @@ INDEX_ENTRY_FORMAT, INDEX_ENTRY_SIZE, TRIGGER_FRAME_FLAG, + RAW_IMAGE_FLAG, + RAW_IMAGE_HEADER_FORMAT, + RAW_IMAGE_HEADER_SIZE, BlobCompletionHandler, NoOpBlobCompletionHandler, CompositeBlobCapture, @@ -111,7 +116,7 @@ def test_flags_stored_in_index(self, tmp_path: pathlib.Path) -> None: _img_off, _img_sz, flags, _fn_len, _res, _fn_raw = struct.unpack_from(INDEX_ENTRY_FORMAT, data, SUPERBLOCK_SIZE) assert flags == 0xCAFEBABE - def test_update_last_flags_overwrites_flags_field(self, tmp_path: pathlib.Path) -> None: + def test_update_last_flags_sets_flags_field(self, tmp_path: pathlib.Path) -> None: blob = BlobFsCapture(tmp_path / "cap.blob", num_slots=4, max_image_bytes=512 * 1024) blob.capture(_make_image(), "frame.png", flags=0) blob.update_last_flags(TRIGGER_FRAME_FLAG) @@ -120,6 +125,35 @@ def test_update_last_flags_overwrites_flags_field(self, tmp_path: pathlib.Path) _img_off, _img_sz, flags, _fn_len, _res, _fn_raw = struct.unpack_from(INDEX_ENTRY_FORMAT, data, SUPERBLOCK_SIZE) assert flags == TRIGGER_FRAME_FLAG + def test_raw_capture_sets_raw_image_flag(self, tmp_path: pathlib.Path) -> None: + blob = BlobFsCapture(tmp_path / "cap.blob", num_slots=4, max_image_bytes=512 * 1024, raw_capture=True) + blob.capture(_make_image(), "frame.png") + blob.close() + data = (tmp_path / "cap.blob").read_bytes() + _img_off, _img_sz, flags, _fn_len, _res, _fn_raw = struct.unpack_from(INDEX_ENTRY_FORMAT, data, SUPERBLOCK_SIZE) + assert flags & RAW_IMAGE_FLAG + + def test_raw_capture_stores_header_and_pixels(self, tmp_path: pathlib.Path) -> None: + image = _make_image() + blob = BlobFsCapture(tmp_path / "cap.blob", num_slots=4, max_image_bytes=512 * 1024, raw_capture=True) + blob.capture(image, "frame.png") + blob.close() + data = (tmp_path / "cap.blob").read_bytes() + img_off, img_sz, *_ = struct.unpack_from(INDEX_ENTRY_FORMAT, data, SUPERBLOCK_SIZE) + raw = data[img_off : img_off + img_sz] + assert len(raw) == RAW_IMAGE_HEADER_SIZE + image.size + h, w, c = struct.unpack_from(RAW_IMAGE_HEADER_FORMAT, raw, 0) + assert (h, w, c) == (image.shape[0], image.shape[1], image.shape[2]) + + def test_update_last_flags_preserves_raw_image_flag(self, tmp_path: pathlib.Path) -> None: + blob = BlobFsCapture(tmp_path / "cap.blob", num_slots=4, max_image_bytes=512 * 1024, raw_capture=True) + blob.capture(_make_image(), "frame.png") + blob.update_last_flags(TRIGGER_FRAME_FLAG) + blob.close() + data = (tmp_path / "cap.blob").read_bytes() + _img_off, _img_sz, flags, *_ = struct.unpack_from(INDEX_ENTRY_FORMAT, data, SUPERBLOCK_SIZE) + assert flags == (RAW_IMAGE_FLAG | TRIGGER_FRAME_FLAG) + def test_open_on_size_mismatch_creates_fresh_blob(self, tmp_path: pathlib.Path) -> None: path = tmp_path / "cap.blob" blob = BlobFsCapture(path, num_slots=4, max_image_bytes=256 * 1024) @@ -354,6 +388,50 @@ def test_no_siblings_when_none_exist(self, tmp_path: pathlib.Path) -> None: extracted = BlobExtractor(path).extract(tmp_path / "out") assert {p.name for p in extracted} == {"only.png"} + def test_raw_capture_extractor_produces_png(self, tmp_path: pathlib.Path) -> None: + names = ["a.png", "b.png"] + path = tmp_path / "cap.blob" + blob = BlobFsCapture(path, num_slots=4, max_image_bytes=512 * 1024, raw_capture=True) + for name in names: + blob.capture(_make_image(), name) + blob.close() + dest = tmp_path / "out" + extracted = BlobExtractor(path).extract(dest) + assert len(extracted) == 2 + assert {p.name for p in extracted} == set(names) + for p in extracted: + arr = cv2.imdecode(np.frombuffer(p.read_bytes(), dtype=np.uint8), cv2.IMREAD_COLOR) + assert arr is not None + + def test_raw_capture_round_trip_pixel_equality(self, tmp_path: pathlib.Path) -> None: + image = _make_image() + path = tmp_path / "cap.blob" + blob = BlobFsCapture(path, num_slots=4, max_image_bytes=512 * 1024, raw_capture=True) + blob.capture(image, "frame.png") + blob.close() + extracted = BlobExtractor(path).extract(tmp_path / "out") + assert len(extracted) == 1 + decoded = cv2.imdecode(np.frombuffer(extracted[0].read_bytes(), dtype=np.uint8), cv2.IMREAD_COLOR) + assert decoded is not None + # cv2 stores BGR; compare channel-by-channel (lossless PNG round-trip must be identical) + assert np.array_equal(image, decoded) + + def test_raw_and_png_blobs_extracted_transparently(self, tmp_path: pathlib.Path) -> None: + """Primary blob uses PNG, post sibling uses raw; both extract to valid PNGs.""" + primary = tmp_path / "event.blob" + post = tmp_path / "event-post.blob" + b1 = BlobFsCapture(primary, num_slots=4, max_image_bytes=512 * 1024, raw_capture=False) + b1.capture(_make_image(), "png.png") + b1.close() + b2 = BlobFsCapture(post, num_slots=4, max_image_bytes=512 * 1024, raw_capture=True) + b2.capture(_make_image(), "raw.png") + b2.close() + extracted = BlobExtractor(primary).extract(tmp_path / "out") + assert {p.name for p in extracted} == {"png.png", "raw.png"} + for p in extracted: + arr = cv2.imdecode(np.frombuffer(p.read_bytes(), dtype=np.uint8), cv2.IMREAD_COLOR) + assert arr is not None + # --------------------------------------------------------------------------- # BlobExtractor — tar/tar.gz archive support @@ -730,6 +808,35 @@ def test_get_last_filepath_delegates_to_primary(self, tmp_path: pathlib.Path) -> assert comp.get_last_filepath() is None comp.close() + def test_composite_raw_capture_post_blob_also_raw(self, tmp_path: pathlib.Path) -> None: + follow_up = 2 + cap_folder = tmp_path / "rotated" + comp = CompositeBlobCapture( + tmp_path / "active.blob", + cap_folder, + num_slots=8, + max_image_bytes=512 * 1024, + follow_up_count=follow_up, + completion_handler=NoOpBlobCompletionHandler(), + raw_capture=True, + ) + comp.capture(_make_image(), "trigger.png") + comp.rotate("ev-raw") + for i in range(follow_up): + comp.capture(_make_image(), f"followup{i}.png") + comp.close() + + for blob_path in [cap_folder / "ev-raw.blob", cap_folder / "ev-raw-post.blob"]: + data = blob_path.read_bytes() + _, _, flags, *_ = struct.unpack_from(INDEX_ENTRY_FORMAT, data, SUPERBLOCK_SIZE) + assert flags & RAW_IMAGE_FLAG + + all_extracted = BlobExtractor(cap_folder / "ev-raw.blob").extract(tmp_path / "out") + assert len(all_extracted) == 1 + follow_up + for p in all_extracted: + arr = cv2.imdecode(np.frombuffer(p.read_bytes(), dtype=np.uint8), cv2.IMREAD_COLOR) + assert arr is not None + # --------------------------------------------------------------------------- # add_args / make_capture_backend @@ -746,6 +853,19 @@ def test_add_args_registers_defaults(self) -> None: assert args.blob_num_slots == 60 assert args.blob_max_image_bytes == 4 * 1024 * 1024 assert args.blob_png_compression == 1 + assert args.blob_raw_capture is False + + def test_add_args_raw_capture_flag(self) -> None: + parser = argparse.ArgumentParser() + add_args(parser) + args = parser.parse_args(["--blob-raw-capture"]) + assert args.blob_raw_capture is True + + def test_add_args_raw_capture_default_injectable(self) -> None: + parser = argparse.ArgumentParser() + add_args(parser, blob_raw_capture=True) + args = parser.parse_args([]) + assert args.blob_raw_capture is True def test_add_args_allows_overrides(self) -> None: parser = argparse.ArgumentParser()