Spaces:
Running on Zero
Running on Zero
| """ | |
| Resumable chunked writer for Part 2 analysis results, one row per | |
| image_id. Same chunk-then-consolidate pattern as ManifestWriter (see | |
| manifest.py) so a crashed/interrupted analysis run can resume without | |
| re-calling Gemini or re-computing embeddings for images already done. | |
| """ | |
| import glob | |
| import os | |
| from datetime import datetime, timezone | |
| import pandas as pd | |
| ANALYSIS_COLUMNS = [ | |
| "image_id", | |
| "race_classification", | |
| "L", "a", "b", "h_star", "skin_tone_error", | |
| "clip_emb", "face_emb", "face_det_score", "face_bbox_xywh", | |
| "n_faces_detected", "embed_status", | |
| "error", | |
| "analyzed_at", | |
| ] | |
| class AnalysisWriter: | |
| def __init__(self, output_dir, checkpoint_every=25): | |
| self.output_dir = output_dir | |
| self.chunk_dir = os.path.join(output_dir, "analysis_chunks") | |
| os.makedirs(self.chunk_dir, exist_ok=True) | |
| self.checkpoint_every = checkpoint_every | |
| self._buffer = [] | |
| self._chunk_idx = len(glob.glob(os.path.join(self.chunk_dir, "chunk_*.parquet"))) | |
| self.done_image_ids = self._load_done_ids() | |
| def _load_done_ids(self): | |
| chunks = sorted(glob.glob(os.path.join(self.chunk_dir, "chunk_*.parquet"))) | |
| if not chunks: | |
| return set() | |
| dfs = [pd.read_parquet(c, columns=["image_id"]) for c in chunks] | |
| return set(pd.concat(dfs, ignore_index=True)["image_id"]) | |
| def is_done(self, image_id): | |
| return image_id in self.done_image_ids | |
| def add_row(self, image_id, **fields): | |
| row = {"image_id": image_id, "analyzed_at": datetime.now(timezone.utc).isoformat()} | |
| row.update({col: fields.get(col) for col in ANALYSIS_COLUMNS if col not in row}) | |
| self._buffer.append(row) | |
| self.done_image_ids.add(image_id) | |
| if len(self._buffer) >= self.checkpoint_every: | |
| self.flush() | |
| return image_id | |
| def flush(self): | |
| if not self._buffer: | |
| return | |
| df = pd.DataFrame(self._buffer, columns=ANALYSIS_COLUMNS) | |
| path = os.path.join(self.chunk_dir, f"chunk_{self._chunk_idx:05d}.parquet") | |
| df.to_parquet(path, index=False) | |
| self._chunk_idx += 1 | |
| self._buffer = [] | |
| def finalize(self, out_path=None): | |
| self.flush() | |
| out_path = out_path or os.path.join(self.output_dir, "analysis.parquet") | |
| chunks = sorted(glob.glob(os.path.join(self.chunk_dir, "chunk_*.parquet"))) | |
| if not chunks: | |
| raise RuntimeError("No chunks to consolidate -- did analysis produce any rows?") | |
| df = pd.concat([pd.read_parquet(c) for c in chunks], ignore_index=True) | |
| df = df.drop_duplicates(subset=["image_id"], keep="last").reset_index(drop=True) | |
| df.to_parquet(out_path, index=False) | |
| return out_path, df | |