Spaces:
Running on Zero
Running on Zero
| """ | |
| Part 2 entry point: run race classification, skin tone extraction, and | |
| CLIP/face embeddings over every image in a Part 1 manifest, one image at | |
| a time (resumable, same philosophy as Part 1's run_generation). | |
| Which analyses run is selectable via `components` -- e.g. run only skin | |
| tone + face embeddings and skip CLIP (useful when CLIP's ~4GB download | |
| won't fit on disk). Default is all four. | |
| Usage: | |
| import pandas as pd | |
| from src.analyze import run_analysis | |
| manifest_df = pd.read_parquet("runs/my_audit/manifest.parquet") | |
| analysis_path, enriched_df = run_analysis(manifest_df, output_dir="runs/my_audit") | |
| # or a subset: | |
| run_analysis(manifest_df, "runs/my_audit", components={"skin", "face"}) | |
| """ | |
| import glob | |
| import os | |
| from PIL import Image | |
| from .zerogpu import GPU | |
| from .analysis_writer import AnalysisWriter | |
| from .race_classification import classify_race | |
| from .skin_tone import analyze_skin_tone_from_img | |
| from .embeddings import embed_clip, embed_face | |
| ALL_COMPONENTS = {"race", "skin", "clip", "face"} | |
| def _zerogpu_startup_probe(): | |
| """ZeroGPU requires at least one @spaces.GPU function to exist at startup. | |
| We deliberately DON'T route the real analysis through a per-image GPU call: | |
| ZeroGPU forks a fresh process per call, so lazily-loaded models (the 4GB | |
| CLIP + InsightFace) reloaded on every image and blew past the GPU token's | |
| lifetime ("expired GPU proxy token"). Instead the model steps run in the | |
| main process (below), where models load once and are reused. This probe | |
| exists only to satisfy the startup requirement.""" | |
| return True | |
| def _run_model_analyses(pil_image, components): | |
| """Skin tone (BiSeNet), CLIP, and face embeddings for one image. Runs in | |
| the main process (CPU on ZeroGPU) so the heavy models load once and are | |
| reused across images. Race classification is separate -- it's a Gemini API | |
| call handled in _analyze_one. | |
| Collects step failures in fields["_errors"] (a list) so a bad step doesn't | |
| lose the others; _analyze_one folds those into the final `error` column.""" | |
| fields = {"_errors": []} | |
| if "skin" in components: | |
| skin = analyze_skin_tone_from_img(pil_image) | |
| if "error" in skin: | |
| fields["skin_tone_error"] = skin["error"] | |
| else: | |
| fields["L"], fields["a"], fields["b"], fields["h_star"] = ( | |
| skin["L"], skin["a"], skin["b"], skin["h_star"] | |
| ) | |
| if "clip" in components: | |
| try: | |
| fields["clip_emb"] = embed_clip(pil_image).tolist() | |
| except Exception as e: | |
| fields["clip_emb"] = None | |
| fields["_errors"].append(f"clip:{type(e).__name__}:{e}") | |
| if "face" in components: | |
| try: | |
| face_vec, det_score, bbox, n_faces, status = embed_face(pil_image) | |
| fields["face_emb"] = face_vec.tolist() if face_vec is not None else None | |
| fields["face_det_score"] = det_score | |
| fields["face_bbox_xywh"] = bbox | |
| fields["n_faces_detected"] = n_faces | |
| fields["embed_status"] = status | |
| except Exception as e: | |
| fields["embed_status"] = f"error:{type(e).__name__}" | |
| fields["_errors"].append(f"face:{type(e).__name__}:{e}") | |
| return fields | |
| def _analyze_one(pil_image, components, gemini_api_key=None): | |
| """Runs the selected analyses for one image. Never raises -- a failure | |
| of an individual step is recorded in the `error` field so one bad step | |
| doesn't lose the others. GPU-bound steps run in _gpu_analyze_one (which | |
| ZeroGPU attaches a GPU to); the Gemini race call stays on CPU. | |
| gemini_api_key: passed to the classifier so callers can supply a | |
| per-request key (bring-your-own-key UI) instead of relying on the | |
| GOOGLE_API_KEY env var.""" | |
| fields = {} | |
| errors = [] | |
| if "race" in components: | |
| try: | |
| fields["race_classification"] = classify_race(pil_image, api_key=gemini_api_key) | |
| except Exception as e: | |
| fields["race_classification"] = None | |
| errors.append(f"race_classification:{type(e).__name__}:{e}") | |
| model_fields = _run_model_analyses(pil_image, components) | |
| errors.extend(model_fields.pop("_errors", [])) | |
| fields.update(model_fields) | |
| if errors: | |
| fields["error"] = " | ".join(errors) | |
| return fields | |
| def _clear_prior_analysis(output_dir): | |
| """Remove any existing analysis chunks/outputs so a run starts fresh. | |
| Used when `force=True` (e.g. the user changed which components to run -- | |
| resumability is image-level, not component-level, so re-running with a | |
| different component set requires starting over).""" | |
| chunk_dir = os.path.join(output_dir, "analysis_chunks") | |
| for f in glob.glob(os.path.join(chunk_dir, "chunk_*.parquet")): | |
| os.remove(f) | |
| for name in ("analysis.parquet", "analyzed_manifest.parquet"): | |
| p = os.path.join(output_dir, name) | |
| if os.path.exists(p): | |
| os.remove(p) | |
| def run_analysis(manifest_df, output_dir, components=None, force=False, | |
| on_progress=None, gemini_api_key=None): | |
| """ | |
| manifest_df: the Part 1 manifest (must have image_id, file_path columns) | |
| output_dir: same run directory used for generation -- analysis_chunks/ | |
| and analysis.parquet are written alongside images/ and manifest.parquet | |
| components: which analyses to run, subset of {"race","skin","clip","face"}. | |
| None => all four. | |
| force: if True, discard any prior analysis for this run and start fresh. | |
| (Resumability is image-level; re-running with a different component | |
| set needs force=True to actually re-analyze already-seen images.) | |
| on_progress: optional callback(done, total) | |
| gemini_api_key: per-request Google API key for race classification | |
| (bring-your-own-key); falls back to the GOOGLE_API_KEY env var. | |
| Returns: (analysis_path, enriched_df) where enriched_df is manifest_df | |
| left-joined with the analysis results on image_id. | |
| """ | |
| components = ALL_COMPONENTS if components is None else set(components) | |
| if not components: | |
| raise ValueError("No analysis components selected.") | |
| if force: | |
| _clear_prior_analysis(output_dir) | |
| writer = AnalysisWriter(output_dir) | |
| total = len(manifest_df) | |
| done = 0 | |
| for _, row in manifest_df.iterrows(): | |
| image_id = row["image_id"] | |
| if writer.is_done(image_id): | |
| done += 1 | |
| if on_progress: | |
| on_progress(done, total) | |
| continue | |
| img = Image.open(row["file_path"]).convert("RGB") | |
| fields = _analyze_one(img, components, gemini_api_key=gemini_api_key) | |
| writer.add_row(image_id, **fields) | |
| done += 1 | |
| if on_progress: | |
| on_progress(done, total) | |
| analysis_path, analysis_df = writer.finalize() | |
| enriched_df = manifest_df.merge(analysis_df, on="image_id", how="left") | |
| enriched_path = os.path.join(output_dir, "analyzed_manifest.parquet") | |
| enriched_df.to_parquet(enriched_path, index=False) | |
| return enriched_path, enriched_df | |