From 612fdd5789fd570d87a63f247110e838b672cc8b Mon Sep 17 00:00:00 2001 From: Ryan Shpeherd Date: Mon, 10 Aug 2026 15:22:22 -0400 Subject: [PATCH] After story 4 --- src/batcher.py | 68 +++++++++ src/export_face_detector.py | 104 +++++++++++++ src/face_detector.py | 295 ++++++++++++++++++++++++++++++++++++ src/gpu_manager.py | 101 ++++++++++++ src/orchestrator.py | 41 +++++ tests/test_story_04.py | 98 ++++++++++++ 6 files changed, 707 insertions(+) create mode 100644 src/batcher.py create mode 100644 src/export_face_detector.py create mode 100644 src/face_detector.py create mode 100644 src/gpu_manager.py create mode 100644 tests/test_story_04.py diff --git a/src/batcher.py b/src/batcher.py new file mode 100644 index 0000000..66a9a44 --- /dev/null +++ b/src/batcher.py @@ -0,0 +1,68 @@ +"""Dynamic batching utilities for face detection inference.""" + +import logging +import time +from collections import deque +from typing import Callable, Deque, Generic, List, Optional, TypeVar + +from gpu_manager import GPUMemoryManager + +logger = logging.getLogger(__name__) + +T = TypeVar("T") +R = TypeVar("R") + + +class DynamicBatcher(Generic[T, R]): + """Accumulate items into batches and flush based on size or timeout.""" + + def __init__( + self, + process_fn: Callable[[List[T]], List[R]], + gpu_manager: Optional[GPUMemoryManager] = None, + max_batch_size: int = 16, + batch_timeout_ms: float = 100.0, + min_batch_size: int = 1, + ): + self.process_fn = process_fn + self.gpu_manager = gpu_manager + self.max_batch_size = max_batch_size + self.batch_timeout_ms = batch_timeout_ms + self.min_batch_size = min_batch_size + self._queue: Deque[T] = deque() + self._last_flush = time.monotonic() + + def add(self, item: T) -> List[R]: + """Add an item and return any results if a batch was flushed.""" + self._queue.append(item) + if len(self._queue) >= self._current_batch_size(): + return self.flush() + + elapsed_ms = (time.monotonic() - self._last_flush) * 1000 + if elapsed_ms >= self.batch_timeout_ms and len(self._queue) >= self.min_batch_size: + return self.flush() + return [] + + def flush(self) -> List[R]: + """Flush all queued items through the processor.""" + if not self._queue: + return [] + + batch = [self._queue.popleft() for _ in range(min(len(self._queue), self._current_batch_size()))] + results = self.process_fn(batch) + self._last_flush = time.monotonic() + + if self.gpu_manager is not None: + self.gpu_manager.adjust_batch_size() + self.gpu_manager.empty_cache() + + return results + + def _current_batch_size(self) -> int: + if self.gpu_manager is not None: + return min(self.gpu_manager.current_batch_size, self.max_batch_size) + return self.max_batch_size + + @property + def queued_count(self) -> int: + return len(self._queue) diff --git a/src/export_face_detector.py b/src/export_face_detector.py new file mode 100644 index 0000000..04591c9 --- /dev/null +++ b/src/export_face_detector.py @@ -0,0 +1,104 @@ +"""Export a YOLOv8n face-detection model to ONNX and build a TensorRT engine.""" + +import argparse +import json +import logging +from pathlib import Path + +logger = logging.getLogger(__name__) + + +def export_onnx(output_dir: Path, input_size: int = 640) -> Path: + """Export YOLOv8n to ONNX with the required input shape.""" + output_dir.mkdir(parents=True, exist_ok=True) + onnx_path = output_dir / "face_detector.onnx" + metadata_path = output_dir / "model.json" + + try: + from ultralytics import YOLO + + model = YOLO("yolov8n.pt") + model.export( + format="onnx", + imgsz=input_size, + half=False, + simplify=True, + dynamic=False, + ) + exported = Path("yolov8n.onnx") + if exported.exists(): + exported.rename(onnx_path) + logger.info("Exported ONNX model to %s", onnx_path) + except Exception as exc: + logger.error("Failed to export ONNX model: %s", exc) + raise + + metadata = { + "input_shape": [1, 3, input_size, input_size], + "mean": [0.485, 0.456, 0.406], + "std": [0.229, 0.224, 0.225], + "confidence_threshold": 0.25, + "iou_threshold": 0.45, + } + metadata_path.write_text(json.dumps(metadata, indent=2)) + return onnx_path + + +def build_tensorrt_engine(onnx_path: Path, output_dir: Path, max_batch_size: int = 32) -> Path: + """Build and serialize a TensorRT FP32 engine from an ONNX file.""" + engine_path = output_dir / "face_detector.trt" + + try: + import tensorrt as trt + + logger = trt.Logger(trt.Logger.WARNING) + builder = trt.Builder(logger) + network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) + parser = trt.OnnxParser(network, logger) + + onnx_data = onnx_path.read_bytes() + if not parser.parse(onnx_data): + for error in range(parser.num_errors): + logger.error(parser.get_error(error)) + raise RuntimeError("ONNX parsing failed") + + config = builder.create_builder_config() + config.max_workspace_size = 4 << 30 # 4GB + config.set_flag(trt.BuilderFlag.FP32) + + profile = builder.create_optimization_profile() + input_name = network.get_input(0).name + profile.set_shape( + input_name, + (1, 3, 640, 640), + (max_batch_size // 2, 3, 640, 640), + (max_batch_size, 3, 640, 640), + ) + config.add_optimization_profile(profile) + + engine = builder.build_engine(network, config) + if engine is None: + raise RuntimeError("TensorRT engine build failed") + + engine_path.write_bytes(engine.serialize()) + logger.info("Serialized TensorRT engine to %s", engine_path) + return engine_path + except Exception as exc: + logger.error("Failed to build TensorRT engine: %s", exc) + raise + + +def main(): + logging.basicConfig(level=logging.INFO) + parser = argparse.ArgumentParser(description="Export YOLOv8n face detector to ONNX/TensorRT") + parser.add_argument("--output-dir", type=Path, default=Path("models/face_detector")) + parser.add_argument("--input-size", type=int, default=640) + parser.add_argument("--max-batch-size", type=int, default=32) + args = parser.parse_args() + + onnx_path = export_onnx(args.output_dir, args.input_size) + build_tensorrt_engine(onnx_path, args.output_dir, args.max_batch_size) + + +if __name__ == "__main__": + main() diff --git a/src/face_detector.py b/src/face_detector.py new file mode 100644 index 0000000..027c8df --- /dev/null +++ b/src/face_detector.py @@ -0,0 +1,295 @@ +""" +Face detection runner. + +Wraps a YOLOv8n model exported to ONNX/TensorRT, runs batched inference on +sampled frames, applies NMS, and produces 224×224 face crops for the classifier. +""" + +import logging +import time +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple + +import numpy as np + +logger = logging.getLogger(__name__) + + +class Detection: + """A single face detection result.""" + + def __init__( + self, + frame_path: str, + x1: float, + y1: float, + x2: float, + y2: float, + confidence: float, + crop_path: Optional[str] = None, + ): + self.frame_path = frame_path + self.x1 = x1 + self.y1 = y1 + self.x2 = x2 + self.y2 = y2 + self.confidence = confidence + self.crop_path = crop_path + + def to_dict(self) -> Dict[str, Any]: + return { + "frame_path": self.frame_path, + "x1": self.x1, + "y1": self.y1, + "x2": self.x2, + "y2": self.y2, + "confidence": self.confidence, + "crop_path": self.crop_path, + } + + +class FaceDetector: + """Run batched face detection on video frames.""" + + def __init__( + self, + engine_path: str, + input_size: int = 640, + confidence_threshold: float = 0.25, + iou_threshold: float = 0.45, + max_faces_per_frame: int = 10, + max_faces_per_video: int = 100, + device: str = "cuda", + ): + self.engine_path = engine_path + self.input_size = input_size + self.confidence_threshold = confidence_threshold + self.iou_threshold = iou_threshold + self.max_faces_per_frame = max_faces_per_frame + self.max_faces_per_video = max_faces_per_video + self.device = device + self._session: Optional[Any] = None + self._load_model() + + def _load_model(self): + """Load the inference backend.""" + path = Path(self.engine_path) + if not path.exists(): + logger.warning("Face detector engine not found at %s; using placeholder", self.engine_path) + return + + suffix = path.suffix.lower() + if suffix == ".onnx": + try: + import onnxruntime as ort + + providers = ["CUDAExecutionProvider"] if self.device.startswith("cuda") else ["CPUExecutionProvider"] + self._session = ort.InferenceSession(str(path), providers=providers) + logger.info("Loaded ONNX face detector from %s", self.engine_path) + except Exception as exc: + logger.warning("Failed to load ONNX face detector: %s", exc) + elif suffix in (".trt", ".engine", ".plan"): + try: + import tensorrt as trt + + with trt.Logger() as trt_logger, open(path, "rb") as f: + runtime = trt.Runtime(trt_logger) + self._session = runtime.deserialize_cuda_engine(f.read()) + logger.info("Loaded TensorRT face detector from %s", self.engine_path) + except Exception as exc: + logger.warning("Failed to load TensorRT face detector: %s", exc) + else: + logger.warning("Unsupported face detector format: %s", suffix) + + def detect_faces( + self, + frame_paths: List[str], + batch_size: int = 16, + ) -> List[List[Detection]]: + """Detect faces in a list of frame image paths.""" + all_results: List[List[Detection]] = [] + for i in range(0, len(frame_paths), batch_size): + batch = frame_paths[i : i + batch_size] + batch_results = self._detect_batch(batch) + all_results.extend(batch_results) + return all_results + + def _detect_batch(self, frame_paths: List[str]) -> List[List[Detection]]: + """Run detection on one batch and return detections per frame.""" + preprocessed = [] + for path in frame_paths: + try: + preprocessed.append(self._preprocess(path)) + except Exception as exc: + logger.warning("Failed to preprocess %s: %s", path, exc) + preprocessed.append(np.zeros((3, self.input_size, self.input_size), dtype=np.float32)) + + batch_input = np.stack(preprocessed, axis=0) + + outputs = self._infer(batch_input) + + results = [] + for idx, path in enumerate(frame_paths): + try: + detections = self._parse_outputs(outputs, idx, path) + detections = self._nms(detections) + detections = detections[: self.max_faces_per_frame] + results.append(detections) + except Exception as exc: + logger.warning("Failed to parse detections for %s: %s", path, exc) + results.append([]) + return results + + def _preprocess(self, frame_path: str) -> np.ndarray: + """Load and normalize a frame for the detector.""" + from PIL import Image + + image = Image.open(frame_path).convert("RGB") + image = image.resize((self.input_size, self.input_size), Image.Resampling.BILINEAR) + arr = np.array(image, dtype=np.float32) / 255.0 + mean = np.array([0.485, 0.456, 0.406], dtype=np.float32) + std = np.array([0.229, 0.224, 0.225], dtype=np.float32) + arr = (arr - mean) / std + return np.transpose(arr, (2, 0, 1)) # HWC -> CHW + + def _infer(self, batch_input: np.ndarray) -> Any: + """Run inference on the preprocessed batch.""" + if self._session is None: + return self._placeholder_outputs(batch_input.shape[0]) + + try: + if hasattr(self._session, "run"): + input_name = self._session.get_inputs()[0].name + return self._session.run(None, {input_name: batch_input}) + + # TensorRT execution + import pycuda.driver as cuda + import pycuda.autoinit # noqa: F401 + + context = self._session.create_execution_context() + output_shape = (batch_input.shape[0], 84, 8400) # YOLOv8n default shape + d_input = cuda.mem_alloc(batch_input.nbytes) + d_output = cuda.mem_alloc(np.prod(output_shape) * np.dtype(np.float32).itemsize) + stream = cuda.Stream() + cuda.memcpy_htod_async(d_input, batch_input, stream) + context.execute_async_v2(bindings=[int(d_input), int(d_output)], stream_handle=stream.handle) + output = np.empty(output_shape, dtype=np.float32) + cuda.memcpy_dtoh_async(output, d_output, stream) + stream.synchronize() + return output + except Exception as exc: + logger.warning("Face detector inference failed: %s", exc) + return self._placeholder_outputs(batch_input.shape[0]) + + def _placeholder_outputs(self, batch_size: int) -> np.ndarray: + """Return an empty output tensor when no model is loaded.""" + return np.zeros((batch_size, 84, 8400), dtype=np.float32) + + def _parse_outputs(self, outputs: Any, batch_index: int, frame_path: str) -> List[Detection]: + """Parse raw inference outputs into Detection objects.""" + if isinstance(outputs, list): + raw = outputs[0][batch_index] # (84, 8400) + else: + raw = outputs[batch_index] + + # YOLOv8 output layout: (cx, cy, w, h, cls scores...) + raw = raw.T # (8400, 84) + scores = raw[:, 4:].max(axis=1) + mask = scores >= self.confidence_threshold + candidates = raw[mask] + scores = scores[mask] + + detections = [] + for row, score in zip(candidates, scores): + cx, cy, w, h = row[:4] + x1 = cx - w / 2 + y1 = cy - h / 2 + x2 = cx + w / 2 + y2 = cy + h / 2 + detections.append( + Detection( + frame_path=frame_path, + x1=float(x1), + y1=float(y1), + x2=float(x2), + y2=float(y2), + confidence=float(score), + ) + ) + return detections + + def _nms(self, detections: List[Detection]) -> List[Detection]: + """Apply greedy Non-Maximum Suppression.""" + if not detections: + return [] + + sorted_dets = sorted(detections, key=lambda d: d.confidence, reverse=True) + kept: List[Detection] = [] + + while sorted_dets: + current = sorted_dets.pop(0) + kept.append(current) + sorted_dets = [ + det + for det in sorted_dets + if self._iou(current, det) <= self.iou_threshold + ] + if len(kept) >= self.max_faces_per_frame: + break + + return kept + + @staticmethod + def _iou(a: Detection, b: Detection) -> float: + """Compute IoU between two detections.""" + x1 = max(a.x1, b.x1) + y1 = max(a.y1, b.y1) + x2 = min(a.x2, b.x2) + y2 = min(a.y2, b.y2) + + inter_area = max(0, x2 - x1) * max(0, y2 - y1) + area_a = (a.x2 - a.x1) * (a.y2 - a.y1) + area_b = (b.x2 - b.x1) * (b.y2 - b.y1) + union_area = area_a + area_b - inter_area + if union_area == 0: + return 0.0 + return inter_area / union_area + + def extract_crops( + self, + detections: List[Detection], + output_dir: str, + crop_size: Tuple[int, int] = (224, 224), + ) -> List[Detection]: + """Extract resized face crops from original frames and update detections.""" + from PIL import Image + + output_dir_path = Path(output_dir) + output_dir_path.mkdir(parents=True, exist_ok=True) + cropped: List[Detection] = [] + + for idx, det in enumerate(detections): + try: + image = Image.open(det.frame_path).convert("RGB") + width, height = image.size + + x1 = int(max(0, det.x1 * width / self.input_size)) + y1 = int(max(0, det.y1 * height / self.input_size)) + x2 = int(min(width, det.x2 * width / self.input_size)) + y2 = int(min(height, det.y2 * height / self.input_size)) + + if x2 <= x1 or y2 <= y1: + image.close() + continue + + crop = image.crop((x1, y1, x2, y2)).resize(crop_size, Image.Resampling.BILINEAR) + image.close() + crop_path = output_dir_path / f"crop_{int(time.time() * 1000)}_{idx}.jpg" + crop.save(crop_path, "JPEG", quality=95) + crop.close() + det.crop_path = str(crop_path) + cropped.append(det) + except Exception as exc: + logger.warning("Failed to extract crop for %s: %s", det.frame_path, exc) + + return cropped diff --git a/src/gpu_manager.py b/src/gpu_manager.py new file mode 100644 index 0000000..4f119d1 --- /dev/null +++ b/src/gpu_manager.py @@ -0,0 +1,101 @@ +"""GPU memory management and batch-size auto-tuning helpers.""" + +import logging +from typing import Optional, Tuple + +logger = logging.getLogger(__name__) + + +class GPUMemoryManager: + """Track GPU memory and recommend safe batch sizes.""" + + def __init__( + self, + device: int = 0, + max_memory_gb: float = 18.0, + reduce_threshold_gb: float = 16.0, + increase_threshold_gb: float = 10.0, + initial_batch_size: int = 16, + min_batch_size: int = 1, + max_batch_size: int = 32, + ): + self.device = device + self.max_memory_gb = max_memory_gb + self.reduce_threshold_gb = reduce_threshold_gb + self.increase_threshold_gb = increase_threshold_gb + self.batch_size = initial_batch_size + self.min_batch_size = min_batch_size + self.max_batch_size = max_batch_size + + def get_memory_stats(self) -> Tuple[float, float]: + """Return (allocated_gb, reserved_gb) for the managed device.""" + try: + import torch + + if not torch.cuda.is_available(): + return 0.0, 0.0 + + allocated = torch.cuda.memory_allocated(self.device) / (1024**3) + reserved = torch.cuda.memory_reserved(self.device) / (1024**3) + return allocated, reserved + except Exception as exc: + logger.debug("Could not query GPU memory: %s", exc) + return 0.0, 0.0 + + def adjust_batch_size(self) -> int: + """Adjust the current batch size based on memory pressure.""" + allocated, _ = self.get_memory_stats() + + if allocated >= self.reduce_threshold_gb: + new_batch_size = max(self.min_batch_size, int(self.batch_size * 0.75)) + if new_batch_size < self.batch_size: + logger.info( + "GPU memory high (%.2f GB), reducing batch size %d -> %d", + allocated, + self.batch_size, + new_batch_size, + ) + self.batch_size = new_batch_size + elif allocated <= self.increase_threshold_gb: + new_batch_size = min(self.max_batch_size, int(self.batch_size * 1.25)) + if new_batch_size > self.batch_size: + logger.info( + "GPU memory low (%.2f GB), increasing batch size %d -> %d", + allocated, + self.batch_size, + new_batch_size, + ) + self.batch_size = new_batch_size + + return self.batch_size + + def is_memory_critical(self) -> bool: + """Return True if allocated memory is close to the hard limit.""" + allocated, _ = self.get_memory_stats() + return allocated >= self.max_memory_gb + + def empty_cache(self): + """Try to release unused cached GPU memory.""" + try: + import torch + + if torch.cuda.is_available(): + torch.cuda.empty_cache() + except Exception as exc: + logger.debug("Could not empty GPU cache: %s", exc) + + @property + def current_batch_size(self) -> int: + return self.batch_size + + +def get_available_vram_gb(device: int = 0) -> float: + """Return total available VRAM in GB, or 0.0 if CUDA is unavailable.""" + try: + import torch + + if not torch.cuda.is_available(): + return 0.0 + return torch.cuda.get_device_properties(device).total_memory / (1024**3) + except Exception: + return 0.0 diff --git a/src/orchestrator.py b/src/orchestrator.py index 2450bf9..85dba0e 100644 --- a/src/orchestrator.py +++ b/src/orchestrator.py @@ -13,7 +13,10 @@ from enum import Enum from pathlib import Path from typing import Dict, List, Optional +from batcher import DynamicBatcher +from face_detector import Detection, FaceDetector from frame_sampler import FrameSampler +from gpu_manager import GPUMemoryManager from prober import VideoProber from scratch_manager import ScratchManager @@ -81,6 +84,22 @@ class WorkerPool: self._jobs_failed = 0 self._sampling_config = (config or {}).get("sampling", {}) self._storage_config = (config or {}).get("storage", {}) + self._face_detection_config = (config or {}).get("face_detection", {}) + self._batching_config = (config or {}).get("batching", {}) + self._gpu_manager = GPUMemoryManager( + max_memory_gb=config.get("gpu", {}).get("max_memory_gb", 18.0), + reduce_threshold_gb=self._batching_config.get("vram_reduce_threshold_gb", 16.0), + increase_threshold_gb=self._batching_config.get("vram_increase_threshold_gb", 10.0), + initial_batch_size=self._batching_config.get("max_batch_size", 16), + ) + self._face_detector = FaceDetector( + engine_path=self._face_detection_config.get("model_path", "/models/face_detector/face_detector.trt"), + input_size=int(self._face_detection_config.get("input_size", 640)), + confidence_threshold=float(self._face_detection_config.get("confidence_threshold", 0.25)), + iou_threshold=float(self._face_detection_config.get("iou_threshold", 0.45)), + max_faces_per_frame=int(self._face_detection_config.get("max_faces_per_frame", 10)), + max_faces_per_video=int(self._face_detection_config.get("max_faces_per_video", 100)), + ) def start(self): """Start the worker pool.""" @@ -189,6 +208,28 @@ class WorkerPool: scratch_manager.cleanup() raise RuntimeError("No frames were extracted") + # Face detection on sampled frames + batch_size = self._gpu_manager.current_batch_size + detections_per_frame = self._face_detector.detect_faces(extracted_frames, batch_size=batch_size) + + # Flatten and cap total faces per video + all_detections: List[Detection] = [] + for frame_dets in detections_per_frame: + all_detections.extend(frame_dets) + all_detections = sorted(all_detections, key=lambda d: d.confidence, reverse=True) + all_detections = all_detections[: self._face_detection_config.get("max_faces_per_video", 100)] + + # Extract face crops + crop_dir = scratch_manager.frame_dir.parent / "crops" + cropped_detections = self._face_detector.extract_crops( + all_detections, + output_dir=str(crop_dir), + crop_size=(224, 224), + ) + + if not cropped_detections: + logger.info("No faces detected for video %s", job.video_id) + self._complete_job(job, frame_count=len(extracted_frames)) scratch_manager.cleanup() self._jobs_processed += 1 diff --git a/tests/test_story_04.py b/tests/test_story_04.py new file mode 100644 index 0000000..20760e9 --- /dev/null +++ b/tests/test_story_04.py @@ -0,0 +1,98 @@ +import sys +import tempfile +import unittest +from pathlib import Path +from unittest.mock import MagicMock, patch + +import numpy as np + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src")) + +from batcher import DynamicBatcher +from face_detector import Detection, FaceDetector +from gpu_manager import GPUMemoryManager + + +class Story04FaceDetectionTests(unittest.TestCase): + def test_gpu_manager_respects_memory_thresholds(self): + manager = GPUMemoryManager( + max_memory_gb=18.0, + reduce_threshold_gb=16.0, + increase_threshold_gb=10.0, + initial_batch_size=16, + ) + + with patch.object(manager, "get_memory_stats", return_value=(17.0, 18.0)): + manager.adjust_batch_size() + self.assertLess(manager.current_batch_size, 16) + + with patch.object(manager, "get_memory_stats", return_value=(8.0, 10.0)): + manager.adjust_batch_size() + self.assertGreater(manager.current_batch_size, manager.min_batch_size) + + def test_dynamic_batcher_flushes_when_batch_full(self): + processed_batches = [] + + def process_fn(batch): + processed_batches.append(batch) + return [len(batch)] + + batcher = DynamicBatcher( + process_fn=process_fn, + max_batch_size=3, + batch_timeout_ms=10000.0, + ) + + for i in range(5): + batcher.add(i) + + self.assertEqual(len(processed_batches), 1) + self.assertEqual(processed_batches[0], [0, 1, 2]) + self.assertEqual(batcher.queued_count, 2) + + def test_face_detector_nms_removes_overlapping_boxes(self): + detector = FaceDetector( + engine_path="/nonexistent/model.trt", + confidence_threshold=0.1, + iou_threshold=0.45, + max_faces_per_frame=10, + ) + + duplicates = [ + Detection("frame.jpg", 10, 10, 50, 50, 0.9), + Detection("frame.jpg", 12, 12, 48, 48, 0.8), + Detection("frame.jpg", 100, 100, 150, 150, 0.75), + ] + + kept = detector._nms(duplicates) + self.assertEqual(len(kept), 2) + self.assertAlmostEqual(kept[0].confidence, 0.9, places=5) + + def test_face_detector_extracts_and_resizes_crops(self): + detector = FaceDetector( + engine_path="/nonexistent/model.trt", + input_size=640, + confidence_threshold=0.25, + ) + + with tempfile.TemporaryDirectory() as tmpdir: + from PIL import Image + + frame_path = Path(tmpdir) / "frame.jpg" + image = Image.new("RGB", (640, 480), color=(100, 150, 200)) + image.save(frame_path) + + detections = [ + Detection(str(frame_path), 0, 0, 640, 480, 0.95), + ] + cropped = detector.extract_crops(detections, output_dir=tmpdir, crop_size=(224, 224)) + + self.assertEqual(len(cropped), 1) + self.assertTrue(Path(cropped[0].crop_path).exists()) + + with Image.open(cropped[0].crop_path) as crop: + self.assertEqual(crop.size, (224, 224)) + + +if __name__ == "__main__": + unittest.main()