Spaces:
Running
on
Zero
Running
on
Zero
| #!/usr/bin/env python | |
| from __future__ import annotations | |
| import os | |
| import pathlib | |
| import sys | |
| import cv2 | |
| import gradio as gr | |
| import huggingface_hub | |
| import numpy as np | |
| import spaces | |
| import torch | |
| from huggingface_hub import hf_hub_download | |
| torch.jit.script = lambda f: f | |
| sys.path.insert(0, "face_detection") | |
| sys.path.insert(0, "face_parsing") | |
| sys.path.insert(0, "fpage") | |
| sys.path.insert(0, "roi_tanh_warping") | |
| from ibug.age_estimation import AgeEstimator | |
| from ibug.face_detection import RetinaFacePredictor | |
| from ibug.face_parsing.utils import label_colormap | |
| DESCRIPTION = "# [FP-Age](https://github.com/ibug-group/fpage)" | |
| def is_lfs_pointer_file(path: pathlib.Path) -> bool: | |
| try: | |
| with open(path, "r") as f: | |
| # Git LFS pointer files usually start with version line | |
| version_line = f.readline() | |
| if version_line.startswith("version https://git-lfs.github.com/spec/"): | |
| # Check for the presence of oid and size lines | |
| oid_line = f.readline() | |
| size_line = f.readline() | |
| if oid_line.startswith("oid sha256:") and size_line.startswith("size "): | |
| return True | |
| except Exception as e: | |
| print(f"Error reading file {path}: {e}") | |
| return False | |
| lfs_model_paths = sorted(pathlib.Path("face_parsing").rglob("*.torch")) | |
| for lfs_model_path in lfs_model_paths: | |
| if is_lfs_pointer_file(lfs_model_path): | |
| os.remove(lfs_model_path) | |
| out_path = hf_hub_download( | |
| "public-data/ibug-face-parsing", | |
| filename=lfs_model_path.name, | |
| repo_type="model", | |
| subfolder=lfs_model_path.parts[-3], | |
| ) | |
| os.symlink(out_path, lfs_model_path) | |
| device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") | |
| detector = RetinaFacePredictor(threshold=0.8, device="cpu", model=RetinaFacePredictor.get_model("mobilenet0.25")) | |
| detector.device = device | |
| detector.net.to(device) | |
| model = AgeEstimator( | |
| device=device, | |
| ckpt=huggingface_hub.hf_hub_download("hysts/ibug", "fpage/models/fpage-resnet50-fcn-14-97.torch"), | |
| encoder="resnet50", | |
| decoder="fcn", | |
| age_classes=97, | |
| face_classes=14, | |
| ) | |
| def predict(image: np.ndarray, max_num_faces: int) -> np.ndarray: | |
| colormap = label_colormap(14) | |
| # RGB -> BGR | |
| image = image[:, :, ::-1] | |
| faces = detector(image, rgb=False) | |
| if len(faces) == 0: | |
| raise RuntimeError("No face was found.") | |
| faces = sorted(list(faces), key=lambda x: -x[4])[:max_num_faces][::-1] | |
| ages, masks = model.predict_img(image, faces, rgb=False) | |
| mask_image = np.zeros_like(image) | |
| for mask in masks: | |
| temp = colormap[mask] | |
| mask_image[temp > 0] = temp[temp > 0] | |
| res = image.astype(float) * 0.5 + mask_image[:, :, ::-1] * 0.5 | |
| res = np.clip(np.round(res), 0, 255).astype(np.uint8) | |
| for face, age in zip(faces, ages): | |
| bbox = np.round(face[:4]).astype(int) | |
| cv2.rectangle( | |
| res, | |
| (bbox[0], bbox[1]), | |
| (bbox[2], bbox[3]), | |
| color=(0, 255, 0), | |
| thickness=2, | |
| ) | |
| text_content = f"Age: ({age: .1f})" | |
| cv2.putText( | |
| res, | |
| text_content, | |
| (bbox[0], bbox[1] - 10), | |
| cv2.FONT_HERSHEY_DUPLEX, | |
| 1, | |
| (255, 255, 255), | |
| lineType=cv2.LINE_AA, | |
| ) | |
| return res[:, :, ::-1] | |
| with gr.Blocks(css="style.css") as demo: | |
| gr.Markdown(DESCRIPTION) | |
| with gr.Row(): | |
| with gr.Column(): | |
| image = gr.Image(label="Input") | |
| max_num_faces = gr.Slider(label="Max Number of Faces", minimum=1, maximum=20, step=1, value=5) | |
| run_button = gr.Button() | |
| with gr.Column(): | |
| result = gr.Image(label="Output") | |
| gr.Examples( | |
| examples=[[path.as_posix(), 5] for path in sorted(pathlib.Path("images").rglob("*.jpg"))], | |
| inputs=[image, max_num_faces], | |
| outputs=result, | |
| fn=predict, | |
| ) | |
| run_button.click( | |
| fn=predict, | |
| inputs=[image, max_num_faces], | |
| outputs=result, | |
| api_name="predict", | |
| ) | |
| if __name__ == "__main__": | |
| demo.queue(max_size=20).launch() | |