Spaces:
Running on Zero
Running on Zero
File size: 2,759 Bytes
e967c28 | 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 | """
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
|