Spaces:
Sleeping
Sleeping
| """ | |
| Image Management Module | |
| Responsible for indexing, validating, and selecting facial images required for the game. | |
| """ | |
| import os | |
| import random | |
| from pathlib import Path | |
| from typing import List, Dict, Literal | |
| class ImagePool: | |
| """Manages indexing and selection of AI-generated and real human face images""" | |
| def __init__(self, ai_dir: str = "image/ai", real_dir: str = "image/real"): | |
| """ | |
| Initialize image pool | |
| Args: | |
| ai_dir: Directory path for AI-generated images | |
| real_dir: Directory path for real human images | |
| Raises: | |
| ValueError: If either directory contains less than 10 images | |
| """ | |
| self.ai_dir = ai_dir | |
| self.real_dir = real_dir | |
| self.ai_images = self._scan_directory(ai_dir) | |
| self.real_images = self._scan_directory(real_dir) | |
| self.validate() | |
| def _scan_directory(self, directory: str) -> List[str]: | |
| """ | |
| Scan directory and return paths of all valid image files | |
| Args: | |
| directory: Directory path to scan | |
| Returns: | |
| List of image file paths | |
| """ | |
| valid_extensions = {'.jpg', '.jpeg', '.png', '.webp', '.bmp'} | |
| image_paths = [] | |
| if not os.path.exists(directory): | |
| return image_paths | |
| for file in Path(directory).rglob('*'): | |
| if file.is_file() and file.suffix.lower() in valid_extensions: | |
| image_paths.append(str(file)) | |
| return sorted(image_paths) # Sort for consistency | |
| def total_ai(self) -> int: | |
| """Return total number of AI-generated images""" | |
| return len(self.ai_images) | |
| def total_real(self) -> int: | |
| """Return total number of real human images""" | |
| return len(self.real_images) | |
| def validate(self) -> None: | |
| """ | |
| Validate that image count meets requirements | |
| Raises: | |
| ValueError: If either category has less than 10 images | |
| """ | |
| if self.total_ai < 10: | |
| raise ValueError( | |
| f"Insufficient AI images: need at least 10, found {self.total_ai} in {self.ai_dir}" | |
| ) | |
| if self.total_real < 10: | |
| raise ValueError( | |
| f"Insufficient real images: need at least 10, found {self.total_real} in {self.real_dir}" | |
| ) | |
| def create_game_set(self) -> List[Dict[str, str]]: | |
| """ | |
| Create a random set of 20 images for a new game (10 AI + 10 real) | |
| Returns: | |
| List of 20 ImageItem dictionaries, each containing 'path' and 'label' fields | |
| Images are in random order | |
| """ | |
| # Randomly select 10 AI images | |
| selected_ai = random.sample(self.ai_images, 10) | |
| # Randomly select 10 real images | |
| selected_real = random.sample(self.real_images, 10) | |
| # Create ImageItem dictionary list | |
| all_images = [ | |
| {"path": path, "label": "AI"} for path in selected_ai | |
| ] + [ | |
| {"path": path, "label": "Human"} for path in selected_real | |
| ] | |
| # Shuffle order | |
| random.shuffle(all_images) | |
| return all_images | |
| def get_random_image(self, label: Literal["AI", "Human"]) -> str: | |
| """ | |
| Get a random image path from specified category | |
| Args: | |
| label: Image category, "AI" or "Human" | |
| Returns: | |
| Randomly selected image path | |
| """ | |
| if label == "AI": | |
| return random.choice(self.ai_images) | |
| else: | |
| return random.choice(self.real_images) | |