From b77abb6e548b9a09b5c96b1592d71332b45d050e Mon Sep 17 00:00:00 2001 From: Luan Pham Date: Tue, 13 Jul 2021 13:14:33 +0700 Subject: [PATCH] Add docstring for some rmn utils --- README.md | 2 +- rmn/__init__.py | 28 ++++++++++++++++++++++++++++ 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 0bc5076..c8799bc 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ # Facial Expression Recognition using Residual Masking Network, in PyTorch -[![pypi package](https://img.shields.io/badge/version-v3.0.1-blue)](https://pypi.org/project/rmn) +[![pypi package](https://img.shields.io/badge/version-v3.0.2-blue)](https://pypi.org/project/rmn) [![circleci](https://circleci.com/gh/phamquiluan/ResidualMaskingNetwork.svg?style=shield&circle-token=3ca7f15b6bd362b625bec536d57c45fe5ef6f2c9)](https://app.circleci.com/pipelines/github/phamquiluan/ResidualMaskingNetwork) [![style](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/phamquiluan/residualmaskingnetwork) [![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/challenges-in-representation-learning-a/facial-expression-recognition-on-fer2013)](https://paperswithcode.com/sota/facial-expression-recognition-on-fer2013?p=challenges-in-representation-learning-a) diff --git a/rmn/__init__.py b/rmn/__init__.py index 07ec7b3..9c8392c 100644 --- a/rmn/__init__.py +++ b/rmn/__init__.py @@ -148,6 +148,23 @@ def __init__(self, face_detector=True): @torch.no_grad() def detect_emotion_for_single_face_image(self, face_image): + """ + Params: + ----------- + face_image : np.ndarray + a cropped face image + + Return: + ----------- + emo_label : str + dominant emotion label + + emo_proba : float + dominant emotion proba + + proba_list : list + all emotion label and their proba + """ assert isinstance(face_image, np.ndarray) face_image = ensure_color(face_image) face_image = cv2.resize(face_image, image_size) @@ -204,6 +221,17 @@ def video_demo(self): @staticmethod def draw(frame, results): + """ + Params: + --------- + frame : np.ndarray + + results : list of dict.keys('xmin', 'xmax', 'ymin', 'ymax', 'emo_label', 'emo_proba') + + Returns: + --------- + frame : np.ndarray + """ for r in results: xmin = r["xmin"] xmax = r["xmax"]