Spaces:
Sleeping
Sleeping
File size: 3,593 Bytes
6f6e572 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 |
"""
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
@property
def total_ai(self) -> int:
"""Return total number of AI-generated images"""
return len(self.ai_images)
@property
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)
|