Spaces:
Sleeping
Sleeping
File size: 4,229 Bytes
735a97b | 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 | """
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()
|