Spaces:
Sleeping
Sleeping
| """ | |
| Scene Image Selector - Randomly selects candidate images for each scene | |
| Saves selections to a folder before composition | |
| """ | |
| import os | |
| import json | |
| import random | |
| import shutil | |
| from pathlib import Path | |
| class SceneSelector: | |
| def __init__(self, manifest_path, candidates_dir, output_dir="selected"): | |
| """ | |
| Initialize the scene selector. | |
| Args: | |
| manifest_path: Path to manifest_response.json | |
| candidates_dir: Path to candidates folder with scene_0-6 subfolders | |
| output_dir: Where to save selected images | |
| """ | |
| self.manifest_path = manifest_path | |
| self.candidates_dir = candidates_dir | |
| self.output_dir = output_dir | |
| self.selections = {} | |
| # Create output directory | |
| Path(self.output_dir).mkdir(parents=True, exist_ok=True) | |
| # Load manifest | |
| with open(manifest_path, 'r') as f: | |
| self.manifest = json.load(f) | |
| print(f"[Selector] Initialized with {len(self.manifest.get('scenes', []))} scenes") | |
| print(f"[Selector] Candidates dir: {candidates_dir}") | |
| print(f"[Selector] Output dir: {output_dir}\n") | |
| def select_all_scenes(self): | |
| """Randomly select one image per scene.""" | |
| scenes = self.manifest.get("scenes", []) | |
| print(f"[Selector] Selecting images for {len(scenes)} scenes...\n") | |
| for scene_idx, scene in enumerate(scenes): | |
| self._select_scene(scene_idx, scene) | |
| # Save selections metadata | |
| self._save_selections_metadata() | |
| print(f"\n[Selector] ✓ Selected {len(scenes)} images!") | |
| print(f"[Selector] Check {self.output_dir}/ for selected images") | |
| print(f"[Selector] Check {self.output_dir}/selections.json for metadata") | |
| def _select_scene(self, scene_idx, scene): | |
| """Randomly select one image for a scene.""" | |
| scene_label = scene.get("label", f"Scene {scene_idx}") | |
| candidates_folder = os.path.join(self.candidates_dir, f"scene_{scene_idx}") | |
| # Get list of candidate images | |
| if not os.path.exists(candidates_folder): | |
| print(f"[Scene {scene_idx}] ✗ Folder not found: {candidates_folder}") | |
| return | |
| candidates = [f for f in os.listdir(candidates_folder) if f.endswith(('.jpg', '.png'))] | |
| if not candidates: | |
| print(f"[Scene {scene_idx}] ✗ No images found in {candidates_folder}") | |
| return | |
| # Randomly select one image | |
| selected_image = random.choice(candidates) | |
| source_path = os.path.join(candidates_folder, selected_image) | |
| # Copy to output folder | |
| output_filename = f"scene_{scene_idx:02d}.jpg" | |
| dest_path = os.path.join(self.output_dir, output_filename) | |
| shutil.copy2(source_path, dest_path) | |
| # Store selection metadata | |
| self.selections[scene_idx] = { | |
| "scene_label": scene_label, | |
| "original_file": selected_image, | |
| "source_folder": candidates_folder, | |
| "total_candidates": len(candidates), | |
| "output_file": output_filename | |
| } | |
| print(f"[Scene {scene_idx}] {scene_label}") | |
| print(f" Available: {len(candidates)} candidates") | |
| print(f" Selected: {selected_image}") | |
| print(f" Copied to: {output_filename}\n") | |
| def _save_selections_metadata(self): | |
| """Save metadata about which images were selected.""" | |
| metadata_path = os.path.join(self.output_dir, "selections.json") | |
| with open(metadata_path, 'w') as f: | |
| json.dump(self.selections, f, indent=2) | |
| print(f"[Selector] Saved selections metadata to {metadata_path}") | |
| if __name__ == "__main__": | |
| # Load manifest from content_gen_server | |
| manifest_file = "../content_gen_server/manifest_response.json" | |
| candidates_dir = "../workspace/renders/candidates" | |
| output_dir = "selected" | |
| # Create and run selector | |
| selector = SceneSelector(manifest_file, candidates_dir, output_dir) | |
| selector.select_all_scenes() | |