-
Notifications
You must be signed in to change notification settings - Fork 90
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
See https://github.com/quic/ai-hub-models/releases/v0.19.0 for changelog. Signed-off-by: QAIHM Team <[email protected]>
- Loading branch information
Showing
381 changed files
with
28,918 additions
and
21,043 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,4 @@ | ||
| # --------------------------------------------------------------------- | ||
| # Copyright (c) 2024 Qualcomm Innovation Center, Inc. All rights reserved. | ||
| # SPDX-License-Identifier: BSD-3-Clause | ||
| # --------------------------------------------------------------------- |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,81 @@ | ||
| # --------------------------------------------------------------------- | ||
| # Copyright (c) 2024 Qualcomm Innovation Center, Inc. All rights reserved. | ||
| # SPDX-License-Identifier: BSD-3-Clause | ||
| # --------------------------------------------------------------------- | ||
|
|
||
| import cv2 | ||
| import numpy as np | ||
| from PIL import Image | ||
| from skimage import io | ||
|
|
||
| from qai_hub_models.models._shared.facemap_3dmm.app import FaceMap_3DMMApp | ||
| from qai_hub_models.models._shared.facemap_3dmm.model import ( | ||
| MODEL_ASSET_VERSION, | ||
| MODEL_ID, | ||
| FaceMap_3DMM, | ||
| ) | ||
| from qai_hub_models.utils.args import ( | ||
| demo_model_from_cli_args, | ||
| get_model_cli_parser, | ||
| get_on_device_demo_parser, | ||
| validate_on_device_demo_args, | ||
| ) | ||
| from qai_hub_models.utils.asset_loaders import CachedWebModelAsset | ||
| from qai_hub_models.utils.display import display_or_save_image | ||
|
|
||
| INPUT_IMAGE_PATH = str( | ||
| CachedWebModelAsset.from_asset_store(MODEL_ID, MODEL_ASSET_VERSION, "face_img.jpg") | ||
| ) | ||
|
|
||
|
|
||
| # Run FaceMap_3DMM end-to-end on a sample image. | ||
| # The demo will display a image with the predicted landmark displayed. | ||
| def main( | ||
| model_cls: type[FaceMap_3DMM] = FaceMap_3DMM, | ||
| model_id: str = MODEL_ID, | ||
| is_test: bool = False, | ||
| ): | ||
| # Demo parameters | ||
| parser = get_model_cli_parser(model_cls) | ||
| parser = get_on_device_demo_parser(parser, add_output_dir=True) | ||
| parser.add_argument( | ||
| "--image", | ||
| type=str, | ||
| default=INPUT_IMAGE_PATH, | ||
| help="image file path or URL", | ||
| ) | ||
| args = parser.parse_args([] if is_test else None) | ||
| model = demo_model_from_cli_args(model_cls, model_id, args) | ||
| validate_on_device_demo_args(args, model_id) | ||
|
|
||
| # Load image | ||
| (_, _, height, width) = FaceMap_3DMM.get_input_spec()["image"][0] | ||
| image = io.imread(args.image) | ||
|
|
||
| print("Model Loaded") | ||
|
|
||
| app = FaceMap_3DMMApp(model) | ||
|
|
||
| # Get face bounding box info (from file or face detector) | ||
| fbox = np.loadtxt(INPUT_IMAGE_PATH.replace(".jpg", "_fbox.txt")) | ||
| x0, x1, y0, y1 = int(fbox[0]), int(fbox[1]), int(fbox[2]), int(fbox[3]) | ||
|
|
||
| lmk, output = app.landmark_prediction(image, x0, x1, y0, y1) | ||
|
|
||
| if not is_test: | ||
| # Annotated lmk | ||
| np.savetxt( | ||
| "qai_hub_models/models/facemap_3dmm/demo_output_lmk.txt", | ||
| lmk.detach().numpy(), | ||
| ) | ||
|
|
||
| # Annotated image | ||
| display_or_save_image( | ||
| Image.fromarray(cv2.cvtColor(output, cv2.COLOR_BGR2RGB)), | ||
| "qai_hub_models/models/facemap_3dmm", | ||
| "demo_output_img.png", | ||
| ) | ||
|
|
||
|
|
||
| if __name__ == "__main__": | ||
| main() |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,64 @@ | ||
| # --------------------------------------------------------------------- | ||
| # Copyright (c) 2024 Qualcomm Innovation Center, Inc. All rights reserved. | ||
| # SPDX-License-Identifier: BSD-3-Clause | ||
| # --------------------------------------------------------------------- | ||
| from __future__ import annotations | ||
|
|
||
| import torch | ||
| import torch.nn as nn | ||
|
|
||
| from qai_hub_models.models._shared.facemap_3dmm.resnet_score_rgb import resnet18_wd2 | ||
| from qai_hub_models.utils.asset_loaders import CachedWebModelAsset, load_torch | ||
| from qai_hub_models.utils.base_model import BaseModel | ||
| from qai_hub_models.utils.input_spec import InputSpec | ||
|
|
||
| MODEL_ID = __name__.split(".")[-2] | ||
| DEFAULT_WEIGHTS = "resnet_wd2_weak_score_1202_3ch.pth.tar" | ||
| MODEL_ASSET_VERSION = 1 | ||
|
|
||
|
|
||
| class FaceMap_3DMM(BaseModel): | ||
| """Exportable FaceMap_3DMM, end-to-end.""" | ||
|
|
||
| def __init__(self, model: nn.Module) -> None: | ||
| super().__init__() | ||
| self.model = model | ||
|
|
||
| @classmethod | ||
| def from_pretrained(cls): | ||
|
|
||
| resnet_model = resnet18_wd2(pretrained=False) | ||
|
|
||
| checkpoint_path = CachedWebModelAsset.from_asset_store( | ||
| MODEL_ID, MODEL_ASSET_VERSION, DEFAULT_WEIGHTS | ||
| ) | ||
| pretrained_dict = load_torch(checkpoint_path)["state_dict"] | ||
| resnet_model.load_state_dict(pretrained_dict) | ||
| resnet_model.to(torch.device("cpu")).eval() | ||
|
|
||
| return cls(resnet_model) | ||
|
|
||
| def forward(self, image): | ||
| """ | ||
| Run ResNet18_0.5 3Ch on `image`, and produce 265 outputs | ||
| Parameters: | ||
| image: Pixel values pre-processed for encoder consumption. | ||
| Range: float[0, 255] | ||
| 3-channel Color Space: RGB | ||
| Returns: | ||
| 3DMM model paramaeters for facial landmark reconstruction: Shape [batch, 265] | ||
| """ | ||
| return self.model(image) | ||
|
|
||
| @staticmethod | ||
| def get_input_spec() -> InputSpec: | ||
| """ | ||
| Returns the input specification (name -> (shape, type). | ||
| """ | ||
| return {"image": ((1, 3, 128, 128), "float32")} | ||
|
|
||
| @staticmethod | ||
| def get_output_names() -> list[str]: | ||
| return ["parameters_3dmm"] |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.