Spaces:
Sleeping
Sleeping
| # app.py | |
| import os | |
| os.environ["TRANSFORMERS_NO_TF"] = "1" | |
| os.environ["TRANSFORMERS_NO_FLAX"] = "1" | |
| os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3" | |
| import gradio as gr | |
| from PIL import Image | |
| import numpy as np | |
| import torch | |
| from hf_model import PretrainedAgeEstimator | |
| from face_utils import FaceCropper | |
| # NEW: diffusers for cartoonizer | |
| from diffusers import StableDiffusionImg2ImgPipeline | |
| # ---------- Load models once ---------- | |
| est = PretrainedAgeEstimator() | |
| cropper = FaceCropper(device=est.device) | |
| # A solid, public SD 1.5 img2img pipeline; fast and reliable | |
| SD15_ID = "runwayml/stable-diffusion-v1-5" | |
| sd_pipe = StableDiffusionImg2ImgPipeline.from_pretrained( | |
| SD15_ID, | |
| torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, | |
| safety_checker=None, # rely on prompts; HF Spaces also has a global filter | |
| ).to(est.device) | |
| # ---------- Helpers ---------- | |
| def _ensure_pil(img): | |
| if isinstance(img, Image.Image): | |
| return img | |
| return Image.fromarray(img) | |
| # ----- Age: single image ----- | |
| def predict_single(img, auto_crop=True, topk=5, show_annot=True): | |
| if img is None: | |
| return {}, "No image provided.", None | |
| img = _ensure_pil(img).convert("RGB") | |
| preview = img | |
| face = None | |
| if auto_crop: | |
| face, annotated, _ = cropper.detect_and_crop(img, select="largest") | |
| preview = annotated if show_annot else img | |
| target = face if face is not None else img | |
| age, top = est.predict(target, topk=topk) | |
| probs = {lbl: float(prob) for lbl, prob in top} | |
| summary = f"**Estimated age:** {age:.1f} years" | |
| return probs, summary, preview | |
| # ----- Age: batch ----- | |
| def predict_batch(files, auto_crop=True, topk=5): | |
| if not files: | |
| return "No files uploaded." | |
| rows = ["| File | Estimated Age | Top-1 | p |", "|---|---:|---|---:|"] | |
| for f in files: | |
| try: | |
| img = Image.open(f.name).convert("RGB") | |
| face = None | |
| if auto_crop: | |
| face, _, _ = cropper.detect_and_crop(img, select="largest") | |
| target = face if face is not None else img | |
| age, top = est.predict(target, topk=topk) | |
| top1_lbl, top1_p = top[0] | |
| rows.append(f"| {os.path.basename(f.name)} | {age:.1f} | {top1_lbl} | {top1_p:.3f} |") | |
| except Exception: | |
| rows.append(f"| {os.path.basename(f.name)} | (error) | - | - |") | |
| return "\n".join(rows) | |
| # ----- NEW: Cartoonizer (img2img) ----- | |
| def cartoonize(img, prompt, strength=0.6, guidance=7.5, steps=25, seed=0, use_face_crop=True): | |
| """ | |
| img: PIL or numpy | |
| prompt: text description, e.g. "cute cel-shaded cartoon, soft outlines, vibrant colors" | |
| strength: how much to deviate from the input (0.3 subtle → 0.8 strong) | |
| guidance: prompt strength (5–12 typical) | |
| steps: diffusion steps (20–40 typical) | |
| seed: reproducibility (-1 for random) | |
| """ | |
| if img is None: | |
| return None | |
| img = _ensure_pil(img).convert("RGB") | |
| # optional crop to the largest face for better identity preservation | |
| if use_face_crop: | |
| face, _, _ = cropper.detect_and_crop(img, select="largest") | |
| if face is not None: | |
| img = face | |
| # cartoon-y defaults (you can tweak in UI) | |
| base_prompt = ( | |
| "cartoon, cel-shaded, clean lineart, smooth shading, high contrast, vibrant, studio ghibli style, " | |
| "pixar style, highly detailed, 2D illustration" | |
| ) | |
| full_prompt = f"{base_prompt}, {prompt}".strip().strip(",") | |
| generator = None | |
| if seed and seed >= 0: | |
| generator = torch.Generator(device=est.device).manual_seed(int(seed)) | |
| out = sd_pipe( | |
| prompt=full_prompt, | |
| image=img, | |
| strength=float(strength), | |
| guidance_scale=float(guidance), | |
| num_inference_steps=int(steps), | |
| generator=generator, | |
| ) | |
| result = out.images[0] | |
| return result | |
| # ---------- UI ---------- | |
| with gr.Blocks(title="Pretrained Age Estimator + Cartoonizer") as demo: | |
| gr.Markdown("# Pretrained Age Estimator + Cartoonizer") | |
| gr.Markdown("Detects age from a face and can also generate a cartoonized image guided by your text description.") | |
| with gr.Tabs(): | |
| with gr.Tab("Age (Single)"): | |
| with gr.Row(): | |
| with gr.Column(): | |
| inp = gr.Image(type="pil", label="Upload a face image") | |
| cam = gr.Image(sources=["webcam"], type="pil", label="Webcam (optional)") | |
| auto = gr.Checkbox(True, label="Auto face crop (MTCNN)") | |
| topk = gr.Slider(3, 9, value=5, step=1, label="Top-K age ranges") | |
| annot = gr.Checkbox(True, label="Show detection preview") | |
| btn = gr.Button("Predict Age", variant="primary") | |
| with gr.Column(): | |
| out_label = gr.Label(num_top_classes=5, label="Age Prediction (probabilities)") | |
| out_md = gr.Markdown(label="Summary") | |
| out_prev = gr.Image(label="Preview", visible=True) | |
| def run_single(img, cam_img, auto_crop, topk_val, show_annot): | |
| chosen = cam_img if cam_img is not None else img | |
| return predict_single(chosen, auto_crop, int(topk_val), show_annot) | |
| btn.click(fn=run_single, inputs=[inp, cam, auto, topk, annot], | |
| outputs=[out_label, out_md, out_prev]) | |
| with gr.Tab("Age (Batch)"): | |
| files = gr.Files(label="Upload multiple images") | |
| auto_b = gr.Checkbox(True, label="Auto face crop (MTCNN)") | |
| topk_b = gr.Slider(3, 9, value=5, step=1, label="Top-K age ranges") | |
| btn_b = gr.Button("Run batch") | |
| out_table = gr.Markdown() | |
| btn_b.click(fn=predict_batch, inputs=[files, auto_b, topk_b], outputs=out_table) | |
| with gr.Tab("Cartoonizer"): | |
| src = gr.Image(type="pil", label="Source image (face or any photo)") | |
| prompt = gr.Textbox(label="Your style prompt", | |
| value="cute cel-shaded cartoon, clean lines, soft colors") | |
| with gr.Row(): | |
| strength = gr.Slider(0.2, 0.95, value=0.6, step=0.05, label="Transformation strength") | |
| guidance = gr.Slider(3, 15, value=7.5, step=0.5, label="Guidance scale") | |
| steps = gr.Slider(10, 50, value=25, step=1, label="Steps") | |
| seed = gr.Number(value=0, precision=0, label="Seed (0 or -1 = random)") | |
| use_crop = gr.Checkbox(True, label="Crop to largest face before stylizing") | |
| btn_c = gr.Button("Generate Cartoon", variant="primary") | |
| out_img = gr.Image(label="Cartoon result") | |
| btn_c.click(fn=cartoonize, | |
| inputs=[src, prompt, strength, guidance, steps, seed, use_crop], | |
| outputs=out_img) | |
| if __name__ == "__main__": | |
| demo.launch() | |