Spaces:
Running on Zero
Running on Zero
| """Inference adapters shared by the Gradio app and local verification.""" | |
| import gc | |
| from pathlib import Path | |
| import time | |
| import common.fix_chumpy # noqa: F401 | |
| import numpy as np | |
| import torch | |
| from camera import PerspectiveCamera | |
| from common.config import (FLAMEConfig, FaceBoxConfig, RealDenseFaceInferenceConfig, | |
| GNFlameOptimizerConfig, load_fitting_config) | |
| from tracker import RealDenseFaceInferencer, GNFlameOptimizer | |
| from visualization import Visualizer | |
| class ReconstructionEngine: | |
| def __init__(self, root: Path, model_weights: str, facebox_weights: str, flame_assets: str, flame_model: str): | |
| self.root = root | |
| self.flame = FLAMEConfig(flame_model_path=flame_model, flame_assets_path=flame_assets) | |
| self.inferencer = RealDenseFaceInferencer( | |
| RealDenseFaceInferenceConfig( | |
| model_config_path=str(root / "configs/model/vitb.yaml"), | |
| model_weights_path=model_weights, flame=self.flame, | |
| compile_model=False, defer_native_context=True, | |
| ), FaceBoxConfig(model_weights_path=facebox_weights), | |
| ) | |
| self.fitting_config = load_fitting_config(root / "configs/fitting/single_image.yaml", GNFlameOptimizerConfig, | |
| flame=self.flame, device="cuda") | |
| def reconstruct(self, image: np.ndarray): | |
| started = time.perf_counter() | |
| # Preserve the official pipeline's bbox refinement and fallback: some | |
| # heavily occluded teaser faces are missed by the separate detector. | |
| inference = self.inferencer.infer_image(image, face_bbox=None) | |
| optimizer = None | |
| visualizer = None | |
| try: | |
| optimizer = GNFlameOptimizer(self.fitting_config) | |
| fit = optimizer.fit_frame(inference, camera=None) | |
| vertices, _ = optimizer.decode(fit) | |
| if not torch.isfinite(vertices).all().item(): | |
| raise RuntimeError("Fitting produced non-finite vertices") | |
| camera = PerspectiveCamera( | |
| fov_y=np.radians(float(fit.camera_fov_y)), | |
| rot=self.fitting_config.camera_rotation, pos=self.fitting_config.camera_position, | |
| width=image.shape[1], height=image.shape[0], | |
| ) | |
| visualizer = Visualizer(flame_assets_path=self.flame.flame_assets_path) | |
| dense = visualizer.vis_align(image, inference.vertex_coord) | |
| dense = visualizer.vis_contour(dense, inference.vertex_coord) | |
| overlay, mesh = visualizer.vis_flame_shading_overlay(camera, image, vertices, return_shading=True) | |
| torch.cuda.synchronize() | |
| return dense, overlay, mesh, { | |
| "seconds": round(time.perf_counter() - started, 3), | |
| "gpu": torch.cuda.get_device_name(0), | |
| "fov_y": float(fit.camera_fov_y), | |
| "vertices": int(vertices.shape[0]), | |
| } | |
| finally: | |
| torch.cuda.synchronize() | |
| if optimizer is not None: | |
| ops = optimizer.solver_ops.extension | |
| ops.destroy_cublas_handle(optimizer.solver_context.cublas_handle) | |
| ops.destroy_cusolver_handle(optimizer.solver_context.cusolver_handle) | |
| del optimizer, visualizer | |
| gc.collect() | |