diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..56573ae --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +./.venv \ No newline at end of file diff --git a/README.md b/README.md index ee846b5..dc10d7e 100644 --- a/README.md +++ b/README.md @@ -1,54 +1,4 @@ - +# Projekt - Rozpoznawanie i przetwarzanie obrazów -# FaceX-Zoo -FaceX-Zoo is a PyTorch toolbox for face recognition. It provides a training module with various supervisory heads and backbones towards state-of-the-art face recognition, as well as a standardized evaluation module which enables to evaluate the models in most of the popular benchmarks just by editing a simple configuration. Also, a simple yet fully functional face SDK is provided for the validation and primary application of the trained models. Rather than including as many as possible of the prior techniques, we enable FaceX-Zoo to easilyupgrade and extend along with the development of face related domains. Please refer to the [technical report](https://arxiv.org/pdf/2101.04407.pdf) for more detailed information about this project. - -About the name: -* "Face" - this repo is mainly for face recognition. -* "X" - we also aim to provide something beyond face recognition, e.g. face parsing, face lightning. -* "Zoo" - there include a lot of algorithms and models in this repo. -![image](data/images/arch.jpg) - -# What's New -- [Oct. 2021] [Swin Transformer](https://arxiv.org/pdf/2103.14030.pdf) is supported now! We obtain a quite promising result of 98.17 in MegaFace rank1 protocol with Swin-S. For more results, pretrained models and logs please refer to [3.1 Experiments of SOTA backbones](training_mode/README.md). For training face recognition model with Swin Transformer, please refer to [swin_training](training_mode/swin_training). Note that the input size is 224\*224 but not 112\*112 and we use AdamW+CosineLRScheduler to optimize it instead of SGD+MultiStepLR. -- [Sep. 2021] We provide a [Dockfile](docker/Dockerfile) to buid the docker image of this project. -- [Aug. 2021] [RepVGG](https://arxiv.org/pdf/2101.03697.pdf) has been added to the backbones for face recognition, the performance of RepVGG_A0, B0, B1 can be found in [3.1 Experiments of SOTA backbones](training_mode). -- [Jul. 2021] A method for facial expression recognition named [DMUE](https://openaccess.thecvf.com/content/CVPR2021/papers/She_Dive_Into_Ambiguity_Latent_Distribution_Mining_and_Pairwise_Uncertainty_Estimation_CVPR_2021_paper.pdf) has been accepted by CVPR2021, and all codes have been released [here](addition_module/DMUE). -- [Jun. 2021] We evaluate some knowledge distillation methods on face recognition task, results and codes can be found in [face_lightning](addition_module/face_lightning/KDF) module. -- [May. 2021] Tools to convert a trained model to onnx format and the provided sdk format can be found in [model_convertor](addition_module/model_convertor). -- [Apr. 2021] IJB-C 1:1 protocol has been added to the [evaluation module](test_protocol/test_ijbc.sh). -- [Mar. 2021] [ResNeSt](https://hangzhang.org/files/resnest.pdf) and [ReXNet](https://arxiv.org/pdf/2007.00992.pdf) have been added to the backbones, [MagFace](https://arxiv.org/pdf/2103.06627.pdf) has been added to the heads. -- [Feb. 2021] Distributed training and mixed precision training by [apex](https://github.com/NVIDIA/apex) is supported. Please check [distributed_training](training_mode/distributed_training) and [train_amp.py](training_mode/conventional_training/train_amp.py) -- [Jan. 2021] We commit the initial version of FaceX-Zoo. - -# Requirements -* python >= 3.7.1 -* pytorch >= 1.1.0 -* torchvision >= 0.3.0 - -# Model Training -See [README.md](training_mode/README.md) in [training_mode](training_mode), currently support conventional training and [semi-siamese training](https://arxiv.org/abs/2007.08398). -# Model Evaluation -See [README.md](test_protocol/README.md) in [test_protocol](test_protocol), currently support [LFW](https://people.cs.umass.edu/~elm/papers/lfw.pdf), [CPLFW](http://www.whdeng.cn/CPLFW/Cross-Pose-LFW.pdf), [CALFW](https://arxiv.org/abs/1708.08197), [RFW](https://arxiv.org/abs/1812.00194), [AgeDB30](https://core.ac.uk/download/pdf/83949017.pdf), [IJB-C](http://biometrics.cse.msu.edu/Publications/Face/Mazeetal_IARPAJanusBenchmarkCFaceDatasetAndProtocol_ICB2018.pdf), [MegaFace](https://arxiv.org/abs/1512.00596) and MegaFace-mask. -# Face SDK -See [README.md](face_sdk/README.md) in [face_sdk](face_sdk), currently support face detection, face alignment and face recognition. -# Face Mask Adding -See [README.md](addition_module/face_mask_adding/FMA-3D/README.md) in [FMA-3D](addition_module/face_mask_adding/FMA-3D). - -# License -FaceX-Zoo is released under the [Apache License, Version 2.0](LICENSE). - -# Acknowledgements -This repo is mainly inspired by [InsightFace](https://github.com/deepinsight/insightface), [InsightFace_Pytorch](https://github.com/TreB1eN/InsightFace_Pytorch), [face.evoLVe](https://github.com/ZhaoJ9014/face.evoLVe.PyTorch/blob/master/README.md). We thank the authors a lot for their valuable efforts. - -# Citation -Please consider citing our paper in your publications if the project helps your research. BibTeX reference is as follows. -```BibTeX -@inproceedings{wang2021facex, - author = {Jun Wang, Yinglu Liu, Yibo Hu, Hailin Shi and Tao Mei}, - title = {FaceX-Zoo: A PyTorh Toolbox for Face Recognition}, - journal = {Proceedings of the 29th ACM international conference on Multimedia}, - year = {2021} -} -``` -If you have any questions, please contact with Jun Wang (wangjun492@jd.com), Yinglu Liu (liuyinglu1@jd.com), [Yibo Hu](https://aberhu.github.io/) (huyibo6@jd.com), [Hailin Shi](https://sites.google.com/view/hailin-shi) (shihailin@jd.com) and [Wu Liu](http://drliuwu.com)(liuwu1@jd.com). +- Kamil Majewski 272929 +- Michał Dziamski 272928 diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..fd09efa --- /dev/null +++ b/requirements.txt @@ -0,0 +1,5 @@ +opencv-python +torch +torchvision +scikit-image +PySide6 \ No newline at end of file diff --git a/src/configurator/config.json b/src/configurator/config.json new file mode 100644 index 0000000..af6afe5 --- /dev/null +++ b/src/configurator/config.json @@ -0,0 +1,19 @@ +{ + "frame_crop": { + "x": 0, + "y": 0, + "width": 1400, + "height": 1200 + }, + "detection": { + "confidence_threshold": 0.75 + }, + "alert": { + "show_rectangle": true, + "show_alert": true, + "play_sound": false + }, + "overlay": { + "show_lines": true + } +} \ No newline at end of file diff --git a/src/configurator/configurator.py b/src/configurator/configurator.py new file mode 100644 index 0000000..c0d5810 --- /dev/null +++ b/src/configurator/configurator.py @@ -0,0 +1,118 @@ +import sys +import os +import json +from PySide6.QtWidgets import ( + QApplication, QWidget, QVBoxLayout, QLabel, + QSpinBox, QCheckBox, QPushButton, QHBoxLayout, QDoubleSpinBox +) + +CONFIG_PATH = os.path.join(os.path.dirname(__file__), 'config.json') + + +class Configurator(QWidget): + def __init__(self): + super().__init__() + self.setWindowTitle("Konfigurator") + self.setFixedSize(300, 300) + + self.layout = QVBoxLayout() + self._load_config() + + self.layout.addWidget(QLabel("Kadrowanie (x, y, szer, wys):")) + self.crop_x = QSpinBox() + self.crop_y = QSpinBox() + self.crop_w = QSpinBox() + self.crop_h = QSpinBox() + for box in [self.crop_x, self.crop_y, self.crop_w, self.crop_h]: + box.setRange(0, 1920) + self.crop_x.setValue(self.cfg.get("frame_crop", {}).get("x", 0)) + self.crop_y.setValue(self.cfg.get("frame_crop", {}).get("y", 0)) + self.crop_w.setValue(self.cfg.get("frame_crop", {}).get("width", 0)) + self.crop_h.setValue(self.cfg.get("frame_crop", {}).get("height", 0)) + row = QHBoxLayout() + for widget in [self.crop_x, self.crop_y, self.crop_w, self.crop_h]: + row.addWidget(widget) + self.layout.addLayout(row) + + self.layout.addWidget(QLabel("Confidence threshold:")) + self.confidence = QDoubleSpinBox() + self.confidence.setRange(0.0001, 1.0000) + self.confidence.setDecimals(4) + self.confidence.setSingleStep(0.0001) + default_conf = self.cfg.get("detection", {}).get("confidence_threshold", 0.9991) + self.confidence.setValue(float(default_conf)) + self.layout.addWidget(self.confidence) + + self.rect_cb = QCheckBox("Ramka") + self.alert_cb = QCheckBox("Alert na ekranie") + self.sound_cb = QCheckBox("Dźwięk") + + self.rect_cb.setChecked(self.cfg.get("alert", {}).get("show_rectangle", False)) + self.alert_cb.setChecked(self.cfg.get("alert", {}).get("show_alert", False)) + self.sound_cb.setChecked(self.cfg.get("alert", {}).get("play_sound", False)) + + self.layout.addWidget(self.rect_cb) + self.layout.addWidget(self.alert_cb) + self.layout.addWidget(self.sound_cb) + + self.lines_cb = QCheckBox("Linie pomocnicze") + self.lines_cb.setChecked(self.cfg.get("overlay", {}).get("show_lines", False)) + self.layout.addWidget(self.lines_cb) + + self.save_btn = QPushButton("Zapisz") + self.save_btn.clicked.connect(self.save_config) + self.layout.addWidget(self.save_btn) + + self.setLayout(self.layout) + + def _load_config(self): + if os.path.exists(CONFIG_PATH): + try: + with open(CONFIG_PATH, 'r') as f: + self.cfg = json.load(f) + except Exception: + self.cfg = {} + else: + self.cfg = {} + + def save_config(self): + config = { + "frame_crop": { + "x": self.crop_x.value(), + "y": self.crop_y.value(), + "width": self.crop_w.value(), + "height": self.crop_h.value() + }, + "detection": { + "confidence_threshold": round(self.confidence.value(), 4) + }, + "alert": { + "show_rectangle": self.rect_cb.isChecked(), + "show_alert": self.alert_cb.isChecked(), + "play_sound": self.sound_cb.isChecked() + }, + "overlay": { + "show_lines": self.lines_cb.isChecked() + } + } + with open(CONFIG_PATH, 'w') as f: + json.dump(config, f, indent=4) + self.close() + + @staticmethod + def load_user_config(): + config_path = os.path.abspath(os.path.join(os.path.dirname(__file__), 'config.json')) + if os.path.exists(config_path): + try: + with open(config_path, 'r') as f: + return json.load(f) + except Exception as e: + print(f"Could not load user config: {e}") + return {} + + +if __name__ == "__main__": + app = QApplication(sys.argv) + window = Configurator() + window.show() + sys.exit(app.exec()) diff --git a/src/data/model/kamil1.jpeg b/src/data/model/kamil1.jpeg new file mode 100644 index 0000000..244bf93 Binary files /dev/null and b/src/data/model/kamil1.jpeg differ diff --git a/src/data/model/kamil2.jpeg b/src/data/model/kamil2.jpeg new file mode 100644 index 0000000..c066dbe Binary files /dev/null and b/src/data/model/kamil2.jpeg differ diff --git a/src/data/model/kamil3.jpeg b/src/data/model/kamil3.jpeg new file mode 100644 index 0000000..37c29b2 Binary files /dev/null and b/src/data/model/kamil3.jpeg differ diff --git a/src/data/model/kamil4.jpeg b/src/data/model/kamil4.jpeg new file mode 100644 index 0000000..dfe6ac2 Binary files /dev/null and b/src/data/model/kamil4.jpeg differ diff --git a/src/data/model/michal.jpeg b/src/data/model/michal.jpeg new file mode 100644 index 0000000..8b7100d Binary files /dev/null and b/src/data/model/michal.jpeg differ diff --git a/src/faceDetection/detection.py b/src/faceDetection/detection.py new file mode 100644 index 0000000..135a09b --- /dev/null +++ b/src/faceDetection/detection.py @@ -0,0 +1,277 @@ +import json +import os +import sys +import cv2 +import numpy as np +import glob +import torch + +from src.configurator.configurator import Configurator +from src.faceDetection.face_sdk.face_sdk.core.model_loader.face_detection.FaceDetModelLoader import FaceDetModelLoader +from src.faceDetection.face_sdk.face_sdk.core.model_handler.face_detection.FaceDetModelHandler import \ + FaceDetModelHandler +from src.faceDetection.face_sdk.face_sdk.core.model_loader.face_alignment.FaceAlignModelLoader import \ + FaceAlignModelLoader +from src.faceDetection.face_sdk.face_sdk.core.model_handler.face_alignment.FaceAlignModelHandler import \ + FaceAlignModelHandler +from src.faceDetection.face_sdk.face_sdk.core.model_loader.face_recognition.FaceRecModelLoader import FaceRecModelLoader +from src.faceDetection.face_sdk.face_sdk.core.model_handler.face_recognition.FaceRecModelHandler import \ + FaceRecModelHandler +from src.faceDetection.face_sdk.face_sdk.core.image_cropper.arcface_cropper.FaceRecImageCropper import \ + FaceRecImageCropper + +sdk_root = os.path.join(os.path.dirname(__file__), "face_sdk", "face_sdk") +if sdk_root not in sys.path: + sys.path.insert(0, sdk_root) + + +class FaceRecognition: + def __init__(self): + model_path = os.path.join( + os.path.dirname(__file__), + "face_sdk/face_sdk", + "models" + ) + + det_loader = FaceDetModelLoader(model_path, "face_detection", "face_detection_1.0") + det_model, det_cfg = det_loader.load_model() + self.det_handler = FaceDetModelHandler(det_model, "cpu", det_cfg) + + align_loader = FaceAlignModelLoader(model_path, "face_alignment", "face_alignment_1.0") + align_model, align_cfg = align_loader.load_model() + align_model = align_model.to('cpu') + self.align_handler = FaceAlignModelHandler(align_model, "cpu", align_cfg) + + rec_loader = FaceRecModelLoader(model_path, "face_recognition", "face_recognition_1.0") + rec_model, rec_cfg = rec_loader.load_model() + rec_model = rec_model.to('cpu') + self.rec_handler = FaceRecModelHandler(rec_model, "cpu", rec_cfg) + + self.face_cropper = FaceRecImageCropper() + self.known_faces = {} + + def _preprocess_face(self, frame, box): + try: + landmarks = self.align_handler.inference_on_image(frame, box) + landmarks_list = [] + for (x, y) in landmarks.astype(np.int32): + landmarks_list.extend((x, y)) + + cropped_image = self.face_cropper.crop_image_by_mat(frame, landmarks_list) + + if cropped_image is None or cropped_image.size == 0: + return None + + if len(cropped_image.shape) == 2: + cropped_image = cv2.cvtColor(cropped_image, cv2.COLOR_GRAY2BGR) + + cropped_image = cropped_image.astype('float32') / 255.0 + + + return cropped_image + except Exception as e: + print(f"Face processing error: {e}") + return None + + def register_face(self, frame, name, box): + try: + cropped_face = self._preprocess_face(frame, box) + if cropped_face is None: + return False + + embedding = self.rec_handler.inference_on_image(cropped_face) + + if name not in self.known_faces: + self.known_faces[name] = [] + + if torch.is_tensor(embedding): + embedding = embedding.cpu().numpy() + elif isinstance(embedding, np.ndarray): + pass + else: + raise ValueError("Unknown embedding type") + + self.known_faces[name].append(embedding) + print(f"Correctly registered face: {name}") + return True + except Exception as e: + print(f"Registering face error: {e}") + return False + + def _scale_similarity(self, similarity): + min_val = 0.9965 + max_val = 0.999 + + if similarity <= min_val: + return 0.0 + elif similarity >= max_val: + return similarity + else: + return (similarity - min_val) / (max_val - min_val) * max_val + + def recognize_face(self, frame, box, threshold): + try: + cropped_face = self._preprocess_face(frame, box) + if cropped_face is None: + return {"identity": "Unknown", "similarity": 0} + + unknown_embedding = self.rec_handler.inference_on_image(cropped_face) + unknown_embedding = _normalize(unknown_embedding) + + best_match = "Unknown" + max_similarity = 0 + + for name, embeddings in self.known_faces.items(): + for known_embedding in embeddings: + known_embedding = _normalize(known_embedding) + + similarity = np.dot(unknown_embedding, known_embedding.T) + + if similarity > max_similarity: + max_similarity = similarity + best_match = name + + scaled_similarity = self._scale_similarity(max_similarity) + + if scaled_similarity >= threshold: + identity = best_match + else: + identity = "Unknown" + + print( + f"Best match: {identity}, Similarity: {scaled_similarity:.4f}, Threshold: {threshold}") + + return {"identity": identity, "similarity": float(scaled_similarity)} + + except Exception as e: + print(f"CRITICAL ERROR in recognition: {e}") + return {"identity": "Unknown", "similarity": 0} + + +def _normalize(vec): + norm = np.linalg.norm(vec) + return vec if norm == 0 else vec / norm +class FaceDetector: + def __init__(self): + model_path = os.path.join( + os.path.dirname(__file__), + "face_sdk/face_sdk", + "models" + ) + + det_loader = FaceDetModelLoader(model_path, "face_detection", "face_detection_1.0") + det_model, det_cfg = det_loader.load_model() + self.det_handler = FaceDetModelHandler(det_model, "cpu", det_cfg) + + self.recognizer = FaceRecognition() + + self._load_known_faces() + self.user_config = Configurator.load_user_config() + + def _load_known_faces(self): + data_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'data', 'model')) + if os.path.exists(data_dir): + print(f"Loading known faces from: {data_dir}") + for img_path in glob.glob(os.path.join(data_dir, "*.jpg")): + name = os.path.splitext(os.path.basename(img_path))[0] + print(f"Proccessing: {name}") + + img = cv2.imread(img_path) + if img is not None: + bboxes = self.detect_faces(img) + if bboxes.size > 0: + if self.recognizer.register_face(img, name, bboxes[0]): + print(f"Registered: {name}") + else: + print(f"Could not register: {name}") + else: + print(f"Could not detect face: {name}") + else: + print(f"Cannot load: {img_path}") + else: + print(f"No known_faces folder: {data_dir}") + + def detect_faces(self, frame): + return self.det_handler.inference_on_image(frame) + + def analyze(self, frame): + crop_cfg = self.user_config.get("frame_crop", {}) + x = crop_cfg.get("x", 0) + y = crop_cfg.get("y", 0) + w = crop_cfg.get("width", frame.shape[1]) + h = crop_cfg.get("height", frame.shape[0]) + frame = frame[y:y + h, x:x + w] + + bboxes = self.detect_faces(frame) + if bboxes.size == 0: + return frame + + threshold = float(self.user_config.get("detection", {}).get("confidence_threshold", 0.9991)) + + results = [self.recognizer.recognize_face(frame, box, threshold) for box in bboxes] + + frame = self.draw_faces(frame, bboxes, results) + + if self.user_config.get("overlay", {}).get("show_lines", False): + h, w = frame.shape[:2] + cv2.line(frame, (0, 0), (w, h), (255, 0, 0), 1) + cv2.line(frame, (w, 0), (0, h), (255, 0, 0), 1) + + return frame + + def draw_faces(self, frame, bboxes, recognition_results): + for box, result in zip(bboxes, recognition_results): + x1, y1, x2, y2 = map(int, box[:4]) + color = (0, 255, 0) if result['identity'] != "Unknown" else (0, 0, 255) + label = f"{result['identity']} ({result['similarity']:.4f})" + + if self.user_config.get("alert", {}).get("show_rectangle", True): + cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2) + + if self.user_config.get("alert", {}).get("show_alert", False): + cv2.putText(frame, label, (x1, y1 - 10), + cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1) + + if self.user_config.get("alert", {}).get("play_sound", False): + self._play_sound() + return frame + + def _play_sound(self): + import platform + if platform.system() == "Windows": + import winsound + winsound.Beep(1000, 150) + else: + import os + os.system('printf "\\a"') + +if __name__ == "__main__": + print("System Initialization...") + detector = FaceDetector() + + print("Starting the camera...") + cap = cv2.VideoCapture(0) + if not cap.isOpened(): + print("Error: Cannot open camera!") + exit() + + print("Start processing...") + while True: + ret, frame = cap.read() + if not ret: + print("Error while fetching frame") + break + + try: + result = detector.analyze(frame) + cv2.imshow("Face detection", result) + + if cv2.waitKey(1) & 0xFF == ord('q'): + break + except Exception as e: + print(f"Processing error: {e}") + break + + cap.release() + cv2.destroyAllWindows() + print("System closed") diff --git a/LICENSE b/src/faceDetection/face_sdk/LICENSE similarity index 100% rename from LICENSE rename to src/faceDetection/face_sdk/LICENSE diff --git a/src/faceDetection/face_sdk/README.md b/src/faceDetection/face_sdk/README.md new file mode 100644 index 0000000..21336d8 --- /dev/null +++ b/src/faceDetection/face_sdk/README.md @@ -0,0 +1,54 @@ + + +# FaceX-Zoo +FaceX-Zoo is a PyTorch toolbox for face recognition. It provides a training module with various supervisory heads and backbones towards state-of-the-art face recognition, as well as a standardized evaluation module which enables to evaluate the models in most of the popular benchmarks just by editing a simple configuration. Also, a simple yet fully functional face SDK is provided for the validation and primary application of the trained models. Rather than including as many as possible of the prior techniques, we enable FaceX-Zoo to easilyupgrade and extend along with the development of face related domains. Please refer to the [technical report](https://arxiv.org/pdf/2101.04407.pdf) for more detailed information about this project. + +About the name: +* "Face" - this repo is mainly for face recognition. +* "X" - we also aim to provide something beyond face recognition, e.g. face parsing, face lightning. +* "Zoo" - there include a lot of algorithms and models in this repo. +![image](data/images/arch.jpg) + +# What's New +- [Oct. 2021] [Swin Transformer](https://arxiv.org/pdf/2103.14030.pdf) is supported now! We obtain a quite promising result of 98.17 in MegaFace rank1 protocol with Swin-S. For more results, pretrained models and logs please refer to [3.1 Experiments of SOTA backbones](training_mode/README.md). For training face recognition model with Swin Transformer, please refer to [swin_training](training_mode/swin_training). Note that the input size is 224\*224 but not 112\*112 and we use AdamW+CosineLRScheduler to optimize it instead of SGD+MultiStepLR. +- [Sep. 2021] We provide a [Dockfile](docker/Dockerfile) to buid the docker image of this project. +- [Aug. 2021] [RepVGG](https://arxiv.org/pdf/2101.03697.pdf) has been added to the backbones for face recognition, the performance of RepVGG_A0, B0, B1 can be found in [3.1 Experiments of SOTA backbones](training_mode). +- [Jul. 2021] A method for facial expression recognition named [DMUE](https://openaccess.thecvf.com/content/CVPR2021/papers/She_Dive_Into_Ambiguity_Latent_Distribution_Mining_and_Pairwise_Uncertainty_Estimation_CVPR_2021_paper.pdf) has been accepted by CVPR2021, and all codes have been released [here](addition_module/DMUE). +- [Jun. 2021] We evaluate some knowledge distillation methods on face recognition task, results and codes can be found in [face_lightning](addition_module/face_lightning/KDF) module. +- [May. 2021] Tools to convert a trained model to onnx format and the provided sdk format can be found in [model_convertor](addition_module/model_convertor). +- [Apr. 2021] IJB-C 1:1 protocol has been added to the [evaluation module](test_protocol/test_ijbc.sh). +- [Mar. 2021] [ResNeSt](https://hangzhang.org/files/resnest.pdf) and [ReXNet](https://arxiv.org/pdf/2007.00992.pdf) have been added to the backbones, [MagFace](https://arxiv.org/pdf/2103.06627.pdf) has been added to the heads. +- [Feb. 2021] Distributed training and mixed precision training by [apex](https://github.com/NVIDIA/apex) is supported. Please check [distributed_training](training_mode/distributed_training) and [train_amp.py](training_mode/conventional_training/train_amp.py) +- [Jan. 2021] We commit the initial version of FaceX-Zoo. + +# Requirements +* python >= 3.7.1 +* pytorch >= 1.1.0 +* torchvision >= 0.3.0 + +# Model Training +See [README.md](training_mode/README.md) in [training_mode](training_mode), currently support conventional training and [semi-siamese training](https://arxiv.org/abs/2007.08398). +# Model Evaluation +See [README.md](test_protocol/README.md) in [test_protocol](test_protocol), currently support [LFW](https://people.cs.umass.edu/~elm/papers/lfw.pdf), [CPLFW](http://www.whdeng.cn/CPLFW/Cross-Pose-LFW.pdf), [CALFW](https://arxiv.org/abs/1708.08197), [RFW](https://arxiv.org/abs/1812.00194), [AgeDB30](https://core.ac.uk/download/pdf/83949017.pdf), [IJB-C](http://biometrics.cse.msu.edu/Publications/Face/Mazeetal_IARPAJanusBenchmarkCFaceDatasetAndProtocol_ICB2018.pdf), [MegaFace](https://arxiv.org/abs/1512.00596) and MegaFace-mask. +# Face SDK +See [README.md](face_sdk/README.md) in [face_sdk](face_sdk), currently support face detection, face alignment and face recognition. +# Face Mask Adding +See [README.md](addition_module/face_mask_adding/FMA-3D/README.md) in [FMA-3D](addition_module/face_mask_adding/FMA-3D). + +# License +FaceX-Zoo is released under the [Apache License, Version 2.0](LICENSE). + +# Acknowledgements +This repo is mainly inspired by [InsightFace](https://github.com/deepinsight/insightface), [InsightFace_Pytorch](https://github.com/TreB1eN/InsightFace_Pytorch), [face.evoLVe](https://github.com/ZhaoJ9014/face.evoLVe.PyTorch/blob/master/README.md). We thank the authors a lot for their valuable efforts. + +# Citation +Please consider citing our paper in your publications if the project helps your research. BibTeX reference is as follows. +```BibTeX +@inproceedings{wang2021facex, + author = {Jun Wang, Yinglu Liu, Yibo Hu, Hailin Shi and Tao Mei}, + title = {FaceX-Zoo: A PyTorh Toolbox for Face Recognition}, + journal = {Proceedings of the 29th ACM international conference on Multimedia}, + year = {2021} +} +``` +If you have any questions, please contact with Jun Wang (wangjun492@jd.com), Yinglu Liu (liuyinglu1@jd.com), [Yibo Hu](https://aberhu.github.io/) (huyibo6@jd.com), [Hailin Shi](https://sites.google.com/view/hailin-shi) (shihailin@jd.com) and [Wu Liu](http://drliuwu.com)(liuwu1@jd.com). diff --git a/addition_module/DMUE/preprocess/face_alignmenet/__init__.py b/src/faceDetection/face_sdk/__init__.py similarity index 100% rename from addition_module/DMUE/preprocess/face_alignmenet/__init__.py rename to src/faceDetection/face_sdk/__init__.py diff --git a/addition_module/DMUE/.gitignore b/src/faceDetection/face_sdk/addition_module/DMUE/.gitignore similarity index 100% rename from addition_module/DMUE/.gitignore rename to src/faceDetection/face_sdk/addition_module/DMUE/.gitignore diff --git a/addition_module/DMUE/README.md b/src/faceDetection/face_sdk/addition_module/DMUE/README.md similarity index 100% rename from addition_module/DMUE/README.md rename to src/faceDetection/face_sdk/addition_module/DMUE/README.md diff --git a/addition_module/DMUE/preprocess/face_alignmenet/core/__init__.py b/src/faceDetection/face_sdk/addition_module/DMUE/__init__.py similarity index 100% rename from addition_module/DMUE/preprocess/face_alignmenet/core/__init__.py rename to src/faceDetection/face_sdk/addition_module/DMUE/__init__.py diff --git a/addition_module/DMUE/config.py b/src/faceDetection/face_sdk/addition_module/DMUE/config.py similarity index 100% rename from addition_module/DMUE/config.py rename to src/faceDetection/face_sdk/addition_module/DMUE/config.py diff --git a/addition_module/DMUE/convert_weights.py b/src/faceDetection/face_sdk/addition_module/DMUE/convert_weights.py similarity index 100% rename from addition_module/DMUE/convert_weights.py rename to src/faceDetection/face_sdk/addition_module/DMUE/convert_weights.py diff --git a/addition_module/DMUE/datasets/__init__.py b/src/faceDetection/face_sdk/addition_module/DMUE/datasets/__init__.py similarity index 100% rename from addition_module/DMUE/datasets/__init__.py rename to src/faceDetection/face_sdk/addition_module/DMUE/datasets/__init__.py diff --git a/addition_module/DMUE/datasets/bases.py b/src/faceDetection/face_sdk/addition_module/DMUE/datasets/bases.py similarity index 96% rename from addition_module/DMUE/datasets/bases.py rename to src/faceDetection/face_sdk/addition_module/DMUE/datasets/bases.py index 2370dd8..01e4825 100644 --- a/addition_module/DMUE/datasets/bases.py +++ b/src/faceDetection/face_sdk/addition_module/DMUE/datasets/bases.py @@ -1,47 +1,47 @@ -import os -import lmdb -import random -import numpy as np - -from torch.utils.data import Dataset - -from .image_utils import add_gaussian_noise - - -def read_lmdb(key, txn): - """Keep reading image until succeed. - This can avoid IOError incurred by heavy IO process.""" - got_img = False - while not got_img: - try: - img = np.frombuffer(txn.get(key.encode('utf-8')), dtype=np.uint8).reshape(256, 256, 3) # here fixed - got_img = True - except IOError: - print("IOError incurred when reading '{}'. Will redo. Don't worry. Just chill.".format(key)) - pass - return img - - -class ImageDataset(Dataset): - def __init__(self, dataset, transform=None, lmdb_f=None, mode=None): - self.dataset = dataset - self.img_label_list = [x for _, x in self.dataset] - self.transform = transform - self.txn = lmdb.open(lmdb_f, readonly=True).begin(write=False) - self.mode = mode - - def __len__(self): - return len(self.dataset) - - def __getitem__(self, index): - key, label = self.dataset[index] - img = read_lmdb(key, self.txn) - - if self.mode is None: - if random.uniform(0, 1) > 0.75: - img = add_gaussian_noise(img) - - if self.transform is not None: - img = self.transform(img) - - return img, label +import os +import lmdb +import random +import numpy as np + +from torch.utils.data import Dataset + +from .image_utils import add_gaussian_noise + + +def read_lmdb(key, txn): + """Keep reading image until succeed. + This can avoid IOError incurred by heavy IO process.""" + got_img = False + while not got_img: + try: + img = np.frombuffer(txn.get(key.encode('utf-8')), dtype=np.uint8).reshape(256, 256, 3) # here fixed + got_img = True + except IOError: + print("IOError incurred when reading '{}'. Will redo. Don't worry. Just chill.".format(key)) + pass + return img + + +class ImageDataset(Dataset): + def __init__(self, dataset, transform=None, lmdb_f=None, mode=None): + self.dataset = dataset + self.img_label_list = [x for _, x in self.dataset] + self.transform = transform + self.txn = lmdb.open(lmdb_f, readonly=True).begin(write=False) + self.mode = mode + + def __len__(self): + return len(self.dataset) + + def __getitem__(self, index): + key, label = self.dataset[index] + img = read_lmdb(key, self.txn) + + if self.mode is None: + if random.uniform(0, 1) > 0.75: + img = add_gaussian_noise(img) + + if self.transform is not None: + img = self.transform(img) + + return img, label diff --git a/addition_module/DMUE/datasets/image_utils.py b/src/faceDetection/face_sdk/addition_module/DMUE/datasets/image_utils.py similarity index 96% rename from addition_module/DMUE/datasets/image_utils.py rename to src/faceDetection/face_sdk/addition_module/DMUE/datasets/image_utils.py index 259483a..cb1e903 100644 --- a/addition_module/DMUE/datasets/image_utils.py +++ b/src/faceDetection/face_sdk/addition_module/DMUE/datasets/image_utils.py @@ -1,22 +1,22 @@ -import cv2 -import numpy as np - - -def add_gaussian_noise(image_array, mean=0.0, var=30): - std = var**0.5 - noisy_img = image_array + np.random.normal(mean, std, image_array.shape) - noisy_img_clipped = np.clip(noisy_img, 0, 255).astype(np.uint8) - return noisy_img_clipped - - -def flip_image(image_array): - return cv2.flip(image_array, 1) - - -def color2gray(image_array): - gray = cv2.cvtColor(image_array, cv2.COLOR_RGB2GRAY) - gray_img_3d = image_array.copy() - gray_img_3d[:, :, 0] = gray - gray_img_3d[:, :, 1] = gray - gray_img_3d[:, :, 2] = gray - return gray_img_3d +import cv2 +import numpy as np + + +def add_gaussian_noise(image_array, mean=0.0, var=30): + std = var**0.5 + noisy_img = image_array + np.random.normal(mean, std, image_array.shape) + noisy_img_clipped = np.clip(noisy_img, 0, 255).astype(np.uint8) + return noisy_img_clipped + + +def flip_image(image_array): + return cv2.flip(image_array, 1) + + +def color2gray(image_array): + gray = cv2.cvtColor(image_array, cv2.COLOR_RGB2GRAY) + gray_img_3d = image_array.copy() + gray_img_3d[:, :, 0] = gray + gray_img_3d[:, :, 1] = gray + gray_img_3d[:, :, 2] = gray + return gray_img_3d diff --git a/addition_module/DMUE/datasets/make_dataloader.py b/src/faceDetection/face_sdk/addition_module/DMUE/datasets/make_dataloader.py similarity index 97% rename from addition_module/DMUE/datasets/make_dataloader.py rename to src/faceDetection/face_sdk/addition_module/DMUE/datasets/make_dataloader.py index aabaff6..00025f6 100644 --- a/addition_module/DMUE/datasets/make_dataloader.py +++ b/src/faceDetection/face_sdk/addition_module/DMUE/datasets/make_dataloader.py @@ -1,47 +1,47 @@ -import torch -import torchvision.transforms as T -from torch.utils.data import DataLoader - -from .bases import ImageDataset -from .preprocessing import RandomErasing -from .sampler import ImbalancedDatasetSampler -from .image_utils import add_gaussian_noise, flip_image, color2gray - - -def make_dataloader(cfg): - train_transforms = T.Compose([ - T.ToPILImage(), - T.Resize(cfg.ori_shape), - T.RandomCrop(cfg.image_crop_size), - T.RandomHorizontalFlip(p=cfg.horizontal_flip_p), - T.ToTensor(), - T.Normalize(mean=cfg.normalize_mean, - std=cfg.normalize_std), - T.RandomErasing(scale=(0.02, 0.25)) - ]) - val_transforms = T.Compose([ - T.ToPILImage(), - T.Resize(cfg.ori_shape), - T.CenterCrop(cfg.image_crop_size), - T.ToTensor(), - T.Normalize(mean=cfg.normalize_mean, - std=cfg.normalize_std), - ]) - - train_set = ImageDataset(cfg.train_dataset, transform=train_transforms, lmdb_f=cfg.lmdb_f) - val_set = ImageDataset(cfg.val_dataset, transform=val_transforms, lmdb_f=cfg.lmdb_f, mode='test') - # print('Train set size:', train_set.__len__()) - # print('Test set size:', val_set.__len__()) - - train_loader = DataLoader( - train_set, - sampler=ImbalancedDatasetSampler(train_set), - batch_size=cfg.batch_size, - num_workers=cfg.num_workers, pin_memory=True - ) - - val_loader = DataLoader( - val_set, batch_size=cfg.test_minibatch, shuffle=False, num_workers=cfg.num_workers - ) - - return train_loader, val_loader +import torch +import torchvision.transforms as T +from torch.utils.data import DataLoader + +from .bases import ImageDataset +from .preprocessing import RandomErasing +from .sampler import ImbalancedDatasetSampler +from .image_utils import add_gaussian_noise, flip_image, color2gray + + +def make_dataloader(cfg): + train_transforms = T.Compose([ + T.ToPILImage(), + T.Resize(cfg.ori_shape), + T.RandomCrop(cfg.image_crop_size), + T.RandomHorizontalFlip(p=cfg.horizontal_flip_p), + T.ToTensor(), + T.Normalize(mean=cfg.normalize_mean, + std=cfg.normalize_std), + T.RandomErasing(scale=(0.02, 0.25)) + ]) + val_transforms = T.Compose([ + T.ToPILImage(), + T.Resize(cfg.ori_shape), + T.CenterCrop(cfg.image_crop_size), + T.ToTensor(), + T.Normalize(mean=cfg.normalize_mean, + std=cfg.normalize_std), + ]) + + train_set = ImageDataset(cfg.train_dataset, transform=train_transforms, lmdb_f=cfg.lmdb_f) + val_set = ImageDataset(cfg.val_dataset, transform=val_transforms, lmdb_f=cfg.lmdb_f, mode='test') + # print('Train set size:', train_set.__len__()) + # print('Test set size:', val_set.__len__()) + + train_loader = DataLoader( + train_set, + sampler=ImbalancedDatasetSampler(train_set), + batch_size=cfg.batch_size, + num_workers=cfg.num_workers, pin_memory=True + ) + + val_loader = DataLoader( + val_set, batch_size=cfg.test_minibatch, shuffle=False, num_workers=cfg.num_workers + ) + + return train_loader, val_loader diff --git a/addition_module/DMUE/datasets/preprocessing.py b/src/faceDetection/face_sdk/addition_module/DMUE/datasets/preprocessing.py similarity index 97% rename from addition_module/DMUE/datasets/preprocessing.py rename to src/faceDetection/face_sdk/addition_module/DMUE/datasets/preprocessing.py index e995cad..1135464 100644 --- a/addition_module/DMUE/datasets/preprocessing.py +++ b/src/faceDetection/face_sdk/addition_module/DMUE/datasets/preprocessing.py @@ -1,49 +1,49 @@ -import math -import random - - -class RandomErasing(object): - """ Randomly selects a rectangle region in an image and erases its pixels. - 'Random Erasing Data Augmentation' by Zhong et al. - See https://arxiv.org/pdf/1708.04896.pdf - Args: - probability: The probability that the Random Erasing operation will be performed. - sl: Minimum proportion of erased area against input image. - sh: Maximum proportion of erased area against input image. - r1: Minimum aspect ratio of erased area. - mean: Erasing value. - """ - - def __init__(self, probability=0.5, sl=0.02, sh=0.4, r1=0.3, mean=(0.4914, 0.4822, 0.4465)): - self.probability = probability - self.mean = mean - self.sl = sl - self.sh = sh - self.r1 = r1 - - def __call__(self, img): - - if random.uniform(0, 1) >= self.probability: - return img - - for attempt in range(100): - area = img.size()[1] * img.size()[2] - - target_area = random.uniform(self.sl, self.sh) * area - aspect_ratio = random.uniform(self.r1, 1 / self.r1) - - h = int(round(math.sqrt(target_area * aspect_ratio))) - w = int(round(math.sqrt(target_area / aspect_ratio))) - - if w < img.size()[2] and h < img.size()[1]: - x1 = random.randint(0, img.size()[1] - h) - y1 = random.randint(0, img.size()[2] - w) - if img.size()[0] == 3: - img[0, x1:x1 + h, y1:y1 + w] = self.mean[0] - img[1, x1:x1 + h, y1:y1 + w] = self.mean[1] - img[2, x1:x1 + h, y1:y1 + w] = self.mean[2] - else: - img[0, x1:x1 + h, y1:y1 + w] = self.mean[0] - return img - - return img +import math +import random + + +class RandomErasing(object): + """ Randomly selects a rectangle region in an image and erases its pixels. + 'Random Erasing Data Augmentation' by Zhong et al. + See https://arxiv.org/pdf/1708.04896.pdf + Args: + probability: The probability that the Random Erasing operation will be performed. + sl: Minimum proportion of erased area against input image. + sh: Maximum proportion of erased area against input image. + r1: Minimum aspect ratio of erased area. + mean: Erasing value. + """ + + def __init__(self, probability=0.5, sl=0.02, sh=0.4, r1=0.3, mean=(0.4914, 0.4822, 0.4465)): + self.probability = probability + self.mean = mean + self.sl = sl + self.sh = sh + self.r1 = r1 + + def __call__(self, img): + + if random.uniform(0, 1) >= self.probability: + return img + + for attempt in range(100): + area = img.size()[1] * img.size()[2] + + target_area = random.uniform(self.sl, self.sh) * area + aspect_ratio = random.uniform(self.r1, 1 / self.r1) + + h = int(round(math.sqrt(target_area * aspect_ratio))) + w = int(round(math.sqrt(target_area / aspect_ratio))) + + if w < img.size()[2] and h < img.size()[1]: + x1 = random.randint(0, img.size()[1] - h) + y1 = random.randint(0, img.size()[2] - w) + if img.size()[0] == 3: + img[0, x1:x1 + h, y1:y1 + w] = self.mean[0] + img[1, x1:x1 + h, y1:y1 + w] = self.mean[1] + img[2, x1:x1 + h, y1:y1 + w] = self.mean[2] + else: + img[0, x1:x1 + h, y1:y1 + w] = self.mean[0] + return img + + return img diff --git a/addition_module/DMUE/datasets/sampler.py b/src/faceDetection/face_sdk/addition_module/DMUE/datasets/sampler.py similarity index 97% rename from addition_module/DMUE/datasets/sampler.py rename to src/faceDetection/face_sdk/addition_module/DMUE/datasets/sampler.py index c19877d..26e100a 100644 --- a/addition_module/DMUE/datasets/sampler.py +++ b/src/faceDetection/face_sdk/addition_module/DMUE/datasets/sampler.py @@ -1,46 +1,46 @@ -import torch - - -class ImbalancedDatasetSampler(torch.utils.data.sampler.Sampler): - """Samples elements randomly from a given list of indices for imbalanced dataset - Arguments: - indices (list, optional): a list of indices - num_samples (int, optional): number of samples to draw - """ - - def __init__(self, dataset, indices=None, num_samples=None): - - # if indices is not provided, - # all elements in the dataset will be considered - self.indices = list(range(len(dataset))) \ - if indices is None else indices - - # if num_samples is not provided, - # draw `len(indices)` samples in each iteration - self.num_samples = len(self.indices) \ - if num_samples is None else num_samples - - # distribution of classes in the dataset - label_to_count = {} - for idx in self.indices: - label = self._get_label(dataset, idx) - if label in label_to_count: - label_to_count[label] += 1 - else: - label_to_count[label] = 1 - - # weight for each sample - weights = [1.0 / label_to_count[self._get_label(dataset, idx)] - for idx in self.indices] - self.weights = torch.DoubleTensor(weights) - - def _get_label(self, dataset, idx): - dataset_type = type(dataset) - return dataset.img_label_list[idx] - - def __iter__(self): - return (self.indices[i] for i in torch.multinomial( - self.weights, self.num_samples, replacement=True)) - - def __len__(self): +import torch + + +class ImbalancedDatasetSampler(torch.utils.data.sampler.Sampler): + """Samples elements randomly from a given list of indices for imbalanced dataset + Arguments: + indices (list, optional): a list of indices + num_samples (int, optional): number of samples to draw + """ + + def __init__(self, dataset, indices=None, num_samples=None): + + # if indices is not provided, + # all elements in the dataset will be considered + self.indices = list(range(len(dataset))) \ + if indices is None else indices + + # if num_samples is not provided, + # draw `len(indices)` samples in each iteration + self.num_samples = len(self.indices) \ + if num_samples is None else num_samples + + # distribution of classes in the dataset + label_to_count = {} + for idx in self.indices: + label = self._get_label(dataset, idx) + if label in label_to_count: + label_to_count[label] += 1 + else: + label_to_count[label] = 1 + + # weight for each sample + weights = [1.0 / label_to_count[self._get_label(dataset, idx)] + for idx in self.indices] + self.weights = torch.DoubleTensor(weights) + + def _get_label(self, dataset, idx): + dataset_type = type(dataset) + return dataset.img_label_list[idx] + + def __iter__(self): + return (self.indices[i] for i in torch.multinomial( + self.weights, self.num_samples, replacement=True)) + + def __len__(self): return self.num_samples \ No newline at end of file diff --git a/addition_module/DMUE/images/framework.png b/src/faceDetection/face_sdk/addition_module/DMUE/images/framework.png similarity index 100% rename from addition_module/DMUE/images/framework.png rename to src/faceDetection/face_sdk/addition_module/DMUE/images/framework.png diff --git a/addition_module/DMUE/images/test1.jpg b/src/faceDetection/face_sdk/addition_module/DMUE/images/test1.jpg similarity index 100% rename from addition_module/DMUE/images/test1.jpg rename to src/faceDetection/face_sdk/addition_module/DMUE/images/test1.jpg diff --git a/addition_module/DMUE/inference.py b/src/faceDetection/face_sdk/addition_module/DMUE/inference.py similarity index 98% rename from addition_module/DMUE/inference.py rename to src/faceDetection/face_sdk/addition_module/DMUE/inference.py index 0752e41..11f5742 100644 --- a/addition_module/DMUE/inference.py +++ b/src/faceDetection/face_sdk/addition_module/DMUE/inference.py @@ -47,7 +47,7 @@ def inference(model, img_path, transform, is_cuda=True): if __name__ == '__main__': - img_path = './images/test1.jpg' + img_path = 'images/test1.jpg' transform = T.Compose([ T.Resize(cfg.ori_shape), diff --git a/addition_module/DMUE/losses/SP_KD_Loss.py b/src/faceDetection/face_sdk/addition_module/DMUE/losses/SP_KD_Loss.py similarity index 96% rename from addition_module/DMUE/losses/SP_KD_Loss.py rename to src/faceDetection/face_sdk/addition_module/DMUE/losses/SP_KD_Loss.py index b307243..5f8ea90 100644 --- a/addition_module/DMUE/losses/SP_KD_Loss.py +++ b/src/faceDetection/face_sdk/addition_module/DMUE/losses/SP_KD_Loss.py @@ -1,23 +1,23 @@ -import torch -import torch.nn.functional as F -import numpy as np - - -class SP_KD_Loss(object): - def __call__(self, G_matrixs, G_main_matrixs): - # G_matrixs: List - # G_main_matrixs: List - G_err = [F.mse_loss(G_aux, G_main) for G_aux, G_main in zip(G_matrixs, G_main_matrixs)] - G_err = sum(G_err) / len(G_main_matrixs) - - return G_err - - -if __name__ == '__main__': - loss = SoftLoss() - x = torch.randn([2, 8]) - y = torch.from_numpy(np.array([7,2])) - soft_y = torch.randn([2,8]) - mask = torch.ones(2,8).scatter_(1, y.view(-1,1).long(),0) - soft_y = soft_y * mask - Lsoft = loss(x, y, soft_y) +import torch +import torch.nn.functional as F +import numpy as np + + +class SP_KD_Loss(object): + def __call__(self, G_matrixs, G_main_matrixs): + # G_matrixs: List + # G_main_matrixs: List + G_err = [F.mse_loss(G_aux, G_main) for G_aux, G_main in zip(G_matrixs, G_main_matrixs)] + G_err = sum(G_err) / len(G_main_matrixs) + + return G_err + + +if __name__ == '__main__': + loss = SoftLoss() + x = torch.randn([2, 8]) + y = torch.from_numpy(np.array([7,2])) + soft_y = torch.randn([2,8]) + mask = torch.ones(2,8).scatter_(1, y.view(-1,1).long(),0) + soft_y = soft_y * mask + Lsoft = loss(x, y, soft_y) diff --git a/addition_module/DMUE/losses/__init__.py b/src/faceDetection/face_sdk/addition_module/DMUE/losses/__init__.py similarity index 100% rename from addition_module/DMUE/losses/__init__.py rename to src/faceDetection/face_sdk/addition_module/DMUE/losses/__init__.py diff --git a/addition_module/DMUE/losses/softLoss.py b/src/faceDetection/face_sdk/addition_module/DMUE/losses/softLoss.py similarity index 97% rename from addition_module/DMUE/losses/softLoss.py rename to src/faceDetection/face_sdk/addition_module/DMUE/losses/softLoss.py index 8f71a8a..aed3966 100644 --- a/addition_module/DMUE/losses/softLoss.py +++ b/src/faceDetection/face_sdk/addition_module/DMUE/losses/softLoss.py @@ -1,36 +1,36 @@ -import torch -import torch.nn.functional as F -import numpy as np - - -class SoftLoss(object): - def __call__(self, outputs_x, targets_x, softlabel_x, epoch=None): - # output_x: tensor - # softlabel_x: tensor , size like output_x - probs_x = torch.softmax(outputs_x, dim=1).cuda() - mask = torch.ones_like(probs_x).scatter_(1, targets_x.view(-1, 1).long(), 0).cuda() - probs_x = probs_x * mask - Lsoft = torch.mean((probs_x - softlabel_x)**2) - return Lsoft - -class SoftLoss_normal(object): - def __call__(self, outputs_x, targets_x, softlabel_x, epoch=None): - # output_x: tensor - # softlabel_x: list, each element size like output_x - # it is the normal KD loss, used only in debug - print('The normal KD loss, used only in debug ') - probs_x = torch.softmax(outputs_x, dim=1).cuda() - kd_loss = [torch.mean((probs_x - aux_label)**2) for aux_label in softlabel_x] - kd_loss = sum(kd_loss) / len(softlabel_x) - - return kd_loss - -if __name__ == '__main__': - loss = SoftLoss() - x = torch.randn([2, 8]) - y = torch.from_numpy(np.array([7,2])) - soft_y = torch.randn([2,8]) - mask = torch.ones(2,8).scatter_(1, y.view(-1,1).long(),0) - soft_y = soft_y * mask - Lsoft = loss(x, y, soft_y) +import torch +import torch.nn.functional as F +import numpy as np + + +class SoftLoss(object): + def __call__(self, outputs_x, targets_x, softlabel_x, epoch=None): + # output_x: tensor + # softlabel_x: tensor , size like output_x + probs_x = torch.softmax(outputs_x, dim=1).cuda() + mask = torch.ones_like(probs_x).scatter_(1, targets_x.view(-1, 1).long(), 0).cuda() + probs_x = probs_x * mask + Lsoft = torch.mean((probs_x - softlabel_x)**2) + return Lsoft + +class SoftLoss_normal(object): + def __call__(self, outputs_x, targets_x, softlabel_x, epoch=None): + # output_x: tensor + # softlabel_x: list, each element size like output_x + # it is the normal KD loss, used only in debug + print('The normal KD loss, used only in debug ') + probs_x = torch.softmax(outputs_x, dim=1).cuda() + kd_loss = [torch.mean((probs_x - aux_label)**2) for aux_label in softlabel_x] + kd_loss = sum(kd_loss) / len(softlabel_x) + + return kd_loss + +if __name__ == '__main__': + loss = SoftLoss() + x = torch.randn([2, 8]) + y = torch.from_numpy(np.array([7,2])) + soft_y = torch.randn([2,8]) + mask = torch.ones(2,8).scatter_(1, y.view(-1,1).long(),0) + soft_y = soft_y * mask + Lsoft = loss(x, y, soft_y) \ No newline at end of file diff --git a/addition_module/DMUE/models/__init__.py b/src/faceDetection/face_sdk/addition_module/DMUE/models/__init__.py similarity index 98% rename from addition_module/DMUE/models/__init__.py rename to src/faceDetection/face_sdk/addition_module/DMUE/models/__init__.py index 64794c5..2669fe5 100644 --- a/addition_module/DMUE/models/__init__.py +++ b/src/faceDetection/face_sdk/addition_module/DMUE/models/__init__.py @@ -1,2 +1,2 @@ -from .make_model import make_model +from .make_model import make_model from .make_target_model import make_target_model \ No newline at end of file diff --git a/addition_module/DMUE/models/make_model.py b/src/faceDetection/face_sdk/addition_module/DMUE/models/make_model.py similarity index 97% rename from addition_module/DMUE/models/make_model.py rename to src/faceDetection/face_sdk/addition_module/DMUE/models/make_model.py index d1e98bf..85d21ab 100644 --- a/addition_module/DMUE/models/make_model.py +++ b/src/faceDetection/face_sdk/addition_module/DMUE/models/make_model.py @@ -1,236 +1,236 @@ -import torch -import torch.nn as nn -import torch.nn.functional as F -from torch.nn.parameter import Parameter - -from .resnet_multibranch import ResNet, BasicBlock, Bottleneck -from .resnet_ibn_multibranch import resnet50_ibn_a - - -class GeM(nn.Module): - def __init__(self, p=3.0, eps=1e-6, freeze_p=True): - super(GeM, self).__init__() - self.p = p if freeze_p else Parameter(torch.ones(1) * p) - self.eps = eps - - def forward(self, x): - return F.adaptive_avg_pool2d(x.clamp(min=self.eps).pow(self.p), - (1, 1)).pow(1. / self.p) - - def __repr__(self): - if isinstance(self.p, float): - p = self.p - else: - p = self.p.data.tolist()[0] - return self.__class__.__name__ +\ - '(' + 'p=' + '{:.4f}'.format(p) +\ - ', ' + 'eps=' + str(self.eps) + ')' - - -def weights_init_kaiming(m): - classname = m.__class__.__name__ - if classname.find('Linear') != -1: - nn.init.kaiming_normal_(m.weight, a=0, mode='fan_out') - nn.init.constant_(m.bias, 0.0) - - elif classname.find('Conv') != -1: - nn.init.kaiming_normal_(m.weight, a=0, mode='fan_in') - if m.bias is not None: - nn.init.constant_(m.bias, 0.0) - elif classname.find('BatchNorm') != -1: - if m.affine: - nn.init.constant_(m.weight, 1.0) - nn.init.constant_(m.bias, 0.0) - - -def weights_init_classifier(m): - classname = m.__class__.__name__ - if classname.find('Linear') != -1: - nn.init.normal_(m.weight, std=0.001) - if m.bias: # need to be comment when BiasInCls is True - nn.init.constant_(m.bias, 0.0) - - -class Backbone(nn.Module): - def __init__(self, cfg): - super(Backbone, self).__init__() - last_stride = cfg.last_stride - model_name = cfg.backbone - self.num_classes = cfg.num_classes - self.num_branches = cfg.num_branches - - # base model - if model_name == 'resnet18': - self.in_planes = 512 - self.base = ResNet(num_branches=self.num_branches, - num_classes = self.num_classes, - last_stride=last_stride, - block=BasicBlock, frozen_stages=cfg.frozen_stages, - layers=[2, 2, 2, 2]) - # print('Using resnet18 as a backbone') - elif model_name == 'resnet50_ibn': - self.in_planes = 2048 - self.base = resnet50_ibn_a(last_stride, self.num_branches, self.num_classes) - # print('Using resnet50_ibn as a backbone') - else: - raise ValueError('Unsupported backbone! but got {}'.format(model_name)) - - # pooling after base - if cfg.pooling_method == 'GeM': - # print('Using GeM pooling') - self.gap = GeM() - elif cfg.pooling_method == 'GAP': - self.gap = nn.AdaptiveAvgPool2d(1) - - # loss - self.classifiers = [] - self.KD_whole_cls = False - for i in range(self.num_branches): - output_classes = self.num_classes if i is self.num_branches-1 else self.num_classes -1 - self.classifiers.append(nn.Linear(self.in_planes, output_classes, bias=cfg.BiasInCls)) - self.classifiers[i].apply(weights_init_classifier) - self.classifiers = nn.ModuleList(self.classifiers) - - self.dropout = nn.Dropout(p=0.5) - self.neck = cfg.bnneck - self.second_order_statics = cfg.second_order_statics - if cfg.train_mode == 'sp_confidence': - self.sigmoid = nn.Sigmoid() - input_dim = 2 * self.num_classes - self.project_w = nn.Sequential(nn.Linear(input_dim, cfg.num_classes),nn.PReLU(), nn.Linear(cfg.num_classes, 1)) - - if cfg.bnneck: - # print('Using bnneck') - self.bottleneck = nn.BatchNorm1d(self.in_planes) - self.bottleneck.bias.requires_grad_(False) - self.bottleneck.apply(weights_init_kaiming) - - def forward(self, x, label=None, training_phase=None, c=None): - x_final, x_list, softlabel_list, targets, inds, inds_softlabel = self.base(x, label, training_phase=training_phase, c=c) - if training_phase == 'normal': - x_final = self.gap(x_final).squeeze(2).squeeze(2) - x_final = self.dropout(x_final) - x_final = self.bottleneck(x_final) if self.neck else x_final - x_final = self.classifiers[self.num_branches - 1](x_final) - - return x_final, x_list, targets, None - - # train auxiliary & main branch - x_list = [self.gap(_x).squeeze(2).squeeze(2) for _x in x_list] - x_list = [self.dropout(_x) for _x in x_list] - x_list = [self.bottleneck(_x) for _x in x_list] if self.neck else x_list - x_list = [self.classifiers[i](x_list[i]) for i in range(self.num_branches - 1)] - x_final = self.gap(x_final).squeeze(2).squeeze(2) - - # before dropout cal cos_simlarity in batch - if training_phase == 'sp_confidence': - cos_dot_product_matrix = x_final.mm(x_final.t()) - x_final_norm = torch.norm(x_final, dim=1).unsqueeze(1) - x_norm_product = x_final_norm.mm(x_final_norm.t()) - cos_similarity = cos_dot_product_matrix / x_norm_product - - x_final = self.dropout(x_final) - x_final = self.bottleneck(x_final) if self.neck else x_final - - # make distribution - softlabel_list = [self.gap(_x).squeeze(2).squeeze(2) for _x in softlabel_list] - softlabel_list = [self.bottleneck(_x) for _x in softlabel_list] if self.neck else softlabel_list - - if training_phase == 'sp_confidence': - # sp kd - G_matrixs = [fm.mm(fm.t()) for fm in x_list] - G_main = x_final.mm(x_final.t()) - # construct sub-matrix of G_main before l2-norm - G_main_matrixs = [G_main[ind[:,0], :][:, ind[:,0]] for ind in inds] - # l2 norm - G_matrixs = [F.normalize(G, p=2, dim=1) for G in G_matrixs] - G_main_matrixs = [F.normalize(G, p=2, dim=1) for G in G_main_matrixs] - - softlabel_list = [self.classifiers[i](softlabel_list[i]) for i in range(self.num_branches - 1)] - x_final = self.classifiers[self.num_branches - 1](x_final) - - # convert label into one-hot vector - label = label.unsqueeze(1) - one_hot = torch.zeros(label.shape[0], self.num_classes).cuda().scatter_(1, label, 1) - # cal mean according to different classes - cos_matrixs = [cos_similarity[:, ind[:, 0]] for ind in inds_softlabel] # may be change by torch.Tensor.indexadd(dim, index, tensor) - cos_matrixs_mean = [torch.mean(cos_m, dim=1, keepdim=True) for cos_m in cos_matrixs] - cos_mean = torch.cat(tuple(cos_matrixs_mean), dim=1)# bs x num_class - # scale to 0-1, the effect is minor - cos_mean = (cos_mean - cos_mean.min(dim=1, keepdim=True)[0]) / (cos_mean.max(dim=1, keepdim=True)[0] - cos_mean.min(dim=1, keepdim=True)[0]) - statics = torch.cat((cos_mean, one_hot), dim=1) - - score = self.project_w(statics) - score = self.sigmoid(score) - atten_x_final = x_final * score - - return x_final, x_list, targets, softlabel_list, G_matrixs, G_main_matrixs, score, atten_x_final - else: - raise Exception("Specific training phase") - - - def load_param(self, cfg): - pretrained = cfg.pretrained - param_dict = torch.load(pretrained, map_location=lambda storage,loc: storage.cpu()) - # print('Pretrained choice ', cfg.pretrained_choice) - - if cfg.pretrained_choice == 'msra': - if 'resnet' in cfg.backbone: - for i in param_dict['state_dict']: - if ('loss_layer' in i) or ('fc' in i): - continue - else: - layer_name = i.split('.')[1] - if 'layer4' == layer_name: - op_name = i[16:] - for branch_idx in range(self.num_branches): - j = 'base.'+layer_name +'_'+str(branch_idx)+'.'+op_name - assert j in self.state_dict() - self.state_dict()[j].copy_(param_dict['state_dict'][i]) - else: - j = 'base'+i.split('feat_net')[1] - assert j in self.state_dict() - self.state_dict()[j].copy_(param_dict['state_dict'][i]) - elif 'ShuffleNet' in cfg.backbone: - for i in param_dict['state_dict']: - if ('loss_layer' in i) or ('fc' in i): - continue - else: - k = int(i.split('.')[2]) - if k > 11: - for branch_idx in range(self.num_branches): - j = i.replace('feat_net', 'base') - _pre, _post = j[:4], j[16:] - j = _pre + '.last_stage_{}.'.format(branch_idx) + str(k-12) + _post - assert j in self.state_dict() - self.state_dict()[j].copy_(param_dict['state_dict'][i]) - else: - j = i.replace('feat_net', 'base') - assert j in self.state_dict() - self.state_dict()[j].copy_(param_dict['state_dict'][i]) - elif 'MobileNetV2' in cfg.backbone: - for i in param_dict['state_dict']: - if ('loss_layer' in i) or ('fc' in i): - continue - else: - k = int(i.split('.')[2]) - if k > 13: - for branch_idx in range(self.num_branches): - j = i.replace('feat_net', 'base') - _pre, _post = j[:4], j[16:] - j = _pre + '.last_stage_{}.'.format(branch_idx) + str(k-14) + _post - assert j in self.state_dict() - self.state_dict()[j].copy_(param_dict['state_dict'][i]) - else: - j = i.replace('feat_net', 'base') - assert j in self.state_dict() - self.state_dict()[j].copy_(param_dict['state_dict'][i]) - else: - for i in param_dict: - j = i.replace('module.', '') - self.state_dict()[j].copy_(param_dict[i]) - - -def make_model(cfg): - model = Backbone(cfg) - return model +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.nn.parameter import Parameter + +from .resnet_multibranch import ResNet, BasicBlock, Bottleneck +from .resnet_ibn_multibranch import resnet50_ibn_a + + +class GeM(nn.Module): + def __init__(self, p=3.0, eps=1e-6, freeze_p=True): + super(GeM, self).__init__() + self.p = p if freeze_p else Parameter(torch.ones(1) * p) + self.eps = eps + + def forward(self, x): + return F.adaptive_avg_pool2d(x.clamp(min=self.eps).pow(self.p), + (1, 1)).pow(1. / self.p) + + def __repr__(self): + if isinstance(self.p, float): + p = self.p + else: + p = self.p.data.tolist()[0] + return self.__class__.__name__ +\ + '(' + 'p=' + '{:.4f}'.format(p) +\ + ', ' + 'eps=' + str(self.eps) + ')' + + +def weights_init_kaiming(m): + classname = m.__class__.__name__ + if classname.find('Linear') != -1: + nn.init.kaiming_normal_(m.weight, a=0, mode='fan_out') + nn.init.constant_(m.bias, 0.0) + + elif classname.find('Conv') != -1: + nn.init.kaiming_normal_(m.weight, a=0, mode='fan_in') + if m.bias is not None: + nn.init.constant_(m.bias, 0.0) + elif classname.find('BatchNorm') != -1: + if m.affine: + nn.init.constant_(m.weight, 1.0) + nn.init.constant_(m.bias, 0.0) + + +def weights_init_classifier(m): + classname = m.__class__.__name__ + if classname.find('Linear') != -1: + nn.init.normal_(m.weight, std=0.001) + if m.bias: # need to be comment when BiasInCls is True + nn.init.constant_(m.bias, 0.0) + + +class Backbone(nn.Module): + def __init__(self, cfg): + super(Backbone, self).__init__() + last_stride = cfg.last_stride + model_name = cfg.backbone + self.num_classes = cfg.num_classes + self.num_branches = cfg.num_branches + + # base model + if model_name == 'resnet18': + self.in_planes = 512 + self.base = ResNet(num_branches=self.num_branches, + num_classes = self.num_classes, + last_stride=last_stride, + block=BasicBlock, frozen_stages=cfg.frozen_stages, + layers=[2, 2, 2, 2]) + # print('Using resnet18 as a backbone') + elif model_name == 'resnet50_ibn': + self.in_planes = 2048 + self.base = resnet50_ibn_a(last_stride, self.num_branches, self.num_classes) + # print('Using resnet50_ibn as a backbone') + else: + raise ValueError('Unsupported backbone! but got {}'.format(model_name)) + + # pooling after base + if cfg.pooling_method == 'GeM': + # print('Using GeM pooling') + self.gap = GeM() + elif cfg.pooling_method == 'GAP': + self.gap = nn.AdaptiveAvgPool2d(1) + + # loss + self.classifiers = [] + self.KD_whole_cls = False + for i in range(self.num_branches): + output_classes = self.num_classes if i is self.num_branches-1 else self.num_classes -1 + self.classifiers.append(nn.Linear(self.in_planes, output_classes, bias=cfg.BiasInCls)) + self.classifiers[i].apply(weights_init_classifier) + self.classifiers = nn.ModuleList(self.classifiers) + + self.dropout = nn.Dropout(p=0.5) + self.neck = cfg.bnneck + self.second_order_statics = cfg.second_order_statics + if cfg.train_mode == 'sp_confidence': + self.sigmoid = nn.Sigmoid() + input_dim = 2 * self.num_classes + self.project_w = nn.Sequential(nn.Linear(input_dim, cfg.num_classes),nn.PReLU(), nn.Linear(cfg.num_classes, 1)) + + if cfg.bnneck: + # print('Using bnneck') + self.bottleneck = nn.BatchNorm1d(self.in_planes) + self.bottleneck.bias.requires_grad_(False) + self.bottleneck.apply(weights_init_kaiming) + + def forward(self, x, label=None, training_phase=None, c=None): + x_final, x_list, softlabel_list, targets, inds, inds_softlabel = self.base(x, label, training_phase=training_phase, c=c) + if training_phase == 'normal': + x_final = self.gap(x_final).squeeze(2).squeeze(2) + x_final = self.dropout(x_final) + x_final = self.bottleneck(x_final) if self.neck else x_final + x_final = self.classifiers[self.num_branches - 1](x_final) + + return x_final, x_list, targets, None + + # train auxiliary & main branch + x_list = [self.gap(_x).squeeze(2).squeeze(2) for _x in x_list] + x_list = [self.dropout(_x) for _x in x_list] + x_list = [self.bottleneck(_x) for _x in x_list] if self.neck else x_list + x_list = [self.classifiers[i](x_list[i]) for i in range(self.num_branches - 1)] + x_final = self.gap(x_final).squeeze(2).squeeze(2) + + # before dropout cal cos_simlarity in batch + if training_phase == 'sp_confidence': + cos_dot_product_matrix = x_final.mm(x_final.t()) + x_final_norm = torch.norm(x_final, dim=1).unsqueeze(1) + x_norm_product = x_final_norm.mm(x_final_norm.t()) + cos_similarity = cos_dot_product_matrix / x_norm_product + + x_final = self.dropout(x_final) + x_final = self.bottleneck(x_final) if self.neck else x_final + + # make distribution + softlabel_list = [self.gap(_x).squeeze(2).squeeze(2) for _x in softlabel_list] + softlabel_list = [self.bottleneck(_x) for _x in softlabel_list] if self.neck else softlabel_list + + if training_phase == 'sp_confidence': + # sp kd + G_matrixs = [fm.mm(fm.t()) for fm in x_list] + G_main = x_final.mm(x_final.t()) + # construct sub-matrix of G_main before l2-norm + G_main_matrixs = [G_main[ind[:,0], :][:, ind[:,0]] for ind in inds] + # l2 norm + G_matrixs = [F.normalize(G, p=2, dim=1) for G in G_matrixs] + G_main_matrixs = [F.normalize(G, p=2, dim=1) for G in G_main_matrixs] + + softlabel_list = [self.classifiers[i](softlabel_list[i]) for i in range(self.num_branches - 1)] + x_final = self.classifiers[self.num_branches - 1](x_final) + + # convert label into one-hot vector + label = label.unsqueeze(1) + one_hot = torch.zeros(label.shape[0], self.num_classes).cuda().scatter_(1, label, 1) + # cal mean according to different classes + cos_matrixs = [cos_similarity[:, ind[:, 0]] for ind in inds_softlabel] # may be change by torch.Tensor.indexadd(dim, index, tensor) + cos_matrixs_mean = [torch.mean(cos_m, dim=1, keepdim=True) for cos_m in cos_matrixs] + cos_mean = torch.cat(tuple(cos_matrixs_mean), dim=1)# bs x num_class + # scale to 0-1, the effect is minor + cos_mean = (cos_mean - cos_mean.min(dim=1, keepdim=True)[0]) / (cos_mean.max(dim=1, keepdim=True)[0] - cos_mean.min(dim=1, keepdim=True)[0]) + statics = torch.cat((cos_mean, one_hot), dim=1) + + score = self.project_w(statics) + score = self.sigmoid(score) + atten_x_final = x_final * score + + return x_final, x_list, targets, softlabel_list, G_matrixs, G_main_matrixs, score, atten_x_final + else: + raise Exception("Specific training phase") + + + def load_param(self, cfg): + pretrained = cfg.pretrained + param_dict = torch.load(pretrained, map_location=lambda storage,loc: storage.cpu()) + # print('Pretrained choice ', cfg.pretrained_choice) + + if cfg.pretrained_choice == 'msra': + if 'resnet' in cfg.backbone: + for i in param_dict['state_dict']: + if ('loss_layer' in i) or ('fc' in i): + continue + else: + layer_name = i.split('.')[1] + if 'layer4' == layer_name: + op_name = i[16:] + for branch_idx in range(self.num_branches): + j = 'base.'+layer_name +'_'+str(branch_idx)+'.'+op_name + assert j in self.state_dict() + self.state_dict()[j].copy_(param_dict['state_dict'][i]) + else: + j = 'base'+i.split('feat_net')[1] + assert j in self.state_dict() + self.state_dict()[j].copy_(param_dict['state_dict'][i]) + elif 'ShuffleNet' in cfg.backbone: + for i in param_dict['state_dict']: + if ('loss_layer' in i) or ('fc' in i): + continue + else: + k = int(i.split('.')[2]) + if k > 11: + for branch_idx in range(self.num_branches): + j = i.replace('feat_net', 'base') + _pre, _post = j[:4], j[16:] + j = _pre + '.last_stage_{}.'.format(branch_idx) + str(k-12) + _post + assert j in self.state_dict() + self.state_dict()[j].copy_(param_dict['state_dict'][i]) + else: + j = i.replace('feat_net', 'base') + assert j in self.state_dict() + self.state_dict()[j].copy_(param_dict['state_dict'][i]) + elif 'MobileNetV2' in cfg.backbone: + for i in param_dict['state_dict']: + if ('loss_layer' in i) or ('fc' in i): + continue + else: + k = int(i.split('.')[2]) + if k > 13: + for branch_idx in range(self.num_branches): + j = i.replace('feat_net', 'base') + _pre, _post = j[:4], j[16:] + j = _pre + '.last_stage_{}.'.format(branch_idx) + str(k-14) + _post + assert j in self.state_dict() + self.state_dict()[j].copy_(param_dict['state_dict'][i]) + else: + j = i.replace('feat_net', 'base') + assert j in self.state_dict() + self.state_dict()[j].copy_(param_dict['state_dict'][i]) + else: + for i in param_dict: + j = i.replace('module.', '') + self.state_dict()[j].copy_(param_dict[i]) + + +def make_model(cfg): + model = Backbone(cfg) + return model diff --git a/addition_module/DMUE/models/make_target_model.py b/src/faceDetection/face_sdk/addition_module/DMUE/models/make_target_model.py similarity index 100% rename from addition_module/DMUE/models/make_target_model.py rename to src/faceDetection/face_sdk/addition_module/DMUE/models/make_target_model.py diff --git a/addition_module/DMUE/models/resnet.py b/src/faceDetection/face_sdk/addition_module/DMUE/models/resnet.py similarity index 96% rename from addition_module/DMUE/models/resnet.py rename to src/faceDetection/face_sdk/addition_module/DMUE/models/resnet.py index 21d2d87..9e9d3b0 100644 --- a/addition_module/DMUE/models/resnet.py +++ b/src/faceDetection/face_sdk/addition_module/DMUE/models/resnet.py @@ -1,125 +1,125 @@ -import math -import torch -from torch import nn - - -def conv3x3(in_planes, out_planes, stride=1): - """3x3 convolution with padding""" - return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, - padding=1, bias=False) - - -class BasicBlock(nn.Module): - expansion = 1 - - def __init__(self, inplanes, planes, stride=1, downsample=None): - super(BasicBlock, self).__init__() - self.conv1 = conv3x3(inplanes, planes, stride) - self.bn1 = nn.BatchNorm2d(planes) - self.relu = nn.ReLU(inplace=True) - self.conv2 = conv3x3(planes, planes) - self.bn2 = nn.BatchNorm2d(planes) - self.downsample = downsample - self.stride = stride - - def forward(self, x): - residual = x - - out = self.conv1(x) - out = self.bn1(out) - out = self.relu(out) - - out = self.conv2(out) - out = self.bn2(out) - - if self.downsample is not None: - residual = self.downsample(x) - - out += residual - out = self.relu(out) - - return out - - -class Bottleneck(nn.Module): - expansion = 4 - - def __init__(self, inplanes, planes, stride=1, downsample=None): - super(Bottleneck, self).__init__() - self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) - self.bn1 = nn.BatchNorm2d(planes) - self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, - padding=1, bias=False) - self.bn2 = nn.BatchNorm2d(planes) - self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False) - self.bn3 = nn.BatchNorm2d(planes * 4) - self.relu = nn.ReLU(inplace=True) - self.downsample = downsample - self.stride = stride - - def forward(self, x): - residual = x - - out = self.conv1(x) - out = self.bn1(out) - out = self.relu(out) - - out = self.conv2(out) - out = self.bn2(out) - out = self.relu(out) - - out = self.conv3(out) - out = self.bn3(out) - - if self.downsample is not None: - residual = self.downsample(x) - - out += residual - out = self.relu(out) - - return out - - -class ResNet(nn.Module): - def __init__(self, last_stride=2, block=Bottleneck, layers=[3, 4, 6, 3]): - self.inplanes = 64 - super().__init__() - self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, - bias=False) - self.bn1 = nn.BatchNorm2d(64) - # self.relu = nn.ReLU(inplace=True) # add missed relu - self.maxpool = nn.MaxPool2d(kernel_size=2, stride=None, padding=0) - self.layer1 = self._make_layer(block, 64, layers[0]) - self.layer2 = self._make_layer(block, 128, layers[1], stride=2) - self.layer3 = self._make_layer(block, 256, layers[2], stride=2) - self.layer4 = self._make_layer(block, 512, layers[3], stride=last_stride) - - def _make_layer(self, block, planes, blocks, stride=1): - downsample = None - if stride != 1 or self.inplanes != planes * block.expansion: - downsample = nn.Sequential( - nn.Conv2d(self.inplanes, planes * block.expansion, - kernel_size=1, stride=stride, bias=False), - nn.BatchNorm2d(planes * block.expansion), - ) - - layers = [] - layers.append(block(self.inplanes, planes, stride, downsample)) - self.inplanes = planes * block.expansion - for i in range(1, blocks): - layers.append(block(self.inplanes, planes)) - - return nn.Sequential(*layers) - - def forward(self, x): - x = self.conv1(x) - x = self.bn1(x) - # x = self.relu(x) # add missed relu - x = self.maxpool(x) - - x = self.layer1(x) - x = self.layer2(x) - x = self.layer3(x) - x = self.layer4(x) - - return x +import math +import torch +from torch import nn + + +def conv3x3(in_planes, out_planes, stride=1): + """3x3 convolution with padding""" + return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, + padding=1, bias=False) + + +class BasicBlock(nn.Module): + expansion = 1 + + def __init__(self, inplanes, planes, stride=1, downsample=None): + super(BasicBlock, self).__init__() + self.conv1 = conv3x3(inplanes, planes, stride) + self.bn1 = nn.BatchNorm2d(planes) + self.relu = nn.ReLU(inplace=True) + self.conv2 = conv3x3(planes, planes) + self.bn2 = nn.BatchNorm2d(planes) + self.downsample = downsample + self.stride = stride + + def forward(self, x): + residual = x + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + + out = self.conv2(out) + out = self.bn2(out) + + if self.downsample is not None: + residual = self.downsample(x) + + out += residual + out = self.relu(out) + + return out + + +class Bottleneck(nn.Module): + expansion = 4 + + def __init__(self, inplanes, planes, stride=1, downsample=None): + super(Bottleneck, self).__init__() + self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) + self.bn1 = nn.BatchNorm2d(planes) + self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, + padding=1, bias=False) + self.bn2 = nn.BatchNorm2d(planes) + self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False) + self.bn3 = nn.BatchNorm2d(planes * 4) + self.relu = nn.ReLU(inplace=True) + self.downsample = downsample + self.stride = stride + + def forward(self, x): + residual = x + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + + out = self.conv2(out) + out = self.bn2(out) + out = self.relu(out) + + out = self.conv3(out) + out = self.bn3(out) + + if self.downsample is not None: + residual = self.downsample(x) + + out += residual + out = self.relu(out) + + return out + + +class ResNet(nn.Module): + def __init__(self, last_stride=2, block=Bottleneck, layers=[3, 4, 6, 3]): + self.inplanes = 64 + super().__init__() + self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, + bias=False) + self.bn1 = nn.BatchNorm2d(64) + # self.relu = nn.ReLU(inplace=True) # add missed relu + self.maxpool = nn.MaxPool2d(kernel_size=2, stride=None, padding=0) + self.layer1 = self._make_layer(block, 64, layers[0]) + self.layer2 = self._make_layer(block, 128, layers[1], stride=2) + self.layer3 = self._make_layer(block, 256, layers[2], stride=2) + self.layer4 = self._make_layer(block, 512, layers[3], stride=last_stride) + + def _make_layer(self, block, planes, blocks, stride=1): + downsample = None + if stride != 1 or self.inplanes != planes * block.expansion: + downsample = nn.Sequential( + nn.Conv2d(self.inplanes, planes * block.expansion, + kernel_size=1, stride=stride, bias=False), + nn.BatchNorm2d(planes * block.expansion), + ) + + layers = [] + layers.append(block(self.inplanes, planes, stride, downsample)) + self.inplanes = planes * block.expansion + for i in range(1, blocks): + layers.append(block(self.inplanes, planes)) + + return nn.Sequential(*layers) + + def forward(self, x): + x = self.conv1(x) + x = self.bn1(x) + # x = self.relu(x) # add missed relu + x = self.maxpool(x) + + x = self.layer1(x) + x = self.layer2(x) + x = self.layer3(x) + x = self.layer4(x) + + return x diff --git a/addition_module/DMUE/models/resnet_ibn_a.py b/src/faceDetection/face_sdk/addition_module/DMUE/models/resnet_ibn_a.py similarity index 96% rename from addition_module/DMUE/models/resnet_ibn_a.py rename to src/faceDetection/face_sdk/addition_module/DMUE/models/resnet_ibn_a.py index 46cd1ac..d7d8886 100644 --- a/addition_module/DMUE/models/resnet_ibn_a.py +++ b/src/faceDetection/face_sdk/addition_module/DMUE/models/resnet_ibn_a.py @@ -1,111 +1,111 @@ -import math -import torch -import torch.nn as nn - - -class IBN(nn.Module): - def __init__(self, planes): - super(IBN, self).__init__() - half1 = int(planes/2) - self.half = half1 - half2 = planes - half1 - self.IN = nn.InstanceNorm2d(half1, affine=True) - self.BN = nn.BatchNorm2d(half2) - - def forward(self, x): - split = torch.split(x, self.half, 1) - out1 = self.IN(split[0].contiguous()) - out2 = self.BN(split[1].contiguous()) - out = torch.cat((out1, out2), 1) - return out - - -class Bottleneck_IBN(nn.Module): - expansion = 4 - - def __init__(self, inplanes, planes, ibn=False, stride=1, downsample=None): - super(Bottleneck_IBN, self).__init__() - self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) - if ibn: - self.bn1 = IBN(planes) - else: - self.bn1 = nn.BatchNorm2d(planes) - self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, - padding=1, bias=False) - self.bn2 = nn.BatchNorm2d(planes) - self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1, bias=False) - self.bn3 = nn.BatchNorm2d(planes * self.expansion) - self.relu = nn.ReLU(inplace=True) - self.downsample = downsample - self.stride = stride - - def forward(self, x): - residual = x - - out = self.conv1(x) - out = self.bn1(out) - out = self.relu(out) - - out = self.conv2(out) - out = self.bn2(out) - out = self.relu(out) - - out = self.conv3(out) - out = self.bn3(out) - - if self.downsample is not None: - residual = self.downsample(x) - - out += residual - out = self.relu(out) - - return out - - -class ResNet_IBN(nn.Module): - def __init__(self, last_stride, block, layers): - scale = 64 - self.inplanes = scale - super(ResNet_IBN, self).__init__() - self.conv1 = nn.Conv2d(3, scale, kernel_size=7, stride=2, padding=3, - bias=False) - self.bn1 = nn.BatchNorm2d(scale) - self.relu = nn.ReLU(inplace=True) - self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) - self.layer1 = self._make_layer(block, scale, layers[0]) - self.layer2 = self._make_layer(block, scale*2, layers[1], stride=2) - self.layer3 = self._make_layer(block, scale*4, layers[2], stride=2) - self.layer4 = self._make_layer(block, scale*8, layers[3], stride=last_stride) - - def _make_layer(self, block, planes, blocks, stride=1): - downsample = None - if stride != 1 or self.inplanes != planes * block.expansion: - downsample = nn.Sequential( - nn.Conv2d(self.inplanes, planes * block.expansion, - kernel_size=1, stride=stride, bias=False), - nn.BatchNorm2d(planes * block.expansion), - ) - - layers = [] - ibn = True - if planes == 512: - ibn = False - layers.append(block(self.inplanes, planes, ibn, stride, downsample)) - self.inplanes = planes * block.expansion - for i in range(1, blocks): - layers.append(block(self.inplanes, planes, ibn)) - - return nn.Sequential(*layers) - - def forward(self, x): - x = self.conv1(x) - x = self.bn1(x) - x = self.relu(x) - x = self.maxpool(x) - - x = self.layer1(x) - x = self.layer2(x) - x = self.layer3(x) - x = self.layer4(x) - - return x +import math +import torch +import torch.nn as nn + + +class IBN(nn.Module): + def __init__(self, planes): + super(IBN, self).__init__() + half1 = int(planes/2) + self.half = half1 + half2 = planes - half1 + self.IN = nn.InstanceNorm2d(half1, affine=True) + self.BN = nn.BatchNorm2d(half2) + + def forward(self, x): + split = torch.split(x, self.half, 1) + out1 = self.IN(split[0].contiguous()) + out2 = self.BN(split[1].contiguous()) + out = torch.cat((out1, out2), 1) + return out + + +class Bottleneck_IBN(nn.Module): + expansion = 4 + + def __init__(self, inplanes, planes, ibn=False, stride=1, downsample=None): + super(Bottleneck_IBN, self).__init__() + self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) + if ibn: + self.bn1 = IBN(planes) + else: + self.bn1 = nn.BatchNorm2d(planes) + self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, + padding=1, bias=False) + self.bn2 = nn.BatchNorm2d(planes) + self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1, bias=False) + self.bn3 = nn.BatchNorm2d(planes * self.expansion) + self.relu = nn.ReLU(inplace=True) + self.downsample = downsample + self.stride = stride + + def forward(self, x): + residual = x + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + + out = self.conv2(out) + out = self.bn2(out) + out = self.relu(out) + + out = self.conv3(out) + out = self.bn3(out) + + if self.downsample is not None: + residual = self.downsample(x) + + out += residual + out = self.relu(out) + + return out + + +class ResNet_IBN(nn.Module): + def __init__(self, last_stride, block, layers): + scale = 64 + self.inplanes = scale + super(ResNet_IBN, self).__init__() + self.conv1 = nn.Conv2d(3, scale, kernel_size=7, stride=2, padding=3, + bias=False) + self.bn1 = nn.BatchNorm2d(scale) + self.relu = nn.ReLU(inplace=True) + self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) + self.layer1 = self._make_layer(block, scale, layers[0]) + self.layer2 = self._make_layer(block, scale*2, layers[1], stride=2) + self.layer3 = self._make_layer(block, scale*4, layers[2], stride=2) + self.layer4 = self._make_layer(block, scale*8, layers[3], stride=last_stride) + + def _make_layer(self, block, planes, blocks, stride=1): + downsample = None + if stride != 1 or self.inplanes != planes * block.expansion: + downsample = nn.Sequential( + nn.Conv2d(self.inplanes, planes * block.expansion, + kernel_size=1, stride=stride, bias=False), + nn.BatchNorm2d(planes * block.expansion), + ) + + layers = [] + ibn = True + if planes == 512: + ibn = False + layers.append(block(self.inplanes, planes, ibn, stride, downsample)) + self.inplanes = planes * block.expansion + for i in range(1, blocks): + layers.append(block(self.inplanes, planes, ibn)) + + return nn.Sequential(*layers) + + def forward(self, x): + x = self.conv1(x) + x = self.bn1(x) + x = self.relu(x) + x = self.maxpool(x) + + x = self.layer1(x) + x = self.layer2(x) + x = self.layer3(x) + x = self.layer4(x) + + return x diff --git a/addition_module/DMUE/models/resnet_ibn_multibranch.py b/src/faceDetection/face_sdk/addition_module/DMUE/models/resnet_ibn_multibranch.py similarity index 100% rename from addition_module/DMUE/models/resnet_ibn_multibranch.py rename to src/faceDetection/face_sdk/addition_module/DMUE/models/resnet_ibn_multibranch.py diff --git a/addition_module/DMUE/models/resnet_multibranch.py b/src/faceDetection/face_sdk/addition_module/DMUE/models/resnet_multibranch.py similarity index 97% rename from addition_module/DMUE/models/resnet_multibranch.py rename to src/faceDetection/face_sdk/addition_module/DMUE/models/resnet_multibranch.py index 7091234..fceb806 100644 --- a/addition_module/DMUE/models/resnet_multibranch.py +++ b/src/faceDetection/face_sdk/addition_module/DMUE/models/resnet_multibranch.py @@ -1,188 +1,188 @@ -import math -import torch -from torch import nn - - -def conv3x3(in_planes, out_planes, stride=1): - """3x3 convolution with padding""" - return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, - padding=1, bias=False) - - -class BasicBlock(nn.Module): - expansion = 1 - - def __init__(self, inplanes, planes, stride=1, downsample=None): - super(BasicBlock, self).__init__() - self.conv1 = conv3x3(inplanes, planes, stride) - self.bn1 = nn.BatchNorm2d(planes) - self.relu = nn.ReLU(inplace=True) - self.conv2 = conv3x3(planes, planes) - self.bn2 = nn.BatchNorm2d(planes) - self.downsample = downsample - self.stride = stride - - def forward(self, x): - residual = x - - out = self.conv1(x) - out = self.bn1(out) - out = self.relu(out) - - out = self.conv2(out) - out = self.bn2(out) - - if self.downsample is not None: - residual = self.downsample(x) - - out += residual - out = self.relu(out) - - return out - - -class Bottleneck(nn.Module): - expansion = 4 - - def __init__(self, inplanes, planes, stride=1, downsample=None): - super(Bottleneck, self).__init__() - self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) - self.bn1 = nn.BatchNorm2d(planes) - self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, - padding=1, bias=False) - self.bn2 = nn.BatchNorm2d(planes) - self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False) - self.bn3 = nn.BatchNorm2d(planes * 4) - self.relu = nn.ReLU(inplace=True) - self.downsample = downsample - self.stride = stride - - def forward(self, x): - residual = x - - out = self.conv1(x) - out = self.bn1(out) - out = self.relu(out) - - out = self.conv2(out) - out = self.bn2(out) - out = self.relu(out) - - out = self.conv3(out) - out = self.bn3(out) - - if self.downsample is not None: - residual = self.downsample(x) - - out += residual - out = self.relu(out) - - return out - - -class ResNet(nn.Module): - def __init__(self, num_branches, num_classes, - last_stride=2, block=Bottleneck, frozen_stages=-1,layers=[3, 4, 6, 3]): - self.inplanes = 64 - super().__init__() - self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, - bias=False) - self.bn1 = nn.BatchNorm2d(64) - self.maxpool = nn.MaxPool2d(kernel_size=2, stride=None, padding=0) - self.frozen_stages = frozen_stages - self.layer1 = self._make_layer(block, 64, layers[0]) - self.layer2 = self._make_layer(block, 128, layers[1], stride=2) - self.layer3 = self._make_layer(block, 256, layers[2], stride=2) - - fix_inplanes = self.inplanes - self.num_branches = num_branches - self.num_classes = num_classes - for i in range(self.num_branches): - setattr(self, 'layer4_'+str(i), self._make_layer(block, 512, layers[3], stride=last_stride)) - self.inplanes = fix_inplanes - - def _make_layer(self, block, planes, blocks, stride=1): - downsample = None - if stride != 1 or self.inplanes != planes * block.expansion: - downsample = nn.Sequential( - nn.Conv2d(self.inplanes, planes * block.expansion, - kernel_size=1, stride=stride, bias=False), - nn.BatchNorm2d(planes * block.expansion), - ) - - layers = [] - layers.append(block(self.inplanes, planes, stride, downsample)) - self.inplanes = planes * block.expansion - for i in range(1, blocks): - layers.append(block(self.inplanes, planes)) - - return nn.Sequential(*layers) - - def forward(self, x, target, training_phase, c=None): - x = self.conv1(x) - x = self.bn1(x) - x = self.maxpool(x) - - x = self.layer1(x) - x = self.layer2(x) - x = self.layer3(x) - x_list = [] - softlabel_list = [] - x_final = None - targets = None - - # auxiliary & main branch forward - features, targets, inds = self.split_feature(x, target) - for i in range(self.num_branches-1): - x_list.append(getattr(self, 'layer4_'+str(i))(features[i])) - x_final = getattr(self, 'layer4_'+str(self.num_branches-1))(x) - # make label distribution - features2, inds_softlabel = self.split_feature_makeLD(x, target) - for i in range(self.num_branches-1): - softlabel_list.append(getattr(self, 'layer4_'+str(i))(features2[i])) - - return x_final, x_list, softlabel_list, targets, inds, inds_softlabel - - def split_feature(self, x, target): - x_parts = [] - t_parts = [] - inds = [] - for c in range(self.num_classes): - ind = (target != c).nonzero() - inds.append(ind) - _t = target[ind[:,0]] - t_parts.append(_t.where(_t < c, _t-1)) - x_parts.append(x[ind[:, 0], :, :, :]) - return x_parts, t_parts, inds - - def split_feature_makeLD(self, x, target): - x_parts = [] - inds = [] - for c in range(self.num_classes): - ind = (target == c).nonzero() - inds.append(ind) - x_parts.append(x[ind[:, 0], :, :, :]) - return x_parts, inds - - def random_init(self): - for m in self.modules(): - if isinstance(m, nn.Conv2d): - n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels - m.weight.data.normal_(0, math.sqrt(2. / n)) - elif isinstance(m, nn.BatchNorm2d): - m.weight.data.fill_(1) - m.bias.data.zero_() - - def _freeze_stages(self): - if self.frozen_stages >= 0: - self.bn1.eval() - for m in [self.conv1, self.bn1]: - for param in m.parameters(): - param.requires_grad = False - - for i in range(1, self.frozen_stages + 1): - m = getattr(self, 'layer{}'.format(i)) - # print('layer{}'.format(i)) - m.eval() - for param in m.parameters(): - param.requires_grad = False +import math +import torch +from torch import nn + + +def conv3x3(in_planes, out_planes, stride=1): + """3x3 convolution with padding""" + return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, + padding=1, bias=False) + + +class BasicBlock(nn.Module): + expansion = 1 + + def __init__(self, inplanes, planes, stride=1, downsample=None): + super(BasicBlock, self).__init__() + self.conv1 = conv3x3(inplanes, planes, stride) + self.bn1 = nn.BatchNorm2d(planes) + self.relu = nn.ReLU(inplace=True) + self.conv2 = conv3x3(planes, planes) + self.bn2 = nn.BatchNorm2d(planes) + self.downsample = downsample + self.stride = stride + + def forward(self, x): + residual = x + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + + out = self.conv2(out) + out = self.bn2(out) + + if self.downsample is not None: + residual = self.downsample(x) + + out += residual + out = self.relu(out) + + return out + + +class Bottleneck(nn.Module): + expansion = 4 + + def __init__(self, inplanes, planes, stride=1, downsample=None): + super(Bottleneck, self).__init__() + self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) + self.bn1 = nn.BatchNorm2d(planes) + self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, + padding=1, bias=False) + self.bn2 = nn.BatchNorm2d(planes) + self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False) + self.bn3 = nn.BatchNorm2d(planes * 4) + self.relu = nn.ReLU(inplace=True) + self.downsample = downsample + self.stride = stride + + def forward(self, x): + residual = x + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + + out = self.conv2(out) + out = self.bn2(out) + out = self.relu(out) + + out = self.conv3(out) + out = self.bn3(out) + + if self.downsample is not None: + residual = self.downsample(x) + + out += residual + out = self.relu(out) + + return out + + +class ResNet(nn.Module): + def __init__(self, num_branches, num_classes, + last_stride=2, block=Bottleneck, frozen_stages=-1,layers=[3, 4, 6, 3]): + self.inplanes = 64 + super().__init__() + self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, + bias=False) + self.bn1 = nn.BatchNorm2d(64) + self.maxpool = nn.MaxPool2d(kernel_size=2, stride=None, padding=0) + self.frozen_stages = frozen_stages + self.layer1 = self._make_layer(block, 64, layers[0]) + self.layer2 = self._make_layer(block, 128, layers[1], stride=2) + self.layer3 = self._make_layer(block, 256, layers[2], stride=2) + + fix_inplanes = self.inplanes + self.num_branches = num_branches + self.num_classes = num_classes + for i in range(self.num_branches): + setattr(self, 'layer4_'+str(i), self._make_layer(block, 512, layers[3], stride=last_stride)) + self.inplanes = fix_inplanes + + def _make_layer(self, block, planes, blocks, stride=1): + downsample = None + if stride != 1 or self.inplanes != planes * block.expansion: + downsample = nn.Sequential( + nn.Conv2d(self.inplanes, planes * block.expansion, + kernel_size=1, stride=stride, bias=False), + nn.BatchNorm2d(planes * block.expansion), + ) + + layers = [] + layers.append(block(self.inplanes, planes, stride, downsample)) + self.inplanes = planes * block.expansion + for i in range(1, blocks): + layers.append(block(self.inplanes, planes)) + + return nn.Sequential(*layers) + + def forward(self, x, target, training_phase, c=None): + x = self.conv1(x) + x = self.bn1(x) + x = self.maxpool(x) + + x = self.layer1(x) + x = self.layer2(x) + x = self.layer3(x) + x_list = [] + softlabel_list = [] + x_final = None + targets = None + + # auxiliary & main branch forward + features, targets, inds = self.split_feature(x, target) + for i in range(self.num_branches-1): + x_list.append(getattr(self, 'layer4_'+str(i))(features[i])) + x_final = getattr(self, 'layer4_'+str(self.num_branches-1))(x) + # make label distribution + features2, inds_softlabel = self.split_feature_makeLD(x, target) + for i in range(self.num_branches-1): + softlabel_list.append(getattr(self, 'layer4_'+str(i))(features2[i])) + + return x_final, x_list, softlabel_list, targets, inds, inds_softlabel + + def split_feature(self, x, target): + x_parts = [] + t_parts = [] + inds = [] + for c in range(self.num_classes): + ind = (target != c).nonzero() + inds.append(ind) + _t = target[ind[:,0]] + t_parts.append(_t.where(_t < c, _t-1)) + x_parts.append(x[ind[:, 0], :, :, :]) + return x_parts, t_parts, inds + + def split_feature_makeLD(self, x, target): + x_parts = [] + inds = [] + for c in range(self.num_classes): + ind = (target == c).nonzero() + inds.append(ind) + x_parts.append(x[ind[:, 0], :, :, :]) + return x_parts, inds + + def random_init(self): + for m in self.modules(): + if isinstance(m, nn.Conv2d): + n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + m.weight.data.normal_(0, math.sqrt(2. / n)) + elif isinstance(m, nn.BatchNorm2d): + m.weight.data.fill_(1) + m.bias.data.zero_() + + def _freeze_stages(self): + if self.frozen_stages >= 0: + self.bn1.eval() + for m in [self.conv1, self.bn1]: + for param in m.parameters(): + param.requires_grad = False + + for i in range(1, self.frozen_stages + 1): + m = getattr(self, 'layer{}'.format(i)) + # print('layer{}'.format(i)) + m.eval() + for param in m.parameters(): + param.requires_grad = False diff --git a/addition_module/DMUE/preprocess/face_alignmenet/utils/__init__.py b/src/faceDetection/face_sdk/addition_module/DMUE/preprocess/__init__.py similarity index 100% rename from addition_module/DMUE/preprocess/face_alignmenet/utils/__init__.py rename to src/faceDetection/face_sdk/addition_module/DMUE/preprocess/__init__.py diff --git a/addition_module/DMUE/preprocess/crop_align.py b/src/faceDetection/face_sdk/addition_module/DMUE/preprocess/crop_align.py similarity index 100% rename from addition_module/DMUE/preprocess/crop_align.py rename to src/faceDetection/face_sdk/addition_module/DMUE/preprocess/crop_align.py diff --git a/addition_module/DMUE/preprocess/crop_msra.py b/src/faceDetection/face_sdk/addition_module/DMUE/preprocess/crop_msra.py similarity index 100% rename from addition_module/DMUE/preprocess/crop_msra.py rename to src/faceDetection/face_sdk/addition_module/DMUE/preprocess/crop_msra.py diff --git a/addition_module/DMUE/pretrain/backbones/__init__.py b/src/faceDetection/face_sdk/addition_module/DMUE/preprocess/face_alignmenet/__init__.py similarity index 100% rename from addition_module/DMUE/pretrain/backbones/__init__.py rename to src/faceDetection/face_sdk/addition_module/DMUE/preprocess/face_alignmenet/__init__.py diff --git a/addition_module/DMUE/pretrain/utils/__init__.py b/src/faceDetection/face_sdk/addition_module/DMUE/preprocess/face_alignmenet/core/__init__.py similarity index 100% rename from addition_module/DMUE/pretrain/utils/__init__.py rename to src/faceDetection/face_sdk/addition_module/DMUE/preprocess/face_alignmenet/core/__init__.py diff --git a/addition_module/DMUE/preprocess/face_alignmenet/core/coord_conv.py b/src/faceDetection/face_sdk/addition_module/DMUE/preprocess/face_alignmenet/core/coord_conv.py similarity index 100% rename from addition_module/DMUE/preprocess/face_alignmenet/core/coord_conv.py rename to src/faceDetection/face_sdk/addition_module/DMUE/preprocess/face_alignmenet/core/coord_conv.py diff --git a/addition_module/DMUE/preprocess/face_alignmenet/core/dataloader.py b/src/faceDetection/face_sdk/addition_module/DMUE/preprocess/face_alignmenet/core/dataloader.py similarity index 100% rename from addition_module/DMUE/preprocess/face_alignmenet/core/dataloader.py rename to src/faceDetection/face_sdk/addition_module/DMUE/preprocess/face_alignmenet/core/dataloader.py diff --git a/addition_module/DMUE/preprocess/face_alignmenet/core/evaler.py b/src/faceDetection/face_sdk/addition_module/DMUE/preprocess/face_alignmenet/core/evaler.py similarity index 100% rename from addition_module/DMUE/preprocess/face_alignmenet/core/evaler.py rename to src/faceDetection/face_sdk/addition_module/DMUE/preprocess/face_alignmenet/core/evaler.py diff --git a/addition_module/DMUE/preprocess/face_alignmenet/core/models.py b/src/faceDetection/face_sdk/addition_module/DMUE/preprocess/face_alignmenet/core/models.py similarity index 100% rename from addition_module/DMUE/preprocess/face_alignmenet/core/models.py rename to src/faceDetection/face_sdk/addition_module/DMUE/preprocess/face_alignmenet/core/models.py diff --git a/addition_module/DMUE/preprocess/face_alignmenet/dataloader.py b/src/faceDetection/face_sdk/addition_module/DMUE/preprocess/face_alignmenet/dataloader.py similarity index 100% rename from addition_module/DMUE/preprocess/face_alignmenet/dataloader.py rename to src/faceDetection/face_sdk/addition_module/DMUE/preprocess/face_alignmenet/dataloader.py diff --git a/addition_module/DMUE/preprocess/face_alignmenet/get_landmark.py b/src/faceDetection/face_sdk/addition_module/DMUE/preprocess/face_alignmenet/get_landmark.py similarity index 78% rename from addition_module/DMUE/preprocess/face_alignmenet/get_landmark.py rename to src/faceDetection/face_sdk/addition_module/DMUE/preprocess/face_alignmenet/get_landmark.py index a7dc5b8..04bcd86 100644 --- a/addition_module/DMUE/preprocess/face_alignmenet/get_landmark.py +++ b/src/faceDetection/face_sdk/addition_module/DMUE/preprocess/face_alignmenet/get_landmark.py @@ -1,73 +1,56 @@ -#coding:utf-8 -import matplotlib -import math -import torch -import copy -import time -from torch.autograd import Variable -import shutil -from skimage import io -import numpy as np -from .utils.utils import fan_NME, show_landmarks, get_preds_fromhm -from PIL import Image, ImageDraw -from pylab import * -import os -import sys -import cv2 -import matplotlib.pyplot as plt -import torch -import argparse -import numpy as np -import torch.nn as nn -import time -import os -from .core import models - - -def load_model(): - NUM_LANDMARKS = 98 - - PRETRAINED_WEIGHTS = os.path.join(os.path.dirname(__file__), './ckpt/WFLW_4HG.pth') - GRAY_SCALE = False - HG_BLOCKS = 4 - END_RELU = False - - device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") - # print(device) - - model_ft = models.FAN(HG_BLOCKS, END_RELU, GRAY_SCALE, NUM_LANDMARKS) - - checkpoint = torch.load(PRETRAINED_WEIGHTS, map_location=torch.device('cpu')) - pretrained_weights = checkpoint['state_dict'] - model_weights = model_ft.state_dict() - pretrained_weights = {k: v for k, v in pretrained_weights.items() \ - if k in model_weights} - model_weights.update(pretrained_weights) - model_ft.load_state_dict(model_weights, device) - - model_ft = model_ft.to(device) - model_ft.eval() - - return model_ft - - -def get_lmd(model_ft, img): - model_ft.eval() - # Iterate over data. - with torch.no_grad(): - img = cv2.resize(img, (256, 256), interpolation=cv2.INTER_CUBIC) - img = img.transpose(2, 0, 1)[np.newaxis, :, :, :].astype(np.float32) / 255 - ''' - 这里输入的取值范围在0-1 - - ''' - inputs = torch.from_numpy(img).cuda() - outputs, boundary_channels = model_ft(inputs) - - pred_heatmap = outputs[-1][:, :-1, :, :][0].detach().cpu() # （98， 64， 64） - pred_landmarks, _ = get_preds_fromhm(pred_heatmap.unsqueeze(0)) - pred_landmarks = pred_landmarks.squeeze().numpy() # 98 x 2 list - pred_landmarks = (pred_landmarks * 256 / 64).astype(np.int) - - return pred_landmarks - +#coding:utf-8 +from .utils.utils import get_preds_fromhm +import cv2 +import torch +import numpy as np +import os +from .core import models + + +def load_model(): + NUM_LANDMARKS = 98 + + PRETRAINED_WEIGHTS = os.path.join(os.path.dirname(__file__), './ckpt/WFLW_4HG.pth') + GRAY_SCALE = False + HG_BLOCKS = 4 + END_RELU = False + + device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + # print(device) + + model_ft = models.FAN(HG_BLOCKS, END_RELU, GRAY_SCALE, NUM_LANDMARKS) + + checkpoint = torch.load(PRETRAINED_WEIGHTS, map_location=torch.device('cpu')) + pretrained_weights = checkpoint['state_dict'] + model_weights = model_ft.state_dict() + pretrained_weights = {k: v for k, v in pretrained_weights.items() \ + if k in model_weights} + model_weights.update(pretrained_weights) + model_ft.load_state_dict(model_weights, device) + + model_ft = model_ft.to(device) + model_ft.eval() + + return model_ft + + +def get_lmd(model_ft, img): + model_ft.eval() + # Iterate over data. + with torch.no_grad(): + img = cv2.resize(img, (256, 256), interpolation=cv2.INTER_CUBIC) + img = img.transpose(2, 0, 1)[np.newaxis, :, :, :].astype(np.float32) / 255 + ''' + 这里输入的取值范围在0-1 + + ''' + inputs = torch.from_numpy(img).cuda() + outputs, boundary_channels = model_ft(inputs) + + pred_heatmap = outputs[-1][:, :-1, :, :][0].detach().cpu() # （98， 64， 64） + pred_landmarks, _ = get_preds_fromhm(pred_heatmap.unsqueeze(0)) + pred_landmarks = pred_landmarks.squeeze().numpy() # 98 x 2 list + pred_landmarks = (pred_landmarks * 256 / 64).astype(np.int) + + return pred_landmarks + diff --git a/src/faceDetection/face_sdk/addition_module/DMUE/preprocess/face_alignmenet/utils/__init__.py b/src/faceDetection/face_sdk/addition_module/DMUE/preprocess/face_alignmenet/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/addition_module/DMUE/preprocess/face_alignmenet/utils/utils.py b/src/faceDetection/face_sdk/addition_module/DMUE/preprocess/face_alignmenet/utils/utils.py similarity index 100% rename from addition_module/DMUE/preprocess/face_alignmenet/utils/utils.py rename to src/faceDetection/face_sdk/addition_module/DMUE/preprocess/face_alignmenet/utils/utils.py diff --git a/addition_module/DMUE/preprocess/make_lmdb.py b/src/faceDetection/face_sdk/addition_module/DMUE/preprocess/make_lmdb.py similarity index 100% rename from addition_module/DMUE/preprocess/make_lmdb.py rename to src/faceDetection/face_sdk/addition_module/DMUE/preprocess/make_lmdb.py diff --git a/addition_module/DMUE/preprocess/mtcnn/__init__.py b/src/faceDetection/face_sdk/addition_module/DMUE/preprocess/mtcnn/__init__.py similarity index 100% rename from addition_module/DMUE/preprocess/mtcnn/__init__.py rename to src/faceDetection/face_sdk/addition_module/DMUE/preprocess/mtcnn/__init__.py diff --git a/addition_module/DMUE/preprocess/mtcnn/detect_face.py b/src/faceDetection/face_sdk/addition_module/DMUE/preprocess/mtcnn/detect_face.py similarity index 100% rename from addition_module/DMUE/preprocess/mtcnn/detect_face.py rename to src/faceDetection/face_sdk/addition_module/DMUE/preprocess/mtcnn/detect_face.py diff --git a/addition_module/DMUE/preprocess/mtcnn/mtcnn.py b/src/faceDetection/face_sdk/addition_module/DMUE/preprocess/mtcnn/mtcnn.py similarity index 100% rename from addition_module/DMUE/preprocess/mtcnn/mtcnn.py rename to src/faceDetection/face_sdk/addition_module/DMUE/preprocess/mtcnn/mtcnn.py diff --git a/addition_module/DMUE/preprocess/mtcnn/onet.pt b/src/faceDetection/face_sdk/addition_module/DMUE/preprocess/mtcnn/onet.pt similarity index 100% rename from addition_module/DMUE/preprocess/mtcnn/onet.pt rename to src/faceDetection/face_sdk/addition_module/DMUE/preprocess/mtcnn/onet.pt diff --git a/addition_module/DMUE/preprocess/mtcnn/pnet.pt b/src/faceDetection/face_sdk/addition_module/DMUE/preprocess/mtcnn/pnet.pt similarity index 100% rename from addition_module/DMUE/preprocess/mtcnn/pnet.pt rename to src/faceDetection/face_sdk/addition_module/DMUE/preprocess/mtcnn/pnet.pt diff --git a/addition_module/DMUE/preprocess/mtcnn/rnet.pt b/src/faceDetection/face_sdk/addition_module/DMUE/preprocess/mtcnn/rnet.pt similarity index 100% rename from addition_module/DMUE/preprocess/mtcnn/rnet.pt rename to src/faceDetection/face_sdk/addition_module/DMUE/preprocess/mtcnn/rnet.pt diff --git a/addition_module/DMUE/preprocess/template.npy b/src/faceDetection/face_sdk/addition_module/DMUE/preprocess/template.npy similarity index 100% rename from addition_module/DMUE/preprocess/template.npy rename to src/faceDetection/face_sdk/addition_module/DMUE/preprocess/template.npy diff --git a/src/faceDetection/face_sdk/addition_module/DMUE/pretrain/__init__.py b/src/faceDetection/face_sdk/addition_module/DMUE/pretrain/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/faceDetection/face_sdk/addition_module/DMUE/pretrain/backbones/__init__.py b/src/faceDetection/face_sdk/addition_module/DMUE/pretrain/backbones/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/addition_module/DMUE/pretrain/backbones/resnet.py b/src/faceDetection/face_sdk/addition_module/DMUE/pretrain/backbones/resnet.py similarity index 96% rename from addition_module/DMUE/pretrain/backbones/resnet.py rename to src/faceDetection/face_sdk/addition_module/DMUE/pretrain/backbones/resnet.py index a83abea..53a5b79 100644 --- a/addition_module/DMUE/pretrain/backbones/resnet.py +++ b/src/faceDetection/face_sdk/addition_module/DMUE/pretrain/backbones/resnet.py @@ -1,158 +1,158 @@ -# encoding: utf-8 -import math - -import torch -from torch import nn - - -def conv3x3(in_planes, out_planes, stride=1): - """3x3 convolution with padding""" - return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, - padding=1, bias=False) - - -class BasicBlock(nn.Module): - expansion = 1 - - def __init__(self, inplanes, planes, stride=1, downsample=None): - super(BasicBlock, self).__init__() - self.conv1 = conv3x3(inplanes, planes, stride) - self.bn1 = nn.BatchNorm2d(planes) - self.relu = nn.ReLU(inplace=True) - self.conv2 = conv3x3(planes, planes) - self.bn2 = nn.BatchNorm2d(planes) - self.downsample = downsample - self.stride = stride - - def forward(self, x): - residual = x - - out = self.conv1(x) - out = self.bn1(out) - out = self.relu(out) - - out = self.conv2(out) - out = self.bn2(out) - - if self.downsample is not None: - residual = self.downsample(x) - - out += residual - out = self.relu(out) - - return out - - -class Bottleneck(nn.Module): - expansion = 4 - - def __init__(self, inplanes, planes, stride=1, downsample=None): - super(Bottleneck, self).__init__() - self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) - self.bn1 = nn.BatchNorm2d(planes) - self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, - padding=1, bias=False) - self.bn2 = nn.BatchNorm2d(planes) - self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False) - self.bn3 = nn.BatchNorm2d(planes * 4) - self.relu = nn.ReLU(inplace=True) - self.downsample = downsample - self.stride = stride - - def forward(self, x): - residual = x - - out = self.conv1(x) - out = self.bn1(out) - out = self.relu(out) - - out = self.conv2(out) - out = self.bn2(out) - out = self.relu(out) - - out = self.conv3(out) - out = self.bn3(out) - - if self.downsample is not None: - residual = self.downsample(x) - - out += residual - out = self.relu(out) - - return out - - -class ResNet(nn.Module): - def __init__(self, last_stride=2, block=Bottleneck, frozen_stages=-1,layers=[3, 4, 6, 3]): - self.inplanes = 64 - super().__init__() - self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, - bias=False) - self.bn1 = nn.BatchNorm2d(64) - # self.relu = nn.ReLU(inplace=True) # add missed relu - self.maxpool = nn.MaxPool2d(kernel_size=2, stride=None, padding=0) - self.frozen_stages = frozen_stages - self.layer1 = self._make_layer(block, 64, layers[0]) - self.layer2 = self._make_layer(block, 128, layers[1], stride=2) - self.layer3 = self._make_layer(block, 256, layers[2], stride=2) - self.layer4 = self._make_layer(block, 512, layers[3], stride=last_stride) - - def _make_layer(self, block, planes, blocks, stride=1): - downsample = None - if stride != 1 or self.inplanes != planes * block.expansion: - downsample = nn.Sequential( - nn.Conv2d(self.inplanes, planes * block.expansion, - kernel_size=1, stride=stride, bias=False), - nn.BatchNorm2d(planes * block.expansion), - ) - - layers = [] - layers.append(block(self.inplanes, planes, stride, downsample)) - self.inplanes = planes * block.expansion - for i in range(1, blocks): - layers.append(block(self.inplanes, planes)) - - return nn.Sequential(*layers) - - def _freeze_stages(self): - if self.frozen_stages >= 0: - self.bn1.eval() - for m in [self.conv1, self.bn1]: - for param in m.parameters(): - param.requires_grad = False - - for i in range(1, self.frozen_stages + 1): - m = getattr(self, 'layer{}'.format(i)) - print('layer{}'.format(i)) - m.eval() - for param in m.parameters(): - param.requires_grad = False - - def forward(self, x): - x = self.conv1(x) - x = self.bn1(x) - # x = self.relu(x) # add missed relu - x = self.maxpool(x) - - x = self.layer1(x) - x = self.layer2(x) - x = self.layer3(x) - x = self.layer4(x) - - return x - - def load_param(self, model_path): - param_dict = torch.load(model_path) - for i in param_dict: - if 'fc' in i: - continue - self.state_dict()[i].copy_(param_dict[i]) - - def random_init(self): - for m in self.modules(): - if isinstance(m, nn.Conv2d): - n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels - m.weight.data.normal_(0, math.sqrt(2. / n)) - elif isinstance(m, nn.BatchNorm2d): - m.weight.data.fill_(1) +# encoding: utf-8 +import math + +import torch +from torch import nn + + +def conv3x3(in_planes, out_planes, stride=1): + """3x3 convolution with padding""" + return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, + padding=1, bias=False) + + +class BasicBlock(nn.Module): + expansion = 1 + + def __init__(self, inplanes, planes, stride=1, downsample=None): + super(BasicBlock, self).__init__() + self.conv1 = conv3x3(inplanes, planes, stride) + self.bn1 = nn.BatchNorm2d(planes) + self.relu = nn.ReLU(inplace=True) + self.conv2 = conv3x3(planes, planes) + self.bn2 = nn.BatchNorm2d(planes) + self.downsample = downsample + self.stride = stride + + def forward(self, x): + residual = x + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + + out = self.conv2(out) + out = self.bn2(out) + + if self.downsample is not None: + residual = self.downsample(x) + + out += residual + out = self.relu(out) + + return out + + +class Bottleneck(nn.Module): + expansion = 4 + + def __init__(self, inplanes, planes, stride=1, downsample=None): + super(Bottleneck, self).__init__() + self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) + self.bn1 = nn.BatchNorm2d(planes) + self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, + padding=1, bias=False) + self.bn2 = nn.BatchNorm2d(planes) + self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False) + self.bn3 = nn.BatchNorm2d(planes * 4) + self.relu = nn.ReLU(inplace=True) + self.downsample = downsample + self.stride = stride + + def forward(self, x): + residual = x + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + + out = self.conv2(out) + out = self.bn2(out) + out = self.relu(out) + + out = self.conv3(out) + out = self.bn3(out) + + if self.downsample is not None: + residual = self.downsample(x) + + out += residual + out = self.relu(out) + + return out + + +class ResNet(nn.Module): + def __init__(self, last_stride=2, block=Bottleneck, frozen_stages=-1,layers=[3, 4, 6, 3]): + self.inplanes = 64 + super().__init__() + self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, + bias=False) + self.bn1 = nn.BatchNorm2d(64) + # self.relu = nn.ReLU(inplace=True) # add missed relu + self.maxpool = nn.MaxPool2d(kernel_size=2, stride=None, padding=0) + self.frozen_stages = frozen_stages + self.layer1 = self._make_layer(block, 64, layers[0]) + self.layer2 = self._make_layer(block, 128, layers[1], stride=2) + self.layer3 = self._make_layer(block, 256, layers[2], stride=2) + self.layer4 = self._make_layer(block, 512, layers[3], stride=last_stride) + + def _make_layer(self, block, planes, blocks, stride=1): + downsample = None + if stride != 1 or self.inplanes != planes * block.expansion: + downsample = nn.Sequential( + nn.Conv2d(self.inplanes, planes * block.expansion, + kernel_size=1, stride=stride, bias=False), + nn.BatchNorm2d(planes * block.expansion), + ) + + layers = [] + layers.append(block(self.inplanes, planes, stride, downsample)) + self.inplanes = planes * block.expansion + for i in range(1, blocks): + layers.append(block(self.inplanes, planes)) + + return nn.Sequential(*layers) + + def _freeze_stages(self): + if self.frozen_stages >= 0: + self.bn1.eval() + for m in [self.conv1, self.bn1]: + for param in m.parameters(): + param.requires_grad = False + + for i in range(1, self.frozen_stages + 1): + m = getattr(self, 'layer{}'.format(i)) + print('layer{}'.format(i)) + m.eval() + for param in m.parameters(): + param.requires_grad = False + + def forward(self, x): + x = self.conv1(x) + x = self.bn1(x) + # x = self.relu(x) # add missed relu + x = self.maxpool(x) + + x = self.layer1(x) + x = self.layer2(x) + x = self.layer3(x) + x = self.layer4(x) + + return x + + def load_param(self, model_path): + param_dict = torch.load(model_path) + for i in param_dict: + if 'fc' in i: + continue + self.state_dict()[i].copy_(param_dict[i]) + + def random_init(self): + for m in self.modules(): + if isinstance(m, nn.Conv2d): + n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + m.weight.data.normal_(0, math.sqrt(2. / n)) + elif isinstance(m, nn.BatchNorm2d): + m.weight.data.fill_(1) m.bias.data.zero_() \ No newline at end of file diff --git a/addition_module/DMUE/pretrain/backbones/resnet_ibn_a.py b/src/faceDetection/face_sdk/addition_module/DMUE/pretrain/backbones/resnet_ibn_a.py similarity index 97% rename from addition_module/DMUE/pretrain/backbones/resnet_ibn_a.py rename to src/faceDetection/face_sdk/addition_module/DMUE/pretrain/backbones/resnet_ibn_a.py index b8503af..c956dfc 100644 --- a/addition_module/DMUE/pretrain/backbones/resnet_ibn_a.py +++ b/src/faceDetection/face_sdk/addition_module/DMUE/pretrain/backbones/resnet_ibn_a.py @@ -1,197 +1,197 @@ -import torch -import torch.nn as nn -import math -import torch.utils.model_zoo as model_zoo - -__all__ = ['ResNet_IBN', 'resnet50_ibn_a', 'resnet101_ibn_a', - 'resnet152_ibn_a'] - - -model_urls = { - 'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth', - 'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth', - 'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth', -} - - -class IBN(nn.Module): - def __init__(self, planes): - super(IBN, self).__init__() - half1 = int(planes/2) - self.half = half1 - half2 = planes - half1 - self.IN = nn.InstanceNorm2d(half1, affine=True) - self.BN = nn.BatchNorm2d(half2) - - def forward(self, x): - split = torch.split(x, self.half, 1) - out1 = self.IN(split[0].contiguous()) - out2 = self.BN(split[1].contiguous()) - out = torch.cat((out1, out2), 1) - return out - - -class Bottleneck_IBN(nn.Module): - expansion = 4 - - def __init__(self, inplanes, planes, ibn=False, stride=1, downsample=None): - super(Bottleneck_IBN, self).__init__() - self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) - if ibn: - self.bn1 = IBN(planes) - else: - self.bn1 = nn.BatchNorm2d(planes) - self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, - padding=1, bias=False) - self.bn2 = nn.BatchNorm2d(planes) - self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1, bias=False) - self.bn3 = nn.BatchNorm2d(planes * self.expansion) - self.relu = nn.ReLU(inplace=True) - self.downsample = downsample - self.stride = stride - - def forward(self, x): - residual = x - - out = self.conv1(x) - out = self.bn1(out) - out = self.relu(out) - - out = self.conv2(out) - out = self.bn2(out) - out = self.relu(out) - - out = self.conv3(out) - out = self.bn3(out) - - if self.downsample is not None: - residual = self.downsample(x) - - out += residual - out = self.relu(out) - - return out - - -class ResNet_IBN(nn.Module): - - def __init__(self, last_stride, block, layers, frozen_stages=-1,num_classes=1000): - scale = 64 - self.inplanes = scale - super(ResNet_IBN, self).__init__() - self.conv1 = nn.Conv2d(3, scale, kernel_size=7, stride=2, padding=3, - bias=False) - self.bn1 = nn.BatchNorm2d(scale) - self.relu = nn.ReLU(inplace=True) - self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) - self.frozen_stages = frozen_stages - self.layer1 = self._make_layer(block, scale, layers[0]) - self.layer2 = self._make_layer(block, scale*2, layers[1], stride=2) - self.layer3 = self._make_layer(block, scale*4, layers[2], stride=2) - self.layer4 = self._make_layer(block, scale*8, layers[3], stride=last_stride) - self.avgpool = nn.AvgPool2d(7) - self.fc = nn.Linear(scale * 8 * block.expansion, num_classes) - - for m in self.modules(): - if isinstance(m, nn.Conv2d): - n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels - m.weight.data.normal_(0, math.sqrt(2. / n)) - elif isinstance(m, nn.BatchNorm2d): - m.weight.data.fill_(1) - m.bias.data.zero_() - elif isinstance(m, nn.InstanceNorm2d): - m.weight.data.fill_(1) - m.bias.data.zero_() - - def _make_layer(self, block, planes, blocks, stride=1): - downsample = None - if stride != 1 or self.inplanes != planes * block.expansion: - downsample = nn.Sequential( - nn.Conv2d(self.inplanes, planes * block.expansion, - kernel_size=1, stride=stride, bias=False), - nn.BatchNorm2d(planes * block.expansion), - ) - - layers = [] - ibn = True - if planes == 512: - ibn = False - layers.append(block(self.inplanes, planes, ibn, stride, downsample)) - self.inplanes = planes * block.expansion - for i in range(1, blocks): - layers.append(block(self.inplanes, planes, ibn)) - - return nn.Sequential(*layers) - - def _freeze_stages(self): - if self.frozen_stages >= 0: - self.bn1.eval() - for m in [self.conv1, self.bn1]: - for param in m.parameters(): - param.requires_grad = False - - for i in range(1, self.frozen_stages + 1): - m = getattr(self, 'layer{}'.format(i)) - print('layer{}'.format(i)) - m.eval() - for param in m.parameters(): - param.requires_grad = False - def forward(self, x): - x = self.conv1(x) - x = self.bn1(x) - x = self.relu(x) - x = self.maxpool(x) - - x = self.layer1(x) - x = self.layer2(x) - x = self.layer3(x) - x = self.layer4(x) - - # x = self.avgpool(x) - # x = x.view(x.size(0), -1) - # x = self.fc(x) - - return x - - def load_param(self, model_path): - param_dict = torch.load(model_path) - if 'state_dict' in param_dict: - param_dict = param_dict['state_dict'] - for i in param_dict: - if 'fc' in i: - continue - self.state_dict()[i.replace('module.','')].copy_(param_dict[i]) - - - -def resnet50_ibn_a(last_stride, pretrained=False, **kwargs): - """Constructs a ResNet-50 model. - Args: - pretrained (bool): If True, returns a model pre-trained on ImageNet - """ - model = ResNet_IBN(last_stride, Bottleneck_IBN, [3, 4, 6, 3], **kwargs) - if pretrained: - model.load_state_dict(model_zoo.load_url(model_urls['resnet50'])) - return model - - -def resnet101_ibn_a(last_stride, pretrained=False, **kwargs): - """Constructs a ResNet-101 model. - Args: - pretrained (bool): If True, returns a model pre-trained on ImageNet - """ - model = ResNet_IBN(last_stride, Bottleneck_IBN, [3, 4, 23, 3], **kwargs) - if pretrained: - model.load_state_dict(model_zoo.load_url(model_urls['resnet101'])) - return model - - -def resnet152_ibn_a(last_stride, pretrained=False, **kwargs): - """Constructs a ResNet-152 model. - Args: - pretrained (bool): If True, returns a model pre-trained on ImageNet - """ - model = ResNet_IBN(last_stride, Bottleneck_IBN, [3, 8, 36, 3], **kwargs) - if pretrained: - model.load_state_dict(model_zoo.load_url(model_urls['resnet152'])) +import torch +import torch.nn as nn +import math +import torch.utils.model_zoo as model_zoo + +__all__ = ['ResNet_IBN', 'resnet50_ibn_a', 'resnet101_ibn_a', + 'resnet152_ibn_a'] + + +model_urls = { + 'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth', + 'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth', + 'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth', +} + + +class IBN(nn.Module): + def __init__(self, planes): + super(IBN, self).__init__() + half1 = int(planes/2) + self.half = half1 + half2 = planes - half1 + self.IN = nn.InstanceNorm2d(half1, affine=True) + self.BN = nn.BatchNorm2d(half2) + + def forward(self, x): + split = torch.split(x, self.half, 1) + out1 = self.IN(split[0].contiguous()) + out2 = self.BN(split[1].contiguous()) + out = torch.cat((out1, out2), 1) + return out + + +class Bottleneck_IBN(nn.Module): + expansion = 4 + + def __init__(self, inplanes, planes, ibn=False, stride=1, downsample=None): + super(Bottleneck_IBN, self).__init__() + self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) + if ibn: + self.bn1 = IBN(planes) + else: + self.bn1 = nn.BatchNorm2d(planes) + self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, + padding=1, bias=False) + self.bn2 = nn.BatchNorm2d(planes) + self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1, bias=False) + self.bn3 = nn.BatchNorm2d(planes * self.expansion) + self.relu = nn.ReLU(inplace=True) + self.downsample = downsample + self.stride = stride + + def forward(self, x): + residual = x + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + + out = self.conv2(out) + out = self.bn2(out) + out = self.relu(out) + + out = self.conv3(out) + out = self.bn3(out) + + if self.downsample is not None: + residual = self.downsample(x) + + out += residual + out = self.relu(out) + + return out + + +class ResNet_IBN(nn.Module): + + def __init__(self, last_stride, block, layers, frozen_stages=-1,num_classes=1000): + scale = 64 + self.inplanes = scale + super(ResNet_IBN, self).__init__() + self.conv1 = nn.Conv2d(3, scale, kernel_size=7, stride=2, padding=3, + bias=False) + self.bn1 = nn.BatchNorm2d(scale) + self.relu = nn.ReLU(inplace=True) + self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) + self.frozen_stages = frozen_stages + self.layer1 = self._make_layer(block, scale, layers[0]) + self.layer2 = self._make_layer(block, scale*2, layers[1], stride=2) + self.layer3 = self._make_layer(block, scale*4, layers[2], stride=2) + self.layer4 = self._make_layer(block, scale*8, layers[3], stride=last_stride) + self.avgpool = nn.AvgPool2d(7) + self.fc = nn.Linear(scale * 8 * block.expansion, num_classes) + + for m in self.modules(): + if isinstance(m, nn.Conv2d): + n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + m.weight.data.normal_(0, math.sqrt(2. / n)) + elif isinstance(m, nn.BatchNorm2d): + m.weight.data.fill_(1) + m.bias.data.zero_() + elif isinstance(m, nn.InstanceNorm2d): + m.weight.data.fill_(1) + m.bias.data.zero_() + + def _make_layer(self, block, planes, blocks, stride=1): + downsample = None + if stride != 1 or self.inplanes != planes * block.expansion: + downsample = nn.Sequential( + nn.Conv2d(self.inplanes, planes * block.expansion, + kernel_size=1, stride=stride, bias=False), + nn.BatchNorm2d(planes * block.expansion), + ) + + layers = [] + ibn = True + if planes == 512: + ibn = False + layers.append(block(self.inplanes, planes, ibn, stride, downsample)) + self.inplanes = planes * block.expansion + for i in range(1, blocks): + layers.append(block(self.inplanes, planes, ibn)) + + return nn.Sequential(*layers) + + def _freeze_stages(self): + if self.frozen_stages >= 0: + self.bn1.eval() + for m in [self.conv1, self.bn1]: + for param in m.parameters(): + param.requires_grad = False + + for i in range(1, self.frozen_stages + 1): + m = getattr(self, 'layer{}'.format(i)) + print('layer{}'.format(i)) + m.eval() + for param in m.parameters(): + param.requires_grad = False + def forward(self, x): + x = self.conv1(x) + x = self.bn1(x) + x = self.relu(x) + x = self.maxpool(x) + + x = self.layer1(x) + x = self.layer2(x) + x = self.layer3(x) + x = self.layer4(x) + + # x = self.avgpool(x) + # x = x.view(x.size(0), -1) + # x = self.fc(x) + + return x + + def load_param(self, model_path): + param_dict = torch.load(model_path) + if 'state_dict' in param_dict: + param_dict = param_dict['state_dict'] + for i in param_dict: + if 'fc' in i: + continue + self.state_dict()[i.replace('module.','')].copy_(param_dict[i]) + + + +def resnet50_ibn_a(last_stride, pretrained=False, **kwargs): + """Constructs a ResNet-50 model. + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + """ + model = ResNet_IBN(last_stride, Bottleneck_IBN, [3, 4, 6, 3], **kwargs) + if pretrained: + model.load_state_dict(model_zoo.load_url(model_urls['resnet50'])) + return model + + +def resnet101_ibn_a(last_stride, pretrained=False, **kwargs): + """Constructs a ResNet-101 model. + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + """ + model = ResNet_IBN(last_stride, Bottleneck_IBN, [3, 4, 23, 3], **kwargs) + if pretrained: + model.load_state_dict(model_zoo.load_url(model_urls['resnet101'])) + return model + + +def resnet152_ibn_a(last_stride, pretrained=False, **kwargs): + """Constructs a ResNet-152 model. + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + """ + model = ResNet_IBN(last_stride, Bottleneck_IBN, [3, 8, 36, 3], **kwargs) + if pretrained: + model.load_state_dict(model_zoo.load_url(model_urls['resnet152'])) return model \ No newline at end of file diff --git a/addition_module/DMUE/pretrain/gen_train_file.py b/src/faceDetection/face_sdk/addition_module/DMUE/pretrain/gen_train_file.py similarity index 100% rename from addition_module/DMUE/pretrain/gen_train_file.py rename to src/faceDetection/face_sdk/addition_module/DMUE/pretrain/gen_train_file.py diff --git a/addition_module/DMUE/pretrain/loss_def.py b/src/faceDetection/face_sdk/addition_module/DMUE/pretrain/loss_def.py similarity index 100% rename from addition_module/DMUE/pretrain/loss_def.py rename to src/faceDetection/face_sdk/addition_module/DMUE/pretrain/loss_def.py diff --git a/addition_module/DMUE/pretrain/train_mv_softmax.py b/src/faceDetection/face_sdk/addition_module/DMUE/pretrain/train_mv_softmax.py similarity index 100% rename from addition_module/DMUE/pretrain/train_mv_softmax.py rename to src/faceDetection/face_sdk/addition_module/DMUE/pretrain/train_mv_softmax.py diff --git a/addition_module/DMUE/pretrain/train_res18.sh b/src/faceDetection/face_sdk/addition_module/DMUE/pretrain/train_res18.sh similarity index 100% rename from addition_module/DMUE/pretrain/train_res18.sh rename to src/faceDetection/face_sdk/addition_module/DMUE/pretrain/train_res18.sh diff --git a/addition_module/DMUE/pretrain/train_res50_ibn.sh b/src/faceDetection/face_sdk/addition_module/DMUE/pretrain/train_res50_ibn.sh similarity index 100% rename from addition_module/DMUE/pretrain/train_res50_ibn.sh rename to src/faceDetection/face_sdk/addition_module/DMUE/pretrain/train_res50_ibn.sh diff --git a/addition_module/DMUE/pretrain/utils/AverageMeter.py b/src/faceDetection/face_sdk/addition_module/DMUE/pretrain/utils/AverageMeter.py similarity index 100% rename from addition_module/DMUE/pretrain/utils/AverageMeter.py rename to src/faceDetection/face_sdk/addition_module/DMUE/pretrain/utils/AverageMeter.py diff --git a/src/faceDetection/face_sdk/addition_module/DMUE/pretrain/utils/__init__.py b/src/faceDetection/face_sdk/addition_module/DMUE/pretrain/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/addition_module/DMUE/pretrain/utils/dataset.py b/src/faceDetection/face_sdk/addition_module/DMUE/pretrain/utils/dataset.py similarity index 96% rename from addition_module/DMUE/pretrain/utils/dataset.py rename to src/faceDetection/face_sdk/addition_module/DMUE/pretrain/utils/dataset.py index 8e64a17..c668c10 100644 --- a/addition_module/DMUE/pretrain/utils/dataset.py +++ b/src/faceDetection/face_sdk/addition_module/DMUE/pretrain/utils/dataset.py @@ -1,49 +1,49 @@ -import os -import random - -import cv2 -import torch -import numpy as np -from torch.utils.data import Dataset -from PIL import Image, ImageFile - - -def read_image(img_path): - """Keep reading image until succeed. - This can avoid IOError incurred by heavy IO process.""" - got_img = False - if not os.path.exists(img_path): - raise IOError("{} does not exist".format(img_path)) - while not got_img: - try: - img = Image.open(img_path).convert('RGB') - got_img = True - except IOError: - print("IOError incurred when reading '{}'. Will redo. Don't worry. Just chill.".format(img_path)) - pass - return img - - -class ImageDataset(Dataset): - def __init__(self, data_root, train_file, transform): - self.data_root = data_root - self.transform = transform - self.train_list = [] - train_file_buf = open(train_file) - line = train_file_buf.readline().strip() - while line: - image_path, image_label = line.split(' ') - self.train_list.append((os.path.join(self.data_root, image_path), int(image_label))) - line = train_file_buf.readline().strip() - - def __len__(self): - return len(self.train_list) - - def __getitem__(self, index, transform=None): - img_path, label = self.train_list[index] - img = read_image(img_path) - - if self.transform is not None: - img = self.transform(img) - - return img, label, img_path +import os +import random + +import cv2 +import torch +import numpy as np +from torch.utils.data import Dataset +from PIL import Image, ImageFile + + +def read_image(img_path): + """Keep reading image until succeed. + This can avoid IOError incurred by heavy IO process.""" + got_img = False + if not os.path.exists(img_path): + raise IOError("{} does not exist".format(img_path)) + while not got_img: + try: + img = Image.open(img_path).convert('RGB') + got_img = True + except IOError: + print("IOError incurred when reading '{}'. Will redo. Don't worry. Just chill.".format(img_path)) + pass + return img + + +class ImageDataset(Dataset): + def __init__(self, data_root, train_file, transform): + self.data_root = data_root + self.transform = transform + self.train_list = [] + train_file_buf = open(train_file) + line = train_file_buf.readline().strip() + while line: + image_path, image_label = line.split(' ') + self.train_list.append((os.path.join(self.data_root, image_path), int(image_label))) + line = train_file_buf.readline().strip() + + def __len__(self): + return len(self.train_list) + + def __getitem__(self, index, transform=None): + img_path, label = self.train_list[index] + img = read_image(img_path) + + if self.transform is not None: + img = self.transform(img) + + return img, label, img_path diff --git a/addition_module/DMUE/pretrain/utils/make_transform.py b/src/faceDetection/face_sdk/addition_module/DMUE/pretrain/utils/make_transform.py similarity index 96% rename from addition_module/DMUE/pretrain/utils/make_transform.py rename to src/faceDetection/face_sdk/addition_module/DMUE/pretrain/utils/make_transform.py index bda33c3..761756e 100644 --- a/addition_module/DMUE/pretrain/utils/make_transform.py +++ b/src/faceDetection/face_sdk/addition_module/DMUE/pretrain/utils/make_transform.py @@ -1,28 +1,28 @@ -# encoding: utf-8 -import torch -import torchvision.transforms as T -from torch.utils.data import DataLoader -from .dataset import ImageDataset - -def make_transform(): - ori_shape = (256, 256) - image_crop_size = (224, 224) - horizontal_flip_p = 0.5 - normalize_std = [0.5, 0.5, 0.5] - normalize_mean = [0.5, 0.5, 0.5] - - train_transforms = T.Compose([ - T.Resize(ori_shape), - T.RandomCrop(image_crop_size), - T.RandomHorizontalFlip(p=horizontal_flip_p), - T.ToTensor(), - T.Normalize(mean=normalize_mean, std=normalize_std), - ]) - val_transforms = T.Compose([ - T.Resize(ori_shape), - T.CenterCrop(image_crop_size), - T.ToTensor(), - T.Normalize(mean=normalize_mean, std=normalize_std) - ]) - +# encoding: utf-8 +import torch +import torchvision.transforms as T +from torch.utils.data import DataLoader +from .dataset import ImageDataset + +def make_transform(): + ori_shape = (256, 256) + image_crop_size = (224, 224) + horizontal_flip_p = 0.5 + normalize_std = [0.5, 0.5, 0.5] + normalize_mean = [0.5, 0.5, 0.5] + + train_transforms = T.Compose([ + T.Resize(ori_shape), + T.RandomCrop(image_crop_size), + T.RandomHorizontalFlip(p=horizontal_flip_p), + T.ToTensor(), + T.Normalize(mean=normalize_mean, std=normalize_std), + ]) + val_transforms = T.Compose([ + T.Resize(ori_shape), + T.CenterCrop(image_crop_size), + T.ToTensor(), + T.Normalize(mean=normalize_mean, std=normalize_std) + ]) + return train_transforms, val_transforms \ No newline at end of file diff --git a/addition_module/DMUE/train.py b/src/faceDetection/face_sdk/addition_module/DMUE/train.py similarity index 97% rename from addition_module/DMUE/train.py rename to src/faceDetection/face_sdk/addition_module/DMUE/train.py index 9fe0d16..1b13e56 100644 --- a/addition_module/DMUE/train.py +++ b/src/faceDetection/face_sdk/addition_module/DMUE/train.py @@ -1,197 +1,197 @@ -import os -import gc -import argparse -import numpy as np -import logging as logger - -import torch -import torch.nn as nn -import torch.backends.cudnn as cudnn - -from datasets import make_dataloader -from config import config as cfg -from models import make_model -from losses import SoftLoss, SP_KD_Loss -from utils import write_config_into_log, ramp_up, ramp_down - - -cfg.output_dir = os.path.join(cfg.ckpt_root_dir, cfg.output_dir) -cfg.snapshot_dir = os.path.join(cfg.output_dir, 'snapshots') -cfg.train_log = os.path.join(cfg.output_dir, 'train_log.txt') -cfg.test_log = os.path.join(cfg.output_dir, 'test_log.txt') - -print('Experiments dir: {}'.format(cfg.output_dir)) -if not os.path.exists(cfg.output_dir): - os.makedirs(cfg.output_dir) -if not os.path.exists(cfg.snapshot_dir): - os.makedirs(cfg.snapshot_dir) - -log_format = '%(levelname)s %(asctime)s %(filename)s] %(message)s' -logger.basicConfig(level=logger.INFO, format=log_format, datefmt='%Y-%m-%d %H:%M:%S') -fh = logger.FileHandler(cfg.train_log) -fh.setFormatter(logger.Formatter(log_format)) -logger.getLogger().addHandler(fh) - - -def train(): - logger.info('Initializing....') - cudnn.enable = True - cudnn.benchmark = True - # torch.manual_seed(1) - # torch.cuda.manual_seed(1) - write_config_into_log(cfg) - - logger.info('Building model......') - if cfg.pretrained: - model = make_model(cfg) - model.load_param(cfg) - logger.info('Loaded pretrained model from {0}'.format(cfg.pretrained)) - else: - model = make_model(cfg) - - model.cuda() - model = torch.nn.DataParallel(model) - - optimizer = torch.optim.Adam([{'params': model.module.base.parameters(), 'lr': cfg.get_lr(0)[0]}, - {'params': model.module.classifiers.parameters(), 'lr': cfg.get_lr(0)[1]}], - weight_decay=cfg.weight_decay) - - celoss = nn.CrossEntropyLoss().cuda() - softloss = SoftLoss() - sp_kd_loss = SP_KD_Loss() - criterions = [celoss, softloss, sp_kd_loss] - - train_loader, val_loader = make_dataloader(cfg) - logger.info('Begin training......') - for epoch in range(cfg.start_epoch, cfg.max_epoch): - train_one_epoch(train_loader, val_loader, model, criterions, optimizer, epoch, cfg) - - total_acc = test(cfg, val_loader, model, epoch) - with open(cfg.test_log, 'a+') as f: - f.write('Epoch {0}: Acc is {1:.4f}\n'.format(epoch, total_acc)) - torch.save(obj=model.state_dict(), - f=os.path.join(cfg.snapshot_dir, 'ep{}_acc{:.4f}.pth'.format(epoch, total_acc))) - logger.info('Model saved') - - -def train_one_epoch(train_loader, val_loader, model, criterions, optimizer, epoch, cfg): - model.train() - - ramp_up_w, ramp_down_w = ramp_up(epoch, cfg.ramp_a), ramp_down(epoch, cfg.ramp_a) - w, gamma = cfg.w, cfg.gamma - logger.info('ramp_up_w: {:.4f}, ramp_down_w: {:.4f}, w: {}, gamma: {}'.format(ramp_up_w, ramp_down_w, w, gamma)) - adjust_learning_rate(optimizer, cfg.get_lr(epoch)[0], cfg.get_lr(epoch)[1]) - training_phase = cfg.train_mode - - for idx, batch in enumerate(train_loader, start=1): - img, label = batch - - good_batch = check_batch(label, cfg.num_classes) - if not good_batch: - continue - img, label = img.cuda(), label.cuda() - - x_final, output_x_list, targets, softlabel_list, G_matrixs, G_main_matrixs, score, atten_x_final = model(img, label, training_phase) - - softlabel = torch.zeros_like(x_final) - for c in range(cfg.num_classes): - # sharpen - px = torch.softmax(softlabel_list[c], dim=1) - ptx = px ** (1 / cfg.T) # temparature sharpening - targets_x = ptx / ptx.sum(dim=1, keepdim=True) # normalize - targets_x = targets_x.detach() - - _h, _w = px.shape - _targets = torch.zeros([_h, _w+1]).float() - _targets[:, 0:c], _targets[:, c + 1:] = targets_x[:, 0:c], targets_x[:, c:] - ind = (label == c).nonzero() - softlabel[ind[:, 0]] = _targets.cuda() # bs x num_cls - - softLoss = criterions[1](x_final, label, softlabel) - aux_loss = [criterions[0](_pred, _label) for _pred, _label in zip(output_x_list, targets)] # auxiliary branch - aux_loss = sum(aux_loss) / (cfg.num_branches-1) - CEloss = criterions[0](atten_x_final, label) # main - spLoss = criterions[2](G_matrixs, G_main_matrixs) - - loss = ramp_up_w * CEloss + ramp_up_w * (w * softLoss + gamma * spLoss) + ramp_down_w * aux_loss - - if idx % 20 == 0: - logger.info('Epoch {} Batch {}/{}: CEloss {:.6f}, softLoss {:.6f}, spLoss {:.6f}, aux_loss {:.6f}'.format( - epoch, idx, len(train_loader), CEloss, softLoss, spLoss, aux_loss)) - - # compute gradient and do SGD step - optimizer.zero_grad() - loss.backward() - optimizer.step() - - if idx % 200 == 0: - total_acc = test(cfg, val_loader, model, epoch, idx) - with open(cfg.test_log, 'a+') as f: - f.write('Epoch {0} Batch {1}: Acc is {2:.4f}\n'.format(epoch, idx, total_acc)) - torch.save(obj=model.state_dict(), - f=os.path.join(cfg.snapshot_dir, 'ep{}_b{}_acc{:.4f}.pth'.format(epoch, idx, total_acc))) - logger.info('Model saved') - - gc.collect() - - -def test_worker(val_loader, model): - Pred, Label = [], [] - for idx, batch in enumerate(val_loader, 1): - img, label = batch - pred, output_x_list, targets, softlabel_list = model(img.cuda(), label, 'normal') - - for i in range(pred.shape[0]): - x = pred[i].data.cpu().numpy() - y = label[i].item() - Pred.append(x) - Label.append(y) - - return Pred, Label - - -def test(cfg, val_loader, model, epoch, batch=None): - model.eval() - - pred, label = test_worker(val_loader, model) - pred = [np.where(x == np.max(x))[0][0] for x in pred] - - total = len([x for x in range(len(pred)) if pred[x] == label[x]]) - total_acc = total / len(pred) - logger.info('Epoch {}: Acc is {:.4f}'.format(epoch, total_acc)) if batch is None else \ - logger.info('Epoch {} Batch {}: Acc is {:.4f}'.format(epoch, batch, total_acc)) - - return total_acc - - -def adjust_learning_rate(optimizer, lr1, lr2): - optimizer.param_groups[0]['lr'] = lr1 - optimizer.param_groups[1]['lr'] = lr2 - - -def accuracy(output, target, topk=(1,)): - """Computes the precision@k for the specified values of k""" - maxk = max(topk) - batch_size = target.size(0) - - _, pred = output.topk(maxk, 1, True, True) - pred = pred.t() - correct = pred.eq(target.view(1, -1).expand_as(pred)) - - res = [] - for k in topk: - correct_k = correct[:k].view(-1).float().sum(0, keepdim=True) - res.append(correct_k.mul_(100.0 / batch_size)) - return res - - -def check_batch(label, num_classes): - for i in range(num_classes): - cnt = (label == i).nonzero().shape[0] - if cnt < 2: - return False - return True - - -if __name__ == '__main__': - train() +import os +import gc +import argparse +import numpy as np +import logging as logger + +import torch +import torch.nn as nn +import torch.backends.cudnn as cudnn + +from datasets import make_dataloader +from config import config as cfg +from models import make_model +from losses import SoftLoss, SP_KD_Loss +from utils import write_config_into_log, ramp_up, ramp_down + + +cfg.output_dir = os.path.join(cfg.ckpt_root_dir, cfg.output_dir) +cfg.snapshot_dir = os.path.join(cfg.output_dir, 'snapshots') +cfg.train_log = os.path.join(cfg.output_dir, 'train_log.txt') +cfg.test_log = os.path.join(cfg.output_dir, 'test_log.txt') + +print('Experiments dir: {}'.format(cfg.output_dir)) +if not os.path.exists(cfg.output_dir): + os.makedirs(cfg.output_dir) +if not os.path.exists(cfg.snapshot_dir): + os.makedirs(cfg.snapshot_dir) + +log_format = '%(levelname)s %(asctime)s %(filename)s] %(message)s' +logger.basicConfig(level=logger.INFO, format=log_format, datefmt='%Y-%m-%d %H:%M:%S') +fh = logger.FileHandler(cfg.train_log) +fh.setFormatter(logger.Formatter(log_format)) +logger.getLogger().addHandler(fh) + + +def train(): + logger.info('Initializing....') + cudnn.enable = True + cudnn.benchmark = True + # torch.manual_seed(1) + # torch.cuda.manual_seed(1) + write_config_into_log(cfg) + + logger.info('Building model......') + if cfg.pretrained: + model = make_model(cfg) + model.load_param(cfg) + logger.info('Loaded pretrained model from {0}'.format(cfg.pretrained)) + else: + model = make_model(cfg) + + model.cuda() + model = torch.nn.DataParallel(model) + + optimizer = torch.optim.Adam([{'params': model.module.base.parameters(), 'lr': cfg.get_lr(0)[0]}, + {'params': model.module.classifiers.parameters(), 'lr': cfg.get_lr(0)[1]}], + weight_decay=cfg.weight_decay) + + celoss = nn.CrossEntropyLoss().cuda() + softloss = SoftLoss() + sp_kd_loss = SP_KD_Loss() + criterions = [celoss, softloss, sp_kd_loss] + + train_loader, val_loader = make_dataloader(cfg) + logger.info('Begin training......') + for epoch in range(cfg.start_epoch, cfg.max_epoch): + train_one_epoch(train_loader, val_loader, model, criterions, optimizer, epoch, cfg) + + total_acc = test(cfg, val_loader, model, epoch) + with open(cfg.test_log, 'a+') as f: + f.write('Epoch {0}: Acc is {1:.4f}\n'.format(epoch, total_acc)) + torch.save(obj=model.state_dict(), + f=os.path.join(cfg.snapshot_dir, 'ep{}_acc{:.4f}.pth'.format(epoch, total_acc))) + logger.info('Model saved') + + +def train_one_epoch(train_loader, val_loader, model, criterions, optimizer, epoch, cfg): + model.train() + + ramp_up_w, ramp_down_w = ramp_up(epoch, cfg.ramp_a), ramp_down(epoch, cfg.ramp_a) + w, gamma = cfg.w, cfg.gamma + logger.info('ramp_up_w: {:.4f}, ramp_down_w: {:.4f}, w: {}, gamma: {}'.format(ramp_up_w, ramp_down_w, w, gamma)) + adjust_learning_rate(optimizer, cfg.get_lr(epoch)[0], cfg.get_lr(epoch)[1]) + training_phase = cfg.train_mode + + for idx, batch in enumerate(train_loader, start=1): + img, label = batch + + good_batch = check_batch(label, cfg.num_classes) + if not good_batch: + continue + img, label = img.cuda(), label.cuda() + + x_final, output_x_list, targets, softlabel_list, G_matrixs, G_main_matrixs, score, atten_x_final = model(img, label, training_phase) + + softlabel = torch.zeros_like(x_final) + for c in range(cfg.num_classes): + # sharpen + px = torch.softmax(softlabel_list[c], dim=1) + ptx = px ** (1 / cfg.T) # temparature sharpening + targets_x = ptx / ptx.sum(dim=1, keepdim=True) # normalize + targets_x = targets_x.detach() + + _h, _w = px.shape + _targets = torch.zeros([_h, _w+1]).float() + _targets[:, 0:c], _targets[:, c + 1:] = targets_x[:, 0:c], targets_x[:, c:] + ind = (label == c).nonzero() + softlabel[ind[:, 0]] = _targets.cuda() # bs x num_cls + + softLoss = criterions[1](x_final, label, softlabel) + aux_loss = [criterions[0](_pred, _label) for _pred, _label in zip(output_x_list, targets)] # auxiliary branch + aux_loss = sum(aux_loss) / (cfg.num_branches-1) + CEloss = criterions[0](atten_x_final, label) # main + spLoss = criterions[2](G_matrixs, G_main_matrixs) + + loss = ramp_up_w * CEloss + ramp_up_w * (w * softLoss + gamma * spLoss) + ramp_down_w * aux_loss + + if idx % 20 == 0: + logger.info('Epoch {} Batch {}/{}: CEloss {:.6f}, softLoss {:.6f}, spLoss {:.6f}, aux_loss {:.6f}'.format( + epoch, idx, len(train_loader), CEloss, softLoss, spLoss, aux_loss)) + + # compute gradient and do SGD step + optimizer.zero_grad() + loss.backward() + optimizer.step() + + if idx % 200 == 0: + total_acc = test(cfg, val_loader, model, epoch, idx) + with open(cfg.test_log, 'a+') as f: + f.write('Epoch {0} Batch {1}: Acc is {2:.4f}\n'.format(epoch, idx, total_acc)) + torch.save(obj=model.state_dict(), + f=os.path.join(cfg.snapshot_dir, 'ep{}_b{}_acc{:.4f}.pth'.format(epoch, idx, total_acc))) + logger.info('Model saved') + + gc.collect() + + +def test_worker(val_loader, model): + Pred, Label = [], [] + for idx, batch in enumerate(val_loader, 1): + img, label = batch + pred, output_x_list, targets, softlabel_list = model(img.cuda(), label, 'normal') + + for i in range(pred.shape[0]): + x = pred[i].data.cpu().numpy() + y = label[i].item() + Pred.append(x) + Label.append(y) + + return Pred, Label + + +def test(cfg, val_loader, model, epoch, batch=None): + model.eval() + + pred, label = test_worker(val_loader, model) + pred = [np.where(x == np.max(x))[0][0] for x in pred] + + total = len([x for x in range(len(pred)) if pred[x] == label[x]]) + total_acc = total / len(pred) + logger.info('Epoch {}: Acc is {:.4f}'.format(epoch, total_acc)) if batch is None else \ + logger.info('Epoch {} Batch {}: Acc is {:.4f}'.format(epoch, batch, total_acc)) + + return total_acc + + +def adjust_learning_rate(optimizer, lr1, lr2): + optimizer.param_groups[0]['lr'] = lr1 + optimizer.param_groups[1]['lr'] = lr2 + + +def accuracy(output, target, topk=(1,)): + """Computes the precision@k for the specified values of k""" + maxk = max(topk) + batch_size = target.size(0) + + _, pred = output.topk(maxk, 1, True, True) + pred = pred.t() + correct = pred.eq(target.view(1, -1).expand_as(pred)) + + res = [] + for k in topk: + correct_k = correct[:k].view(-1).float().sum(0, keepdim=True) + res.append(correct_k.mul_(100.0 / batch_size)) + return res + + +def check_batch(label, num_classes): + for i in range(num_classes): + cnt = (label == i).nonzero().shape[0] + if cnt < 2: + return False + return True + + +if __name__ == '__main__': + train() diff --git a/addition_module/DMUE/train.sh b/src/faceDetection/face_sdk/addition_module/DMUE/train.sh similarity index 100% rename from addition_module/DMUE/train.sh rename to src/faceDetection/face_sdk/addition_module/DMUE/train.sh diff --git a/addition_module/DMUE/train_ddp.py b/src/faceDetection/face_sdk/addition_module/DMUE/train_ddp.py similarity index 97% rename from addition_module/DMUE/train_ddp.py rename to src/faceDetection/face_sdk/addition_module/DMUE/train_ddp.py index 0db091d..404df2f 100644 --- a/addition_module/DMUE/train_ddp.py +++ b/src/faceDetection/face_sdk/addition_module/DMUE/train_ddp.py @@ -1,269 +1,269 @@ -import os -import gc -import argparse -import numpy as np -import logging as logger -import argparse - -import torch -import torch.nn as nn -import torch.distributed as dist -import torch.utils.data.distributed -import torch.backends.cudnn as cudnn -from apex.parallel import DistributedDataParallel -from apex.parallel import convert_syncbn_model - -from datasets import make_dataloader -from config import config as cfg -from models import make_model -from losses import SoftLoss, SP_KD_Loss -from utils import write_config_into_log, ramp_up, ramp_down - - -parser = argparse.ArgumentParser(description="DDP parameters") -parser.add_argument('--local_rank', type=int, default=0) -parser.add_argument('--sync_bn', action='store_true', help='enabling apex sync BN') - -args = parser.parse_args() - - -cfg.output_dir = os.path.join(cfg.ckpt_root_dir, cfg.output_dir) -cfg.snapshot_dir = os.path.join(cfg.output_dir, 'snapshots') -cfg.train_log = os.path.join(cfg.output_dir, 'train_log.txt') -cfg.test_log = os.path.join(cfg.output_dir, 'test_log.txt') - -if args.local_rank == 0: - print('Experiments dir: {}'.format(cfg.output_dir)) - if not os.path.exists(cfg.output_dir): - os.makedirs(cfg.output_dir) - if not os.path.exists(cfg.snapshot_dir): - os.makedirs(cfg.snapshot_dir) - - log_format = '%(levelname)s %(asctime)s %(filename)s] %(message)s' - logger.basicConfig(level=logger.INFO, format=log_format, datefmt='%Y-%m-%d %H:%M:%S') - fh = logger.FileHandler(cfg.train_log) - fh.setFormatter(logger.Formatter(log_format)) - logger.getLogger().addHandler(fh) - - -def train(): - if args.local_rank == 0: - logger.info('Initializing....') - cudnn.enable = True - cudnn.benchmark = True - # torch.manual_seed(1) - # torch.cuda.manual_seed(1) - args.distributed = False - if 'WORLD_SIZE' in os.environ: - args.distributed = int(os.environ['WORLD_SIZE']) > 1 - args.gpu = 0 - args.world_size = 1 - if args.distributed: - args.gpu = args.local_rank - torch.cuda.set_device(args.gpu) - torch.distributed.init_process_group(backend='nccl', init_method='env://') - args.world_size = torch.distributed.get_world_size() - - if args.local_rank == 0: - write_config_into_log(cfg) - - if args.local_rank == 0: - logger.info('Building model......') - if cfg.pretrained: - model = make_model(cfg) - model.load_param(cfg) - if args.local_rank == 0: - logger.info('Loaded pretrained model from {0}'.format(cfg.pretrained)) - else: - model = make_model(cfg) - - if args.sync_bn: - if args.local_rank == 0: logging.info("using apex synced BN") - model = convert_syncbn_model(model) - model.cuda() - if args.distributed: - # By default, apex.parallel.DistributedDataParallel overlaps communication with - # computation in the backward pass. - # delay_allreduce delays all communication to the end of the backward pass. - model = DistributedDataParallel(model, delay_allreduce=True) - else: - model = torch.nn.DataParallel(model) - - optimizer = torch.optim.Adam([{'params': model.module.base.parameters(), 'lr': cfg.get_lr(0)[0]}, - {'params': model.module.classifiers.parameters(), 'lr': cfg.get_lr(0)[1]}], - weight_decay=cfg.weight_decay) - - celoss = nn.CrossEntropyLoss().cuda() - softloss = SoftLoss() - sp_kd_loss = SP_KD_Loss() - criterions = [celoss, softloss, sp_kd_loss] - - cfg.batch_size = cfg.batch_size // args.world_size - cfg.num_workers = cfg.num_workers // args.world_size - train_loader, val_loader = make_dataloader(cfg) - - if args.local_rank == 0: - logger.info('Begin training......') - for epoch in range(cfg.start_epoch, cfg.max_epoch): - train_one_epoch(train_loader, val_loader, model, criterions, optimizer, epoch, cfg) - - total_acc = test(cfg, val_loader, model, epoch) - if args.local_rank == 0: - with open(cfg.test_log, 'a+') as f: - f.write('Epoch {0}: Acc is {1:.4f}\n'.format(epoch, total_acc)) - torch.save(obj=model.state_dict(), - f=os.path.join(cfg.snapshot_dir, 'ep{}_acc{:.4f}.pth'.format(epoch, total_acc))) - logger.info('Model saved') - - -def train_one_epoch(train_loader, val_loader, model, criterions, optimizer, epoch, cfg): - model.train() - - ramp_up_w, ramp_down_w = ramp_up(epoch, cfg.ramp_a), ramp_down(epoch, cfg.ramp_a) - w, gamma = cfg.w, cfg.gamma - if args.local_rank == 0: - logger.info('ramp_up_w: {:.4f}, ramp_down_w: {:.4f}, w: {}, gamma: {}'.format(ramp_up_w, ramp_down_w, w, gamma)) - adjust_learning_rate(optimizer, cfg.get_lr(epoch)[0], cfg.get_lr(epoch)[1]) - training_phase = cfg.train_mode - - for idx, batch in enumerate(train_loader, start=1): - img, label = batch - - good_batch = check_batch(label, cfg.num_classes) - good_batch = torch.cuda.IntTensor([good_batch]) - dist.all_reduce(good_batch, op=dist.ReduceOp.SUM) - good_batch = good_batch // args.world_size - good_batch = bool(good_batch.item()) - if not good_batch: - continue - img, label = img.cuda(), label.cuda() - - x_final, output_x_list, targets, softlabel_list, G_matrixs, G_main_matrixs, score, atten_x_final = model(img, label, training_phase) - - softlabel = torch.zeros_like(x_final) - for c in range(cfg.num_classes): - # sharpen - px = torch.softmax(softlabel_list[c], dim=1) - ptx = px ** (1 / cfg.T) # temparature sharpening - targets_x = ptx / ptx.sum(dim=1, keepdim=True) # normalize - targets_x = targets_x.detach() - - _h, _w = px.shape - _targets = torch.zeros([_h, _w+1]).float() - _targets[:, 0:c], _targets[:, c + 1:] = targets_x[:, 0:c], targets_x[:, c:] - ind = (label == c).nonzero() - softlabel[ind[:, 0]] = _targets.cuda() # bs x num_cls - - softLoss = criterions[1](x_final, label, softlabel) - aux_loss = [criterions[0](_pred, _label) for _pred, _label in zip(output_x_list, targets)] # auxiliary branch - aux_loss = sum(aux_loss) / (cfg.num_branches-1) - CEloss = criterions[0](atten_x_final, label) # main - spLoss = criterions[2](G_matrixs, G_main_matrixs) - - loss = ramp_up_w * CEloss + ramp_up_w * (w * softLoss + gamma * spLoss) + ramp_down_w * aux_loss - - if idx % 20 == 0: - if args.distributed: - CEloss = reduce_tensor(CEloss) - softLoss = reduce_tensor(softLoss) - spLoss = reduce_tensor(spLoss) - aux_loss = reduce_tensor(aux_loss) - if args.local_rank == 0: - logger.info('Epoch {} Batch {}/{}: CEloss {:.6f}, softLoss {:.6f}, spLoss {:.6f}, aux_loss {:.6f}'.format( - epoch, idx, len(train_loader), CEloss, softLoss, spLoss, aux_loss)) - - # compute gradient and do SGD step - optimizer.zero_grad() - loss.backward() - optimizer.step() - - if idx % 200 == 0: - total_acc = test(cfg, val_loader, model, epoch, idx) - if args.local_rank == 0: - with open(cfg.test_log, 'a+') as f: - f.write('Epoch {0} Batch {1}: Acc is {2:.4f}\n'.format(epoch, idx, total_acc)) - torch.save(obj=model.state_dict(), - f=os.path.join(cfg.snapshot_dir, 'ep{}_b{}_acc{:.4f}.pth'.format(epoch, idx, total_acc))) - logger.info('Model saved') - - gc.collect() - - -def test_worker(val_loader, model): - Pred, Label = [], [] - for idx, batch in enumerate(val_loader, 1): - img, label = batch - pred, output_x_list, targets, softlabel_list = model(img.cuda(), label, 'normal') - - for i in range(pred.shape[0]): - x = pred[i].data.cpu().numpy() - y = label[i].item() - Pred.append(x) - Label.append(y) - - return Pred, Label - - -def test(cfg, val_loader, model, epoch, batch=None): - model.eval() - - pred, label = test_worker(val_loader, model) - pred = [np.where(x == np.max(x))[0][0] for x in pred] - - total = len([x for x in range(len(pred)) if pred[x] == label[x]]) - total_acc = total / len(pred) - if args.distributed: - total_acc = reduce_float(total_acc) - if args.local_rank == 0: - logger.info('Epoch {}: Acc is {:.4f}'.format(epoch, total_acc)) if batch is None else \ - logger.info('Epoch {} Batch {}: Acc is {:.4f}'.format(epoch, batch, total_acc)) - - return total_acc - - -def adjust_learning_rate(optimizer, lr1, lr2): - optimizer.param_groups[0]['lr'] = lr1 - optimizer.param_groups[1]['lr'] = lr2 - - -def accuracy(output, target, topk=(1,)): - """Computes the precision@k for the specified values of k""" - maxk = max(topk) - batch_size = target.size(0) - - _, pred = output.topk(maxk, 1, True, True) - pred = pred.t() - correct = pred.eq(target.view(1, -1).expand_as(pred)) - - res = [] - for k in topk: - correct_k = correct[:k].view(-1).float().sum(0, keepdim=True) - res.append(correct_k.mul_(100.0 / batch_size)) - return res - - -def check_batch(label, num_classes): - for i in range(num_classes): - cnt = (label == i).nonzero().shape[0] - if cnt < 2: - return False - return True - - -def reduce_tensor(tensor): - rt = tensor.clone() - dist.all_reduce(rt, op=dist.ReduceOp.SUM) - rt /= args.world_size - return rt - - -def reduce_float(f): - tensor = torch.cuda.FloatTensor([f]) - rt = tensor.clone() - dist.all_reduce(rt, op=dist.ReduceOp.SUM) - rt /= args.world_size - return rt.item() - - -if __name__ == '__main__': - train() +import os +import gc +import argparse +import numpy as np +import logging as logger +import argparse + +import torch +import torch.nn as nn +import torch.distributed as dist +import torch.utils.data.distributed +import torch.backends.cudnn as cudnn +from apex.parallel import DistributedDataParallel +from apex.parallel import convert_syncbn_model + +from datasets import make_dataloader +from config import config as cfg +from models import make_model +from losses import SoftLoss, SP_KD_Loss +from utils import write_config_into_log, ramp_up, ramp_down + + +parser = argparse.ArgumentParser(description="DDP parameters") +parser.add_argument('--local_rank', type=int, default=0) +parser.add_argument('--sync_bn', action='store_true', help='enabling apex sync BN') + +args = parser.parse_args() + + +cfg.output_dir = os.path.join(cfg.ckpt_root_dir, cfg.output_dir) +cfg.snapshot_dir = os.path.join(cfg.output_dir, 'snapshots') +cfg.train_log = os.path.join(cfg.output_dir, 'train_log.txt') +cfg.test_log = os.path.join(cfg.output_dir, 'test_log.txt') + +if args.local_rank == 0: + print('Experiments dir: {}'.format(cfg.output_dir)) + if not os.path.exists(cfg.output_dir): + os.makedirs(cfg.output_dir) + if not os.path.exists(cfg.snapshot_dir): + os.makedirs(cfg.snapshot_dir) + + log_format = '%(levelname)s %(asctime)s %(filename)s] %(message)s' + logger.basicConfig(level=logger.INFO, format=log_format, datefmt='%Y-%m-%d %H:%M:%S') + fh = logger.FileHandler(cfg.train_log) + fh.setFormatter(logger.Formatter(log_format)) + logger.getLogger().addHandler(fh) + + +def train(): + if args.local_rank == 0: + logger.info('Initializing....') + cudnn.enable = True + cudnn.benchmark = True + # torch.manual_seed(1) + # torch.cuda.manual_seed(1) + args.distributed = False + if 'WORLD_SIZE' in os.environ: + args.distributed = int(os.environ['WORLD_SIZE']) > 1 + args.gpu = 0 + args.world_size = 1 + if args.distributed: + args.gpu = args.local_rank + torch.cuda.set_device(args.gpu) + torch.distributed.init_process_group(backend='nccl', init_method='env://') + args.world_size = torch.distributed.get_world_size() + + if args.local_rank == 0: + write_config_into_log(cfg) + + if args.local_rank == 0: + logger.info('Building model......') + if cfg.pretrained: + model = make_model(cfg) + model.load_param(cfg) + if args.local_rank == 0: + logger.info('Loaded pretrained model from {0}'.format(cfg.pretrained)) + else: + model = make_model(cfg) + + if args.sync_bn: + if args.local_rank == 0: logging.info("using apex synced BN") + model = convert_syncbn_model(model) + model.cuda() + if args.distributed: + # By default, apex.parallel.DistributedDataParallel overlaps communication with + # computation in the backward pass. + # delay_allreduce delays all communication to the end of the backward pass. + model = DistributedDataParallel(model, delay_allreduce=True) + else: + model = torch.nn.DataParallel(model) + + optimizer = torch.optim.Adam([{'params': model.module.base.parameters(), 'lr': cfg.get_lr(0)[0]}, + {'params': model.module.classifiers.parameters(), 'lr': cfg.get_lr(0)[1]}], + weight_decay=cfg.weight_decay) + + celoss = nn.CrossEntropyLoss().cuda() + softloss = SoftLoss() + sp_kd_loss = SP_KD_Loss() + criterions = [celoss, softloss, sp_kd_loss] + + cfg.batch_size = cfg.batch_size // args.world_size + cfg.num_workers = cfg.num_workers // args.world_size + train_loader, val_loader = make_dataloader(cfg) + + if args.local_rank == 0: + logger.info('Begin training......') + for epoch in range(cfg.start_epoch, cfg.max_epoch): + train_one_epoch(train_loader, val_loader, model, criterions, optimizer, epoch, cfg) + + total_acc = test(cfg, val_loader, model, epoch) + if args.local_rank == 0: + with open(cfg.test_log, 'a+') as f: + f.write('Epoch {0}: Acc is {1:.4f}\n'.format(epoch, total_acc)) + torch.save(obj=model.state_dict(), + f=os.path.join(cfg.snapshot_dir, 'ep{}_acc{:.4f}.pth'.format(epoch, total_acc))) + logger.info('Model saved') + + +def train_one_epoch(train_loader, val_loader, model, criterions, optimizer, epoch, cfg): + model.train() + + ramp_up_w, ramp_down_w = ramp_up(epoch, cfg.ramp_a), ramp_down(epoch, cfg.ramp_a) + w, gamma = cfg.w, cfg.gamma + if args.local_rank == 0: + logger.info('ramp_up_w: {:.4f}, ramp_down_w: {:.4f}, w: {}, gamma: {}'.format(ramp_up_w, ramp_down_w, w, gamma)) + adjust_learning_rate(optimizer, cfg.get_lr(epoch)[0], cfg.get_lr(epoch)[1]) + training_phase = cfg.train_mode + + for idx, batch in enumerate(train_loader, start=1): + img, label = batch + + good_batch = check_batch(label, cfg.num_classes) + good_batch = torch.cuda.IntTensor([good_batch]) + dist.all_reduce(good_batch, op=dist.ReduceOp.SUM) + good_batch = good_batch // args.world_size + good_batch = bool(good_batch.item()) + if not good_batch: + continue + img, label = img.cuda(), label.cuda() + + x_final, output_x_list, targets, softlabel_list, G_matrixs, G_main_matrixs, score, atten_x_final = model(img, label, training_phase) + + softlabel = torch.zeros_like(x_final) + for c in range(cfg.num_classes): + # sharpen + px = torch.softmax(softlabel_list[c], dim=1) + ptx = px ** (1 / cfg.T) # temparature sharpening + targets_x = ptx / ptx.sum(dim=1, keepdim=True) # normalize + targets_x = targets_x.detach() + + _h, _w = px.shape + _targets = torch.zeros([_h, _w+1]).float() + _targets[:, 0:c], _targets[:, c + 1:] = targets_x[:, 0:c], targets_x[:, c:] + ind = (label == c).nonzero() + softlabel[ind[:, 0]] = _targets.cuda() # bs x num_cls + + softLoss = criterions[1](x_final, label, softlabel) + aux_loss = [criterions[0](_pred, _label) for _pred, _label in zip(output_x_list, targets)] # auxiliary branch + aux_loss = sum(aux_loss) / (cfg.num_branches-1) + CEloss = criterions[0](atten_x_final, label) # main + spLoss = criterions[2](G_matrixs, G_main_matrixs) + + loss = ramp_up_w * CEloss + ramp_up_w * (w * softLoss + gamma * spLoss) + ramp_down_w * aux_loss + + if idx % 20 == 0: + if args.distributed: + CEloss = reduce_tensor(CEloss) + softLoss = reduce_tensor(softLoss) + spLoss = reduce_tensor(spLoss) + aux_loss = reduce_tensor(aux_loss) + if args.local_rank == 0: + logger.info('Epoch {} Batch {}/{}: CEloss {:.6f}, softLoss {:.6f}, spLoss {:.6f}, aux_loss {:.6f}'.format( + epoch, idx, len(train_loader), CEloss, softLoss, spLoss, aux_loss)) + + # compute gradient and do SGD step + optimizer.zero_grad() + loss.backward() + optimizer.step() + + if idx % 200 == 0: + total_acc = test(cfg, val_loader, model, epoch, idx) + if args.local_rank == 0: + with open(cfg.test_log, 'a+') as f: + f.write('Epoch {0} Batch {1}: Acc is {2:.4f}\n'.format(epoch, idx, total_acc)) + torch.save(obj=model.state_dict(), + f=os.path.join(cfg.snapshot_dir, 'ep{}_b{}_acc{:.4f}.pth'.format(epoch, idx, total_acc))) + logger.info('Model saved') + + gc.collect() + + +def test_worker(val_loader, model): + Pred, Label = [], [] + for idx, batch in enumerate(val_loader, 1): + img, label = batch + pred, output_x_list, targets, softlabel_list = model(img.cuda(), label, 'normal') + + for i in range(pred.shape[0]): + x = pred[i].data.cpu().numpy() + y = label[i].item() + Pred.append(x) + Label.append(y) + + return Pred, Label + + +def test(cfg, val_loader, model, epoch, batch=None): + model.eval() + + pred, label = test_worker(val_loader, model) + pred = [np.where(x == np.max(x))[0][0] for x in pred] + + total = len([x for x in range(len(pred)) if pred[x] == label[x]]) + total_acc = total / len(pred) + if args.distributed: + total_acc = reduce_float(total_acc) + if args.local_rank == 0: + logger.info('Epoch {}: Acc is {:.4f}'.format(epoch, total_acc)) if batch is None else \ + logger.info('Epoch {} Batch {}: Acc is {:.4f}'.format(epoch, batch, total_acc)) + + return total_acc + + +def adjust_learning_rate(optimizer, lr1, lr2): + optimizer.param_groups[0]['lr'] = lr1 + optimizer.param_groups[1]['lr'] = lr2 + + +def accuracy(output, target, topk=(1,)): + """Computes the precision@k for the specified values of k""" + maxk = max(topk) + batch_size = target.size(0) + + _, pred = output.topk(maxk, 1, True, True) + pred = pred.t() + correct = pred.eq(target.view(1, -1).expand_as(pred)) + + res = [] + for k in topk: + correct_k = correct[:k].view(-1).float().sum(0, keepdim=True) + res.append(correct_k.mul_(100.0 / batch_size)) + return res + + +def check_batch(label, num_classes): + for i in range(num_classes): + cnt = (label == i).nonzero().shape[0] + if cnt < 2: + return False + return True + + +def reduce_tensor(tensor): + rt = tensor.clone() + dist.all_reduce(rt, op=dist.ReduceOp.SUM) + rt /= args.world_size + return rt + + +def reduce_float(f): + tensor = torch.cuda.FloatTensor([f]) + rt = tensor.clone() + dist.all_reduce(rt, op=dist.ReduceOp.SUM) + rt /= args.world_size + return rt.item() + + +if __name__ == '__main__': + train() diff --git a/addition_module/DMUE/train_ddp.sh b/src/faceDetection/face_sdk/addition_module/DMUE/train_ddp.sh similarity index 100% rename from addition_module/DMUE/train_ddp.sh rename to src/faceDetection/face_sdk/addition_module/DMUE/train_ddp.sh diff --git a/addition_module/DMUE/utils/__init__.py b/src/faceDetection/face_sdk/addition_module/DMUE/utils/__init__.py similarity index 96% rename from addition_module/DMUE/utils/__init__.py rename to src/faceDetection/face_sdk/addition_module/DMUE/utils/__init__.py index d5bb509..dd95341 100644 --- a/addition_module/DMUE/utils/__init__.py +++ b/src/faceDetection/face_sdk/addition_module/DMUE/utils/__init__.py @@ -1,3 +1,3 @@ -from .easylog import write_config_into_log -from .ramp import ramp_up, ramp_down - +from .easylog import write_config_into_log +from .ramp import ramp_up, ramp_down + diff --git a/addition_module/DMUE/utils/easylog.py b/src/faceDetection/face_sdk/addition_module/DMUE/utils/easylog.py similarity index 97% rename from addition_module/DMUE/utils/easylog.py rename to src/faceDetection/face_sdk/addition_module/DMUE/utils/easylog.py index 47db4b4..37b73bf 100644 --- a/addition_module/DMUE/utils/easylog.py +++ b/src/faceDetection/face_sdk/addition_module/DMUE/utils/easylog.py @@ -1,13 +1,13 @@ -import os - -def write_config_into_log(cfg): - attrs = dir(cfg) - no_print_attributes = ['ckpt_root_dir', 'default_seed', 'get_lr', - 'tb_dump_interval', 'iter_per_epoch', 'program_name', - 'this_model_dir', 'train_data_num_thread', 'train_dp_name', - 'workspace_limit', 'train_dataset', 'val_dataset'] - for attr in attrs: - if not (('__' in attr) or (attr in no_print_attributes)): - with open( os.path.join(cfg.output_dir, 'config.txt'), 'a+') as f: - f.write('{0},{1}\n'.format(attr, getattr(cfg, attr))) - +import os + +def write_config_into_log(cfg): + attrs = dir(cfg) + no_print_attributes = ['ckpt_root_dir', 'default_seed', 'get_lr', + 'tb_dump_interval', 'iter_per_epoch', 'program_name', + 'this_model_dir', 'train_data_num_thread', 'train_dp_name', + 'workspace_limit', 'train_dataset', 'val_dataset'] + for attr in attrs: + if not (('__' in attr) or (attr in no_print_attributes)): + with open( os.path.join(cfg.output_dir, 'config.txt'), 'a+') as f: + f.write('{0},{1}\n'.format(attr, getattr(cfg, attr))) + diff --git a/addition_module/DMUE/utils/ramp.py b/src/faceDetection/face_sdk/addition_module/DMUE/utils/ramp.py similarity index 96% rename from addition_module/DMUE/utils/ramp.py rename to src/faceDetection/face_sdk/addition_module/DMUE/utils/ramp.py index 08838cc..35c65b4 100644 --- a/addition_module/DMUE/utils/ramp.py +++ b/src/faceDetection/face_sdk/addition_module/DMUE/utils/ramp.py @@ -1,15 +1,15 @@ -import math - -def ramp_up(epoch, alpha, lamda=1): - if epoch > alpha: - return lamda - else: - w = lamda * math.exp(-5*math.pow((1-epoch/alpha), 2)) - return w - -def ramp_down(epoch, alpha, lamda=1): - if epoch < alpha: - return lamda - else: - w = lamda * math.exp(-1*math.pow((1-alpha/epoch), 2)) +import math + +def ramp_up(epoch, alpha, lamda=1): + if epoch > alpha: + return lamda + else: + w = lamda * math.exp(-5*math.pow((1-epoch/alpha), 2)) + return w + +def ramp_down(epoch, alpha, lamda=1): + if epoch < alpha: + return lamda + else: + w = lamda * math.exp(-1*math.pow((1-alpha/epoch), 2)) return w \ No newline at end of file diff --git a/addition_module/DSDG/.gitignore b/src/faceDetection/face_sdk/addition_module/DSDG/.gitignore similarity index 100% rename from addition_module/DSDG/.gitignore rename to src/faceDetection/face_sdk/addition_module/DSDG/.gitignore diff --git a/addition_module/DSDG/DUM/Load_OULUNPUcrop_train.py b/src/faceDetection/face_sdk/addition_module/DSDG/DUM/Load_OULUNPUcrop_train.py similarity index 100% rename from addition_module/DSDG/DUM/Load_OULUNPUcrop_train.py rename to src/faceDetection/face_sdk/addition_module/DSDG/DUM/Load_OULUNPUcrop_train.py diff --git a/addition_module/DSDG/DUM/Load_OULUNPUcrop_valtest.py b/src/faceDetection/face_sdk/addition_module/DSDG/DUM/Load_OULUNPUcrop_valtest.py similarity index 100% rename from addition_module/DSDG/DUM/Load_OULUNPUcrop_valtest.py rename to src/faceDetection/face_sdk/addition_module/DSDG/DUM/Load_OULUNPUcrop_valtest.py diff --git a/src/faceDetection/face_sdk/addition_module/DSDG/DUM/__init__.py b/src/faceDetection/face_sdk/addition_module/DSDG/DUM/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/addition_module/DSDG/DUM/checkpoint/CDCN_U_P1.pkl b/src/faceDetection/face_sdk/addition_module/DSDG/DUM/checkpoint/CDCN_U_P1.pkl similarity index 100% rename from addition_module/DSDG/DUM/checkpoint/CDCN_U_P1.pkl rename to src/faceDetection/face_sdk/addition_module/DSDG/DUM/checkpoint/CDCN_U_P1.pkl diff --git a/src/faceDetection/face_sdk/addition_module/DSDG/DUM/make_dataset/__init__.py b/src/faceDetection/face_sdk/addition_module/DSDG/DUM/make_dataset/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/addition_module/DSDG/DUM/make_dataset/crop_dataset.py b/src/faceDetection/face_sdk/addition_module/DSDG/DUM/make_dataset/crop_dataset.py similarity index 96% rename from addition_module/DSDG/DUM/make_dataset/crop_dataset.py rename to src/faceDetection/face_sdk/addition_module/DSDG/DUM/make_dataset/crop_dataset.py index 5edc22f..2e7a8bf 100644 --- a/addition_module/DSDG/DUM/make_dataset/crop_dataset.py +++ b/src/faceDetection/face_sdk/addition_module/DSDG/DUM/make_dataset/crop_dataset.py @@ -1,116 +1,116 @@ -import math -import os -import cv2 -import numpy as np - - -root_dir = '/export2/home/wht/oulu_images_crop/' - -img_root = '/export2/home/wht/oulu_images/train_img_flod/' -map_root = '/export2/home/wht/oulu_images/train_depth_flod/' -bbox_root = '/export2/home/wht/oulu_images/train_bbox_flod/' - - -def crop_face_from_scene(image, face_name_full, scale): - f = open(face_name_full, 'r') - lines = f.readlines() - lines = lines[0].split(' ') - y1, x1, w, h = [int(ele) for ele in lines[:4]] - f.close() - y2 = y1 + w - x2 = x1 + h - - y_mid = (y1 + y2) / 2.0 - x_mid = (x1 + x2) / 2.0 - h_img, w_img = image.shape[0], image.shape[1] - # w_img,h_img=image.size - w_scale = scale * w - h_scale = scale * h - y1 = y_mid - w_scale / 2.0 - x1 = x_mid - h_scale / 2.0 - y2 = y_mid + w_scale / 2.0 - x2 = x_mid + h_scale / 2.0 - y1 = max(math.floor(y1), 0) - x1 = max(math.floor(x1), 0) - y2 = min(math.floor(y2), w_img) - x2 = min(math.floor(x2), h_img) - - # region=image[y1:y2,x1:x2] - region = image[x1:x2, y1:y2] - return region - - -def crop_face_from_scene_prnet(image, face_name_full, scale): - h_img, w_img = image.shape[0], image.shape[1] - f = open(face_name_full, 'r') - lines = f.readlines() - lines = lines[0].split(' ') - l, r, t, b = [int(ele) for ele in lines[:4]] - if l < 0: - l = 0 - if r > w_img: - r = w_img - if t < 0: - t = 0 - if b > h_img: - b = h_img - y1 = l - x1 = t - w = r - l - h = b - t - f.close() - y2 = y1 + w - x2 = x1 + h - - y_mid = (y1 + y2) / 2.0 - x_mid = (x1 + x2) / 2.0 - # w_img,h_img=image.size - w_scale = scale * w - h_scale = scale * h - y1 = y_mid - w_scale / 2.0 - x1 = x_mid - h_scale / 2.0 - y2 = y_mid + w_scale / 2.0 - x2 = x_mid + h_scale / 2.0 - y1 = max(math.floor(y1), 0) - x1 = max(math.floor(x1), 0) - y2 = min(math.floor(y2), w_img) - x2 = min(math.floor(x2), h_img) - - # region=image[y1:y2,x1:x2] - region = image[x1:x2, y1:y2] - return region - - -vedio_list = os.listdir(bbox_root) -for i, vedio_name in enumerate(vedio_list): - print(i) - bbox_list = os.listdir(os.path.join(bbox_root, vedio_name)) - for bbox_name in bbox_list: - - face_scale = np.random.randint(12, 15) - face_scale = face_scale / 10.0 - - # face_scale = 1.3 - - bbox_path = os.path.join(bbox_root, vedio_name, bbox_name) - - img_path = os.path.join(img_root, vedio_name, bbox_name[:-4] + '.jpg') - - img = cv2.imread(img_path) - img_crop = cv2.resize(crop_face_from_scene_prnet(img, bbox_path, face_scale), (256, 256)) - - img_crop_path = os.path.join(root_dir, 'train_img_flod') - - if not os.path.exists(os.path.join(img_crop_path, vedio_name)): - os.makedirs(os.path.join(img_crop_path, vedio_name)) - - cv2.imwrite(os.path.join(img_crop_path, vedio_name, bbox_name[:-4] + '.jpg'), img_crop) - - map_path = os.path.join(map_root, vedio_name, bbox_name[:-9] + 'depth1D.jpg') - - map = cv2.imread(map_path, 0) - map_crop = cv2.resize(crop_face_from_scene_prnet(map, bbox_path, face_scale), (32, 32)) - map_crop_path = os.path.join(root_dir, 'train_depth_flod') - if not os.path.exists(os.path.join(map_crop_path, vedio_name)): - os.makedirs(os.path.join(map_crop_path, vedio_name)) - cv2.imwrite(os.path.join(map_crop_path, vedio_name, bbox_name[:-9] + 'depth1D.jpg'), map_crop) +import math +import os +import cv2 +import numpy as np + + +root_dir = '/export2/home/wht/oulu_images_crop/' + +img_root = '/export2/home/wht/oulu_images/train_img_flod/' +map_root = '/export2/home/wht/oulu_images/train_depth_flod/' +bbox_root = '/export2/home/wht/oulu_images/train_bbox_flod/' + + +def crop_face_from_scene(image, face_name_full, scale): + f = open(face_name_full, 'r') + lines = f.readlines() + lines = lines[0].split(' ') + y1, x1, w, h = [int(ele) for ele in lines[:4]] + f.close() + y2 = y1 + w + x2 = x1 + h + + y_mid = (y1 + y2) / 2.0 + x_mid = (x1 + x2) / 2.0 + h_img, w_img = image.shape[0], image.shape[1] + # w_img,h_img=image.size + w_scale = scale * w + h_scale = scale * h + y1 = y_mid - w_scale / 2.0 + x1 = x_mid - h_scale / 2.0 + y2 = y_mid + w_scale / 2.0 + x2 = x_mid + h_scale / 2.0 + y1 = max(math.floor(y1), 0) + x1 = max(math.floor(x1), 0) + y2 = min(math.floor(y2), w_img) + x2 = min(math.floor(x2), h_img) + + # region=image[y1:y2,x1:x2] + region = image[x1:x2, y1:y2] + return region + + +def crop_face_from_scene_prnet(image, face_name_full, scale): + h_img, w_img = image.shape[0], image.shape[1] + f = open(face_name_full, 'r') + lines = f.readlines() + lines = lines[0].split(' ') + l, r, t, b = [int(ele) for ele in lines[:4]] + if l < 0: + l = 0 + if r > w_img: + r = w_img + if t < 0: + t = 0 + if b > h_img: + b = h_img + y1 = l + x1 = t + w = r - l + h = b - t + f.close() + y2 = y1 + w + x2 = x1 + h + + y_mid = (y1 + y2) / 2.0 + x_mid = (x1 + x2) / 2.0 + # w_img,h_img=image.size + w_scale = scale * w + h_scale = scale * h + y1 = y_mid - w_scale / 2.0 + x1 = x_mid - h_scale / 2.0 + y2 = y_mid + w_scale / 2.0 + x2 = x_mid + h_scale / 2.0 + y1 = max(math.floor(y1), 0) + x1 = max(math.floor(x1), 0) + y2 = min(math.floor(y2), w_img) + x2 = min(math.floor(x2), h_img) + + # region=image[y1:y2,x1:x2] + region = image[x1:x2, y1:y2] + return region + + +vedio_list = os.listdir(bbox_root) +for i, vedio_name in enumerate(vedio_list): + print(i) + bbox_list = os.listdir(os.path.join(bbox_root, vedio_name)) + for bbox_name in bbox_list: + + face_scale = np.random.randint(12, 15) + face_scale = face_scale / 10.0 + + # face_scale = 1.3 + + bbox_path = os.path.join(bbox_root, vedio_name, bbox_name) + + img_path = os.path.join(img_root, vedio_name, bbox_name[:-4] + '.jpg') + + img = cv2.imread(img_path) + img_crop = cv2.resize(crop_face_from_scene_prnet(img, bbox_path, face_scale), (256, 256)) + + img_crop_path = os.path.join(root_dir, 'train_img_flod') + + if not os.path.exists(os.path.join(img_crop_path, vedio_name)): + os.makedirs(os.path.join(img_crop_path, vedio_name)) + + cv2.imwrite(os.path.join(img_crop_path, vedio_name, bbox_name[:-4] + '.jpg'), img_crop) + + map_path = os.path.join(map_root, vedio_name, bbox_name[:-9] + 'depth1D.jpg') + + map = cv2.imread(map_path, 0) + map_crop = cv2.resize(crop_face_from_scene_prnet(map, bbox_path, face_scale), (32, 32)) + map_crop_path = os.path.join(root_dir, 'train_depth_flod') + if not os.path.exists(os.path.join(map_crop_path, vedio_name)): + os.makedirs(os.path.join(map_crop_path, vedio_name)) + cv2.imwrite(os.path.join(map_crop_path, vedio_name, bbox_name[:-9] + 'depth1D.jpg'), map_crop) diff --git a/addition_module/DSDG/DUM/models/CDCNs_u.py b/src/faceDetection/face_sdk/addition_module/DSDG/DUM/models/CDCNs_u.py similarity index 100% rename from addition_module/DSDG/DUM/models/CDCNs_u.py rename to src/faceDetection/face_sdk/addition_module/DSDG/DUM/models/CDCNs_u.py diff --git a/addition_module/DSDG/DUM/models/ResNet_u.py b/src/faceDetection/face_sdk/addition_module/DSDG/DUM/models/ResNet_u.py similarity index 97% rename from addition_module/DSDG/DUM/models/ResNet_u.py rename to src/faceDetection/face_sdk/addition_module/DSDG/DUM/models/ResNet_u.py index 4b133b4..5354565 100644 --- a/addition_module/DSDG/DUM/models/ResNet_u.py +++ b/src/faceDetection/face_sdk/addition_module/DSDG/DUM/models/ResNet_u.py @@ -1,114 +1,114 @@ -import math - -import torch -import torch.nn.functional as F -import torch.utils.model_zoo as model_zoo -from torch import nn -from torch.nn import Parameter -import pdb -import numpy as np - - -class ResidualBlock(nn.Module): - def __init__(self, inchannel, outchannel, stride=1): - super(ResidualBlock, self).__init__() - self.left = nn.Sequential( - nn.Conv2d(inchannel, outchannel, kernel_size=3, stride=stride, padding=1, bias=False), - nn.BatchNorm2d(outchannel), - nn.ReLU(inplace=True), - nn.Conv2d(outchannel, outchannel, kernel_size=3, stride=1, padding=1, bias=False), - nn.BatchNorm2d(outchannel) - ) - self.shortcut = nn.Sequential() - if stride != 1 or inchannel != outchannel: - self.shortcut = nn.Sequential( - nn.Conv2d(inchannel, outchannel, kernel_size=1, stride=stride, bias=False), - nn.BatchNorm2d(outchannel) - ) - - def forward(self, x): - out = self.left(x) - out += self.shortcut(x) - out = F.relu(out) - return out - - -class ResNet(nn.Module): - def __init__(self, ResidualBlock = ResidualBlock): - super(ResNet, self).__init__() - self.inchannel = 64 - self.conv1 = nn.Sequential( - nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False), - nn.BatchNorm2d(64), - nn.ReLU(), - ) - self.layer1 = self.make_layer(ResidualBlock, 64, 2, stride=1) - self.layer2 = self.make_layer(ResidualBlock, 128, 2, stride=2) - self.layer3 = self.make_layer(ResidualBlock, 196, 2, stride=2) - self.layer4 = self.make_layer(ResidualBlock, 256, 2, stride=2) - self.lastconv1 = nn.Sequential( - nn.Conv2d(128 + 196 + 256, 128, kernel_size=3, stride=1, padding=1, bias=False), - nn.BatchNorm2d(128), - nn.ReLU(), - ) - self.lastconv2 = nn.Sequential( - nn.Conv2d(128, 64, kernel_size=3, stride=1, padding=1, bias=False), - nn.BatchNorm2d(64), - nn.ReLU(), - ) - self.lastconv3 = nn.Sequential( - nn.Conv2d(64, 1, kernel_size=3, stride=1, padding=1, bias=False), - nn.ReLU(), - ) - self.downsample32x32 = nn.Upsample(size=(32, 32), mode='bilinear') - - self.mu_head = nn.Sequential( - nn.Conv2d(64, 1, kernel_size=3, stride=1, padding=1, bias=False), - - ) - self.logvar_head = nn.Sequential( - nn.Conv2d(64, 1, kernel_size=3, stride=1, padding=1, bias=False), - - ) - - def make_layer(self, block, channels, num_blocks, stride): - strides = [stride] + [1] * (num_blocks - 1) #strides=[1,1] - layers = [] - for stride in strides: - layers.append(block(self.inchannel, channels, stride)) - self.inchannel = channels - return nn.Sequential(*layers) - - def _reparameterize(self, mu, logvar): - std = torch.exp(logvar).sqrt() - epsilon = torch.randn_like(std) - return mu + epsilon * std - - def forward(self, x): - x_input = x - x = self.conv1(x) - x_block1 = self.layer1(x) - x_block2 = self.layer2(x_block1) - x_block2_32 = self.downsample32x32(x_block2) - x_block3 = self.layer3(x_block2) - x_block3_32 = self.downsample32x32(x_block3) - x_block4 = self.layer4(x_block3) - x_block4_32 = self.downsample32x32(x_block4) - - x_concat = torch.cat((x_block2_32, x_block3_32, x_block4_32), dim=1) - - x = self.lastconv1(x_concat) - x = self.lastconv2(x) - mu = self.mu_head(x) - mu = mu.squeeze(1) - logvar = self.logvar_head(x) - logvar = logvar.squeeze(1) - embedding = self._reparameterize(mu, logvar) - - return mu, logvar, embedding, x_concat, x_block2, x_block3, x_block4, x_input - - -def ResNet18_u(): - - return ResNet(ResidualBlock) - +import math + +import torch +import torch.nn.functional as F +import torch.utils.model_zoo as model_zoo +from torch import nn +from torch.nn import Parameter +import pdb +import numpy as np + + +class ResidualBlock(nn.Module): + def __init__(self, inchannel, outchannel, stride=1): + super(ResidualBlock, self).__init__() + self.left = nn.Sequential( + nn.Conv2d(inchannel, outchannel, kernel_size=3, stride=stride, padding=1, bias=False), + nn.BatchNorm2d(outchannel), + nn.ReLU(inplace=True), + nn.Conv2d(outchannel, outchannel, kernel_size=3, stride=1, padding=1, bias=False), + nn.BatchNorm2d(outchannel) + ) + self.shortcut = nn.Sequential() + if stride != 1 or inchannel != outchannel: + self.shortcut = nn.Sequential( + nn.Conv2d(inchannel, outchannel, kernel_size=1, stride=stride, bias=False), + nn.BatchNorm2d(outchannel) + ) + + def forward(self, x): + out = self.left(x) + out += self.shortcut(x) + out = F.relu(out) + return out + + +class ResNet(nn.Module): + def __init__(self, ResidualBlock = ResidualBlock): + super(ResNet, self).__init__() + self.inchannel = 64 + self.conv1 = nn.Sequential( + nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False), + nn.BatchNorm2d(64), + nn.ReLU(), + ) + self.layer1 = self.make_layer(ResidualBlock, 64, 2, stride=1) + self.layer2 = self.make_layer(ResidualBlock, 128, 2, stride=2) + self.layer3 = self.make_layer(ResidualBlock, 196, 2, stride=2) + self.layer4 = self.make_layer(ResidualBlock, 256, 2, stride=2) + self.lastconv1 = nn.Sequential( + nn.Conv2d(128 + 196 + 256, 128, kernel_size=3, stride=1, padding=1, bias=False), + nn.BatchNorm2d(128), + nn.ReLU(), + ) + self.lastconv2 = nn.Sequential( + nn.Conv2d(128, 64, kernel_size=3, stride=1, padding=1, bias=False), + nn.BatchNorm2d(64), + nn.ReLU(), + ) + self.lastconv3 = nn.Sequential( + nn.Conv2d(64, 1, kernel_size=3, stride=1, padding=1, bias=False), + nn.ReLU(), + ) + self.downsample32x32 = nn.Upsample(size=(32, 32), mode='bilinear') + + self.mu_head = nn.Sequential( + nn.Conv2d(64, 1, kernel_size=3, stride=1, padding=1, bias=False), + + ) + self.logvar_head = nn.Sequential( + nn.Conv2d(64, 1, kernel_size=3, stride=1, padding=1, bias=False), + + ) + + def make_layer(self, block, channels, num_blocks, stride): + strides = [stride] + [1] * (num_blocks - 1) #strides=[1,1] + layers = [] + for stride in strides: + layers.append(block(self.inchannel, channels, stride)) + self.inchannel = channels + return nn.Sequential(*layers) + + def _reparameterize(self, mu, logvar): + std = torch.exp(logvar).sqrt() + epsilon = torch.randn_like(std) + return mu + epsilon * std + + def forward(self, x): + x_input = x + x = self.conv1(x) + x_block1 = self.layer1(x) + x_block2 = self.layer2(x_block1) + x_block2_32 = self.downsample32x32(x_block2) + x_block3 = self.layer3(x_block2) + x_block3_32 = self.downsample32x32(x_block3) + x_block4 = self.layer4(x_block3) + x_block4_32 = self.downsample32x32(x_block4) + + x_concat = torch.cat((x_block2_32, x_block3_32, x_block4_32), dim=1) + + x = self.lastconv1(x_concat) + x = self.lastconv2(x) + mu = self.mu_head(x) + mu = mu.squeeze(1) + logvar = self.logvar_head(x) + logvar = logvar.squeeze(1) + embedding = self._reparameterize(mu, logvar) + + return mu, logvar, embedding, x_concat, x_block2, x_block3, x_block4, x_input + + +def ResNet18_u(): + + return ResNet(ResidualBlock) + diff --git a/src/faceDetection/face_sdk/addition_module/DSDG/DUM/models/__init__.py b/src/faceDetection/face_sdk/addition_module/DSDG/DUM/models/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/addition_module/DSDG/DUM/test.py b/src/faceDetection/face_sdk/addition_module/DSDG/DUM/test.py similarity index 97% rename from addition_module/DSDG/DUM/test.py rename to src/faceDetection/face_sdk/addition_module/DSDG/DUM/test.py index 7beb098..fb94a21 100644 --- a/addition_module/DSDG/DUM/test.py +++ b/src/faceDetection/face_sdk/addition_module/DSDG/DUM/test.py @@ -1,168 +1,168 @@ -from __future__ import print_function, division -import torch - -torch.set_printoptions(profile="full") -import matplotlib.pyplot as plt -import argparse, os -import numpy as np -import shutil -from torch.utils.data import DataLoader -from torchvision import transforms - -from models.CDCNs_u import Conv2d_cd, CDCN_u - -from Load_OULUNPUcrop_valtest import Spoofing_valtest, Normaliztion_valtest, ToTensor_valtest - -import torch.optim as optim - -from utils import performances - -# Dataset root - -val_image_dir = '/export2/home/wht/oulu_img_crop/dev_file_flod/' -test_image_dir = '/export2/home/wht/oulu_img_crop/test_file_flod/' - -val_map_dir = '/export2/home/wht/oulu_img_crop/dev_depth_flod/' -test_map_dir = '/export2/home/wht/oulu_img_crop/test_depth_flod/' - -val_list = '/export2/home/wht/oulu_img_crop/protocols/Protocol_1/Dev.txt' -test_list = '/export2/home/wht/oulu_img_crop/protocols/Protocol_1/Test.txt' - - -# main function -def test(): - # GPU & log file --> if use DataParallel, please comment this command - os.environ["CUDA_VISIBLE_DEVICES"] = '0, 1, 2, 3' - - isExists = os.path.exists(args.log) - if not isExists: - os.makedirs(args.log) - log_file = open(args.log + '/' + args.log + '_log_P1.txt', 'w') - - log_file.write('Oulu-NPU, P1:\n ') - log_file.flush() - - print('test!\n') - log_file.write('test!\n') - log_file.flush() - - model = CDCN_u(basic_conv=Conv2d_cd, theta=0.7) - # model = ResNet18_u() - - model = model.cuda() - model = torch.nn.DataParallel(model) - model.load_state_dict(torch.load('./DUM/checkpoint/CDCN_U_P1.pkl', map_location='cuda:0')) - - print(model) - - optimizer = optim.Adam(model.parameters(), lr=0.001, weight_decay=0.00005) - - for epoch in range(args.epochs): - - model.eval() - - with torch.no_grad(): - ########################################### - ''' val ''' - ########################################### - # val for threshold - val_data = Spoofing_valtest(val_list, val_image_dir, val_map_dir, - transform=transforms.Compose([Normaliztion_valtest(), ToTensor_valtest()])) - dataloader_val = DataLoader(val_data, batch_size=1, shuffle=False, num_workers=4) - - map_score_list = [] - - for i, sample_batched in enumerate(dataloader_val): - # get the inputs - inputs, spoof_label = sample_batched['image_x'].cuda(), sample_batched['spoofing_label'].cuda() - val_maps = sample_batched['val_map_x'].cuda() # binary map from PRNet - - optimizer.zero_grad() - - # pdb.set_trace() - map_score = 0.0 - for frame_t in range(inputs.shape[1]): - mu, logvar, map_x, x_concat, x_Block1, x_Block2, x_Block3, x_input = model( - inputs[:, frame_t, :, :, :]) - score_norm = torch.sum(mu) / torch.sum(val_maps[:, frame_t, :, :]) - map_score += score_norm - - map_score = map_score / inputs.shape[1] - - map_score_list.append('{} {}\n'.format(map_score, spoof_label[0][0])) - # pdb.set_trace() - map_score_val_filename = args.log + '/' + args.log + '_map_score_val.txt' - with open(map_score_val_filename, 'w') as file: - file.writelines(map_score_list) - - ########################################### - ''' test ''' - ########################################## - # test for ACC - test_data = Spoofing_valtest(test_list, test_image_dir, test_map_dir, - transform=transforms.Compose([Normaliztion_valtest(), ToTensor_valtest()])) - dataloader_test = DataLoader(test_data, batch_size=1, shuffle=False, num_workers=4) - - map_score_list = [] - - for i, sample_batched in enumerate(dataloader_test): - # get the inputs - inputs, spoof_label = sample_batched['image_x'].cuda(), sample_batched['spoofing_label'].cuda() - test_maps = sample_batched['val_map_x'].cuda() - - optimizer.zero_grad() - - # pdb.set_trace() - map_score = 0.0 - for frame_t in range(inputs.shape[1]): - mu, logvar, map_x, x_concat, x_Block1, x_Block2, x_Block3, x_input = model( - inputs[:, frame_t, :, :, :]) - - score_norm = torch.sum(mu) / torch.sum(test_maps[:, frame_t, :, :]) - map_score += score_norm - - map_score = map_score / inputs.shape[1] - - map_score_list.append('{} {}\n'.format(map_score, spoof_label[0][0])) - - map_score_test_filename = args.log + '/' + args.log + '_map_score_test.txt' - with open(map_score_test_filename, 'w') as file: - file.writelines(map_score_list) - - ############################################################# - # performance measurement both val and test - ############################################################# - val_threshold, test_threshold, val_ACC, val_ACER, test_ACC, test_APCER, test_BPCER, test_ACER, test_ACER_test_threshold = performances( - map_score_val_filename, map_score_test_filename) - - print('epoch:%d, Val: val_threshold= %.4f, val_ACC= %.4f, val_ACER= %.4f' % ( - epoch + 1, val_threshold, val_ACC, val_ACER)) - log_file.write('\n epoch:%d, Val: val_threshold= %.4f, val_ACC= %.4f, val_ACER= %.4f \n' % ( - epoch + 1, val_threshold, val_ACC, val_ACER)) - - print('epoch:%d, Test: ACC= %.4f, APCER= %.4f, BPCER= %.4f, ACER= %.4f' % ( - epoch + 1, test_ACC, test_APCER, test_BPCER, test_ACER)) - log_file.write('epoch:%d, Test: ACC= %.4f, APCER= %.4f, BPCER= %.4f, ACER= %.4f \n' % ( - epoch + 1, test_ACC, test_APCER, test_BPCER, test_ACER)) - log_file.flush() - - print('Finished Training') - log_file.close() - - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="save quality using landmarkpose model") - parser.add_argument('--gpus', type=str, default='0,1,2,3', help='the gpu id used for predict') - parser.add_argument('--lr', type=float, default=0.001, help='initial learning rate') - parser.add_argument('--batchsize', type=int, default=32, help='initial batchsize') - parser.add_argument('--step_size', type=int, default=500, help='how many epochs lr decays once') # 500 - parser.add_argument('--gamma', type=float, default=0.5, help='gamma of optim.lr_scheduler.StepLR, decay of lr') - parser.add_argument('--kl_lambda', type=float, default=0.001, help='') - parser.add_argument('--echo_batches', type=int, default=50, help='how many batches display once') # 50 - parser.add_argument('--epochs', type=int, default=1, help='total training epochs') - parser.add_argument('--log', type=str, default="CDCN_U_P1_test", help='log and save model name') - parser.add_argument('--finetune', action='store_true', default=False, help='whether finetune other models') - parser.add_argument('--test', action='store_true', default=True, help='') - - args = parser.parse_args() - test() +from __future__ import print_function, division +import torch + +torch.set_printoptions(profile="full") +import matplotlib.pyplot as plt +import argparse, os +import numpy as np +import shutil +from torch.utils.data import DataLoader +from torchvision import transforms + +from models.CDCNs_u import Conv2d_cd, CDCN_u + +from Load_OULUNPUcrop_valtest import Spoofing_valtest, Normaliztion_valtest, ToTensor_valtest + +import torch.optim as optim + +from utils import performances + +# Dataset root + +val_image_dir = '/export2/home/wht/oulu_img_crop/dev_file_flod/' +test_image_dir = '/export2/home/wht/oulu_img_crop/test_file_flod/' + +val_map_dir = '/export2/home/wht/oulu_img_crop/dev_depth_flod/' +test_map_dir = '/export2/home/wht/oulu_img_crop/test_depth_flod/' + +val_list = '/export2/home/wht/oulu_img_crop/protocols/Protocol_1/Dev.txt' +test_list = '/export2/home/wht/oulu_img_crop/protocols/Protocol_1/Test.txt' + + +# main function +def test(): + # GPU & log file --> if use DataParallel, please comment this command + os.environ["CUDA_VISIBLE_DEVICES"] = '0, 1, 2, 3' + + isExists = os.path.exists(args.log) + if not isExists: + os.makedirs(args.log) + log_file = open(args.log + '/' + args.log + '_log_P1.txt', 'w') + + log_file.write('Oulu-NPU, P1:\n ') + log_file.flush() + + print('test!\n') + log_file.write('test!\n') + log_file.flush() + + model = CDCN_u(basic_conv=Conv2d_cd, theta=0.7) + # model = ResNet18_u() + + model = model.cuda() + model = torch.nn.DataParallel(model) + model.load_state_dict(torch.load('./DUM/checkpoint/CDCN_U_P1.pkl', map_location='cuda:0')) + + print(model) + + optimizer = optim.Adam(model.parameters(), lr=0.001, weight_decay=0.00005) + + for epoch in range(args.epochs): + + model.eval() + + with torch.no_grad(): + ########################################### + ''' val ''' + ########################################### + # val for threshold + val_data = Spoofing_valtest(val_list, val_image_dir, val_map_dir, + transform=transforms.Compose([Normaliztion_valtest(), ToTensor_valtest()])) + dataloader_val = DataLoader(val_data, batch_size=1, shuffle=False, num_workers=4) + + map_score_list = [] + + for i, sample_batched in enumerate(dataloader_val): + # get the inputs + inputs, spoof_label = sample_batched['image_x'].cuda(), sample_batched['spoofing_label'].cuda() + val_maps = sample_batched['val_map_x'].cuda() # binary map from PRNet + + optimizer.zero_grad() + + # pdb.set_trace() + map_score = 0.0 + for frame_t in range(inputs.shape[1]): + mu, logvar, map_x, x_concat, x_Block1, x_Block2, x_Block3, x_input = model( + inputs[:, frame_t, :, :, :]) + score_norm = torch.sum(mu) / torch.sum(val_maps[:, frame_t, :, :]) + map_score += score_norm + + map_score = map_score / inputs.shape[1] + + map_score_list.append('{} {}\n'.format(map_score, spoof_label[0][0])) + # pdb.set_trace() + map_score_val_filename = args.log + '/' + args.log + '_map_score_val.txt' + with open(map_score_val_filename, 'w') as file: + file.writelines(map_score_list) + + ########################################### + ''' test ''' + ########################################## + # test for ACC + test_data = Spoofing_valtest(test_list, test_image_dir, test_map_dir, + transform=transforms.Compose([Normaliztion_valtest(), ToTensor_valtest()])) + dataloader_test = DataLoader(test_data, batch_size=1, shuffle=False, num_workers=4) + + map_score_list = [] + + for i, sample_batched in enumerate(dataloader_test): + # get the inputs + inputs, spoof_label = sample_batched['image_x'].cuda(), sample_batched['spoofing_label'].cuda() + test_maps = sample_batched['val_map_x'].cuda() + + optimizer.zero_grad() + + # pdb.set_trace() + map_score = 0.0 + for frame_t in range(inputs.shape[1]): + mu, logvar, map_x, x_concat, x_Block1, x_Block2, x_Block3, x_input = model( + inputs[:, frame_t, :, :, :]) + + score_norm = torch.sum(mu) / torch.sum(test_maps[:, frame_t, :, :]) + map_score += score_norm + + map_score = map_score / inputs.shape[1] + + map_score_list.append('{} {}\n'.format(map_score, spoof_label[0][0])) + + map_score_test_filename = args.log + '/' + args.log + '_map_score_test.txt' + with open(map_score_test_filename, 'w') as file: + file.writelines(map_score_list) + + ############################################################# + # performance measurement both val and test + ############################################################# + val_threshold, test_threshold, val_ACC, val_ACER, test_ACC, test_APCER, test_BPCER, test_ACER, test_ACER_test_threshold = performances( + map_score_val_filename, map_score_test_filename) + + print('epoch:%d, Val: val_threshold= %.4f, val_ACC= %.4f, val_ACER= %.4f' % ( + epoch + 1, val_threshold, val_ACC, val_ACER)) + log_file.write('\n epoch:%d, Val: val_threshold= %.4f, val_ACC= %.4f, val_ACER= %.4f \n' % ( + epoch + 1, val_threshold, val_ACC, val_ACER)) + + print('epoch:%d, Test: ACC= %.4f, APCER= %.4f, BPCER= %.4f, ACER= %.4f' % ( + epoch + 1, test_ACC, test_APCER, test_BPCER, test_ACER)) + log_file.write('epoch:%d, Test: ACC= %.4f, APCER= %.4f, BPCER= %.4f, ACER= %.4f \n' % ( + epoch + 1, test_ACC, test_APCER, test_BPCER, test_ACER)) + log_file.flush() + + print('Finished Training') + log_file.close() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="save quality using landmarkpose model") + parser.add_argument('--gpus', type=str, default='0,1,2,3', help='the gpu id used for predict') + parser.add_argument('--lr', type=float, default=0.001, help='initial learning rate') + parser.add_argument('--batchsize', type=int, default=32, help='initial batchsize') + parser.add_argument('--step_size', type=int, default=500, help='how many epochs lr decays once') # 500 + parser.add_argument('--gamma', type=float, default=0.5, help='gamma of optim.lr_scheduler.StepLR, decay of lr') + parser.add_argument('--kl_lambda', type=float, default=0.001, help='') + parser.add_argument('--echo_batches', type=int, default=50, help='how many batches display once') # 50 + parser.add_argument('--epochs', type=int, default=1, help='total training epochs') + parser.add_argument('--log', type=str, default="CDCN_U_P1_test", help='log and save model name') + parser.add_argument('--finetune', action='store_true', default=False, help='whether finetune other models') + parser.add_argument('--test', action='store_true', default=True, help='') + + args = parser.parse_args() + test() diff --git a/addition_module/DSDG/DUM/train.py b/src/faceDetection/face_sdk/addition_module/DSDG/DUM/train.py similarity index 97% rename from addition_module/DSDG/DUM/train.py rename to src/faceDetection/face_sdk/addition_module/DSDG/DUM/train.py index c3c622c..259d953 100644 --- a/addition_module/DSDG/DUM/train.py +++ b/src/faceDetection/face_sdk/addition_module/DSDG/DUM/train.py @@ -1,289 +1,289 @@ -from __future__ import print_function, division -import torch -import matplotlib.pyplot as plt -import argparse, os -import numpy as np -from torch.utils.data import DataLoader -from torchvision import transforms - -from models.CDCNs_u import Conv2d_cd, CDCN_u - -from Load_OULUNPUcrop_train import Spoofing_train_g, SeparateBatchSampler, Normaliztion, ToTensor, \ - RandomHorizontalFlip, Cutout, RandomErasing -from Load_OULUNPUcrop_valtest import Spoofing_valtest, Normaliztion_valtest, ToTensor_valtest - -import torch.nn.functional as F -import torch.nn as nn -import torch.optim as optim - -from utils import AvgrageMeter, performances - -# Dataset root -train_image_dir = '/export2/home/wht/oulu_img_crop/train_file_flod/' -val_image_dir = '/export2/home/wht/oulu_img_crop/dev_file_flod/' -test_image_dir = '/export2/home/wht/oulu_img_crop/test_file_flod/' - -train_map_dir = '/export2/home/wht/oulu_img_crop/train_depth_flod/' -val_map_dir = '/export2/home/wht/oulu_img_crop/dev_depth_flod/' -test_map_dir = '/export2/home/wht/oulu_img_crop/test_depth_flod/' - -train_list = '/export2/home/wht/oulu_img_crop/protocols/Protocol_1/Train_g.txt' -val_list = '/export2/home/wht/oulu_img_crop/protocols/Protocol_1/Dev.txt' -test_list = '/export2/home/wht/oulu_img_crop/protocols/Protocol_1/Test.txt' - - -def contrast_depth_conv(input): - ''' compute contrast depth in both of (out, label) ''' - ''' - input 32x32 - output 8x32x32 - ''' - - kernel_filter_list = [ - [[1, 0, 0], [0, -1, 0], [0, 0, 0]], [[0, 1, 0], [0, -1, 0], [0, 0, 0]], [[0, 0, 1], [0, -1, 0], [0, 0, 0]], - [[0, 0, 0], [1, -1, 0], [0, 0, 0]], [[0, 0, 0], [0, -1, 1], [0, 0, 0]], - [[0, 0, 0], [0, -1, 0], [1, 0, 0]], [[0, 0, 0], [0, -1, 0], [0, 1, 0]], [[0, 0, 0], [0, -1, 0], [0, 0, 1]] - ] - - kernel_filter = np.array(kernel_filter_list, np.float32) - - kernel_filter = torch.from_numpy(kernel_filter.astype(np.float)).float().cuda() - # weights (in_channel, out_channel, kernel, kernel) - kernel_filter = kernel_filter.unsqueeze(dim=1) - - input = input.unsqueeze(dim=1).expand(input.shape[0], 8, input.shape[1], input.shape[2]) - - contrast_depth = F.conv2d(input, weight=kernel_filter, groups=8) - - return contrast_depth - - -class Contrast_depth_loss(nn.Module): - def __init__(self): - super(Contrast_depth_loss, self).__init__() - return - - def forward(self, out, label): - contrast_out = contrast_depth_conv(out) - contrast_label = contrast_depth_conv(label) - - criterion_MSE = nn.MSELoss().cuda() - - loss = criterion_MSE(contrast_out, contrast_label) - - return loss - - -def train_test(): - isExists = os.path.exists(args.log) - if not isExists: - os.makedirs(args.log) - log_file = open(args.log + '/' + args.log + '_log_P1.txt', 'a') - - log_file.write('Oulu-NPU, P1:\n ') - log_file.flush() - - print('train from scratch!\n') - log_file.write('train from scratch!\n') - log_file.write('lr:%.6f, lamda_kl:%.6f , batchsize:%d\n' % (args.lr, args.kl_lambda, args.batchsize)) - log_file.flush() - - model = CDCN_u(basic_conv=Conv2d_cd, theta=0.7) - # model = ResNet18_u() - - model = model.cuda() - model = torch.nn.DataParallel(model) - - lr = args.lr - optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=0.00005) - scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=args.step_size, gamma=args.gamma) - - print(model) - - criterion_absolute_loss = nn.MSELoss().cuda() - criterion_contrastive_loss = Contrast_depth_loss().cuda() - - for epoch in range(args.epochs): - if (epoch + 1) % args.step_size == 0: - lr *= args.gamma - - loss_absolute_real = AvgrageMeter() - loss_absolute_fake = AvgrageMeter() - loss_contra_real = AvgrageMeter() - loss_contra_fake = AvgrageMeter() - loss_kl_real = AvgrageMeter() - loss_kl_fake = AvgrageMeter() - - ########################################### - ''' train ''' - ########################################### - model.train() - - # load random 16-frame clip data every epoch - train_data = Spoofing_train_g(train_list, train_image_dir, train_map_dir, - transform=transforms.Compose( - [RandomErasing(), RandomHorizontalFlip(), ToTensor(), Cutout(), - Normaliztion()])) - train_real_idx, train_fake_idx = train_data.get_idx() - batch_sampler = SeparateBatchSampler(train_real_idx, train_fake_idx, batch_size=args.batchsize, ratio=args.ratio) - dataloader_train = DataLoader(train_data, num_workers=8, batch_sampler=batch_sampler) - - for i, sample_batched in enumerate(dataloader_train): - # get the inputs - inputs, map_label, spoof_label = sample_batched['image_x'].cuda(), sample_batched['map_x'].cuda(), \ - sample_batched['spoofing_label'].cuda() - - optimizer.zero_grad() - - # forward + backward + optimize - mu, logvar, map_x, x_concat, x_Block1, x_Block2, x_Block3, x_input = model(inputs) - - mu_real = mu[:int(args.batchsize * args.ratio), :, :] - logvar_real = logvar[:int(args.batchsize * args.ratio), :, :] - map_x_real = map_x[:int(args.batchsize * args.ratio), :, :] - map_label_real = map_label[:int(args.batchsize * args.ratio), :, :] - - absolute_loss_real = criterion_absolute_loss(map_x_real, map_label_real) - contrastive_loss_real = criterion_contrastive_loss(map_x_real, map_label_real) - kl_loss_real = -(1 + logvar_real - (mu_real - map_label_real).pow(2) - logvar_real.exp()) / 2 - kl_loss_real = kl_loss_real.sum(dim=1).sum(dim=1).mean() - kl_loss_real = args.kl_lambda * kl_loss_real - - mu_fake = mu[int(args.batchsize * args.ratio):, :, :] - logvar_fake = logvar[int(args.batchsize * args.ratio):, :, :] - map_x_fake = map_x[int(args.batchsize * args.ratio):, :, :] - map_label_fake = map_label[int(args.batchsize * args.ratio):, :, :] - - absolute_loss_fake = 0.1 * criterion_absolute_loss(map_x_fake, map_label_fake) - contrastive_loss_fake = 0.1 * criterion_contrastive_loss(map_x_fake, map_label_fake) - kl_loss_fake = -(1 + logvar_fake - (mu_fake - map_label_fake).pow(2) - logvar_fake.exp()) / 2 - kl_loss_fake = kl_loss_fake.sum(dim=1).sum(dim=1).mean() - kl_loss_fake = 0.1 * args.kl_lambda * kl_loss_fake - - absolute_loss = absolute_loss_real + absolute_loss_fake - contrastive_loss = contrastive_loss_real + contrastive_loss_fake - kl_loss = kl_loss_real + kl_loss_fake - - loss = absolute_loss + contrastive_loss + kl_loss - - loss.backward() - - optimizer.step() - - n = inputs.size(0) - loss_absolute_real.update(absolute_loss_real.data, n) - loss_absolute_fake.update(absolute_loss_fake.data, n) - loss_contra_real.update(contrastive_loss_real.data, n) - loss_contra_fake.update(contrastive_loss_fake.data, n) - loss_kl_real.update(kl_loss_real.data, n) - loss_kl_fake.update(kl_loss_fake.data, n) - - scheduler.step() - # whole epoch average - print( - 'epoch:%d, Train: Absolute_loss: real=%.4f,fake=%.4f, ' - 'Contrastive_loss: real=%.4f,fake=%.4f, kl_loss: real=%.4f,fake=%.4f' % ( - epoch + 1, loss_absolute_real.avg, loss_absolute_fake.avg, loss_contra_real.avg, loss_contra_fake.avg, - loss_kl_real.avg, loss_kl_fake.avg)) - - # validation/test - if epoch < 200: - epoch_test = 200 - else: - epoch_test = 50 - # epoch_test = 1 - if epoch % epoch_test == epoch_test - 1: - model.eval() - - with torch.no_grad(): - ########################################### - ''' val ''' - ########################################### - # val for threshold - val_data = Spoofing_valtest(val_list, val_image_dir, val_map_dir, - transform=transforms.Compose([Normaliztion_valtest(), ToTensor_valtest()])) - dataloader_val = DataLoader(val_data, batch_size=1, shuffle=False, num_workers=4) - - map_score_list = [] - - for i, sample_batched in enumerate(dataloader_val): - # get the inputs - inputs, spoof_label = sample_batched['image_x'].cuda(), sample_batched['spoofing_label'].cuda() - val_maps = sample_batched['val_map_x'].cuda() # binary map from PRNet - - optimizer.zero_grad() - - mu, logvar, map_x, x_concat, x_Block1, x_Block2, x_Block3, x_input = model(inputs.squeeze(0)) - score_norm = mu.sum(dim=1).sum(dim=1) / val_maps.squeeze(0).sum(dim=1).sum(dim=1) - map_score = score_norm.mean() - map_score_list.append('{} {}\n'.format(map_score, spoof_label[0][0])) - - map_score_val_filename = args.log + '/' + args.log + '_map_score_val.txt' - with open(map_score_val_filename, 'w') as file: - file.writelines(map_score_list) - - ########################################### - ''' test ''' - ########################################## - # test for ACC - test_data = Spoofing_valtest(test_list, test_image_dir, test_map_dir, - transform=transforms.Compose([Normaliztion_valtest(), ToTensor_valtest()])) - dataloader_test = DataLoader(test_data, batch_size=1, shuffle=False, num_workers=4) - - map_score_list = [] - - for i, sample_batched in enumerate(dataloader_test): - # get the inputs - inputs, spoof_label = sample_batched['image_x'].cuda(), sample_batched['spoofing_label'].cuda() - test_maps = sample_batched['val_map_x'].cuda() - - optimizer.zero_grad() - mu, logvar, map_x, x_concat, x_Block1, x_Block2, x_Block3, x_input = model(inputs.squeeze(0)) - score_norm = mu.sum(dim=1).sum(dim=1) / test_maps.squeeze(0).sum(dim=1).sum(dim=1) - map_score = score_norm.mean() - map_score_list.append('{} {}\n'.format(map_score, spoof_label[0][0])) - - map_score_test_filename = args.log + '/' + args.log + '_map_score_test.txt' - with open(map_score_test_filename, 'w') as file: - file.writelines(map_score_list) - - ############################################################# - # performance measurement both val and test - ############################################################# - val_threshold, test_threshold, val_ACC, val_ACER, test_ACC, test_APCER, test_BPCER, test_ACER, test_ACER_test_threshold = performances( - map_score_val_filename, map_score_test_filename) - - print('epoch:%d, Val: val_threshold= %.4f, val_ACC= %.4f, val_ACER= %.4f' % ( - epoch + 1, val_threshold, val_ACC, val_ACER)) - log_file.write('\n epoch:%d, Val: val_threshold= %.4f, val_ACC= %.4f, val_ACER= %.4f \n' % ( - epoch + 1, val_threshold, val_ACC, val_ACER)) - - print('epoch:%d, Test: ACC= %.4f, APCER= %.4f, BPCER= %.4f, ACER= %.4f' % ( - epoch + 1, test_ACC, test_APCER, test_BPCER, test_ACER)) - log_file.write('epoch:%d, Test: ACC= %.4f, APCER= %.4f, BPCER= %.4f, ACER= %.4f \n' % ( - epoch + 1, test_ACC, test_APCER, test_BPCER, test_ACER)) - log_file.flush() - - if epoch % epoch_test == epoch_test - 1: - # save the model until the next improvement - torch.save(model.state_dict(), args.log + '/' + args.log + '_%d.pkl' % (epoch + 1)) - - print('Finished Training') - log_file.close() - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="save quality using landmarkpose model") - parser.add_argument('--gpus', type=str, default='0, 1, 2, 3', help='the gpu id used for predict') - parser.add_argument('--lr', type=float, default=0.0001, help='initial learning rate') - parser.add_argument('--batchsize', type=int, default=64, help='initial batchsize') - parser.add_argument('--step_size', type=int, default=500, help='how many epochs lr decays once') # 500 - parser.add_argument('--gamma', type=float, default=0.5, help='gamma of optim.lr_scheduler.StepLR, decay of lr') - parser.add_argument('--kl_lambda', type=float, default=0.001, help='') - parser.add_argument('--ratio', type=float, default=0.75, help='real and fake in batchsize ') - parser.add_argument('--echo_batches', type=int, default=50, help='how many batches display once') # 50 - parser.add_argument('--epochs', type=int, default=1600, help='total training epochs') - parser.add_argument('--log', type=str, default="CDCN_U_P1", help='log and save model name') - parser.add_argument('--finetune', action='store_true', default=False, help='whether finetune other models') - - args = parser.parse_args() - train_test() +from __future__ import print_function, division +import torch +import matplotlib.pyplot as plt +import argparse, os +import numpy as np +from torch.utils.data import DataLoader +from torchvision import transforms + +from models.CDCNs_u import Conv2d_cd, CDCN_u + +from Load_OULUNPUcrop_train import Spoofing_train_g, SeparateBatchSampler, Normaliztion, ToTensor, \ + RandomHorizontalFlip, Cutout, RandomErasing +from Load_OULUNPUcrop_valtest import Spoofing_valtest, Normaliztion_valtest, ToTensor_valtest + +import torch.nn.functional as F +import torch.nn as nn +import torch.optim as optim + +from utils import AvgrageMeter, performances + +# Dataset root +train_image_dir = '/export2/home/wht/oulu_img_crop/train_file_flod/' +val_image_dir = '/export2/home/wht/oulu_img_crop/dev_file_flod/' +test_image_dir = '/export2/home/wht/oulu_img_crop/test_file_flod/' + +train_map_dir = '/export2/home/wht/oulu_img_crop/train_depth_flod/' +val_map_dir = '/export2/home/wht/oulu_img_crop/dev_depth_flod/' +test_map_dir = '/export2/home/wht/oulu_img_crop/test_depth_flod/' + +train_list = '/export2/home/wht/oulu_img_crop/protocols/Protocol_1/Train_g.txt' +val_list = '/export2/home/wht/oulu_img_crop/protocols/Protocol_1/Dev.txt' +test_list = '/export2/home/wht/oulu_img_crop/protocols/Protocol_1/Test.txt' + + +def contrast_depth_conv(input): + ''' compute contrast depth in both of (out, label) ''' + ''' + input 32x32 + output 8x32x32 + ''' + + kernel_filter_list = [ + [[1, 0, 0], [0, -1, 0], [0, 0, 0]], [[0, 1, 0], [0, -1, 0], [0, 0, 0]], [[0, 0, 1], [0, -1, 0], [0, 0, 0]], + [[0, 0, 0], [1, -1, 0], [0, 0, 0]], [[0, 0, 0], [0, -1, 1], [0, 0, 0]], + [[0, 0, 0], [0, -1, 0], [1, 0, 0]], [[0, 0, 0], [0, -1, 0], [0, 1, 0]], [[0, 0, 0], [0, -1, 0], [0, 0, 1]] + ] + + kernel_filter = np.array(kernel_filter_list, np.float32) + + kernel_filter = torch.from_numpy(kernel_filter.astype(np.float)).float().cuda() + # weights (in_channel, out_channel, kernel, kernel) + kernel_filter = kernel_filter.unsqueeze(dim=1) + + input = input.unsqueeze(dim=1).expand(input.shape[0], 8, input.shape[1], input.shape[2]) + + contrast_depth = F.conv2d(input, weight=kernel_filter, groups=8) + + return contrast_depth + + +class Contrast_depth_loss(nn.Module): + def __init__(self): + super(Contrast_depth_loss, self).__init__() + return + + def forward(self, out, label): + contrast_out = contrast_depth_conv(out) + contrast_label = contrast_depth_conv(label) + + criterion_MSE = nn.MSELoss().cuda() + + loss = criterion_MSE(contrast_out, contrast_label) + + return loss + + +def train_test(): + isExists = os.path.exists(args.log) + if not isExists: + os.makedirs(args.log) + log_file = open(args.log + '/' + args.log + '_log_P1.txt', 'a') + + log_file.write('Oulu-NPU, P1:\n ') + log_file.flush() + + print('train from scratch!\n') + log_file.write('train from scratch!\n') + log_file.write('lr:%.6f, lamda_kl:%.6f , batchsize:%d\n' % (args.lr, args.kl_lambda, args.batchsize)) + log_file.flush() + + model = CDCN_u(basic_conv=Conv2d_cd, theta=0.7) + # model = ResNet18_u() + + model = model.cuda() + model = torch.nn.DataParallel(model) + + lr = args.lr + optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=0.00005) + scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=args.step_size, gamma=args.gamma) + + print(model) + + criterion_absolute_loss = nn.MSELoss().cuda() + criterion_contrastive_loss = Contrast_depth_loss().cuda() + + for epoch in range(args.epochs): + if (epoch + 1) % args.step_size == 0: + lr *= args.gamma + + loss_absolute_real = AvgrageMeter() + loss_absolute_fake = AvgrageMeter() + loss_contra_real = AvgrageMeter() + loss_contra_fake = AvgrageMeter() + loss_kl_real = AvgrageMeter() + loss_kl_fake = AvgrageMeter() + + ########################################### + ''' train ''' + ########################################### + model.train() + + # load random 16-frame clip data every epoch + train_data = Spoofing_train_g(train_list, train_image_dir, train_map_dir, + transform=transforms.Compose( + [RandomErasing(), RandomHorizontalFlip(), ToTensor(), Cutout(), + Normaliztion()])) + train_real_idx, train_fake_idx = train_data.get_idx() + batch_sampler = SeparateBatchSampler(train_real_idx, train_fake_idx, batch_size=args.batchsize, ratio=args.ratio) + dataloader_train = DataLoader(train_data, num_workers=8, batch_sampler=batch_sampler) + + for i, sample_batched in enumerate(dataloader_train): + # get the inputs + inputs, map_label, spoof_label = sample_batched['image_x'].cuda(), sample_batched['map_x'].cuda(), \ + sample_batched['spoofing_label'].cuda() + + optimizer.zero_grad() + + # forward + backward + optimize + mu, logvar, map_x, x_concat, x_Block1, x_Block2, x_Block3, x_input = model(inputs) + + mu_real = mu[:int(args.batchsize * args.ratio), :, :] + logvar_real = logvar[:int(args.batchsize * args.ratio), :, :] + map_x_real = map_x[:int(args.batchsize * args.ratio), :, :] + map_label_real = map_label[:int(args.batchsize * args.ratio), :, :] + + absolute_loss_real = criterion_absolute_loss(map_x_real, map_label_real) + contrastive_loss_real = criterion_contrastive_loss(map_x_real, map_label_real) + kl_loss_real = -(1 + logvar_real - (mu_real - map_label_real).pow(2) - logvar_real.exp()) / 2 + kl_loss_real = kl_loss_real.sum(dim=1).sum(dim=1).mean() + kl_loss_real = args.kl_lambda * kl_loss_real + + mu_fake = mu[int(args.batchsize * args.ratio):, :, :] + logvar_fake = logvar[int(args.batchsize * args.ratio):, :, :] + map_x_fake = map_x[int(args.batchsize * args.ratio):, :, :] + map_label_fake = map_label[int(args.batchsize * args.ratio):, :, :] + + absolute_loss_fake = 0.1 * criterion_absolute_loss(map_x_fake, map_label_fake) + contrastive_loss_fake = 0.1 * criterion_contrastive_loss(map_x_fake, map_label_fake) + kl_loss_fake = -(1 + logvar_fake - (mu_fake - map_label_fake).pow(2) - logvar_fake.exp()) / 2 + kl_loss_fake = kl_loss_fake.sum(dim=1).sum(dim=1).mean() + kl_loss_fake = 0.1 * args.kl_lambda * kl_loss_fake + + absolute_loss = absolute_loss_real + absolute_loss_fake + contrastive_loss = contrastive_loss_real + contrastive_loss_fake + kl_loss = kl_loss_real + kl_loss_fake + + loss = absolute_loss + contrastive_loss + kl_loss + + loss.backward() + + optimizer.step() + + n = inputs.size(0) + loss_absolute_real.update(absolute_loss_real.data, n) + loss_absolute_fake.update(absolute_loss_fake.data, n) + loss_contra_real.update(contrastive_loss_real.data, n) + loss_contra_fake.update(contrastive_loss_fake.data, n) + loss_kl_real.update(kl_loss_real.data, n) + loss_kl_fake.update(kl_loss_fake.data, n) + + scheduler.step() + # whole epoch average + print( + 'epoch:%d, Train: Absolute_loss: real=%.4f,fake=%.4f, ' + 'Contrastive_loss: real=%.4f,fake=%.4f, kl_loss: real=%.4f,fake=%.4f' % ( + epoch + 1, loss_absolute_real.avg, loss_absolute_fake.avg, loss_contra_real.avg, loss_contra_fake.avg, + loss_kl_real.avg, loss_kl_fake.avg)) + + # validation/test + if epoch < 200: + epoch_test = 200 + else: + epoch_test = 50 + # epoch_test = 1 + if epoch % epoch_test == epoch_test - 1: + model.eval() + + with torch.no_grad(): + ########################################### + ''' val ''' + ########################################### + # val for threshold + val_data = Spoofing_valtest(val_list, val_image_dir, val_map_dir, + transform=transforms.Compose([Normaliztion_valtest(), ToTensor_valtest()])) + dataloader_val = DataLoader(val_data, batch_size=1, shuffle=False, num_workers=4) + + map_score_list = [] + + for i, sample_batched in enumerate(dataloader_val): + # get the inputs + inputs, spoof_label = sample_batched['image_x'].cuda(), sample_batched['spoofing_label'].cuda() + val_maps = sample_batched['val_map_x'].cuda() # binary map from PRNet + + optimizer.zero_grad() + + mu, logvar, map_x, x_concat, x_Block1, x_Block2, x_Block3, x_input = model(inputs.squeeze(0)) + score_norm = mu.sum(dim=1).sum(dim=1) / val_maps.squeeze(0).sum(dim=1).sum(dim=1) + map_score = score_norm.mean() + map_score_list.append('{} {}\n'.format(map_score, spoof_label[0][0])) + + map_score_val_filename = args.log + '/' + args.log + '_map_score_val.txt' + with open(map_score_val_filename, 'w') as file: + file.writelines(map_score_list) + + ########################################### + ''' test ''' + ########################################## + # test for ACC + test_data = Spoofing_valtest(test_list, test_image_dir, test_map_dir, + transform=transforms.Compose([Normaliztion_valtest(), ToTensor_valtest()])) + dataloader_test = DataLoader(test_data, batch_size=1, shuffle=False, num_workers=4) + + map_score_list = [] + + for i, sample_batched in enumerate(dataloader_test): + # get the inputs + inputs, spoof_label = sample_batched['image_x'].cuda(), sample_batched['spoofing_label'].cuda() + test_maps = sample_batched['val_map_x'].cuda() + + optimizer.zero_grad() + mu, logvar, map_x, x_concat, x_Block1, x_Block2, x_Block3, x_input = model(inputs.squeeze(0)) + score_norm = mu.sum(dim=1).sum(dim=1) / test_maps.squeeze(0).sum(dim=1).sum(dim=1) + map_score = score_norm.mean() + map_score_list.append('{} {}\n'.format(map_score, spoof_label[0][0])) + + map_score_test_filename = args.log + '/' + args.log + '_map_score_test.txt' + with open(map_score_test_filename, 'w') as file: + file.writelines(map_score_list) + + ############################################################# + # performance measurement both val and test + ############################################################# + val_threshold, test_threshold, val_ACC, val_ACER, test_ACC, test_APCER, test_BPCER, test_ACER, test_ACER_test_threshold = performances( + map_score_val_filename, map_score_test_filename) + + print('epoch:%d, Val: val_threshold= %.4f, val_ACC= %.4f, val_ACER= %.4f' % ( + epoch + 1, val_threshold, val_ACC, val_ACER)) + log_file.write('\n epoch:%d, Val: val_threshold= %.4f, val_ACC= %.4f, val_ACER= %.4f \n' % ( + epoch + 1, val_threshold, val_ACC, val_ACER)) + + print('epoch:%d, Test: ACC= %.4f, APCER= %.4f, BPCER= %.4f, ACER= %.4f' % ( + epoch + 1, test_ACC, test_APCER, test_BPCER, test_ACER)) + log_file.write('epoch:%d, Test: ACC= %.4f, APCER= %.4f, BPCER= %.4f, ACER= %.4f \n' % ( + epoch + 1, test_ACC, test_APCER, test_BPCER, test_ACER)) + log_file.flush() + + if epoch % epoch_test == epoch_test - 1: + # save the model until the next improvement + torch.save(model.state_dict(), args.log + '/' + args.log + '_%d.pkl' % (epoch + 1)) + + print('Finished Training') + log_file.close() + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="save quality using landmarkpose model") + parser.add_argument('--gpus', type=str, default='0, 1, 2, 3', help='the gpu id used for predict') + parser.add_argument('--lr', type=float, default=0.0001, help='initial learning rate') + parser.add_argument('--batchsize', type=int, default=64, help='initial batchsize') + parser.add_argument('--step_size', type=int, default=500, help='how many epochs lr decays once') # 500 + parser.add_argument('--gamma', type=float, default=0.5, help='gamma of optim.lr_scheduler.StepLR, decay of lr') + parser.add_argument('--kl_lambda', type=float, default=0.001, help='') + parser.add_argument('--ratio', type=float, default=0.75, help='real and fake in batchsize ') + parser.add_argument('--echo_batches', type=int, default=50, help='how many batches display once') # 50 + parser.add_argument('--epochs', type=int, default=1600, help='total training epochs') + parser.add_argument('--log', type=str, default="CDCN_U_P1", help='log and save model name') + parser.add_argument('--finetune', action='store_true', default=False, help='whether finetune other models') + + args = parser.parse_args() + train_test() diff --git a/addition_module/DSDG/DUM/utils.py b/src/faceDetection/face_sdk/addition_module/DSDG/DUM/utils.py similarity index 100% rename from addition_module/DSDG/DUM/utils.py rename to src/faceDetection/face_sdk/addition_module/DSDG/DUM/utils.py diff --git a/addition_module/DSDG/README.md b/src/faceDetection/face_sdk/addition_module/DSDG/README.md similarity index 99% rename from addition_module/DSDG/README.md rename to src/faceDetection/face_sdk/addition_module/DSDG/README.md index f299eb1..6d230d9 100644 --- a/addition_module/DSDG/README.md +++ b/src/faceDetection/face_sdk/addition_module/DSDG/README.md @@ -1,45 +1,45 @@ -# DSDG - -Official Pytorch code of paper [Dual Spoof Disentanglement Generation for Face Anti-spoofing with Depth Uncertainty Learning](https://arxiv.org/pdf/2112.00568.pdf) in IEEE Transactions on Circuits and Systems for Video Technology. - -Face anti-spoofing (FAS) plays a vital role in preventing face recognition systems from presentation attacks. Existing face anti-spoofing datasets lack diversity due to the insufficient identity and insignificant variance, which limits the generalization ability of FAS model. In this paper, we propose Dual Spoof Disentanglement Generation (DSDG) framework to tackle this challenge by "anti-spoofing via generation". Depending on the interpretable factorized latent disentanglement in Variational Autoencoder (VAE), DSDG learns a joint distribution of the identity representation and the spoofing pattern representation in the latent space. Then, large-scale paired live and spoofing images can be generated from random noise to boost the diversity of the training set. However, some generated face images are partially distorted due to the inherent defect of VAE. Such noisy samples are hard to predict precise depth values, thus may obstruct the widely-used depth supervised optimization. To tackle this issue, we further introduce a lightweight Depth Uncertainty Module (DUM), which alleviates the adverse effects of noisy samples by depth uncertainty learning. DUM is developed without extra-dependency, thus can be flexibly integrated with any depth supervised network for face anti-spoofing. We evaluate the effectiveness of the proposed method on five popular benchmarks and achieve state-of-the-art results under both intra- and inter- test settings. - -## Requirements -Our experiments are conducted under the following environments: -- Python == 3.8 -- Pytorch == 1.6.0 -- torchvision == 0.7.0 - -## Training -Before training, we need to extract frame images for some video data sets. Then, we use [MTCNN](https://github.com/ipazc/mtcnn) for face detection and [PRNet](https://github.com/YadiraF/PRNet) for face depth map prediction. We give an example of the OULU-NPU dataset: -* Configure the paths in [./data/extract_frame.py](https://github.com/JDAI-CV/FaceX-Zoo/blob/main/addition_module/DSDG/data/extract_frame.py) to extract frames from videos. -* Configure the paths in [./data/bbox.py](https://github.com/JDAI-CV/FaceX-Zoo/blob/main/addition_module/DSDG/data/bbox.py) to get the location of face in each frame with MTCNN. -* Utilize the PRNet to get the depth map of face in each frame. -* Save the processed data in `./oulu_images/ `. - -#### DSDG -* Download the LightCNN-29 model from this [link](https://drive.google.com/file/d/1Jn6aXtQ84WY-7J3Tpr2_j6sX0ch9yucS/view) and put it to `./ip_checkpoint`. -* Run [./data/make_train_list.py](https://github.com/JDAI-CV/FaceX-Zoo/blob/main/addition_module/DSDG/data/make_train_list.py) to generate training list. -* Run [./train_generator.sh](https://github.com/JDAI-CV/FaceX-Zoo/blob/main/addition_module/DSDG/train_generator.sh) to train the generator. -* Run [./generated.py](https://github.com/JDAI-CV/FaceX-Zoo/blob/main/addition_module/DSDG/generated.py) to generate paired face anti-spoofing data. -* Save the processed data in `./fake_images/` and utilize the PRNet to get the depth map. - -#### DUM -* Run the [./DUM/make_dataset/crop_dataset.py](https://github.com/JDAI-CV/FaceX-Zoo/blob/main/addition_module/DSDG/DUM/make_dataset/crop_dataset.py) to crop the original data. Save the cropped data in `./oulu_images_crop`. -* Move the generated data in `./fake_images/` to `./oulu_images_crop/` and upgrade the protocol. -* Run [./DUM/train.py](https://github.com/JDAI-CV/FaceX-Zoo/blob/main/addition_module/DSDG/DUM/train.py) to train the model with original data and generated data. - -## Testing -We provide a CDCN model with DUM trained on OULU-NPU Protocol-1, and the following shows how to test it. -* The trained model is released in `./DUM/checkpoint/`. -* Run [./DUM/test.py](https://github.com/JDAI-CV/FaceX-Zoo/blob/main/addition_module/DSDG/DUM/test.py) to test on OULU-NPU Protocol-1. - -## Citation -Please consider citing our paper in your publications if the project helps your research. - -## Acknowledgements -This repo is based on the following projects, thank the authors a lot. -- [BradyFU/DVG](https://github.com/BradyFU/DVG) -- [ZitongYu/CDCN](https://github.com/ZitongYu/CDCN) -- [YadiraF/PRNet](https://github.com/YadiraF/PRNet) +# DSDG + +Official Pytorch code of paper [Dual Spoof Disentanglement Generation for Face Anti-spoofing with Depth Uncertainty Learning](https://arxiv.org/pdf/2112.00568.pdf) in IEEE Transactions on Circuits and Systems for Video Technology. + +Face anti-spoofing (FAS) plays a vital role in preventing face recognition systems from presentation attacks. Existing face anti-spoofing datasets lack diversity due to the insufficient identity and insignificant variance, which limits the generalization ability of FAS model. In this paper, we propose Dual Spoof Disentanglement Generation (DSDG) framework to tackle this challenge by "anti-spoofing via generation". Depending on the interpretable factorized latent disentanglement in Variational Autoencoder (VAE), DSDG learns a joint distribution of the identity representation and the spoofing pattern representation in the latent space. Then, large-scale paired live and spoofing images can be generated from random noise to boost the diversity of the training set. However, some generated face images are partially distorted due to the inherent defect of VAE. Such noisy samples are hard to predict precise depth values, thus may obstruct the widely-used depth supervised optimization. To tackle this issue, we further introduce a lightweight Depth Uncertainty Module (DUM), which alleviates the adverse effects of noisy samples by depth uncertainty learning. DUM is developed without extra-dependency, thus can be flexibly integrated with any depth supervised network for face anti-spoofing. We evaluate the effectiveness of the proposed method on five popular benchmarks and achieve state-of-the-art results under both intra- and inter- test settings. + +## Requirements +Our experiments are conducted under the following environments: +- Python == 3.8 +- Pytorch == 1.6.0 +- torchvision == 0.7.0 + +## Training +Before training, we need to extract frame images for some video data sets. Then, we use [MTCNN](https://github.com/ipazc/mtcnn) for face detection and [PRNet](https://github.com/YadiraF/PRNet) for face depth map prediction. We give an example of the OULU-NPU dataset: +* Configure the paths in [./data/extract_frame.py](https://github.com/JDAI-CV/FaceX-Zoo/blob/main/addition_module/DSDG/data/extract_frame.py) to extract frames from videos. +* Configure the paths in [./data/bbox.py](https://github.com/JDAI-CV/FaceX-Zoo/blob/main/addition_module/DSDG/data/bbox.py) to get the location of face in each frame with MTCNN. +* Utilize the PRNet to get the depth map of face in each frame. +* Save the processed data in `./oulu_images/ `. + +#### DSDG +* Download the LightCNN-29 model from this [link](https://drive.google.com/file/d/1Jn6aXtQ84WY-7J3Tpr2_j6sX0ch9yucS/view) and put it to `./ip_checkpoint`. +* Run [./data/make_train_list.py](https://github.com/JDAI-CV/FaceX-Zoo/blob/main/addition_module/DSDG/data/make_train_list.py) to generate training list. +* Run [./train_generator.sh](https://github.com/JDAI-CV/FaceX-Zoo/blob/main/addition_module/DSDG/train_generator.sh) to train the generator. +* Run [./generated.py](https://github.com/JDAI-CV/FaceX-Zoo/blob/main/addition_module/DSDG/generated.py) to generate paired face anti-spoofing data. +* Save the processed data in `./fake_images/` and utilize the PRNet to get the depth map. + +#### DUM +* Run the [./DUM/make_dataset/crop_dataset.py](https://github.com/JDAI-CV/FaceX-Zoo/blob/main/addition_module/DSDG/DUM/make_dataset/crop_dataset.py) to crop the original data. Save the cropped data in `./oulu_images_crop`. +* Move the generated data in `./fake_images/` to `./oulu_images_crop/` and upgrade the protocol. +* Run [./DUM/train.py](https://github.com/JDAI-CV/FaceX-Zoo/blob/main/addition_module/DSDG/DUM/train.py) to train the model with original data and generated data. + +## Testing +We provide a CDCN model with DUM trained on OULU-NPU Protocol-1, and the following shows how to test it. +* The trained model is released in `./DUM/checkpoint/`. +* Run [./DUM/test.py](https://github.com/JDAI-CV/FaceX-Zoo/blob/main/addition_module/DSDG/DUM/test.py) to test on OULU-NPU Protocol-1. + +## Citation +Please consider citing our paper in your publications if the project helps your research. + +## Acknowledgements +This repo is based on the following projects, thank the authors a lot. +- [BradyFU/DVG](https://github.com/BradyFU/DVG) +- [ZitongYu/CDCN](https://github.com/ZitongYu/CDCN) +- [YadiraF/PRNet](https://github.com/YadiraF/PRNet) - [ipazc/mtcnn](https://github.com/ipazc/mtcnn) \ No newline at end of file diff --git a/src/faceDetection/face_sdk/addition_module/DSDG/__init__.py b/src/faceDetection/face_sdk/addition_module/DSDG/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/addition_module/DSDG/data/__init__.py b/src/faceDetection/face_sdk/addition_module/DSDG/data/__init__.py similarity index 98% rename from addition_module/DSDG/data/__init__.py rename to src/faceDetection/face_sdk/addition_module/DSDG/data/__init__.py index 54f70c1..6e228d8 100644 --- a/addition_module/DSDG/data/__init__.py +++ b/src/faceDetection/face_sdk/addition_module/DSDG/data/__init__.py @@ -1,2 +1,2 @@ -from .generation_dataset import GenDataset_s -from .recognition_dataset import ImageList, SeparateBatchSampler, SeparateImageList +from .generation_dataset import GenDataset_s +from .recognition_dataset import ImageList, SeparateBatchSampler, SeparateImageList diff --git a/addition_module/DSDG/data/bbox.py b/src/faceDetection/face_sdk/addition_module/DSDG/data/bbox.py similarity index 96% rename from addition_module/DSDG/data/bbox.py rename to src/faceDetection/face_sdk/addition_module/DSDG/data/bbox.py index 7c39d71..b063527 100644 --- a/addition_module/DSDG/data/bbox.py +++ b/src/faceDetection/face_sdk/addition_module/DSDG/data/bbox.py @@ -1,36 +1,36 @@ -from mtcnn import MTCNN -import cv2 -import os - -os.environ['CUDA_VISIBLE_DEVICES'] = '3' - -detector = MTCNN() - -root_dir = '/export2/home/wht/oulu_images/train_img_flod/' -bbox_dir = '/export2/home/wht/oulu_images/train_bbox_flod/' - -if not os.path.exists(bbox_dir): - os.makedirs(bbox_dir) - -img_list = [] -f = open('/export2/home/wht/oulu_images/images_list.txt', 'r') - -for line in f.readlines(): - img_list.append(str(line[:-1])) - -for i, img in enumerate(img_list): - imgpath = img.split(',')[1] - imgdir = img.split(',')[0] - image = cv2.cvtColor(cv2.imread(root_dir + imgpath), cv2.COLOR_BGR2RGB) - result = detector.detect_faces(image) - if len(result) < 0.5: - continue - bounding_box = result[0]['box'] - bounding_box = ' '.join(str(x) for x in bounding_box) - dat_path = os.path.join(bbox_dir + imgdir) - if not os.path.exists(dat_path): - os.makedirs(dat_path) - fp = open(bbox_dir + imgpath[:-4] + '.dat', 'w') - fp.write(bounding_box) - fp.close() - +from mtcnn import MTCNN +import cv2 +import os + +os.environ['CUDA_VISIBLE_DEVICES'] = '3' + +detector = MTCNN() + +root_dir = '/export2/home/wht/oulu_images/train_img_flod/' +bbox_dir = '/export2/home/wht/oulu_images/train_bbox_flod/' + +if not os.path.exists(bbox_dir): + os.makedirs(bbox_dir) + +img_list = [] +f = open('/export2/home/wht/oulu_images/images_list.txt', 'r') + +for line in f.readlines(): + img_list.append(str(line[:-1])) + +for i, img in enumerate(img_list): + imgpath = img.split(',')[1] + imgdir = img.split(',')[0] + image = cv2.cvtColor(cv2.imread(root_dir + imgpath), cv2.COLOR_BGR2RGB) + result = detector.detect_faces(image) + if len(result) < 0.5: + continue + bounding_box = result[0]['box'] + bounding_box = ' '.join(str(x) for x in bounding_box) + dat_path = os.path.join(bbox_dir + imgdir) + if not os.path.exists(dat_path): + os.makedirs(dat_path) + fp = open(bbox_dir + imgpath[:-4] + '.dat', 'w') + fp.write(bounding_box) + fp.close() + diff --git a/addition_module/DSDG/data/extract_frame.py b/src/faceDetection/face_sdk/addition_module/DSDG/data/extract_frame.py similarity index 96% rename from addition_module/DSDG/data/extract_frame.py rename to src/faceDetection/face_sdk/addition_module/DSDG/data/extract_frame.py index 88384dd..3f01404 100644 --- a/addition_module/DSDG/data/extract_frame.py +++ b/src/faceDetection/face_sdk/addition_module/DSDG/data/extract_frame.py @@ -1,37 +1,37 @@ -import os -import cv2 - - -def save_image(image, img_dir, vedio_name, num): - flod_path = img_dir + vedio_name[:-4] + '/' - if not os.path.exists(flod_path): - os.makedirs(flod_path) - address = flod_path + vedio_name[:-4] + '_' + str(num) + '_scene.jpg' - cv2.imwrite(address, image) - - -root_dir = '/export2/home/wht/Oulu_Origin/Train_files/' - -img_dir = '/export2/home/wht/oulu_images/train_img_flod/' - -if not os.path.exists(img_dir): - os.makedirs(img_dir) - -id_list = os.listdir(root_dir) -num = len(id_list) - -for id in id_list[:num]: - vedio_dir = os.path.join(root_dir, id) - vedio_lists = os.listdir(vedio_dir) - vedioname_list = [] - for name in vedio_lists: - if name[-4:] == '.avi': - vedioname_list.append(name) - for j, vedio_name in enumerate(vedioname_list): - videoCapture = cv2.VideoCapture(vedio_dir + '/' + vedio_name) - success, frame = videoCapture.read() - i = 1 - while success: - success, frame = videoCapture.read() - if success == True: - save_image(frame, img_dir, vedio_name, i) +import os +import cv2 + + +def save_image(image, img_dir, vedio_name, num): + flod_path = img_dir + vedio_name[:-4] + '/' + if not os.path.exists(flod_path): + os.makedirs(flod_path) + address = flod_path + vedio_name[:-4] + '_' + str(num) + '_scene.jpg' + cv2.imwrite(address, image) + + +root_dir = '/export2/home/wht/Oulu_Origin/Train_files/' + +img_dir = '/export2/home/wht/oulu_images/train_img_flod/' + +if not os.path.exists(img_dir): + os.makedirs(img_dir) + +id_list = os.listdir(root_dir) +num = len(id_list) + +for id in id_list[:num]: + vedio_dir = os.path.join(root_dir, id) + vedio_lists = os.listdir(vedio_dir) + vedioname_list = [] + for name in vedio_lists: + if name[-4:] == '.avi': + vedioname_list.append(name) + for j, vedio_name in enumerate(vedioname_list): + videoCapture = cv2.VideoCapture(vedio_dir + '/' + vedio_name) + success, frame = videoCapture.read() + i = 1 + while success: + success, frame = videoCapture.read() + if success == True: + save_image(frame, img_dir, vedio_name, i) diff --git a/addition_module/DSDG/data/generation_dataset.py b/src/faceDetection/face_sdk/addition_module/DSDG/data/generation_dataset.py similarity index 96% rename from addition_module/DSDG/data/generation_dataset.py rename to src/faceDetection/face_sdk/addition_module/DSDG/data/generation_dataset.py index 9825506..97a080a 100644 --- a/addition_module/DSDG/data/generation_dataset.py +++ b/src/faceDetection/face_sdk/addition_module/DSDG/data/generation_dataset.py @@ -1,121 +1,121 @@ -import os -import math -import random -import numpy as np -import cv2 -from PIL import Image -from collections import defaultdict - -import torch -import torch.utils.data as data -import torchvision.transforms as transforms -""" -domain0 => spoof -domain1 => live -""" - - -class GenDataset_s(data.Dataset): - def __init__(self, img_root, list_file, attack_type): - super(GenDataset_s, self).__init__() - - self.img_root = img_root - self.list_file = list_file - self.attack_type = attack_type - - self.transform = transforms.Compose([ - transforms.ToTensor(), - ]) - - self.img_spoof_list, self.make_pair_dict = self.file_reader() - - def __getitem__(self, index): - img_line = self.img_spoof_list[index] - img_name_spoof, label, domain_flag, spoof_label = img_line.strip().split(' ') - if int(spoof_label) == 2: - spoof_type = 0 - elif int(spoof_label) == 3: - spoof_type = 1 - elif int(spoof_label) == 4: - spoof_type = 2 - elif int(spoof_label) == 5: - spoof_type = 3 - else: - print('spoof type out of except') - spoof_type = -1 - exit() - - img_name_live = self.get_pair(label, '1') - - img_domain0 = cv2.imread(os.path.join(self.img_root, 'train_img_flod/', img_name_spoof)) - img_domain0 = cv2.cvtColor(img_domain0, cv2.COLOR_BGR2RGB) - img_domain1 = cv2.imread(os.path.join(self.img_root, 'train_img_flod/', img_name_live)) - img_domain1 = cv2.cvtColor(img_domain1, cv2.COLOR_BGR2RGB) - - bbox_domain0_path = os.path.join(self.img_root, 'train_bbox_flod/', img_name_spoof[:-4] + '.dat') - bbox_domain1_path = os.path.join(self.img_root, 'train_bbox_flod/', img_name_live[:-4] + '.dat') - - face_scale = 1.0 - - img_domain0_crop = cv2.resize(self.crop_face_from_scene(img_domain0, bbox_domain0_path, face_scale), (256, 256)) - img_domain1_crop = cv2.resize(self.crop_face_from_scene(img_domain1, bbox_domain1_path, face_scale), (256, 256)) - - img_domain0 = self.transform(img_domain0_crop) - img_domain1 = self.transform(img_domain1_crop) - - return {'0': img_domain0, '1': img_domain1, 'type': spoof_type} - - def __len__(self): - return len(self.img_spoof_list) - - def file_reader(self): - - def dict_profile(): - return {'0': [], '1': [], 'type': []} - - with open(self.list_file) as file: - img_list = file.readlines() - img_list = [x.strip() for x in img_list] - - make_pair_dict = defaultdict(dict_profile) - img_spoof_list = [] - - for line in img_list: - img_name, label, domain_flag, spoof_label = line.strip().split(' ') - make_pair_dict[label][domain_flag].append(img_name) - - if domain_flag == '0': - img_spoof_list.append(line) - - return img_spoof_list, make_pair_dict - - def get_pair(self, label, domain_flag): - img_name = random.choice(self.make_pair_dict[label][domain_flag]) - return img_name - - def crop_face_from_scene(self, image, face_name_full, scale): - f = open(face_name_full, 'r') - lines = f.readlines() - lines = lines[0].split(' ') - y1, x1, w, h = [int(ele) for ele in lines[:4]] - f.close() - y2 = y1 + w - x2 = x1 + h - - y_mid = (y1 + y2) / 2.0 - x_mid = (x1 + x2) / 2.0 - h_img, w_img = image.shape[0], image.shape[1] - # w_img,h_img=image.size - w_scale = scale * w - h_scale = scale * h - y1 = y_mid - w_scale / 2.0 - x1 = x_mid - h_scale / 2.0 - y2 = y_mid + w_scale / 2.0 - x2 = x_mid + h_scale / 2.0 - y1 = max(math.floor(y1), 0) - x1 = max(math.floor(x1), 0) - y2 = min(math.floor(y2), w_img) - x2 = min(math.floor(x2), h_img) - - region = image[x1:x2, y1:y2] - return region +import os +import math +import random +import numpy as np +import cv2 +from PIL import Image +from collections import defaultdict + +import torch +import torch.utils.data as data +import torchvision.transforms as transforms +""" +domain0 => spoof +domain1 => live +""" + + +class GenDataset_s(data.Dataset): + def __init__(self, img_root, list_file, attack_type): + super(GenDataset_s, self).__init__() + + self.img_root = img_root + self.list_file = list_file + self.attack_type = attack_type + + self.transform = transforms.Compose([ + transforms.ToTensor(), + ]) + + self.img_spoof_list, self.make_pair_dict = self.file_reader() + + def __getitem__(self, index): + img_line = self.img_spoof_list[index] + img_name_spoof, label, domain_flag, spoof_label = img_line.strip().split(' ') + if int(spoof_label) == 2: + spoof_type = 0 + elif int(spoof_label) == 3: + spoof_type = 1 + elif int(spoof_label) == 4: + spoof_type = 2 + elif int(spoof_label) == 5: + spoof_type = 3 + else: + print('spoof type out of except') + spoof_type = -1 + exit() + + img_name_live = self.get_pair(label, '1') + + img_domain0 = cv2.imread(os.path.join(self.img_root, 'train_img_flod/', img_name_spoof)) + img_domain0 = cv2.cvtColor(img_domain0, cv2.COLOR_BGR2RGB) + img_domain1 = cv2.imread(os.path.join(self.img_root, 'train_img_flod/', img_name_live)) + img_domain1 = cv2.cvtColor(img_domain1, cv2.COLOR_BGR2RGB) + + bbox_domain0_path = os.path.join(self.img_root, 'train_bbox_flod/', img_name_spoof[:-4] + '.dat') + bbox_domain1_path = os.path.join(self.img_root, 'train_bbox_flod/', img_name_live[:-4] + '.dat') + + face_scale = 1.0 + + img_domain0_crop = cv2.resize(self.crop_face_from_scene(img_domain0, bbox_domain0_path, face_scale), (256, 256)) + img_domain1_crop = cv2.resize(self.crop_face_from_scene(img_domain1, bbox_domain1_path, face_scale), (256, 256)) + + img_domain0 = self.transform(img_domain0_crop) + img_domain1 = self.transform(img_domain1_crop) + + return {'0': img_domain0, '1': img_domain1, 'type': spoof_type} + + def __len__(self): + return len(self.img_spoof_list) + + def file_reader(self): + + def dict_profile(): + return {'0': [], '1': [], 'type': []} + + with open(self.list_file) as file: + img_list = file.readlines() + img_list = [x.strip() for x in img_list] + + make_pair_dict = defaultdict(dict_profile) + img_spoof_list = [] + + for line in img_list: + img_name, label, domain_flag, spoof_label = line.strip().split(' ') + make_pair_dict[label][domain_flag].append(img_name) + + if domain_flag == '0': + img_spoof_list.append(line) + + return img_spoof_list, make_pair_dict + + def get_pair(self, label, domain_flag): + img_name = random.choice(self.make_pair_dict[label][domain_flag]) + return img_name + + def crop_face_from_scene(self, image, face_name_full, scale): + f = open(face_name_full, 'r') + lines = f.readlines() + lines = lines[0].split(' ') + y1, x1, w, h = [int(ele) for ele in lines[:4]] + f.close() + y2 = y1 + w + x2 = x1 + h + + y_mid = (y1 + y2) / 2.0 + x_mid = (x1 + x2) / 2.0 + h_img, w_img = image.shape[0], image.shape[1] + # w_img,h_img=image.size + w_scale = scale * w + h_scale = scale * h + y1 = y_mid - w_scale / 2.0 + x1 = x_mid - h_scale / 2.0 + y2 = y_mid + w_scale / 2.0 + x2 = x_mid + h_scale / 2.0 + y1 = max(math.floor(y1), 0) + x1 = max(math.floor(x1), 0) + y2 = min(math.floor(y2), w_img) + x2 = min(math.floor(x2), h_img) + + region = image[x1:x2, y1:y2] + return region diff --git a/addition_module/DSDG/data/make_train_list.py b/src/faceDetection/face_sdk/addition_module/DSDG/data/make_train_list.py similarity index 96% rename from addition_module/DSDG/data/make_train_list.py rename to src/faceDetection/face_sdk/addition_module/DSDG/data/make_train_list.py index ec4661f..fa9509e 100644 --- a/addition_module/DSDG/data/make_train_list.py +++ b/src/faceDetection/face_sdk/addition_module/DSDG/data/make_train_list.py @@ -1,29 +1,29 @@ -import os -import random - -root_dir = '/export2/home/wht/oulu_images/train_img_flod/' - -origin_txt = '/export2/home/wht/oulu_images/protocols/Protocol_1/Train.txt' -f = open('./train_list/train_list_oulu_p1.txt', 'w') -ff = open(origin_txt, 'r') -video_names = ff.readlines() - -for video_name in video_names: - video_name = video_name.split(',')[1][:-1] - img_list = os.listdir(os.path.join(root_dir, video_name)) - img_num = len(img_list) - if img_num > 50: - random_img = random.sample(range(1, img_num), 50) - else: - random_img = random.sample(range(1, img_num), img_num) - if int(video_name[-1:]) == 1: - domain_flag = str(1) - else: - domain_flag = str(0) - spoof_type = video_name[-1:] - for j in random_img: - img_name = video_name + '/' + img_list[j] - label = video_name[4:6] - - f.write(img_name + ' ' + label + ' ' + domain_flag + ' ' + spoof_type + '\n') -f.close() +import os +import random + +root_dir = '/export2/home/wht/oulu_images/train_img_flod/' + +origin_txt = '/export2/home/wht/oulu_images/protocols/Protocol_1/Train.txt' +f = open('./train_list/train_list_oulu_p1.txt', 'w') +ff = open(origin_txt, 'r') +video_names = ff.readlines() + +for video_name in video_names: + video_name = video_name.split(',')[1][:-1] + img_list = os.listdir(os.path.join(root_dir, video_name)) + img_num = len(img_list) + if img_num > 50: + random_img = random.sample(range(1, img_num), 50) + else: + random_img = random.sample(range(1, img_num), img_num) + if int(video_name[-1:]) == 1: + domain_flag = str(1) + else: + domain_flag = str(0) + spoof_type = video_name[-1:] + for j in random_img: + img_name = video_name + '/' + img_list[j] + label = video_name[4:6] + + f.write(img_name + ' ' + label + ' ' + domain_flag + ' ' + spoof_type + '\n') +f.close() diff --git a/addition_module/DSDG/data/recognition_dataset.py b/src/faceDetection/face_sdk/addition_module/DSDG/data/recognition_dataset.py similarity index 97% rename from addition_module/DSDG/data/recognition_dataset.py rename to src/faceDetection/face_sdk/addition_module/DSDG/data/recognition_dataset.py index 30aa313..c9d2ccf 100644 --- a/addition_module/DSDG/data/recognition_dataset.py +++ b/src/faceDetection/face_sdk/addition_module/DSDG/data/recognition_dataset.py @@ -1,157 +1,157 @@ -import os -import copy -import math -import time -import random -import numpy as np -from PIL import Image - -import torch -import torch.utils.data as data -import torchvision.transforms as transforms - - -def default_loader(path): - img = Image.open(path).convert('L') - return img - - -def default_list_reader(fileList): - imgList = [] - with open(fileList, 'r') as file: - for line in file.readlines(): - imgPath, label, domain = line.strip().split(' ') - imgList.append((imgPath, int(label), int(domain))) - return imgList - - -class ImageList(data.Dataset): - def __init__(self, root, fileList, list_reader=default_list_reader, loader=default_loader): - self.root = root - self.imgList = list_reader(fileList) - self.loader = loader - - self.transform = transforms.Compose([ - transforms.RandomCrop(128), - transforms.ToTensor(), - ]) - - def __getitem__(self, index): - imgPath, target, domain = self.imgList[index] - img = self.loader(os.path.join(self.root, imgPath)) - - img = self.transform(img) - - return {'img': img, 'label': target, 'domain_flag': domain} - - def __len__(self): - return len(self.imgList) - - -class SeparateBatchSampler(object): - def __init__(self, real_data_idx, fake_data_idx, batch_size=128, ratio=0.5, put_back=False): - self.batch_size = batch_size - self.ratio = ratio - self.real_data_num = len(real_data_idx) - self.fake_data_num = len(fake_data_idx) - self.max_num_image = max(self.real_data_num, self.fake_data_num) - - self.real_data_idx = real_data_idx - self.fake_data_idx = fake_data_idx - - self.processed_idx = copy.deepcopy(self.real_data_idx) - - def __len__(self): - return self.max_num_image // (int(self.batch_size * self.ratio)) - - def __iter__(self): - batch_size_real_data = int(math.floor(self.ratio * self.batch_size)) - batch_size_fake_data = self.batch_size - batch_size_real_data - - self.processed_idx = copy.deepcopy(self.real_data_idx) - rand_real_data_idx = np.random.permutation(len(self.real_data_idx) // 2) - for i in range(self.__len__()): - batch = [] - idx_fake_data = random.sample(self.fake_data_idx, batch_size_fake_data // 2) - - for j in range(batch_size_real_data // 2): - idx = rand_real_data_idx[(i * batch_size_real_data + j) % (self.real_data_num // 2)] - batch.append(self.processed_idx[2 * idx]) - batch.append(self.processed_idx[2 * idx + 1]) - - for idx in idx_fake_data: - batch.append(2 * idx + self.real_data_num) - batch.append(2 * idx + 1 + self.real_data_num) - yield batch - - -class SeparateImageList(data.Dataset): - def __init__(self, real_data_path, real_list_path, fake_data_path, fake_total_num, ratio=0.5): - - self.transform = transforms.Compose([ - transforms.RandomCrop(128), - transforms.RandomHorizontalFlip(), - transforms.ToTensor() - ]) - - # load real nir/vis data - real_data_list, real_data_idx = self.list_reader(real_data_path, real_list_path) - - # load fake nir/vis data from noise - _idx = np.random.permutation(fake_total_num) - fake_data_list = [] - fake_data_idx = [] - for i in range(0, fake_total_num): - _fake_img_name = str(_idx[i] + 1) + '.jpg' - - # nir_noise and vis_noise are the path of the fake data - _fake_img_nir_name = 'nir_noise/' + _fake_img_name - _fake_img_vis_name = 'vis_noise/' + _fake_img_name - - _fake_img_nir_path = os.path.join(fake_data_path, _fake_img_nir_name) - _fake_img_vis_path = os.path.join(fake_data_path, _fake_img_vis_name) - fake_data_list.append((_fake_img_nir_path, -1, 0)) - fake_data_list.append((_fake_img_vis_path, -1, 1)) - fake_data_idx.append(i) - - self.real_data_idx = real_data_idx - self.fake_data_idx = fake_data_idx - - real_data_list.extend(fake_data_list) - self.all_list = real_data_list - - self.ratio = ratio - - print('real: {}, fake: {}, total: {}, ratio: {}\n'.format(len(self.real_data_idx), len(self.fake_data_idx), - len(self.all_list), self.ratio)) - - def get_idx(self): - return self.real_data_idx, self.fake_data_idx - - def list_reader(self, root_path, fileList): - imgList = [] - imgIdx = [] - img_index = 0 - with open(fileList, 'r') as file: - for line in file.readlines(): - img_name, label, domain = line.strip().split(' ') - img_path = os.path.join(root_path, img_name) - imgList.append((img_path, int(label), int(domain))) - imgIdx.append(img_index) - img_index += 1 - return imgList, imgIdx - - def loader(self, path): - img = Image.open(path).convert('L') - return img - - def __getitem__(self, index): - imgPath, target, domain = self.all_list[index] - img = self.loader(imgPath) - - img = self.transform(img) - - return {'img': img, 'label': target, 'domain_flag': domain} - - def __len__(self): - return len(self.all_list) +import os +import copy +import math +import time +import random +import numpy as np +from PIL import Image + +import torch +import torch.utils.data as data +import torchvision.transforms as transforms + + +def default_loader(path): + img = Image.open(path).convert('L') + return img + + +def default_list_reader(fileList): + imgList = [] + with open(fileList, 'r') as file: + for line in file.readlines(): + imgPath, label, domain = line.strip().split(' ') + imgList.append((imgPath, int(label), int(domain))) + return imgList + + +class ImageList(data.Dataset): + def __init__(self, root, fileList, list_reader=default_list_reader, loader=default_loader): + self.root = root + self.imgList = list_reader(fileList) + self.loader = loader + + self.transform = transforms.Compose([ + transforms.RandomCrop(128), + transforms.ToTensor(), + ]) + + def __getitem__(self, index): + imgPath, target, domain = self.imgList[index] + img = self.loader(os.path.join(self.root, imgPath)) + + img = self.transform(img) + + return {'img': img, 'label': target, 'domain_flag': domain} + + def __len__(self): + return len(self.imgList) + + +class SeparateBatchSampler(object): + def __init__(self, real_data_idx, fake_data_idx, batch_size=128, ratio=0.5, put_back=False): + self.batch_size = batch_size + self.ratio = ratio + self.real_data_num = len(real_data_idx) + self.fake_data_num = len(fake_data_idx) + self.max_num_image = max(self.real_data_num, self.fake_data_num) + + self.real_data_idx = real_data_idx + self.fake_data_idx = fake_data_idx + + self.processed_idx = copy.deepcopy(self.real_data_idx) + + def __len__(self): + return self.max_num_image // (int(self.batch_size * self.ratio)) + + def __iter__(self): + batch_size_real_data = int(math.floor(self.ratio * self.batch_size)) + batch_size_fake_data = self.batch_size - batch_size_real_data + + self.processed_idx = copy.deepcopy(self.real_data_idx) + rand_real_data_idx = np.random.permutation(len(self.real_data_idx) // 2) + for i in range(self.__len__()): + batch = [] + idx_fake_data = random.sample(self.fake_data_idx, batch_size_fake_data // 2) + + for j in range(batch_size_real_data // 2): + idx = rand_real_data_idx[(i * batch_size_real_data + j) % (self.real_data_num // 2)] + batch.append(self.processed_idx[2 * idx]) + batch.append(self.processed_idx[2 * idx + 1]) + + for idx in idx_fake_data: + batch.append(2 * idx + self.real_data_num) + batch.append(2 * idx + 1 + self.real_data_num) + yield batch + + +class SeparateImageList(data.Dataset): + def __init__(self, real_data_path, real_list_path, fake_data_path, fake_total_num, ratio=0.5): + + self.transform = transforms.Compose([ + transforms.RandomCrop(128), + transforms.RandomHorizontalFlip(), + transforms.ToTensor() + ]) + + # load real nir/vis data + real_data_list, real_data_idx = self.list_reader(real_data_path, real_list_path) + + # load fake nir/vis data from noise + _idx = np.random.permutation(fake_total_num) + fake_data_list = [] + fake_data_idx = [] + for i in range(0, fake_total_num): + _fake_img_name = str(_idx[i] + 1) + '.jpg' + + # nir_noise and vis_noise are the path of the fake data + _fake_img_nir_name = 'nir_noise/' + _fake_img_name + _fake_img_vis_name = 'vis_noise/' + _fake_img_name + + _fake_img_nir_path = os.path.join(fake_data_path, _fake_img_nir_name) + _fake_img_vis_path = os.path.join(fake_data_path, _fake_img_vis_name) + fake_data_list.append((_fake_img_nir_path, -1, 0)) + fake_data_list.append((_fake_img_vis_path, -1, 1)) + fake_data_idx.append(i) + + self.real_data_idx = real_data_idx + self.fake_data_idx = fake_data_idx + + real_data_list.extend(fake_data_list) + self.all_list = real_data_list + + self.ratio = ratio + + print('real: {}, fake: {}, total: {}, ratio: {}\n'.format(len(self.real_data_idx), len(self.fake_data_idx), + len(self.all_list), self.ratio)) + + def get_idx(self): + return self.real_data_idx, self.fake_data_idx + + def list_reader(self, root_path, fileList): + imgList = [] + imgIdx = [] + img_index = 0 + with open(fileList, 'r') as file: + for line in file.readlines(): + img_name, label, domain = line.strip().split(' ') + img_path = os.path.join(root_path, img_name) + imgList.append((img_path, int(label), int(domain))) + imgIdx.append(img_index) + img_index += 1 + return imgList, imgIdx + + def loader(self, path): + img = Image.open(path).convert('L') + return img + + def __getitem__(self, index): + imgPath, target, domain = self.all_list[index] + img = self.loader(imgPath) + + img = self.transform(img) + + return {'img': img, 'label': target, 'domain_flag': domain} + + def __len__(self): + return len(self.all_list) diff --git a/addition_module/DSDG/generated.py b/src/faceDetection/face_sdk/addition_module/DSDG/generated.py similarity index 96% rename from addition_module/DSDG/generated.py rename to src/faceDetection/face_sdk/addition_module/DSDG/generated.py index e84f896..b88b7f8 100644 --- a/addition_module/DSDG/generated.py +++ b/src/faceDetection/face_sdk/addition_module/DSDG/generated.py @@ -1,78 +1,78 @@ -import os -import time -import argparse -import numpy as np -from PIL import Image - -import torch -import torch.nn as nn -import torch.optim as optim -import torch.nn.functional as F -import torch.backends.cudnn as cudnn -import torchvision.utils as vutils - -from data import * -from networks import * -from misc import * - -parser = argparse.ArgumentParser() - -parser.add_argument('--gpu_ids', default='0,1,2,3', type=str) -parser.add_argument('--hdim', default=128, type=int) -parser.add_argument('--attack_type', default=4, type=int, help='attack type number') -parser.add_argument('--pre_model', default='./result/model_oulu_p1/netG_model_epoch_200_iter_0.pth', type=str) -parser.add_argument('--out_path', default='./fake_images/oulu_p1/', type=str) - -spoof_img_name = 'p1_0_' -live_img_name = 'p1_1_' - - -def main(): - global opt, model - args = parser.parse_args() - print(args) - - os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_ids - cudnn.benchmark = True - - if not os.path.exists(args.out_path): - os.makedirs(args.out_path) - - netE_nir, netE_vis, netG, netCls = define_G(hdim=args.hdim, attack_type=args.attack_type) - load_model(netG, args.pre_model) - netG.eval() - - num = 0 - for n in range(200): - noise = torch.zeros(100, args.hdim).normal_(0, 1) - noise_s = torch.zeros(100, args.hdim).normal_(0, 1) - noise = torch.cat((noise_s, noise, noise), dim=1) - noise = noise.cuda() - - fake = netG(noise) - - spoof = fake[:, 0:3, :, :].data.cpu().numpy() - live = fake[:, 3:6, :, :].data.cpu().numpy() - - os.makedirs(os.path.join(args.out_path, spoof_img_name + str(n))) - os.makedirs(os.path.join(args.out_path, live_img_name + str(n))) - - for i in range(spoof.shape[0]): - num = num + 1 - save_img = spoof[i, :, :, :] - save_img = np.transpose((255 * save_img).astype('uint8'), (1, 2, 0)) - output = Image.fromarray(save_img) - save_name = str(i) + '_scene.jpg' - output.save(os.path.join(args.out_path, spoof_img_name + str(n), save_name)) - - save_img = live[i, :, :, :] - save_img = np.transpose((255 * save_img).astype('uint8'), (1, 2, 0)) - output = Image.fromarray(save_img) - save_name = str(i) + '_scene.jpg' - output.save(os.path.join(args.out_path, live_img_name + str(n), save_name)) - - print(num) - - -if __name__ == "__main__": - main() +import os +import time +import argparse +import numpy as np +from PIL import Image + +import torch +import torch.nn as nn +import torch.optim as optim +import torch.nn.functional as F +import torch.backends.cudnn as cudnn +import torchvision.utils as vutils + +from data import * +from networks import * +from misc import * + +parser = argparse.ArgumentParser() + +parser.add_argument('--gpu_ids', default='0,1,2,3', type=str) +parser.add_argument('--hdim', default=128, type=int) +parser.add_argument('--attack_type', default=4, type=int, help='attack type number') +parser.add_argument('--pre_model', default='./result/model_oulu_p1/netG_model_epoch_200_iter_0.pth', type=str) +parser.add_argument('--out_path', default='./fake_images/oulu_p1/', type=str) + +spoof_img_name = 'p1_0_' +live_img_name = 'p1_1_' + + +def main(): + global opt, model + args = parser.parse_args() + print(args) + + os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_ids + cudnn.benchmark = True + + if not os.path.exists(args.out_path): + os.makedirs(args.out_path) + + netE_nir, netE_vis, netG, netCls = define_G(hdim=args.hdim, attack_type=args.attack_type) + load_model(netG, args.pre_model) + netG.eval() + + num = 0 + for n in range(200): + noise = torch.zeros(100, args.hdim).normal_(0, 1) + noise_s = torch.zeros(100, args.hdim).normal_(0, 1) + noise = torch.cat((noise_s, noise, noise), dim=1) + noise = noise.cuda() + + fake = netG(noise) + + spoof = fake[:, 0:3, :, :].data.cpu().numpy() + live = fake[:, 3:6, :, :].data.cpu().numpy() + + os.makedirs(os.path.join(args.out_path, spoof_img_name + str(n))) + os.makedirs(os.path.join(args.out_path, live_img_name + str(n))) + + for i in range(spoof.shape[0]): + num = num + 1 + save_img = spoof[i, :, :, :] + save_img = np.transpose((255 * save_img).astype('uint8'), (1, 2, 0)) + output = Image.fromarray(save_img) + save_name = str(i) + '_scene.jpg' + output.save(os.path.join(args.out_path, spoof_img_name + str(n), save_name)) + + save_img = live[i, :, :, :] + save_img = np.transpose((255 * save_img).astype('uint8'), (1, 2, 0)) + output = Image.fromarray(save_img) + save_name = str(i) + '_scene.jpg' + output.save(os.path.join(args.out_path, live_img_name + str(n), save_name)) + + print(num) + + +if __name__ == "__main__": + main() diff --git a/addition_module/DSDG/misc/__init__.py b/src/faceDetection/face_sdk/addition_module/DSDG/misc/__init__.py similarity index 91% rename from addition_module/DSDG/misc/__init__.py rename to src/faceDetection/face_sdk/addition_module/DSDG/misc/__init__.py index 57de4a3..7b5491b 100644 --- a/addition_module/DSDG/misc/__init__.py +++ b/src/faceDetection/face_sdk/addition_module/DSDG/misc/__init__.py @@ -1,2 +1,2 @@ - -from .util import * + +from .util import * diff --git a/addition_module/DSDG/misc/util.py b/src/faceDetection/face_sdk/addition_module/DSDG/misc/util.py similarity index 96% rename from addition_module/DSDG/misc/util.py rename to src/faceDetection/face_sdk/addition_module/DSDG/misc/util.py index cfa059b..545fec0 100644 --- a/addition_module/DSDG/misc/util.py +++ b/src/faceDetection/face_sdk/addition_module/DSDG/misc/util.py @@ -1,109 +1,109 @@ -import os -import time -import argparse -import numpy as np - -import torch -import torch.nn.functional as F - - -def rgb2gray(img): - r, g, b = torch.split(img, 1, dim=1) - return torch.mul(r, 0.299) + torch.mul(g, 0.587) + torch.mul(b, 0.114) - - -def reparameterize(mu, logvar): - std = logvar.mul(0.5).exp_() - eps = torch.cuda.FloatTensor(std.size()).normal_() - return eps.mul(std).add_(mu) - - -def kl_loss(mu, logvar, prior_mu=0): - v_kl = mu.add(-prior_mu).pow(2).add_(logvar.exp()).mul_(-1).add_(1).add_(logvar) - v_kl = v_kl.sum(dim=-1).mul_(-0.5) # (batch, 2) - return v_kl - - -def reconstruction_loss(prediction, target, size_average=False): - error = (prediction - target).view(prediction.size(0), -1) - error = error ** 2 - error = torch.sum(error, dim=-1) - - if size_average: - error = error.mean() - else: - error = error.sum() - return error - - -def load_model(model, pretrained): - weights = torch.load(pretrained) - pretrained_dict = weights['model'].state_dict() - model_dict = model.state_dict() - # 1. filter out unnecessary keys - pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict} - # 2. overwrite entries in the existing state dict - model_dict.update(pretrained_dict) - # 3. load the new state dict - model.load_state_dict(model_dict) - - -def save_checkpoint(model_path, model, epoch, iteration, name): - model_out_path = model_path + name + "model_epoch_{}_iter_{}.pth".format(epoch, iteration) - state = {"epoch": epoch, "model": model} - if not os.path.exists(model_path): - os.makedirs(model_path) - torch.save(state, model_out_path) - print("Checkpoint saved to {}".format(model_out_path)) - - -class AverageMeter(object): - """Computes and stores the average and current value""" - - def __init__(self): - self.reset() - - def reset(self): - self.val = 0 - self.avg = 0 - self.sum = 0 - self.count = 0 - - def update(self, val, n=1): - self.val = val - self.sum += val * n - self.count += n - self.avg = self.sum / self.count - - -def MMD_Loss(fc_nir, fc_vis): - mean_fc_nir = torch.mean(fc_nir, 0) - mean_fc_vis = torch.mean(fc_vis, 0) - loss_mmd = F.mse_loss(mean_fc_nir, mean_fc_vis) - return loss_mmd - - -def adjust_learning_rate(lr, step, optimizer, epoch): - scale = 0.457305051927326 - lr = lr * (scale ** (epoch // step)) - print('lr: {}'.format(lr)) - if (epoch != 0) & (epoch % step == 0): - print('Change lr') - for param_group in optimizer.param_groups: - param_group['lr'] = param_group['lr'] * scale - - -def accuracy(output, target, topk=(1,)): - """Computes the precision@k for the specified values of k""" - maxk = max(topk) - batch_size = target.size(0) - - _, pred = output.topk(maxk, 1, True, True) - pred = pred.t() - correct = pred.eq(target.view(1, -1).expand_as(pred)) - - res = [] - for k in topk: - correct_k = correct[:k].view(-1).float().sum(0) - res.append(correct_k.mul_(100.0 / batch_size)) - return res +import os +import time +import argparse +import numpy as np + +import torch +import torch.nn.functional as F + + +def rgb2gray(img): + r, g, b = torch.split(img, 1, dim=1) + return torch.mul(r, 0.299) + torch.mul(g, 0.587) + torch.mul(b, 0.114) + + +def reparameterize(mu, logvar): + std = logvar.mul(0.5).exp_() + eps = torch.cuda.FloatTensor(std.size()).normal_() + return eps.mul(std).add_(mu) + + +def kl_loss(mu, logvar, prior_mu=0): + v_kl = mu.add(-prior_mu).pow(2).add_(logvar.exp()).mul_(-1).add_(1).add_(logvar) + v_kl = v_kl.sum(dim=-1).mul_(-0.5) # (batch, 2) + return v_kl + + +def reconstruction_loss(prediction, target, size_average=False): + error = (prediction - target).view(prediction.size(0), -1) + error = error ** 2 + error = torch.sum(error, dim=-1) + + if size_average: + error = error.mean() + else: + error = error.sum() + return error + + +def load_model(model, pretrained): + weights = torch.load(pretrained) + pretrained_dict = weights['model'].state_dict() + model_dict = model.state_dict() + # 1. filter out unnecessary keys + pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict} + # 2. overwrite entries in the existing state dict + model_dict.update(pretrained_dict) + # 3. load the new state dict + model.load_state_dict(model_dict) + + +def save_checkpoint(model_path, model, epoch, iteration, name): + model_out_path = model_path + name + "model_epoch_{}_iter_{}.pth".format(epoch, iteration) + state = {"epoch": epoch, "model": model} + if not os.path.exists(model_path): + os.makedirs(model_path) + torch.save(state, model_out_path) + print("Checkpoint saved to {}".format(model_out_path)) + + +class AverageMeter(object): + """Computes and stores the average and current value""" + + def __init__(self): + self.reset() + + def reset(self): + self.val = 0 + self.avg = 0 + self.sum = 0 + self.count = 0 + + def update(self, val, n=1): + self.val = val + self.sum += val * n + self.count += n + self.avg = self.sum / self.count + + +def MMD_Loss(fc_nir, fc_vis): + mean_fc_nir = torch.mean(fc_nir, 0) + mean_fc_vis = torch.mean(fc_vis, 0) + loss_mmd = F.mse_loss(mean_fc_nir, mean_fc_vis) + return loss_mmd + + +def adjust_learning_rate(lr, step, optimizer, epoch): + scale = 0.457305051927326 + lr = lr * (scale ** (epoch // step)) + print('lr: {}'.format(lr)) + if (epoch != 0) & (epoch % step == 0): + print('Change lr') + for param_group in optimizer.param_groups: + param_group['lr'] = param_group['lr'] * scale + + +def accuracy(output, target, topk=(1,)): + """Computes the precision@k for the specified values of k""" + maxk = max(topk) + batch_size = target.size(0) + + _, pred = output.topk(maxk, 1, True, True) + pred = pred.t() + correct = pred.eq(target.view(1, -1).expand_as(pred)) + + res = [] + for k in topk: + correct_k = correct[:k].view(-1).float().sum(0) + res.append(correct_k.mul_(100.0 / batch_size)) + return res diff --git a/addition_module/DSDG/networks/__init__.py b/src/faceDetection/face_sdk/addition_module/DSDG/networks/__init__.py similarity index 97% rename from addition_module/DSDG/networks/__init__.py rename to src/faceDetection/face_sdk/addition_module/DSDG/networks/__init__.py index 663101e..f3c1fae 100644 --- a/addition_module/DSDG/networks/__init__.py +++ b/src/faceDetection/face_sdk/addition_module/DSDG/networks/__init__.py @@ -1,32 +1,32 @@ -import torch -from .generator import Encoder, Decoder, Encoder_s, Decoder_s, Cls -from .light_cnn import network_29layers_v2, resblock - - -# define generator -def define_G(hdim=256, attack_type=4): - netE_nir = Encoder_s(hdim=hdim) - netCls = Cls(hdim=hdim, attack_type=attack_type) - netE_vis = Encoder(hdim=hdim) - netG = Decoder_s(hdim=hdim) - - netE_nir = torch.nn.DataParallel(netE_nir).cuda() - netE_vis = torch.nn.DataParallel(netE_vis).cuda() - netG = torch.nn.DataParallel(netG).cuda() - netCls = torch.nn.DataParallel(netCls).cuda() - - return netE_nir, netE_vis, netG, netCls - - -# define identity preserving && feature extraction net -def define_IP(is_train=False): - netIP = network_29layers_v2(resblock, [1, 2, 3, 4], is_train) - netIP = torch.nn.DataParallel(netIP).cuda() - return netIP - - -# define recognition network -def LightCNN_29v2(num_classes=10000, is_train=True): - net = network_29layers_v2(resblock, [1, 2, 3, 4], is_train, num_classes=num_classes) - net = torch.nn.DataParallel(net).cuda() - return net +import torch +from .generator import Encoder, Decoder, Encoder_s, Decoder_s, Cls +from .light_cnn import network_29layers_v2, resblock + + +# define generator +def define_G(hdim=256, attack_type=4): + netE_nir = Encoder_s(hdim=hdim) + netCls = Cls(hdim=hdim, attack_type=attack_type) + netE_vis = Encoder(hdim=hdim) + netG = Decoder_s(hdim=hdim) + + netE_nir = torch.nn.DataParallel(netE_nir).cuda() + netE_vis = torch.nn.DataParallel(netE_vis).cuda() + netG = torch.nn.DataParallel(netG).cuda() + netCls = torch.nn.DataParallel(netCls).cuda() + + return netE_nir, netE_vis, netG, netCls + + +# define identity preserving && feature extraction net +def define_IP(is_train=False): + netIP = network_29layers_v2(resblock, [1, 2, 3, 4], is_train) + netIP = torch.nn.DataParallel(netIP).cuda() + return netIP + + +# define recognition network +def LightCNN_29v2(num_classes=10000, is_train=True): + net = network_29layers_v2(resblock, [1, 2, 3, 4], is_train, num_classes=num_classes) + net = torch.nn.DataParallel(net).cuda() + return net diff --git a/addition_module/DSDG/networks/generator.py b/src/faceDetection/face_sdk/addition_module/DSDG/networks/generator.py similarity index 97% rename from addition_module/DSDG/networks/generator.py rename to src/faceDetection/face_sdk/addition_module/DSDG/networks/generator.py index 211cb02..0b373a6 100644 --- a/addition_module/DSDG/networks/generator.py +++ b/src/faceDetection/face_sdk/addition_module/DSDG/networks/generator.py @@ -1,192 +1,192 @@ -import numpy as np - -import torch -import torch.nn as nn -import torch.nn.functional as F -from torch.autograd import Variable - - -class Encoder(nn.Module): - def __init__(self, hdim=256): - super(Encoder, self).__init__() - - self.hdim = hdim - - self.main = nn.Sequential( - nn.Conv2d(3, 32, 5, 2, 2, bias=False), - nn.InstanceNorm2d(32, eps=0.001), - nn.LeakyReLU(0.2), - nn.AvgPool2d(2), - make_layer(_Residual_Block, 1, 32, 64), - nn.AvgPool2d(2), - make_layer(_Residual_Block, 1, 64, 128), - nn.AvgPool2d(2), - make_layer(_Residual_Block, 1, 128, 256), - nn.AvgPool2d(2), - make_layer(_Residual_Block, 1, 256, 512), - nn.AvgPool2d(2), - make_layer(_Residual_Block, 1, 512, 512) - ) - - # mu and logvar - self.fc = nn.Linear(512 * 4 * 4, 2 * hdim) - - def forward(self, x): - z = self.main(x).view(x.size(0), -1) - z = self.fc(z) - mu, logvar = torch.split(z, split_size_or_sections=self.hdim, dim=-1) - - return mu, logvar - - -class Encoder_s(nn.Module): - def __init__(self, hdim=256): - super(Encoder_s, self).__init__() - - self.hdim = hdim - - self.main = nn.Sequential( - nn.Conv2d(3, 32, 5, 2, 2, bias=False), - nn.InstanceNorm2d(32, eps=0.001), - nn.LeakyReLU(0.2), - nn.AvgPool2d(2), - make_layer(_Residual_Block, 1, 32, 64), - nn.AvgPool2d(2), - make_layer(_Residual_Block, 1, 64, 128), - nn.AvgPool2d(2), - make_layer(_Residual_Block, 1, 128, 256), - nn.AvgPool2d(2), - make_layer(_Residual_Block, 1, 256, 512), - nn.AvgPool2d(2), - make_layer(_Residual_Block, 1, 512, 512) - ) - - # mu and logvar - self.fc = nn.Linear(512 * 4 * 4, 4 * hdim) - # self.fc_at = nn.Linear(hdim, attack_type) - - def forward(self, x): - z = self.main(x).view(x.size(0), -1) - z = self.fc(z) - mu, logvar, mu_a, logvar_a = torch.split(z, split_size_or_sections=self.hdim, dim=-1) - # a_type = self.fc_at(a_t) - - return mu, logvar, mu_a, logvar_a - - -class Cls(nn.Module): - def __init__(self, hdim=256, attack_type=4): - super(Cls, self).__init__() - - self.fc = nn.Linear(hdim, attack_type) - - def forward(self, x): - a_cls = self.fc(x) - - return a_cls - -class Decoder(nn.Module): - def __init__(self, hdim=256): - super(Decoder, self).__init__() - - self.fc = nn.Sequential( - nn.Linear(hdim, 512 * 4 * 4), - nn.ReLU(True) - ) - - self.main = nn.Sequential( - make_layer(_Residual_Block, 1, 512, 512), - nn.Upsample(scale_factor=2, mode='nearest'), - make_layer(_Residual_Block, 1, 512, 512), - nn.Upsample(scale_factor=2, mode='nearest'), - make_layer(_Residual_Block, 1, 512, 512), - nn.Upsample(scale_factor=2, mode='nearest'), - make_layer(_Residual_Block, 1, 512, 256), - nn.Upsample(scale_factor=2, mode='nearest'), - make_layer(_Residual_Block, 1, 256, 128), - nn.Upsample(scale_factor=2, mode='nearest'), - make_layer(_Residual_Block, 2, 128, 64), - nn.Upsample(scale_factor=2, mode='nearest'), - nn.Conv2d(64, 3 + 3, 5, 1, 2) - ) - - def forward(self, z): - z = z.view(z.size(0), -1) - y = self.fc(z) - x = y.view(z.size(0), -1, 4, 4) - img = torch.sigmoid(self.main(x)) - return img - - -class Decoder_s(nn.Module): - def __init__(self, hdim=256): - super(Decoder_s, self).__init__() - - self.fc = nn.Sequential( - nn.Linear(3 * hdim, 512 * 4 * 4), - nn.ReLU(True) - ) - - self.main = nn.Sequential( - make_layer(_Residual_Block, 1, 512, 512), - nn.Upsample(scale_factor=2, mode='nearest'), - make_layer(_Residual_Block, 1, 512, 512), - nn.Upsample(scale_factor=2, mode='nearest'), - make_layer(_Residual_Block, 1, 512, 512), - nn.Upsample(scale_factor=2, mode='nearest'), - make_layer(_Residual_Block, 1, 512, 256), - nn.Upsample(scale_factor=2, mode='nearest'), - make_layer(_Residual_Block, 1, 256, 128), - nn.Upsample(scale_factor=2, mode='nearest'), - make_layer(_Residual_Block, 2, 128, 64), - nn.Upsample(scale_factor=2, mode='nearest'), - nn.Conv2d(64, 3 + 3, 5, 1, 2) - ) - - def forward(self, z): - z = z.view(z.size(0), -1) - y = self.fc(z) - x = y.view(z.size(0), -1, 4, 4) - img = torch.sigmoid(self.main(x)) - return img - - -class _Residual_Block(nn.Module): - def __init__(self, inc=64, outc=64, groups=1): - super(_Residual_Block, self).__init__() - - if inc is not outc: - self.conv_expand = nn.Conv2d(in_channels=inc, out_channels=outc, kernel_size=1, stride=1, padding=0, - groups=1, bias=False) - else: - self.conv_expand = None - - self.conv1 = nn.Conv2d(in_channels=inc, out_channels=outc, kernel_size=3, stride=1, padding=1, groups=groups, - bias=False) - self.bn1 = nn.InstanceNorm2d(outc, eps=0.001) - self.relu1 = nn.LeakyReLU(0.2, inplace=True) - self.conv2 = nn.Conv2d(in_channels=outc, out_channels=outc, kernel_size=3, stride=1, padding=1, groups=groups, - bias=False) - self.bn2 = nn.InstanceNorm2d(outc, eps=0.001) - self.relu2 = nn.LeakyReLU(0.2, inplace=True) - - def forward(self, x): - if self.conv_expand is not None: - identity_data = self.conv_expand(x) - else: - identity_data = x - - output = self.relu1(self.bn1(self.conv1(x))) - output = self.conv2(output) - output = self.relu2(self.bn2(torch.add(output, identity_data))) - return output - - -def make_layer(block, num_of_layer, inc=64, outc=64, groups=1): - if num_of_layer < 1: - num_of_layer = 1 - layers = [] - layers.append(block(inc=inc, outc=outc, groups=groups)) - for _ in range(1, num_of_layer): - layers.append(block(inc=outc, outc=outc, groups=groups)) - return nn.Sequential(*layers) +import numpy as np + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.autograd import Variable + + +class Encoder(nn.Module): + def __init__(self, hdim=256): + super(Encoder, self).__init__() + + self.hdim = hdim + + self.main = nn.Sequential( + nn.Conv2d(3, 32, 5, 2, 2, bias=False), + nn.InstanceNorm2d(32, eps=0.001), + nn.LeakyReLU(0.2), + nn.AvgPool2d(2), + make_layer(_Residual_Block, 1, 32, 64), + nn.AvgPool2d(2), + make_layer(_Residual_Block, 1, 64, 128), + nn.AvgPool2d(2), + make_layer(_Residual_Block, 1, 128, 256), + nn.AvgPool2d(2), + make_layer(_Residual_Block, 1, 256, 512), + nn.AvgPool2d(2), + make_layer(_Residual_Block, 1, 512, 512) + ) + + # mu and logvar + self.fc = nn.Linear(512 * 4 * 4, 2 * hdim) + + def forward(self, x): + z = self.main(x).view(x.size(0), -1) + z = self.fc(z) + mu, logvar = torch.split(z, split_size_or_sections=self.hdim, dim=-1) + + return mu, logvar + + +class Encoder_s(nn.Module): + def __init__(self, hdim=256): + super(Encoder_s, self).__init__() + + self.hdim = hdim + + self.main = nn.Sequential( + nn.Conv2d(3, 32, 5, 2, 2, bias=False), + nn.InstanceNorm2d(32, eps=0.001), + nn.LeakyReLU(0.2), + nn.AvgPool2d(2), + make_layer(_Residual_Block, 1, 32, 64), + nn.AvgPool2d(2), + make_layer(_Residual_Block, 1, 64, 128), + nn.AvgPool2d(2), + make_layer(_Residual_Block, 1, 128, 256), + nn.AvgPool2d(2), + make_layer(_Residual_Block, 1, 256, 512), + nn.AvgPool2d(2), + make_layer(_Residual_Block, 1, 512, 512) + ) + + # mu and logvar + self.fc = nn.Linear(512 * 4 * 4, 4 * hdim) + # self.fc_at = nn.Linear(hdim, attack_type) + + def forward(self, x): + z = self.main(x).view(x.size(0), -1) + z = self.fc(z) + mu, logvar, mu_a, logvar_a = torch.split(z, split_size_or_sections=self.hdim, dim=-1) + # a_type = self.fc_at(a_t) + + return mu, logvar, mu_a, logvar_a + + +class Cls(nn.Module): + def __init__(self, hdim=256, attack_type=4): + super(Cls, self).__init__() + + self.fc = nn.Linear(hdim, attack_type) + + def forward(self, x): + a_cls = self.fc(x) + + return a_cls + +class Decoder(nn.Module): + def __init__(self, hdim=256): + super(Decoder, self).__init__() + + self.fc = nn.Sequential( + nn.Linear(hdim, 512 * 4 * 4), + nn.ReLU(True) + ) + + self.main = nn.Sequential( + make_layer(_Residual_Block, 1, 512, 512), + nn.Upsample(scale_factor=2, mode='nearest'), + make_layer(_Residual_Block, 1, 512, 512), + nn.Upsample(scale_factor=2, mode='nearest'), + make_layer(_Residual_Block, 1, 512, 512), + nn.Upsample(scale_factor=2, mode='nearest'), + make_layer(_Residual_Block, 1, 512, 256), + nn.Upsample(scale_factor=2, mode='nearest'), + make_layer(_Residual_Block, 1, 256, 128), + nn.Upsample(scale_factor=2, mode='nearest'), + make_layer(_Residual_Block, 2, 128, 64), + nn.Upsample(scale_factor=2, mode='nearest'), + nn.Conv2d(64, 3 + 3, 5, 1, 2) + ) + + def forward(self, z): + z = z.view(z.size(0), -1) + y = self.fc(z) + x = y.view(z.size(0), -1, 4, 4) + img = torch.sigmoid(self.main(x)) + return img + + +class Decoder_s(nn.Module): + def __init__(self, hdim=256): + super(Decoder_s, self).__init__() + + self.fc = nn.Sequential( + nn.Linear(3 * hdim, 512 * 4 * 4), + nn.ReLU(True) + ) + + self.main = nn.Sequential( + make_layer(_Residual_Block, 1, 512, 512), + nn.Upsample(scale_factor=2, mode='nearest'), + make_layer(_Residual_Block, 1, 512, 512), + nn.Upsample(scale_factor=2, mode='nearest'), + make_layer(_Residual_Block, 1, 512, 512), + nn.Upsample(scale_factor=2, mode='nearest'), + make_layer(_Residual_Block, 1, 512, 256), + nn.Upsample(scale_factor=2, mode='nearest'), + make_layer(_Residual_Block, 1, 256, 128), + nn.Upsample(scale_factor=2, mode='nearest'), + make_layer(_Residual_Block, 2, 128, 64), + nn.Upsample(scale_factor=2, mode='nearest'), + nn.Conv2d(64, 3 + 3, 5, 1, 2) + ) + + def forward(self, z): + z = z.view(z.size(0), -1) + y = self.fc(z) + x = y.view(z.size(0), -1, 4, 4) + img = torch.sigmoid(self.main(x)) + return img + + +class _Residual_Block(nn.Module): + def __init__(self, inc=64, outc=64, groups=1): + super(_Residual_Block, self).__init__() + + if inc is not outc: + self.conv_expand = nn.Conv2d(in_channels=inc, out_channels=outc, kernel_size=1, stride=1, padding=0, + groups=1, bias=False) + else: + self.conv_expand = None + + self.conv1 = nn.Conv2d(in_channels=inc, out_channels=outc, kernel_size=3, stride=1, padding=1, groups=groups, + bias=False) + self.bn1 = nn.InstanceNorm2d(outc, eps=0.001) + self.relu1 = nn.LeakyReLU(0.2, inplace=True) + self.conv2 = nn.Conv2d(in_channels=outc, out_channels=outc, kernel_size=3, stride=1, padding=1, groups=groups, + bias=False) + self.bn2 = nn.InstanceNorm2d(outc, eps=0.001) + self.relu2 = nn.LeakyReLU(0.2, inplace=True) + + def forward(self, x): + if self.conv_expand is not None: + identity_data = self.conv_expand(x) + else: + identity_data = x + + output = self.relu1(self.bn1(self.conv1(x))) + output = self.conv2(output) + output = self.relu2(self.bn2(torch.add(output, identity_data))) + return output + + +def make_layer(block, num_of_layer, inc=64, outc=64, groups=1): + if num_of_layer < 1: + num_of_layer = 1 + layers = [] + layers.append(block(inc=inc, outc=outc, groups=groups)) + for _ in range(1, num_of_layer): + layers.append(block(inc=outc, outc=outc, groups=groups)) + return nn.Sequential(*layers) diff --git a/addition_module/DSDG/networks/light_cnn.py b/src/faceDetection/face_sdk/addition_module/DSDG/networks/light_cnn.py similarity index 97% rename from addition_module/DSDG/networks/light_cnn.py rename to src/faceDetection/face_sdk/addition_module/DSDG/networks/light_cnn.py index 37a7dee..4604727 100644 --- a/addition_module/DSDG/networks/light_cnn.py +++ b/src/faceDetection/face_sdk/addition_module/DSDG/networks/light_cnn.py @@ -1,98 +1,98 @@ -import torch -import torch.nn as nn -import torch.nn.functional as F - - -class mfm(nn.Module): - def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1, type=1): - super(mfm, self).__init__() - self.out_channels = out_channels - if type == 1: - self.filter = nn.Conv2d(in_channels, 2 * out_channels, kernel_size=kernel_size, stride=stride, - padding=padding) - else: - self.filter = nn.Linear(in_channels, 2 * out_channels) - - def forward(self, x): - x = self.filter(x) - out = torch.split(x, self.out_channels, 1) - return torch.max(out[0], out[1]) - - -class group(nn.Module): - def __init__(self, in_channels, out_channels, kernel_size, stride, padding): - super(group, self).__init__() - self.conv_a = mfm(in_channels, in_channels, 1, 1, 0) - self.conv = mfm(in_channels, out_channels, kernel_size, stride, padding) - - def forward(self, x): - x = self.conv_a(x) - x = self.conv(x) - return x - - -class resblock(nn.Module): - def __init__(self, in_channels, out_channels): - super(resblock, self).__init__() - self.conv1 = mfm(in_channels, out_channels, kernel_size=3, stride=1, padding=1) - self.conv2 = mfm(out_channels, out_channels, kernel_size=3, stride=1, padding=1) - - def forward(self, x): - res = x - out = self.conv1(x) - out = self.conv2(out) - out = out + res - return out - - -class network_29layers_v2(nn.Module): - def __init__(self, block, layers, is_train=False, num_classes=80013): - super(network_29layers_v2, self).__init__() - self.is_train = is_train - self.conv1 = mfm(1, 48, 5, 1, 2) - self.block1 = self._make_layer(block, layers[0], 48, 48) - self.group1 = group(48, 96, 3, 1, 1) - self.block2 = self._make_layer(block, layers[1], 96, 96) - self.group2 = group(96, 192, 3, 1, 1) - self.block3 = self._make_layer(block, layers[2], 192, 192) - self.group3 = group(192, 128, 3, 1, 1) - self.block4 = self._make_layer(block, layers[3], 128, 128) - self.group4 = group(128, 128, 3, 1, 1) - self.fc = nn.Linear(8 * 8 * 128, 256) - - if self.is_train: - self.fc2_ = nn.Linear(256, num_classes, bias=False) - - def _make_layer(self, block, num_blocks, in_channels, out_channels): - layers = [] - for i in range(0, num_blocks): - layers.append(block(in_channels, out_channels)) - return nn.Sequential(*layers) - - def forward(self, x): - x = self.conv1(x) - x = F.max_pool2d(x, 2) + F.avg_pool2d(x, 2) - - x = self.block1(x) - x = self.group1(x) - x = F.max_pool2d(x, 2) + F.avg_pool2d(x, 2) - - x = self.block2(x) - x = self.group2(x) - x = F.max_pool2d(x, 2) + F.avg_pool2d(x, 2) - - x = self.block3(x) - x = self.group3(x) - x = self.block4(x) - x = self.group4(x) - x = F.max_pool2d(x, 2) + F.avg_pool2d(x, 2) - - x = x.view(x.size(0), -1) - fc = self.fc(x) - - if self.is_train: - x = F.dropout(fc, training=self.training) - out = self.fc2_(x) - return out, fc - else: - return fc +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class mfm(nn.Module): + def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1, type=1): + super(mfm, self).__init__() + self.out_channels = out_channels + if type == 1: + self.filter = nn.Conv2d(in_channels, 2 * out_channels, kernel_size=kernel_size, stride=stride, + padding=padding) + else: + self.filter = nn.Linear(in_channels, 2 * out_channels) + + def forward(self, x): + x = self.filter(x) + out = torch.split(x, self.out_channels, 1) + return torch.max(out[0], out[1]) + + +class group(nn.Module): + def __init__(self, in_channels, out_channels, kernel_size, stride, padding): + super(group, self).__init__() + self.conv_a = mfm(in_channels, in_channels, 1, 1, 0) + self.conv = mfm(in_channels, out_channels, kernel_size, stride, padding) + + def forward(self, x): + x = self.conv_a(x) + x = self.conv(x) + return x + + +class resblock(nn.Module): + def __init__(self, in_channels, out_channels): + super(resblock, self).__init__() + self.conv1 = mfm(in_channels, out_channels, kernel_size=3, stride=1, padding=1) + self.conv2 = mfm(out_channels, out_channels, kernel_size=3, stride=1, padding=1) + + def forward(self, x): + res = x + out = self.conv1(x) + out = self.conv2(out) + out = out + res + return out + + +class network_29layers_v2(nn.Module): + def __init__(self, block, layers, is_train=False, num_classes=80013): + super(network_29layers_v2, self).__init__() + self.is_train = is_train + self.conv1 = mfm(1, 48, 5, 1, 2) + self.block1 = self._make_layer(block, layers[0], 48, 48) + self.group1 = group(48, 96, 3, 1, 1) + self.block2 = self._make_layer(block, layers[1], 96, 96) + self.group2 = group(96, 192, 3, 1, 1) + self.block3 = self._make_layer(block, layers[2], 192, 192) + self.group3 = group(192, 128, 3, 1, 1) + self.block4 = self._make_layer(block, layers[3], 128, 128) + self.group4 = group(128, 128, 3, 1, 1) + self.fc = nn.Linear(8 * 8 * 128, 256) + + if self.is_train: + self.fc2_ = nn.Linear(256, num_classes, bias=False) + + def _make_layer(self, block, num_blocks, in_channels, out_channels): + layers = [] + for i in range(0, num_blocks): + layers.append(block(in_channels, out_channels)) + return nn.Sequential(*layers) + + def forward(self, x): + x = self.conv1(x) + x = F.max_pool2d(x, 2) + F.avg_pool2d(x, 2) + + x = self.block1(x) + x = self.group1(x) + x = F.max_pool2d(x, 2) + F.avg_pool2d(x, 2) + + x = self.block2(x) + x = self.group2(x) + x = F.max_pool2d(x, 2) + F.avg_pool2d(x, 2) + + x = self.block3(x) + x = self.group3(x) + x = self.block4(x) + x = self.group4(x) + x = F.max_pool2d(x, 2) + F.avg_pool2d(x, 2) + + x = x.view(x.size(0), -1) + fc = self.fc(x) + + if self.is_train: + x = F.dropout(fc, training=self.training) + out = self.fc2_(x) + return out, fc + else: + return fc diff --git a/addition_module/DSDG/oulu_images/protocols/Protocol_1/Dev.txt b/src/faceDetection/face_sdk/addition_module/DSDG/oulu_images/protocols/Protocol_1/Dev.txt similarity index 100% rename from addition_module/DSDG/oulu_images/protocols/Protocol_1/Dev.txt rename to src/faceDetection/face_sdk/addition_module/DSDG/oulu_images/protocols/Protocol_1/Dev.txt diff --git a/addition_module/DSDG/oulu_images/protocols/Protocol_1/Test.txt b/src/faceDetection/face_sdk/addition_module/DSDG/oulu_images/protocols/Protocol_1/Test.txt similarity index 100% rename from addition_module/DSDG/oulu_images/protocols/Protocol_1/Test.txt rename to src/faceDetection/face_sdk/addition_module/DSDG/oulu_images/protocols/Protocol_1/Test.txt diff --git a/addition_module/DSDG/oulu_images/protocols/Protocol_1/Train.txt b/src/faceDetection/face_sdk/addition_module/DSDG/oulu_images/protocols/Protocol_1/Train.txt similarity index 100% rename from addition_module/DSDG/oulu_images/protocols/Protocol_1/Train.txt rename to src/faceDetection/face_sdk/addition_module/DSDG/oulu_images/protocols/Protocol_1/Train.txt diff --git a/addition_module/DSDG/oulu_images_crop/protocols/Protocol_1/Dev.txt b/src/faceDetection/face_sdk/addition_module/DSDG/oulu_images_crop/protocols/Protocol_1/Dev.txt similarity index 100% rename from addition_module/DSDG/oulu_images_crop/protocols/Protocol_1/Dev.txt rename to src/faceDetection/face_sdk/addition_module/DSDG/oulu_images_crop/protocols/Protocol_1/Dev.txt diff --git a/addition_module/DSDG/oulu_images_crop/protocols/Protocol_1/Test.txt b/src/faceDetection/face_sdk/addition_module/DSDG/oulu_images_crop/protocols/Protocol_1/Test.txt similarity index 100% rename from addition_module/DSDG/oulu_images_crop/protocols/Protocol_1/Test.txt rename to src/faceDetection/face_sdk/addition_module/DSDG/oulu_images_crop/protocols/Protocol_1/Test.txt diff --git a/addition_module/DSDG/oulu_images_crop/protocols/Protocol_1/Train_g.txt b/src/faceDetection/face_sdk/addition_module/DSDG/oulu_images_crop/protocols/Protocol_1/Train_g.txt similarity index 100% rename from addition_module/DSDG/oulu_images_crop/protocols/Protocol_1/Train_g.txt rename to src/faceDetection/face_sdk/addition_module/DSDG/oulu_images_crop/protocols/Protocol_1/Train_g.txt diff --git a/addition_module/DSDG/train_generator.py b/src/faceDetection/face_sdk/addition_module/DSDG/train_generator.py similarity index 97% rename from addition_module/DSDG/train_generator.py rename to src/faceDetection/face_sdk/addition_module/DSDG/train_generator.py index d1ba648..6b8f9ea 100644 --- a/addition_module/DSDG/train_generator.py +++ b/src/faceDetection/face_sdk/addition_module/DSDG/train_generator.py @@ -1,221 +1,221 @@ -import os -import time -import argparse -import numpy as np - -import torch -import torch.nn as nn -import torch.optim as optim -import torch.nn.functional as F -from torch.autograd import Variable -import torch.backends.cudnn as cudnn -import torchvision.utils as vutils - -from data import * -from networks import * -from misc import * - -parser = argparse.ArgumentParser() - -parser.add_argument('--gpu_ids', default='0', type=str) -parser.add_argument('--workers', default=8, type=int) -parser.add_argument('--lr', default=0.0002, type=float) -parser.add_argument('--batch_size', default=8, type=int) -parser.add_argument('--all_epochs', default=500, type=int) -parser.add_argument('--pre_epoch', default=0, type=int, help='train from previous model') -parser.add_argument('--hdim', default=128, type=int, help='dim of the latent code') -parser.add_argument('--attack_type', default=1, type=int, help='attack type number') -parser.add_argument('--out_path', default='results/', type=str, help='folder to sive output images') -parser.add_argument('--checkpoint_path', default='results/', type=str, help='folder to save output checkpoints') -parser.add_argument('--print_freq', type=int, default=20, help='print frequency') - -parser.add_argument('--test_epoch', default=100, type=int) -parser.add_argument('--save_epoch', default=1, type=int) - -parser.add_argument('--ip_model', default='', type=str, help='path of identity preserving model') -parser.add_argument('--img_root', default='', type=str, ) -parser.add_argument('--train_list', default='', type=str) - -parser.add_argument('--lambda_mmd', default=10, type=float) -parser.add_argument('--lambda_ip', default=10, type=float) -parser.add_argument('--lambda_pair', default=0.5, type=float) -parser.add_argument('--lambda_type', default=1, type=float) -parser.add_argument('--lambda_ort', default=0.1, type=float) - -""" -nir => spoof -vis => live -""" - - -def main(): - global args - args = parser.parse_args() - print(args) - - os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_ids - cudnn.benchmark = True - - if not os.path.exists(args.out_path): - os.makedirs(args.out_path) - - # generator - netE_nir, netE_vis, netG, netCls = define_G(hdim=args.hdim, attack_type=args.attack_type) - - if args.pre_epoch: - print('load pretrained model %d' % args.pre_epoch) - load_model(netE_nir, './model/netE_nir_model_epoch_%d_iter_0.pth' % args.pre_epoch) - load_model(netE_vis, './model/netE_vis_model_epoch_%d_iter_0.pth' % args.pre_epoch) - load_model(netG, './model/netG_model_epoch_%d_iter_0.pth' % args.pre_epoch) - - # IP net - netIP = define_IP() - - print("=> loading pretrained identity preserving model '{}'".format(args.ip_model)) - checkpoint = torch.load(args.ip_model) - - pretrained_dict = checkpoint['state_dict'] - model_dict = netIP.state_dict() - pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict} - model_dict.update(pretrained_dict) - netIP.load_state_dict(model_dict) - - for param in netIP.parameters(): - param.requires_grad = False - - # optimizer - optimizer = optim.Adam(list(netE_nir.parameters()) + list(netE_vis.parameters()) + list(netG.parameters()), - lr=args.lr) - - # criterion - criterion_type = torch.nn.CrossEntropyLoss().cuda() - criterionL2 = torch.nn.MSELoss().cuda() - - train_loader = torch.utils.data.DataLoader( - GenDataset_s(img_root=args.img_root, list_file=args.train_list, attack_type=args.attack_type), - batch_size=args.batch_size, shuffle=True, - num_workers=args.workers, pin_memory=True) - - # train - start_epoch = args.pre_epoch + 1 - for epoch in range(start_epoch, args.all_epochs + 1): - batch_time = AverageMeter() - data_time = AverageMeter() - - netE_nir.train() - netE_vis.train() - netG.train() - netIP.eval() - - start_time = time.time() - for iteration, batch in enumerate(train_loader): - data_time.update(time.time() - start_time) - - img_nir = Variable(batch['0'].cuda()) - img_vis = Variable(batch['1'].cuda()) - label_spoof = batch['type'].cuda() - - img = torch.cat((img_nir, img_vis), 1) - - # encoder forward - mu_nir, logvar_nir, mu_a, logvar_a = netE_nir(img_nir) - mu_vis, logvar_vis = netE_vis(img_vis) - - # reparameterization - z_cls = reparameterize(mu_a, logvar_a) - z_nir = reparameterize(mu_nir, logvar_nir) - z_vis = reparameterize(mu_vis, logvar_vis) - - # classify - pre_spoof = netCls(z_cls) - - # generator - rec = netG(torch.cat((z_cls, z_nir, z_vis), dim=1)) - - # vae loss - loss_rec = reconstruction_loss(rec, img, True) / 2.0 - loss_kl = (kl_loss(mu_nir, logvar_nir).mean() + kl_loss(mu_vis, logvar_vis).mean() - + kl_loss(mu_a, logvar_a).mean()) / 3.0 - - # mmd loss - loss_mmd = args.lambda_mmd * torch.abs(z_nir.mean(dim=0) - z_vis.mean(dim=0)).mean() - - # classify loss - loss_cls = args.lambda_type * criterion_type(pre_spoof, label_spoof) - - # Angular Orthogonal Loss - loss_ort = args.lambda_ort * torch.abs((z_cls * z_nir).sum(dim=1).mean()) - - # identity preserving loss - rec_nir = rec[:, 0:3, :, :] - rec_vis = rec[:, 3:6, :, :] - - # 256to128 - img_nir = F.interpolate(img_nir, size=(128, 128), mode='bilinear') - img_vis = F.interpolate(img_vis, size=(128, 128), mode='bilinear') - rec_nir = F.interpolate(rec_nir, size=(128, 128), mode='bilinear') - rec_vis = F.interpolate(rec_vis, size=(128, 128), mode='bilinear') - - nir_fc = netIP(rgb2gray(img_nir)) - vis_fc = netIP(rgb2gray(img_vis)) - rec_nir_fc = netIP(rgb2gray(rec_nir)) - rec_vis_fc = netIP(rgb2gray(rec_vis)) - - nir_fc = F.normalize(nir_fc, p=2, dim=1) - vis_fc = F.normalize(vis_fc, p=2, dim=1) - rec_nir_fc = F.normalize(rec_nir_fc, p=2, dim=1) - rec_vis_fc = F.normalize(rec_vis_fc, p=2, dim=1) - - loss_ip = args.lambda_ip * ( - criterionL2(rec_nir_fc, nir_fc.detach()) + criterionL2(rec_vis_fc, vis_fc.detach())) / 2.0 - loss_pair = args.lambda_pair * criterionL2(rec_nir_fc, rec_vis_fc) - - # - if epoch < 2: - loss = loss_rec + 0.01 * loss_kl + 0.01 * loss_mmd + 0.01 * loss_ip + 0.01 * loss_pair + \ - 0.01 * loss_cls + 0.01 * loss_ort - else: - loss = loss_rec + loss_kl + loss_mmd + loss_ip + loss_pair + loss_cls + loss_ort - - optimizer.zero_grad() - loss.backward() - optimizer.step() - - batch_time.update(time.time() - start_time) - start_time = time.time() - - if iteration % args.print_freq == 0: - info = '====> Epoch: [{:0>4d}][{:3d}/{:3d}] Batch_time: {:4.3f} Data_time: {:4.3f} | '.format( - epoch, iteration, len(train_loader), batch_time.avg, data_time.avg) - - info += 'Loss: rec: {:4.3f} kl: {:4.3f} mmd: {:4.3f} ip: {:4.3f} pair: {:4.3f} cls: {:4.3f} ' \ - 'ort: {:4.3f}'.format(loss_rec.item(), loss_kl.item(), loss_mmd.item(), loss_ip.item(), - loss_pair.item(), loss_cls.item(), loss_ort.item()) - - print(info) - - # test - if epoch % args.test_epoch == 0 or epoch == 1: - noise = torch.zeros(args.batch_size, args.hdim).normal_(0, 1) - noise_s = torch.zeros(args.batch_size, args.hdim).normal_(0, 1) - noise = torch.cat((noise_s, noise, noise), dim=1) - noise = noise.cuda() - - fake = netG(noise) - - vutils.save_image(img_nir.data, '{}/Epoch_{:03d}_img_spoof.png'.format(args.out_path, epoch)) - vutils.save_image(img_vis.data, '{}/Epoch_{:03d}_img_live.png'.format(args.out_path, epoch)) - vutils.save_image(rec_nir.data, '{}/Epoch_{:03d}_rec_spoof.png'.format(args.out_path, epoch)) - vutils.save_image(rec_vis.data, '{}/Epoch_{:03d}_rec_live.png'.format(args.out_path, epoch)) - vutils.save_image(fake[:, 0:3, :, :].data, '{}/Epoch_{:03d}_fake_spoof.png'.format(args.out_path, epoch)) - vutils.save_image(fake[:, 3:6, :, :].data, '{}/Epoch_{:03d}_fake_live.png'.format(args.out_path, epoch)) - - # save models - if epoch % args.save_epoch == 0 or epoch == 1: - save_checkpoint(args.checkpoint_path, netE_nir, epoch, 0, 'netE_spoof_') - save_checkpoint(args.checkpoint_path, netE_vis, epoch, 0, 'netE_live_') - save_checkpoint(args.checkpoint_path, netG, epoch, 0, 'netG_') - - -if __name__ == "__main__": - main() +import os +import time +import argparse +import numpy as np + +import torch +import torch.nn as nn +import torch.optim as optim +import torch.nn.functional as F +from torch.autograd import Variable +import torch.backends.cudnn as cudnn +import torchvision.utils as vutils + +from data import * +from networks import * +from misc import * + +parser = argparse.ArgumentParser() + +parser.add_argument('--gpu_ids', default='0', type=str) +parser.add_argument('--workers', default=8, type=int) +parser.add_argument('--lr', default=0.0002, type=float) +parser.add_argument('--batch_size', default=8, type=int) +parser.add_argument('--all_epochs', default=500, type=int) +parser.add_argument('--pre_epoch', default=0, type=int, help='train from previous model') +parser.add_argument('--hdim', default=128, type=int, help='dim of the latent code') +parser.add_argument('--attack_type', default=1, type=int, help='attack type number') +parser.add_argument('--out_path', default='results/', type=str, help='folder to sive output images') +parser.add_argument('--checkpoint_path', default='results/', type=str, help='folder to save output checkpoints') +parser.add_argument('--print_freq', type=int, default=20, help='print frequency') + +parser.add_argument('--test_epoch', default=100, type=int) +parser.add_argument('--save_epoch', default=1, type=int) + +parser.add_argument('--ip_model', default='', type=str, help='path of identity preserving model') +parser.add_argument('--img_root', default='', type=str, ) +parser.add_argument('--train_list', default='', type=str) + +parser.add_argument('--lambda_mmd', default=10, type=float) +parser.add_argument('--lambda_ip', default=10, type=float) +parser.add_argument('--lambda_pair', default=0.5, type=float) +parser.add_argument('--lambda_type', default=1, type=float) +parser.add_argument('--lambda_ort', default=0.1, type=float) + +""" +nir => spoof +vis => live +""" + + +def main(): + global args + args = parser.parse_args() + print(args) + + os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_ids + cudnn.benchmark = True + + if not os.path.exists(args.out_path): + os.makedirs(args.out_path) + + # generator + netE_nir, netE_vis, netG, netCls = define_G(hdim=args.hdim, attack_type=args.attack_type) + + if args.pre_epoch: + print('load pretrained model %d' % args.pre_epoch) + load_model(netE_nir, './model/netE_nir_model_epoch_%d_iter_0.pth' % args.pre_epoch) + load_model(netE_vis, './model/netE_vis_model_epoch_%d_iter_0.pth' % args.pre_epoch) + load_model(netG, './model/netG_model_epoch_%d_iter_0.pth' % args.pre_epoch) + + # IP net + netIP = define_IP() + + print("=> loading pretrained identity preserving model '{}'".format(args.ip_model)) + checkpoint = torch.load(args.ip_model) + + pretrained_dict = checkpoint['state_dict'] + model_dict = netIP.state_dict() + pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict} + model_dict.update(pretrained_dict) + netIP.load_state_dict(model_dict) + + for param in netIP.parameters(): + param.requires_grad = False + + # optimizer + optimizer = optim.Adam(list(netE_nir.parameters()) + list(netE_vis.parameters()) + list(netG.parameters()), + lr=args.lr) + + # criterion + criterion_type = torch.nn.CrossEntropyLoss().cuda() + criterionL2 = torch.nn.MSELoss().cuda() + + train_loader = torch.utils.data.DataLoader( + GenDataset_s(img_root=args.img_root, list_file=args.train_list, attack_type=args.attack_type), + batch_size=args.batch_size, shuffle=True, + num_workers=args.workers, pin_memory=True) + + # train + start_epoch = args.pre_epoch + 1 + for epoch in range(start_epoch, args.all_epochs + 1): + batch_time = AverageMeter() + data_time = AverageMeter() + + netE_nir.train() + netE_vis.train() + netG.train() + netIP.eval() + + start_time = time.time() + for iteration, batch in enumerate(train_loader): + data_time.update(time.time() - start_time) + + img_nir = Variable(batch['0'].cuda()) + img_vis = Variable(batch['1'].cuda()) + label_spoof = batch['type'].cuda() + + img = torch.cat((img_nir, img_vis), 1) + + # encoder forward + mu_nir, logvar_nir, mu_a, logvar_a = netE_nir(img_nir) + mu_vis, logvar_vis = netE_vis(img_vis) + + # reparameterization + z_cls = reparameterize(mu_a, logvar_a) + z_nir = reparameterize(mu_nir, logvar_nir) + z_vis = reparameterize(mu_vis, logvar_vis) + + # classify + pre_spoof = netCls(z_cls) + + # generator + rec = netG(torch.cat((z_cls, z_nir, z_vis), dim=1)) + + # vae loss + loss_rec = reconstruction_loss(rec, img, True) / 2.0 + loss_kl = (kl_loss(mu_nir, logvar_nir).mean() + kl_loss(mu_vis, logvar_vis).mean() + + kl_loss(mu_a, logvar_a).mean()) / 3.0 + + # mmd loss + loss_mmd = args.lambda_mmd * torch.abs(z_nir.mean(dim=0) - z_vis.mean(dim=0)).mean() + + # classify loss + loss_cls = args.lambda_type * criterion_type(pre_spoof, label_spoof) + + # Angular Orthogonal Loss + loss_ort = args.lambda_ort * torch.abs((z_cls * z_nir).sum(dim=1).mean()) + + # identity preserving loss + rec_nir = rec[:, 0:3, :, :] + rec_vis = rec[:, 3:6, :, :] + + # 256to128 + img_nir = F.interpolate(img_nir, size=(128, 128), mode='bilinear') + img_vis = F.interpolate(img_vis, size=(128, 128), mode='bilinear') + rec_nir = F.interpolate(rec_nir, size=(128, 128), mode='bilinear') + rec_vis = F.interpolate(rec_vis, size=(128, 128), mode='bilinear') + + nir_fc = netIP(rgb2gray(img_nir)) + vis_fc = netIP(rgb2gray(img_vis)) + rec_nir_fc = netIP(rgb2gray(rec_nir)) + rec_vis_fc = netIP(rgb2gray(rec_vis)) + + nir_fc = F.normalize(nir_fc, p=2, dim=1) + vis_fc = F.normalize(vis_fc, p=2, dim=1) + rec_nir_fc = F.normalize(rec_nir_fc, p=2, dim=1) + rec_vis_fc = F.normalize(rec_vis_fc, p=2, dim=1) + + loss_ip = args.lambda_ip * ( + criterionL2(rec_nir_fc, nir_fc.detach()) + criterionL2(rec_vis_fc, vis_fc.detach())) / 2.0 + loss_pair = args.lambda_pair * criterionL2(rec_nir_fc, rec_vis_fc) + + # + if epoch < 2: + loss = loss_rec + 0.01 * loss_kl + 0.01 * loss_mmd + 0.01 * loss_ip + 0.01 * loss_pair + \ + 0.01 * loss_cls + 0.01 * loss_ort + else: + loss = loss_rec + loss_kl + loss_mmd + loss_ip + loss_pair + loss_cls + loss_ort + + optimizer.zero_grad() + loss.backward() + optimizer.step() + + batch_time.update(time.time() - start_time) + start_time = time.time() + + if iteration % args.print_freq == 0: + info = '====> Epoch: [{:0>4d}][{:3d}/{:3d}] Batch_time: {:4.3f} Data_time: {:4.3f} | '.format( + epoch, iteration, len(train_loader), batch_time.avg, data_time.avg) + + info += 'Loss: rec: {:4.3f} kl: {:4.3f} mmd: {:4.3f} ip: {:4.3f} pair: {:4.3f} cls: {:4.3f} ' \ + 'ort: {:4.3f}'.format(loss_rec.item(), loss_kl.item(), loss_mmd.item(), loss_ip.item(), + loss_pair.item(), loss_cls.item(), loss_ort.item()) + + print(info) + + # test + if epoch % args.test_epoch == 0 or epoch == 1: + noise = torch.zeros(args.batch_size, args.hdim).normal_(0, 1) + noise_s = torch.zeros(args.batch_size, args.hdim).normal_(0, 1) + noise = torch.cat((noise_s, noise, noise), dim=1) + noise = noise.cuda() + + fake = netG(noise) + + vutils.save_image(img_nir.data, '{}/Epoch_{:03d}_img_spoof.png'.format(args.out_path, epoch)) + vutils.save_image(img_vis.data, '{}/Epoch_{:03d}_img_live.png'.format(args.out_path, epoch)) + vutils.save_image(rec_nir.data, '{}/Epoch_{:03d}_rec_spoof.png'.format(args.out_path, epoch)) + vutils.save_image(rec_vis.data, '{}/Epoch_{:03d}_rec_live.png'.format(args.out_path, epoch)) + vutils.save_image(fake[:, 0:3, :, :].data, '{}/Epoch_{:03d}_fake_spoof.png'.format(args.out_path, epoch)) + vutils.save_image(fake[:, 3:6, :, :].data, '{}/Epoch_{:03d}_fake_live.png'.format(args.out_path, epoch)) + + # save models + if epoch % args.save_epoch == 0 or epoch == 1: + save_checkpoint(args.checkpoint_path, netE_nir, epoch, 0, 'netE_spoof_') + save_checkpoint(args.checkpoint_path, netE_vis, epoch, 0, 'netE_live_') + save_checkpoint(args.checkpoint_path, netG, epoch, 0, 'netG_') + + +if __name__ == "__main__": + main() diff --git a/addition_module/DSDG/train_generator.sh b/src/faceDetection/face_sdk/addition_module/DSDG/train_generator.sh similarity index 100% rename from addition_module/DSDG/train_generator.sh rename to src/faceDetection/face_sdk/addition_module/DSDG/train_generator.sh diff --git a/addition_module/DSDG/train_list/train_list_oulu_p1.txt b/src/faceDetection/face_sdk/addition_module/DSDG/train_list/train_list_oulu_p1.txt similarity index 100% rename from addition_module/DSDG/train_list/train_list_oulu_p1.txt rename to src/faceDetection/face_sdk/addition_module/DSDG/train_list/train_list_oulu_p1.txt diff --git a/addition_module/HSST/HSST_Prototype.py b/src/faceDetection/face_sdk/addition_module/HSST/HSST_Prototype.py similarity index 100% rename from addition_module/HSST/HSST_Prototype.py rename to src/faceDetection/face_sdk/addition_module/HSST/HSST_Prototype.py diff --git a/addition_module/HSST/README.md b/src/faceDetection/face_sdk/addition_module/HSST/README.md similarity index 100% rename from addition_module/HSST/README.md rename to src/faceDetection/face_sdk/addition_module/HSST/README.md diff --git a/src/faceDetection/face_sdk/addition_module/HSST/__init__.py b/src/faceDetection/face_sdk/addition_module/HSST/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/addition_module/HSST/train.py b/src/faceDetection/face_sdk/addition_module/HSST/train.py similarity index 94% rename from addition_module/HSST/train.py rename to src/faceDetection/face_sdk/addition_module/HSST/train.py index b7dbe28..6b4c11d 100644 --- a/addition_module/HSST/train.py +++ b/src/faceDetection/face_sdk/addition_module/HSST/train.py @@ -1,192 +1,192 @@ -""" -@author: Jun Wang -@date: 20201019 -@contact: jun21wangustc@gmail.com -""" -import os -import sys -import shutil -import argparse -import logging as logger - -import torch -from torch import optim -from torch.utils.data import DataLoader -from tensorboardX import SummaryWriter - -sys.path.append('../../') -from utils.AverageMeter import AverageMeter -from data_processor.train_dataset import ImageDataset_HSST -from backbone.backbone_def import BackboneFactory -from head.head_def import HeadFactory -from test_protocol.utils.online_val import Evaluator - -logger.basicConfig(level=logger.INFO, - format='%(levelname)s %(asctime)s %(filename)s: %(lineno)d] %(message)s', - datefmt='%Y-%m-%d %H:%M:%S') - -def get_lr(optimizer): - """Get the current learning rate from optimizer. - """ - for param_group in optimizer.param_groups: - return param_group['lr'] - -def moving_average(probe, gallery, alpha): - """Update the gallery-set network in the momentum way.(MoCo) - """ - for param_probe, param_gallery in zip(probe.parameters(), gallery.parameters()): - param_gallery.data = \ - alpha* param_gallery.data + (1 - alpha) * param_probe.detach().data - -def train_BN(m): - classname = m.__class__.__name__ - if classname.find('BatchNorm') != -1: - m.train() - -def shuffle_BN(batch_size): - """ShuffleBN for batch, the same as MoCo https://arxiv.org/abs/1911.05722 ####### - """ - shuffle_ids = torch.randperm(batch_size).long().cuda() - reshuffle_ids = torch.zeros(batch_size).long().cuda() - reshuffle_ids.index_copy_(0, shuffle_ids, torch.arange(batch_size).long().cuda()) - return shuffle_ids, reshuffle_ids - -def train_one_epoch(data_loader, probe_net, gallery_net, prototype, optimizer, - criterion, cur_epoch, conf, loss_meter): - """Tain one epoch by semi-siamese training. - """ - for batch_idx, (nir_img, vis_img, id_indexes) in enumerate(data_loader): - batch_size = nir_img.size(0) - global_batch_idx = cur_epoch * len(data_loader) + batch_idx - nir_img = nir_img.cuda() - vis_img = vis_img.cuda() - # set inputs as probe or gallery - shuffle_ids, reshuffle_ids = shuffle_BN(batch_size) - nir_img_probe = probe_net(nir_img) - with torch.no_grad(): - vis_img = vis_img[shuffle_ids] - vis_img_gallery = gallery_net(vis_img)[reshuffle_ids] - vis_img = vis_img[reshuffle_ids] - shuffle_ids, reshuffle_ids = shuffle_BN(batch_size) - vis_img_probe = probe_net(vis_img) - with torch.no_grad(): - nir_img = nir_img[shuffle_ids] - nir_img_gallery = gallery_net(nir_img)[reshuffle_ids] - nir_img = nir_img[reshuffle_ids] - output1, output2, label, id_set = prototype( - nir_img_probe, vis_img_gallery, vis_img_probe, nir_img_gallery, id_indexes) - loss = (criterion(output1, label) + criterion(output2, label))/2 - # update - optimizer.zero_grad() - loss.backward() - optimizer.step() - moving_average(probe_net, gallery_net, conf.alpha) - loss_meter.update(loss.item(), batch_size) - if batch_idx % conf.print_freq == 0: - loss_val = loss_meter.avg - lr = get_lr(optimizer) - logger.info('Epoch %d, iter %d, lr %f, loss %f' % - (cur_epoch, batch_idx, lr, loss_val)) - conf.writer.add_scalar('Train_loss', loss_val, global_batch_idx) - conf.writer.add_scalar('Train_lr', lr, global_batch_idx) - if cur_epoch % conf.save_freq == 0 or cur_epoch == conf.epoches - 1: - saved_name = ('Epoch_{}.pt'.format(cur_epoch)) - torch.save(probe_net.state_dict(), os.path.join(conf.out_dir, saved_name)) - logger.info('save checkpoint %s to disk...' % saved_name) - return id_set - -def train(conf): - """Total training procedure. - """ - conf.device = torch.device('cuda:0') - criterion = torch.nn.CrossEntropyLoss().cuda(conf.device) - backbone_factory = BackboneFactory(conf.backbone_type, conf.backbone_conf_file) - probe_net = backbone_factory.get_backbone() - gallery_net = backbone_factory.get_backbone() - head_factory = HeadFactory(conf.head_type, conf.head_conf_file) - prototype = head_factory.get_head().cuda(conf.device) - probe_net = torch.nn.DataParallel(probe_net).cuda() - gallery_net = torch.nn.DataParallel(gallery_net).cuda() - optimizer = optim.SGD(probe_net.parameters(), lr=conf.lr, momentum=conf.momentum, weight_decay=5e-4) - lr_schedule = optim.lr_scheduler.MultiStepLR(optimizer, milestones=conf.milestones, gamma=0.1) - if conf.resume: - probe_net.load_state_dict(torch.load(args.pretrain_model)) - moving_average(probe_net, gallery_net, 0) - probe_net.train() - gallery_net.eval().apply(train_BN) - - exclude_id_set = set() - loss_meter = AverageMeter() - for epoch in range(conf.epoches): - data_loader = DataLoader( - ImageDataset_HSST(conf.data_root, conf.train_file, exclude_id_set), - conf.batch_size, True, num_workers = 4, drop_last = True) - exclude_id_set = train_one_epoch(data_loader, probe_net, gallery_net, - prototype, optimizer, criterion, epoch, conf, loss_meter) - lr_schedule.step() - if conf.evaluate: - conf.evaluator.evaluate(probe_net) - -if __name__ == '__main__': - conf = argparse.ArgumentParser(description='semi-siamese_training for face recognition.') - conf.add_argument("--data_root", type = str, - help = "The root folder of training set.") - conf.add_argument("--train_file", type = str, - help = "The train file path.") - conf.add_argument('--backbone_type', type=str, default='Mobilefacenets', - help='Mobilefacenets, Resnet.') - conf.add_argument('--backbone_conf_file', type=str, - help='the path of backbone_conf.yaml.') - conf.add_argument("--head_type", type = str, - help = "mv-softmax, arcface, npc-face ...") - conf.add_argument("--head_conf_file", type = str, - help = "the path of head_conf.yaml..") - conf.add_argument('--lr', type = float, default = 0.1, - help='The initial learning rate.') - conf.add_argument("--out_dir", type=str, default='out_dir', - help=" The folder to save models.") - conf.add_argument('--epoches', type = int, default = 130, - help = 'The training epoches.') - conf.add_argument('--step', type = str, default = '60,100,120', - help = 'Step for lr.') - conf.add_argument('--print_freq', type = int, default = 10, - help = 'The print frequency for training state.') - conf.add_argument('--save_freq', type=int, default=1, - help='The save frequency for training state.') - conf.add_argument('--batch_size', type=int, default=128, - help='batch size over all gpus.') - conf.add_argument('--momentum', type=float, default=0.9, - help='The momentum for sgd.') - conf.add_argument('--alpha', type=float, default=0.999, - help='weight of moving_average') - conf.add_argument('--log_dir', type = str, default = 'log', - help = 'The directory to save log.log') - conf.add_argument('--tensorboardx_logdir', type = str, - help = 'The directory to save tensorboardx logs') - conf.add_argument('--pretrain_model', type = str, default = 'mv_epoch_8.pt', - help = 'The path of pretrained model') - conf.add_argument('--resume', '-r', action = 'store_true', default = False, - help = 'Resume from checkpoint or not.') - conf.add_argument('--evaluate', '-e', action = 'store_true', default = False, - help = 'Evaluate the training model.') - conf.add_argument('--test_set', type = str, default = 'LFW', - help = 'Test set to evaluate the model.') - conf.add_argument('--test_data_conf_file', type = str, - help = 'The path of test data conf file.') - args = conf.parse_args() - args.milestones = [int(num) for num in args.step.split(',')] - if not os.path.exists(args.out_dir): - os.makedirs(args.out_dir) - if not os.path.exists(args.log_dir): - os.makedirs(args.log_dir) - tensorboardx_logdir = os.path.join(args.log_dir, args.tensorboardx_logdir) - if os.path.exists(tensorboardx_logdir): - shutil.rmtree(tensorboardx_logdir) - writer = SummaryWriter(log_dir=tensorboardx_logdir) - args.writer = writer - if args.evaluate: - args.evaluator = Evaluator(args.test_set, args.test_data_conf_file) - logger.info('Start optimization.') - logger.info(args) - train(args) - logger.info('Optimization done!') +""" +@author: Jun Wang +@date: 20201019 +@contact: jun21wangustc@gmail.com +""" +import os +import sys +import shutil +import argparse +import logging as logger + +import torch +from torch import optim +from torch.utils.data import DataLoader +from tensorboardX import SummaryWriter + +sys.path.append('../../../../../') +from src.faceDetection.face_sdk.utils.AverageMeter import AverageMeter +from src.faceDetection.face_sdk.data_processor.train_dataset import ImageDataset_HSST +from src.faceDetection.face_sdk.backbone.backbone_def import BackboneFactory +from src.faceDetection.face_sdk.head.head_def import HeadFactory +from src.faceDetection.face_sdk.test_protocol.utils.online_val import Evaluator + +logger.basicConfig(level=logger.INFO, + format='%(levelname)s %(asctime)s %(filename)s: %(lineno)d] %(message)s', + datefmt='%Y-%m-%d %H:%M:%S') + +def get_lr(optimizer): + """Get the current learning rate from optimizer. + """ + for param_group in optimizer.param_groups: + return param_group['lr'] + +def moving_average(probe, gallery, alpha): + """Update the gallery-set network in the momentum way.(MoCo) + """ + for param_probe, param_gallery in zip(probe.parameters(), gallery.parameters()): + param_gallery.data = \ + alpha* param_gallery.data + (1 - alpha) * param_probe.detach().data + +def train_BN(m): + classname = m.__class__.__name__ + if classname.find('BatchNorm') != -1: + m.train() + +def shuffle_BN(batch_size): + """ShuffleBN for batch, the same as MoCo https://arxiv.org/abs/1911.05722 ####### + """ + shuffle_ids = torch.randperm(batch_size).long().cuda() + reshuffle_ids = torch.zeros(batch_size).long().cuda() + reshuffle_ids.index_copy_(0, shuffle_ids, torch.arange(batch_size).long().cuda()) + return shuffle_ids, reshuffle_ids + +def train_one_epoch(data_loader, probe_net, gallery_net, prototype, optimizer, + criterion, cur_epoch, conf, loss_meter): + """Tain one epoch by semi-siamese training. + """ + for batch_idx, (nir_img, vis_img, id_indexes) in enumerate(data_loader): + batch_size = nir_img.size(0) + global_batch_idx = cur_epoch * len(data_loader) + batch_idx + nir_img = nir_img.cuda() + vis_img = vis_img.cuda() + # set inputs as probe or gallery + shuffle_ids, reshuffle_ids = shuffle_BN(batch_size) + nir_img_probe = probe_net(nir_img) + with torch.no_grad(): + vis_img = vis_img[shuffle_ids] + vis_img_gallery = gallery_net(vis_img)[reshuffle_ids] + vis_img = vis_img[reshuffle_ids] + shuffle_ids, reshuffle_ids = shuffle_BN(batch_size) + vis_img_probe = probe_net(vis_img) + with torch.no_grad(): + nir_img = nir_img[shuffle_ids] + nir_img_gallery = gallery_net(nir_img)[reshuffle_ids] + nir_img = nir_img[reshuffle_ids] + output1, output2, label, id_set = prototype( + nir_img_probe, vis_img_gallery, vis_img_probe, nir_img_gallery, id_indexes) + loss = (criterion(output1, label) + criterion(output2, label))/2 + # update + optimizer.zero_grad() + loss.backward() + optimizer.step() + moving_average(probe_net, gallery_net, conf.alpha) + loss_meter.update(loss.item(), batch_size) + if batch_idx % conf.print_freq == 0: + loss_val = loss_meter.avg + lr = get_lr(optimizer) + logger.info('Epoch %d, iter %d, lr %f, loss %f' % + (cur_epoch, batch_idx, lr, loss_val)) + conf.writer.add_scalar('Train_loss', loss_val, global_batch_idx) + conf.writer.add_scalar('Train_lr', lr, global_batch_idx) + if cur_epoch % conf.save_freq == 0 or cur_epoch == conf.epoches - 1: + saved_name = ('Epoch_{}.pt'.format(cur_epoch)) + torch.save(probe_net.state_dict(), os.path.join(conf.out_dir, saved_name)) + logger.info('save checkpoint %s to disk...' % saved_name) + return id_set + +def train(conf): + """Total training procedure. + """ + conf.device = torch.device('cuda:0') + criterion = torch.nn.CrossEntropyLoss().cuda(conf.device) + backbone_factory = BackboneFactory(conf.backbone_type, conf.backbone_conf_file) + probe_net = backbone_factory.get_backbone() + gallery_net = backbone_factory.get_backbone() + head_factory = HeadFactory(conf.head_type, conf.head_conf_file) + prototype = head_factory.get_head().cuda(conf.device) + probe_net = torch.nn.DataParallel(probe_net).cuda() + gallery_net = torch.nn.DataParallel(gallery_net).cuda() + optimizer = optim.SGD(probe_net.parameters(), lr=conf.lr, momentum=conf.momentum, weight_decay=5e-4) + lr_schedule = optim.lr_scheduler.MultiStepLR(optimizer, milestones=conf.milestones, gamma=0.1) + if conf.resume: + probe_net.load_state_dict(torch.load(args.pretrain_model)) + moving_average(probe_net, gallery_net, 0) + probe_net.train() + gallery_net.eval().apply(train_BN) + + exclude_id_set = set() + loss_meter = AverageMeter() + for epoch in range(conf.epoches): + data_loader = DataLoader( + ImageDataset_HSST(conf.data_root, conf.train_file, exclude_id_set), + conf.batch_size, True, num_workers = 4, drop_last = True) + exclude_id_set = train_one_epoch(data_loader, probe_net, gallery_net, + prototype, optimizer, criterion, epoch, conf, loss_meter) + lr_schedule.step() + if conf.evaluate: + conf.evaluator.evaluate(probe_net) + +if __name__ == '__main__': + conf = argparse.ArgumentParser(description='semi-siamese_training for face recognition.') + conf.add_argument("--data_root", type = str, + help = "The root folder of training set.") + conf.add_argument("--train_file", type = str, + help = "The train file path.") + conf.add_argument('--backbone_type', type=str, default='Mobilefacenets', + help='Mobilefacenets, Resnet.') + conf.add_argument('--backbone_conf_file', type=str, + help='the path of backbone_conf.yaml.') + conf.add_argument("--head_type", type = str, + help = "mv-softmax, arcface, npc-face ...") + conf.add_argument("--head_conf_file", type = str, + help = "the path of head_conf.yaml..") + conf.add_argument('--lr', type = float, default = 0.1, + help='The initial learning rate.') + conf.add_argument("--out_dir", type=str, default='out_dir', + help=" The folder to save models.") + conf.add_argument('--epoches', type = int, default = 130, + help = 'The training epoches.') + conf.add_argument('--step', type = str, default = '60,100,120', + help = 'Step for lr.') + conf.add_argument('--print_freq', type = int, default = 10, + help = 'The print frequency for training state.') + conf.add_argument('--save_freq', type=int, default=1, + help='The save frequency for training state.') + conf.add_argument('--batch_size', type=int, default=128, + help='batch size over all gpus.') + conf.add_argument('--momentum', type=float, default=0.9, + help='The momentum for sgd.') + conf.add_argument('--alpha', type=float, default=0.999, + help='weight of moving_average') + conf.add_argument('--log_dir', type = str, default = 'log', + help = 'The directory to save log.log') + conf.add_argument('--tensorboardx_logdir', type = str, + help = 'The directory to save tensorboardx logs') + conf.add_argument('--pretrain_model', type = str, default = 'mv_epoch_8.pt', + help = 'The path of pretrained model') + conf.add_argument('--resume', '-r', action = 'store_true', default = False, + help = 'Resume from checkpoint or not.') + conf.add_argument('--evaluate', '-e', action = 'store_true', default = False, + help = 'Evaluate the training model.') + conf.add_argument('--test_set', type = str, default = 'LFW', + help = 'Test set to evaluate the model.') + conf.add_argument('--test_data_conf_file', type = str, + help = 'The path of test data conf file.') + args = conf.parse_args() + args.milestones = [int(num) for num in args.step.split(',')] + if not os.path.exists(args.out_dir): + os.makedirs(args.out_dir) + if not os.path.exists(args.log_dir): + os.makedirs(args.log_dir) + tensorboardx_logdir = os.path.join(args.log_dir, args.tensorboardx_logdir) + if os.path.exists(tensorboardx_logdir): + shutil.rmtree(tensorboardx_logdir) + writer = SummaryWriter(log_dir=tensorboardx_logdir) + args.writer = writer + if args.evaluate: + args.evaluator = Evaluator(args.test_set, args.test_data_conf_file) + logger.info('Start optimization.') + logger.info(args) + train(args) + logger.info('Optimization done!') diff --git a/addition_module/HSST/train_dataset.py b/src/faceDetection/face_sdk/addition_module/HSST/train_dataset.py similarity index 97% rename from addition_module/HSST/train_dataset.py rename to src/faceDetection/face_sdk/addition_module/HSST/train_dataset.py index 9115d32..6c09871 100644 --- a/addition_module/HSST/train_dataset.py +++ b/src/faceDetection/face_sdk/addition_module/HSST/train_dataset.py @@ -1,168 +1,168 @@ -""" -@author: Hang Du, Jun Wang -@date: 20201101 -@contact: jun21wangustc@gmail.com -""" - -import os -import random -import cv2 -import torch -import numpy as np -from torch.utils.data import Dataset - -def transform(image): - """ Transform a image by cv2. - """ - img_size = image.shape[0] - # random crop - if random.random() > 0.5: - crop_size = 9 - x1_offset = np.random.randint(0, crop_size, size=1)[0] - y1_offset = np.random.randint(0, crop_size, size=1)[0] - x2_offset = np.random.randint(img_size-crop_size, img_size, size=1)[0] - y2_offset = np.random.randint(img_size-crop_size, img_size, size=1)[0] - image = image[x1_offset:x2_offset,y1_offset:y2_offset] - image = cv2.resize(image,(img_size,img_size)) - # horizontal flipping - if random.random() > 0.5: - image = cv2.flip(image, 1) - # grayscale conversion - if random.random() > 0.8: - image= cv2.cvtColor(image,cv2.COLOR_BGR2GRAY) - # rotation - if random.random() > 0.5: - theta = (random.randint(-10,10)) * np.pi / 180 - M_rotate = np.array([[np.cos(theta), -np.sin(theta), 0],[np.sin(theta), np.cos(theta), 0]], dtype=np.float32) - image = cv2.warpAffine(image, M_rotate, (img_size, img_size)) - # normalizing - if image.ndim == 2: - image = (image - 127.5) * 0.0078125 - new_image = np.zeros([3,img_size,img_size], np.float32) - new_image[0,:,:] = image - image = torch.from_numpy(new_image.astype(np.float32)) - else: - image = (image.transpose((2, 0, 1)) - 127.5) * 0.0078125 - image = torch.from_numpy(image.astype(np.float32)) - return image - - -class ImageDataset(Dataset): - def __init__(self, data_root, train_file, crop_eye=False): - self.data_root = data_root - self.train_list = [] - train_file_buf = open(train_file) - line = train_file_buf.readline().strip() - while line: - image_path, image_label = line.split(' ') - self.train_list.append((image_path, int(image_label))) - line = train_file_buf.readline().strip() - self.crop_eye = crop_eye - def __len__(self): - return len(self.train_list) - def __getitem__(self, index): - image_path, image_label = self.train_list[index] - image_path = os.path.join(self.data_root, image_path) - image = cv2.imread(image_path) - if self.crop_eye: - image = image[:60, :] - #image = cv2.resize(image, (128, 128)) #128 * 128 - if random.random() > 0.5: - image = cv2.flip(image, 1) - if image.ndim == 2: - image = image[:, :, np.newaxis] - image = (image.transpose((2, 0, 1)) - 127.5) * 0.0078125 - image = torch.from_numpy(image.astype(np.float32)) - return image, image_label - -class ImageDataset_SST(Dataset): - def __init__(self, data_root, train_file, exclude_id_set): - self.data_root = data_root - label_set = set() - # get id2image_path_list - self.id2image_path_list = {} - train_file_buf = open(train_file) - line = train_file_buf.readline().strip() - while line: - image_path, label = line.split(' ') - label = int(label) - if label in exclude_id_set: - line = train_file_buf.readline().strip() - continue - label_set.add(label) - if not label in self.id2image_path_list: - self.id2image_path_list[label] = [] - self.id2image_path_list[label].append(image_path) - line = train_file_buf.readline().strip() - self.train_list = list(label_set) - print('Valid ids: %d.' % len(self.train_list)) - - def __len__(self): - return len(self.train_list) - - def __getitem__(self, index): - cur_id = self.train_list[index] - cur_image_path_list = self.id2image_path_list[cur_id] - if len(cur_image_path_list) == 1: - image_path1 = cur_image_path_list[0] - image_path2 = cur_image_path_list[0] - else: - training_samples = random.sample(cur_image_path_list, 2) - image_path1 = training_samples[0] - image_path2 = training_samples[1] - image_path1 = os.path.join(self.data_root, image_path1) - image_path2 = os.path.join(self.data_root, image_path2) - image1 = cv2.imread(image_path1) - image2 = cv2.imread(image_path2) - image1 = transform(image1) - image2 = transform(image2) - if random.random() > 0.5: - return image2, image1, cur_id - return image1, image2, cur_id - - -class ImageDataset_HSST(Dataset): - def __init__(self, data_root, train_file, exclude_id_set): - self.data_root = data_root - label_set = set() - # get id2image_path_nir_list & id2image_path_vis_list - self.id2image_path_nir_list = {} - self.id2image_path_vis_list = {} - train_file_buf = open(train_file) - line = train_file_buf.readline().strip() - while line: - nir_img_path, vis_img_path, label = line.split(' ') - label = int(label) - if label in exclude_id_set: - line = train_file_buf.readline().strip() - continue - label_set.add(label) - if not label in self.id2image_path_nir_list: - self.id2image_path_nir_list[label] = [] - if not label in self.id2image_path_vis_list: - self.id2image_path_vis_list[label] = [] - self.id2image_path_nir_list[label].append(nir_img_path) - self.id2image_path_vis_list[label].append(vis_img_path) - line = train_file_buf.readline().strip() - - self.train_list = list(label_set) - - print('Valid ids: %d.' % len(self.train_list)) - - def __len__(self): - return len(self.train_list) - - def __getitem__(self, index): - cur_id = self.train_list[index] - cur_image_path_nir_list = list(set(self.id2image_path_nir_list[cur_id])) - cur_image_path_vis_list = list(set(self.id2image_path_vis_list[cur_id])) - nir_path = random.sample(cur_image_path_nir_list, 1)[0] - vis_path = random.sample(cur_image_path_vis_list, 1)[0] - nir_path = os.path.join(self.data_root, nir_path) - vis_path = os.path.join(self.data_root, vis_path) - nir_image = cv2.imread(nir_path) - vis_image = cv2.imread(vis_path) - nir_image = transform(nir_image) - vis_image = transform(vis_image) - - return nir_image, vis_image, cur_id +""" +@author: Hang Du, Jun Wang +@date: 20201101 +@contact: jun21wangustc@gmail.com +""" + +import os +import random +import cv2 +import torch +import numpy as np +from torch.utils.data import Dataset + +def transform(image): + """ Transform a image by cv2. + """ + img_size = image.shape[0] + # random crop + if random.random() > 0.5: + crop_size = 9 + x1_offset = np.random.randint(0, crop_size, size=1)[0] + y1_offset = np.random.randint(0, crop_size, size=1)[0] + x2_offset = np.random.randint(img_size-crop_size, img_size, size=1)[0] + y2_offset = np.random.randint(img_size-crop_size, img_size, size=1)[0] + image = image[x1_offset:x2_offset,y1_offset:y2_offset] + image = cv2.resize(image,(img_size,img_size)) + # horizontal flipping + if random.random() > 0.5: + image = cv2.flip(image, 1) + # grayscale conversion + if random.random() > 0.8: + image= cv2.cvtColor(image,cv2.COLOR_BGR2GRAY) + # rotation + if random.random() > 0.5: + theta = (random.randint(-10,10)) * np.pi / 180 + M_rotate = np.array([[np.cos(theta), -np.sin(theta), 0],[np.sin(theta), np.cos(theta), 0]], dtype=np.float32) + image = cv2.warpAffine(image, M_rotate, (img_size, img_size)) + # normalizing + if image.ndim == 2: + image = (image - 127.5) * 0.0078125 + new_image = np.zeros([3,img_size,img_size], np.float32) + new_image[0,:,:] = image + image = torch.from_numpy(new_image.astype(np.float32)) + else: + image = (image.transpose((2, 0, 1)) - 127.5) * 0.0078125 + image = torch.from_numpy(image.astype(np.float32)) + return image + + +class ImageDataset(Dataset): + def __init__(self, data_root, train_file, crop_eye=False): + self.data_root = data_root + self.train_list = [] + train_file_buf = open(train_file) + line = train_file_buf.readline().strip() + while line: + image_path, image_label = line.split(' ') + self.train_list.append((image_path, int(image_label))) + line = train_file_buf.readline().strip() + self.crop_eye = crop_eye + def __len__(self): + return len(self.train_list) + def __getitem__(self, index): + image_path, image_label = self.train_list[index] + image_path = os.path.join(self.data_root, image_path) + image = cv2.imread(image_path) + if self.crop_eye: + image = image[:60, :] + #image = cv2.resize(image, (128, 128)) #128 * 128 + if random.random() > 0.5: + image = cv2.flip(image, 1) + if image.ndim == 2: + image = image[:, :, np.newaxis] + image = (image.transpose((2, 0, 1)) - 127.5) * 0.0078125 + image = torch.from_numpy(image.astype(np.float32)) + return image, image_label + +class ImageDataset_SST(Dataset): + def __init__(self, data_root, train_file, exclude_id_set): + self.data_root = data_root + label_set = set() + # get id2image_path_list + self.id2image_path_list = {} + train_file_buf = open(train_file) + line = train_file_buf.readline().strip() + while line: + image_path, label = line.split(' ') + label = int(label) + if label in exclude_id_set: + line = train_file_buf.readline().strip() + continue + label_set.add(label) + if not label in self.id2image_path_list: + self.id2image_path_list[label] = [] + self.id2image_path_list[label].append(image_path) + line = train_file_buf.readline().strip() + self.train_list = list(label_set) + print('Valid ids: %d.' % len(self.train_list)) + + def __len__(self): + return len(self.train_list) + + def __getitem__(self, index): + cur_id = self.train_list[index] + cur_image_path_list = self.id2image_path_list[cur_id] + if len(cur_image_path_list) == 1: + image_path1 = cur_image_path_list[0] + image_path2 = cur_image_path_list[0] + else: + training_samples = random.sample(cur_image_path_list, 2) + image_path1 = training_samples[0] + image_path2 = training_samples[1] + image_path1 = os.path.join(self.data_root, image_path1) + image_path2 = os.path.join(self.data_root, image_path2) + image1 = cv2.imread(image_path1) + image2 = cv2.imread(image_path2) + image1 = transform(image1) + image2 = transform(image2) + if random.random() > 0.5: + return image2, image1, cur_id + return image1, image2, cur_id + + +class ImageDataset_HSST(Dataset): + def __init__(self, data_root, train_file, exclude_id_set): + self.data_root = data_root + label_set = set() + # get id2image_path_nir_list & id2image_path_vis_list + self.id2image_path_nir_list = {} + self.id2image_path_vis_list = {} + train_file_buf = open(train_file) + line = train_file_buf.readline().strip() + while line: + nir_img_path, vis_img_path, label = line.split(' ') + label = int(label) + if label in exclude_id_set: + line = train_file_buf.readline().strip() + continue + label_set.add(label) + if not label in self.id2image_path_nir_list: + self.id2image_path_nir_list[label] = [] + if not label in self.id2image_path_vis_list: + self.id2image_path_vis_list[label] = [] + self.id2image_path_nir_list[label].append(nir_img_path) + self.id2image_path_vis_list[label].append(vis_img_path) + line = train_file_buf.readline().strip() + + self.train_list = list(label_set) + + print('Valid ids: %d.' % len(self.train_list)) + + def __len__(self): + return len(self.train_list) + + def __getitem__(self, index): + cur_id = self.train_list[index] + cur_image_path_nir_list = list(set(self.id2image_path_nir_list[cur_id])) + cur_image_path_vis_list = list(set(self.id2image_path_vis_list[cur_id])) + nir_path = random.sample(cur_image_path_nir_list, 1)[0] + vis_path = random.sample(cur_image_path_vis_list, 1)[0] + nir_path = os.path.join(self.data_root, nir_path) + vis_path = os.path.join(self.data_root, vis_path) + nir_image = cv2.imread(nir_path) + vis_image = cv2.imread(vis_path) + nir_image = transform(nir_image) + vis_image = transform(vis_image) + + return nir_image, vis_image, cur_id diff --git a/src/faceDetection/face_sdk/addition_module/__init__.py b/src/faceDetection/face_sdk/addition_module/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/addition_module/face_body_attribute/README.md b/src/faceDetection/face_sdk/addition_module/face_body_attribute/README.md similarity index 100% rename from addition_module/face_body_attribute/README.md rename to src/faceDetection/face_sdk/addition_module/face_body_attribute/README.md diff --git a/addition_module/face_body_attribute/bagpack/README.md b/src/faceDetection/face_sdk/addition_module/face_body_attribute/bagpack/README.md similarity index 100% rename from addition_module/face_body_attribute/bagpack/README.md rename to src/faceDetection/face_sdk/addition_module/face_body_attribute/bagpack/README.md diff --git a/addition_module/face_body_attribute/body/README.md b/src/faceDetection/face_sdk/addition_module/face_body_attribute/body/README.md similarity index 100% rename from addition_module/face_body_attribute/body/README.md rename to src/faceDetection/face_sdk/addition_module/face_body_attribute/body/README.md diff --git a/addition_module/face_body_attribute/call/README.md b/src/faceDetection/face_sdk/addition_module/face_body_attribute/call/README.md similarity index 100% rename from addition_module/face_body_attribute/call/README.md rename to src/faceDetection/face_sdk/addition_module/face_body_attribute/call/README.md diff --git a/addition_module/face_body_attribute/child/README.md b/src/faceDetection/face_sdk/addition_module/face_body_attribute/child/README.md similarity index 100% rename from addition_module/face_body_attribute/child/README.md rename to src/faceDetection/face_sdk/addition_module/face_body_attribute/child/README.md diff --git a/addition_module/face_body_attribute/downbody/README.md b/src/faceDetection/face_sdk/addition_module/face_body_attribute/downbody/README.md similarity index 100% rename from addition_module/face_body_attribute/downbody/README.md rename to src/faceDetection/face_sdk/addition_module/face_body_attribute/downbody/README.md diff --git a/addition_module/face_body_attribute/hair/README.md b/src/faceDetection/face_sdk/addition_module/face_body_attribute/hair/README.md similarity index 100% rename from addition_module/face_body_attribute/hair/README.md rename to src/faceDetection/face_sdk/addition_module/face_body_attribute/hair/README.md diff --git a/addition_module/face_body_attribute/hand/README.md b/src/faceDetection/face_sdk/addition_module/face_body_attribute/hand/README.md similarity index 100% rename from addition_module/face_body_attribute/hand/README.md rename to src/faceDetection/face_sdk/addition_module/face_body_attribute/hand/README.md diff --git a/addition_module/face_body_attribute/hat/README.md b/src/faceDetection/face_sdk/addition_module/face_body_attribute/hat/README.md similarity index 100% rename from addition_module/face_body_attribute/hat/README.md rename to src/faceDetection/face_sdk/addition_module/face_body_attribute/hat/README.md diff --git a/addition_module/face_body_attribute/orient/README.md b/src/faceDetection/face_sdk/addition_module/face_body_attribute/orient/README.md similarity index 100% rename from addition_module/face_body_attribute/orient/README.md rename to src/faceDetection/face_sdk/addition_module/face_body_attribute/orient/README.md diff --git a/addition_module/face_body_attribute/scarf/README.md b/src/faceDetection/face_sdk/addition_module/face_body_attribute/scarf/README.md similarity index 100% rename from addition_module/face_body_attribute/scarf/README.md rename to src/faceDetection/face_sdk/addition_module/face_body_attribute/scarf/README.md diff --git a/addition_module/face_body_attribute/shoe/README.md b/src/faceDetection/face_sdk/addition_module/face_body_attribute/shoe/README.md similarity index 100% rename from addition_module/face_body_attribute/shoe/README.md rename to src/faceDetection/face_sdk/addition_module/face_body_attribute/shoe/README.md diff --git a/addition_module/face_body_attribute/umbrella/README.md b/src/faceDetection/face_sdk/addition_module/face_body_attribute/umbrella/README.md similarity index 100% rename from addition_module/face_body_attribute/umbrella/README.md rename to src/faceDetection/face_sdk/addition_module/face_body_attribute/umbrella/README.md diff --git a/addition_module/face_body_attribute/upbody/README.md b/src/faceDetection/face_sdk/addition_module/face_body_attribute/upbody/README.md similarity index 100% rename from addition_module/face_body_attribute/upbody/README.md rename to src/faceDetection/face_sdk/addition_module/face_body_attribute/upbody/README.md diff --git a/addition_module/face_lightning/KDF/README.md b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/README.md similarity index 100% rename from addition_module/face_lightning/KDF/README.md rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/README.md diff --git a/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/__init__.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/addition_module/face_lightning/KDF/backbone/MobileFaceNets.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/backbone/MobileFaceNets.py similarity index 100% rename from addition_module/face_lightning/KDF/backbone/MobileFaceNets.py rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/backbone/MobileFaceNets.py diff --git a/addition_module/face_lightning/KDF/backbone/ResNets.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/backbone/ResNets.py similarity index 100% rename from addition_module/face_lightning/KDF/backbone/ResNets.py rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/backbone/ResNets.py diff --git a/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/backbone/__init__.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/backbone/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/addition_module/face_lightning/KDF/backbone/backbone_def.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/backbone/backbone_def.py similarity index 93% rename from addition_module/face_lightning/KDF/backbone/backbone_def.py rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/backbone/backbone_def.py index ee70da7..822b49c 100644 --- a/addition_module/face_lightning/KDF/backbone/backbone_def.py +++ b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/backbone/backbone_def.py @@ -7,8 +7,8 @@ import sys import yaml sys.path.append('../../') -from backbone.ResNets import Resnet -from backbone.MobileFaceNets import MobileFaceNet +from src.faceDetection.face_sdk.backbone.ResNets import Resnet +from src.faceDetection.face_sdk.backbone.MobileFaceNets import MobileFaceNet class BackboneFactory: """Factory to produce backbone according the backbone_conf.yaml. diff --git a/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/loss/__init__.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/loss/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/addition_module/face_lightning/KDF/loss/fitnet.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/loss/fitnet.py similarity index 100% rename from addition_module/face_lightning/KDF/loss/fitnet.py rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/loss/fitnet.py diff --git a/addition_module/face_lightning/KDF/loss/fsp.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/loss/fsp.py similarity index 100% rename from addition_module/face_lightning/KDF/loss/fsp.py rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/loss/fsp.py diff --git a/addition_module/face_lightning/KDF/loss/ft.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/loss/ft.py similarity index 100% rename from addition_module/face_lightning/KDF/loss/ft.py rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/loss/ft.py diff --git a/addition_module/face_lightning/KDF/loss/loss_def.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/loss/loss_def.py similarity index 100% rename from addition_module/face_lightning/KDF/loss/loss_def.py rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/loss/loss_def.py diff --git a/addition_module/face_lightning/KDF/loss/pkt.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/loss/pkt.py similarity index 100% rename from addition_module/face_lightning/KDF/loss/pkt.py rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/loss/pkt.py diff --git a/addition_module/face_lightning/KDF/loss/snn_mimic.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/loss/snn_mimic.py similarity index 100% rename from addition_module/face_lightning/KDF/loss/snn_mimic.py rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/loss/snn_mimic.py diff --git a/addition_module/face_lightning/KDF/loss/soft_target.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/loss/soft_target.py similarity index 100% rename from addition_module/face_lightning/KDF/loss/soft_target.py rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/loss/soft_target.py diff --git a/addition_module/face_lightning/KDF/test_protocol/README.md b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/README.md similarity index 100% rename from addition_module/face_lightning/KDF/test_protocol/README.md rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/README.md diff --git a/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/__init__.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/addition_module/face_lightning/KDF/test_protocol/data_conf.yaml b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/data_conf.yaml similarity index 100% rename from addition_module/face_lightning/KDF/test_protocol/data_conf.yaml rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/data_conf.yaml diff --git a/addition_module/face_lightning/KDF/test_protocol/extract_feature.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/extract_feature.py similarity index 91% rename from addition_module/face_lightning/KDF/test_protocol/extract_feature.py rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/extract_feature.py index 1276bd6..1581830 100644 --- a/addition_module/face_lightning/KDF/test_protocol/extract_feature.py +++ b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/extract_feature.py @@ -7,13 +7,12 @@ import sys import yaml import argparse -import torch -from torch.utils.data import Dataset, DataLoader +from torch.utils.data import DataLoader from utils.model_loader import ModelLoader from utils.extractor.feature_extractor import CommonExtractor sys.path.append('..') -from data_processor.test_dataset import CommonTestDataset -from backbone.backbone_def import BackboneFactory +from src.faceDetection.face_sdk.data_processor.test_dataset import CommonTestDataset +from src.faceDetection.face_sdk.backbone.backbone_def import BackboneFactory if __name__ == '__main__': conf = argparse.ArgumentParser(description='extract features for megaface.') diff --git a/addition_module/face_lightning/KDF/test_protocol/extract_feature.sh b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/extract_feature.sh similarity index 100% rename from addition_module/face_lightning/KDF/test_protocol/extract_feature.sh rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/extract_feature.sh diff --git a/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/ijbc/__init__.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/ijbc/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/ijbc/face_cropper/__init__.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/ijbc/face_cropper/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/addition_module/face_lightning/KDF/test_protocol/ijbc/face_cropper/crop_ijbc_by_arcface.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/ijbc/face_cropper/crop_ijbc_by_arcface.py similarity index 100% rename from addition_module/face_lightning/KDF/test_protocol/ijbc/face_cropper/crop_ijbc_by_arcface.py rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/ijbc/face_cropper/crop_ijbc_by_arcface.py diff --git a/addition_module/face_lightning/KDF/test_protocol/ijbc/ijbc_evaluator.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/ijbc/ijbc_evaluator.py similarity index 100% rename from addition_module/face_lightning/KDF/test_protocol/ijbc/ijbc_evaluator.py rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/ijbc/ijbc_evaluator.py diff --git a/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/lfw/__init__.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/lfw/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/lfw/face_cropper/__init__.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/lfw/face_cropper/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/addition_module/face_lightning/KDF/test_protocol/lfw/face_cropper/crop_agedb_by_arcface.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/lfw/face_cropper/crop_agedb_by_arcface.py similarity index 100% rename from addition_module/face_lightning/KDF/test_protocol/lfw/face_cropper/crop_agedb_by_arcface.py rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/lfw/face_cropper/crop_agedb_by_arcface.py diff --git a/addition_module/face_lightning/KDF/test_protocol/lfw/face_cropper/crop_calfw_by_arcface.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/lfw/face_cropper/crop_calfw_by_arcface.py similarity index 100% rename from addition_module/face_lightning/KDF/test_protocol/lfw/face_cropper/crop_calfw_by_arcface.py rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/lfw/face_cropper/crop_calfw_by_arcface.py diff --git a/addition_module/face_lightning/KDF/test_protocol/lfw/face_cropper/crop_cplfw_by_arcface.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/lfw/face_cropper/crop_cplfw_by_arcface.py similarity index 100% rename from addition_module/face_lightning/KDF/test_protocol/lfw/face_cropper/crop_cplfw_by_arcface.py rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/lfw/face_cropper/crop_cplfw_by_arcface.py diff --git a/addition_module/face_lightning/KDF/test_protocol/lfw/face_cropper/crop_lfw_by_arcface.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/lfw/face_cropper/crop_lfw_by_arcface.py similarity index 100% rename from addition_module/face_lightning/KDF/test_protocol/lfw/face_cropper/crop_lfw_by_arcface.py rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/lfw/face_cropper/crop_lfw_by_arcface.py diff --git a/addition_module/face_lightning/KDF/test_protocol/lfw/face_cropper/crop_rfw_by_arcface.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/lfw/face_cropper/crop_rfw_by_arcface.py similarity index 100% rename from addition_module/face_lightning/KDF/test_protocol/lfw/face_cropper/crop_rfw_by_arcface.py rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/lfw/face_cropper/crop_rfw_by_arcface.py diff --git a/addition_module/face_lightning/KDF/test_protocol/lfw/lfw_evaluator.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/lfw/lfw_evaluator.py similarity index 100% rename from addition_module/face_lightning/KDF/test_protocol/lfw/lfw_evaluator.py rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/lfw/lfw_evaluator.py diff --git a/addition_module/face_lightning/KDF/test_protocol/lfw/pairs_parser.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/lfw/pairs_parser.py similarity index 100% rename from addition_module/face_lightning/KDF/test_protocol/lfw/pairs_parser.py rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/lfw/pairs_parser.py diff --git a/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/megaface/__init__.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/megaface/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/megaface/face_cropper/__init__.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/megaface/face_cropper/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/addition_module/face_lightning/KDF/test_protocol/megaface/face_cropper/crop_eye.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/megaface/face_cropper/crop_eye.py similarity index 100% rename from addition_module/face_lightning/KDF/test_protocol/megaface/face_cropper/crop_eye.py rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/megaface/face_cropper/crop_eye.py diff --git a/addition_module/face_lightning/KDF/test_protocol/megaface/face_cropper/crop_facescrub_by_arcface.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/megaface/face_cropper/crop_facescrub_by_arcface.py similarity index 100% rename from addition_module/face_lightning/KDF/test_protocol/megaface/face_cropper/crop_facescrub_by_arcface.py rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/megaface/face_cropper/crop_facescrub_by_arcface.py diff --git a/addition_module/face_lightning/KDF/test_protocol/megaface/face_cropper/crop_megaface_by_arcface.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/megaface/face_cropper/crop_megaface_by_arcface.py similarity index 100% rename from addition_module/face_lightning/KDF/test_protocol/megaface/face_cropper/crop_megaface_by_arcface.py rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/megaface/face_cropper/crop_megaface_by_arcface.py diff --git a/addition_module/face_lightning/KDF/test_protocol/megaface/megaface_evaluator.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/megaface/megaface_evaluator.py similarity index 100% rename from addition_module/face_lightning/KDF/test_protocol/megaface/megaface_evaluator.py rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/megaface/megaface_evaluator.py diff --git a/addition_module/face_lightning/KDF/test_protocol/remove_noises.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/remove_noises.py similarity index 100% rename from addition_module/face_lightning/KDF/test_protocol/remove_noises.py rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/remove_noises.py diff --git a/addition_module/face_lightning/KDF/test_protocol/remove_noises.sh b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/remove_noises.sh similarity index 100% rename from addition_module/face_lightning/KDF/test_protocol/remove_noises.sh rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/remove_noises.sh diff --git a/test_protocol/test_ijbc.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/test_ijbc.py similarity index 95% rename from test_protocol/test_ijbc.py rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/test_ijbc.py index 828b0d6..310fe66 100644 --- a/test_protocol/test_ijbc.py +++ b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/test_ijbc.py @@ -14,8 +14,8 @@ from utils.model_loader import ModelLoader from utils.extractor.feature_extractor import CommonExtractor sys.path.append('..') -from data_processor.test_dataset import CommonTestDataset -from backbone.backbone_def import BackboneFactory +from src.faceDetection.face_sdk.data_processor.test_dataset import CommonTestDataset +from src.faceDetection.face_sdk.backbone.backbone_def import BackboneFactory def accu_key(elem): return elem[1] diff --git a/addition_module/face_lightning/KDF/test_protocol/test_ijbc.sh b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/test_ijbc.sh similarity index 100% rename from addition_module/face_lightning/KDF/test_protocol/test_ijbc.sh rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/test_ijbc.sh diff --git a/test_protocol/test_lfw.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/test_lfw.py similarity index 93% rename from test_protocol/test_lfw.py rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/test_lfw.py index 66ac969..361f5d3 100644 --- a/test_protocol/test_lfw.py +++ b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/test_lfw.py @@ -15,8 +15,9 @@ from utils.model_loader import ModelLoader from utils.extractor.feature_extractor import CommonExtractor sys.path.append('..') -from data_processor.test_dataset import CommonTestDataset -from backbone.backbone_def import BackboneFactory +from src.faceDetection.face_sdk.backbone.backbone_def import BackboneFactory +sys.path.append('../../../../../../../') +from src.faceDetection.face_sdk.data_processor.test_dataset import CommonTestDataset def accu_key(elem): return elem[1] diff --git a/addition_module/face_lightning/KDF/test_protocol/test_lfw.sh b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/test_lfw.sh similarity index 100% rename from addition_module/face_lightning/KDF/test_protocol/test_lfw.sh rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/test_lfw.sh diff --git a/addition_module/face_lightning/KDF/test_protocol/test_megaface.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/test_megaface.py similarity index 100% rename from addition_module/face_lightning/KDF/test_protocol/test_megaface.py rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/test_megaface.py diff --git a/addition_module/face_lightning/KDF/test_protocol/test_megaface.sh b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/test_megaface.sh similarity index 100% rename from addition_module/face_lightning/KDF/test_protocol/test_megaface.sh rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/test_megaface.sh diff --git a/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/utils/__init__.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/utils/extractor/__init__.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/utils/extractor/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/addition_module/face_lightning/KDF/test_protocol/utils/extractor/feature_extract_save.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/utils/extractor/feature_extract_save.py similarity index 100% rename from addition_module/face_lightning/KDF/test_protocol/utils/extractor/feature_extract_save.py rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/utils/extractor/feature_extract_save.py diff --git a/addition_module/face_lightning/KDF/test_protocol/utils/extractor/feature_extractor.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/utils/extractor/feature_extractor.py similarity index 97% rename from addition_module/face_lightning/KDF/test_protocol/utils/extractor/feature_extractor.py rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/utils/extractor/feature_extractor.py index b5ec9b6..c4fbda7 100644 --- a/addition_module/face_lightning/KDF/test_protocol/utils/extractor/feature_extractor.py +++ b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/utils/extractor/feature_extractor.py @@ -1,99 +1,99 @@ -""" -@author: Jun Wang -@date: 20201016 -@contact: jun21wangustc@gmail.com -""" - -import os -import logging as logger -import numpy as np -import torch -import torch.nn.functional as F -logger.basicConfig(level=logger.INFO, - format='%(levelname)s %(asctime)s %(filename)s: %(lineno)d] %(message)s', - datefmt='%Y-%m-%d %H:%M:%S') - -class CommonExtractor: - """Common feature extractor. - - Attributes: - device(object): device to init model. - """ - def __init__(self, device): - self.device = torch.device(device) - - def extract_online(self, model, data_loader): - """Extract and return features. - - Args: - model(object): initialized model. - data_loader(object): load data to be extracted. - - Returns: - image_name2feature(dict): key is the name of image, value is feature of image. - """ - model.eval() - image_name2feature = {} - with torch.no_grad(): - for batch_idx, (images, filenames) in enumerate(data_loader): - images = images.to(self.device) - _, _, _, _, features = model(images) - features = F.normalize(features) - features = features.cpu().numpy() - for filename, feature in zip(filenames, features): - image_name2feature[filename] = feature - return image_name2feature - - def extract_offline(self, feats_root, model, data_loader): - """Extract and save features. - - Args: - feats_root(str): the path to save features. - model(object): initialized model. - data_loader(object): load data to be extracted. - """ - model.eval() - with torch.no_grad(): - for batch_idx, (images, filenames) in enumerate(data_loader): - images = images.to(self.device) - _, _, _, _, features = model(images) - features = F.normalize(features) - features = features.cpu().numpy() - for filename, feature in zip(filenames, features): - feature_name = os.path.splitext(filename)[0] - feature_path = os.path.join(feats_root, feature_name + '.npy') - feature_dir = os.path.dirname(feature_path) - if not os.path.exists(feature_dir): - os.makedirs(feature_dir) - np.save(feature_path, feature) - if (batch_idx + 1) % 10 == 0: - logger.info('Finished batches: %d/%d.' % (batch_idx+1, len(data_loader))) - -class FeatureHandler: - """Some method to deal with features. - - Atributes: - feats_root(str): the directory which the fetures in. - """ - def __init__(self, feats_root): - self.feats_root = feats_root - - def load_feature(self): - """Load features to memory. - - Returns: - image_name2feature(dict): key is the name of image, value is feature of image. - """ - image_name2feature = {} - for root, dirs, files in os.walk(self.feats_root): - for cur_file in files: - if cur_file.endswith('.npy'): - cur_file_path = os.path.join(root, cur_file) - cur_feats = np.load(cur_file_path) - if self.feats_root.endswith('/'): - cur_short_path = cur_file_path[len(self.feats_root) : ] - else: - cur_short_path = cur_file_path[len(self.feats_root) + 1 : ] - cur_key = cur_short_path.replace('.npy', '.jpg') - image_name2feature[cur_key] = cur_feats - return image_name2feature +""" +@author: Jun Wang +@date: 20201016 +@contact: jun21wangustc@gmail.com +""" + +import os +import logging as logger +import numpy as np +import torch +import torch.nn.functional as F +logger.basicConfig(level=logger.INFO, + format='%(levelname)s %(asctime)s %(filename)s: %(lineno)d] %(message)s', + datefmt='%Y-%m-%d %H:%M:%S') + +class CommonExtractor: + """Common feature extractor. + + Attributes: + device(object): device to init model. + """ + def __init__(self, device): + self.device = torch.device(device) + + def extract_online(self, model, data_loader): + """Extract and return features. + + Args: + model(object): initialized model. + data_loader(object): load data to be extracted. + + Returns: + image_name2feature(dict): key is the name of image, value is feature of image. + """ + model.eval() + image_name2feature = {} + with torch.no_grad(): + for batch_idx, (images, filenames) in enumerate(data_loader): + images = images.to(self.device) + _, _, _, _, features = model(images) + features = F.normalize(features) + features = features.cpu().numpy() + for filename, feature in zip(filenames, features): + image_name2feature[filename] = feature + return image_name2feature + + def extract_offline(self, feats_root, model, data_loader): + """Extract and save features. + + Args: + feats_root(str): the path to save features. + model(object): initialized model. + data_loader(object): load data to be extracted. + """ + model.eval() + with torch.no_grad(): + for batch_idx, (images, filenames) in enumerate(data_loader): + images = images.to(self.device) + _, _, _, _, features = model(images) + features = F.normalize(features) + features = features.cpu().numpy() + for filename, feature in zip(filenames, features): + feature_name = os.path.splitext(filename)[0] + feature_path = os.path.join(feats_root, feature_name + '.npy') + feature_dir = os.path.dirname(feature_path) + if not os.path.exists(feature_dir): + os.makedirs(feature_dir) + np.save(feature_path, feature) + if (batch_idx + 1) % 10 == 0: + logger.info('Finished batches: %d/%d.' % (batch_idx+1, len(data_loader))) + +class FeatureHandler: + """Some method to deal with features. + + Atributes: + feats_root(str): the directory which the fetures in. + """ + def __init__(self, feats_root): + self.feats_root = feats_root + + def load_feature(self): + """Load features to memory. + + Returns: + image_name2feature(dict): key is the name of image, value is feature of image. + """ + image_name2feature = {} + for root, dirs, files in os.walk(self.feats_root): + for cur_file in files: + if cur_file.endswith('.npy'): + cur_file_path = os.path.join(root, cur_file) + cur_feats = np.load(cur_file_path) + if self.feats_root.endswith('/'): + cur_short_path = cur_file_path[len(self.feats_root) : ] + else: + cur_short_path = cur_file_path[len(self.feats_root) + 1 : ] + cur_key = cur_short_path.replace('.npy', '.jpg') + image_name2feature[cur_key] = cur_feats + return image_name2feature diff --git a/addition_module/face_lightning/KDF/test_protocol/utils/feat_concat.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/utils/feat_concat.py similarity index 100% rename from addition_module/face_lightning/KDF/test_protocol/utils/feat_concat.py rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/utils/feat_concat.py diff --git a/addition_module/face_lightning/KDF/test_protocol/utils/model_loader.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/utils/model_loader.py similarity index 100% rename from addition_module/face_lightning/KDF/test_protocol/utils/model_loader.py rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/utils/model_loader.py diff --git a/addition_module/face_lightning/KDF/test_protocol/utils/pair_analysis.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/utils/pair_analysis.py similarity index 100% rename from addition_module/face_lightning/KDF/test_protocol/utils/pair_analysis.py rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/test_protocol/utils/pair_analysis.py diff --git a/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/training_mode/__init__.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/training_mode/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/training_mode/cls_training/__init__.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/training_mode/cls_training/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/addition_module/face_lightning/KDF/training_mode/cls_training/train.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/training_mode/cls_training/train.py similarity index 95% rename from addition_module/face_lightning/KDF/training_mode/cls_training/train.py rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/training_mode/cls_training/train.py index 46df8b0..8010ca2 100644 --- a/addition_module/face_lightning/KDF/training_mode/cls_training/train.py +++ b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/training_mode/cls_training/train.py @@ -15,12 +15,12 @@ from tensorboardX import SummaryWriter sys.path.append('../../') -from backbone.backbone_def import BackboneFactory +from src.faceDetection.face_sdk.backbone.backbone_def import BackboneFactory -sys.path.append('../../../../../') -from utils.AverageMeter import AverageMeter -from data_processor.train_dataset import ImageDataset -from head.head_def import HeadFactory +sys.path.append('../../../../../../../../') +from src.faceDetection.face_sdk.utils.AverageMeter import AverageMeter +from src.faceDetection.face_sdk.data_processor.train_dataset import ImageDataset +from src.faceDetection.face_sdk.head.head_def import HeadFactory logger.basicConfig(level=logger.INFO, format='%(levelname)s %(asctime)s %(filename)s: %(lineno)d] %(message)s', diff --git a/addition_module/face_lightning/KDF/training_mode/cls_training/train.sh b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/training_mode/cls_training/train.sh similarity index 100% rename from addition_module/face_lightning/KDF/training_mode/cls_training/train.sh rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/training_mode/cls_training/train.sh diff --git a/addition_module/face_lightning/KDF/training_mode/head_conf.yaml b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/training_mode/head_conf.yaml similarity index 100% rename from addition_module/face_lightning/KDF/training_mode/head_conf.yaml rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/training_mode/head_conf.yaml diff --git a/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/training_mode/kd_training/__init__.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/training_mode/kd_training/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/addition_module/face_lightning/KDF/training_mode/kd_training/train.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/training_mode/kd_training/train.py similarity index 96% rename from addition_module/face_lightning/KDF/training_mode/kd_training/train.py rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/training_mode/kd_training/train.py index 1da760c..c2a63b4 100644 --- a/addition_module/face_lightning/KDF/training_mode/kd_training/train.py +++ b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/training_mode/kd_training/train.py @@ -15,13 +15,13 @@ from tensorboardX import SummaryWriter sys.path.append('../../') -from backbone.backbone_def import BackboneFactory +from src.faceDetection.face_sdk.backbone.backbone_def import BackboneFactory from loss.loss_def import KDLossFactory -sys.path.append('../../../../../') -from utils.AverageMeter import AverageMeter -from data_processor.train_dataset import ImageDataset -from head.head_def import HeadFactory +sys.path.append('../../../../../../../../') +from src.faceDetection.face_sdk.utils.AverageMeter import AverageMeter +from src.faceDetection.face_sdk.data_processor.train_dataset import ImageDataset +from src.faceDetection.face_sdk.head.head_def import HeadFactory logger.basicConfig(level=logger.INFO, format='%(levelname)s %(asctime)s %(filename)s: %(lineno)d] %(message)s', diff --git a/addition_module/face_lightning/KDF/training_mode/kd_training/train.sh b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/training_mode/kd_training/train.sh similarity index 100% rename from addition_module/face_lightning/KDF/training_mode/kd_training/train.sh rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/training_mode/kd_training/train.sh diff --git a/addition_module/face_lightning/KDF/training_mode/kd_training/train_ft.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/training_mode/kd_training/train_ft.py similarity index 97% rename from addition_module/face_lightning/KDF/training_mode/kd_training/train_ft.py rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/training_mode/kd_training/train_ft.py index 04cb042..3b6fa24 100644 --- a/addition_module/face_lightning/KDF/training_mode/kd_training/train_ft.py +++ b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/training_mode/kd_training/train_ft.py @@ -16,14 +16,14 @@ from tensorboardX import SummaryWriter sys.path.append('../../') -from backbone.backbone_def import BackboneFactory +from src.faceDetection.face_sdk.backbone.backbone_def import BackboneFactory from loss.loss_def import KDLossFactory from utils.net_util import define_paraphraser, define_translator -sys.path.append('../../../../../') -from utils.AverageMeter import AverageMeter -from data_processor.train_dataset import ImageDataset -from head.head_def import HeadFactory +sys.path.append('../../../../../../../../') +from src.faceDetection.face_sdk.utils.AverageMeter import AverageMeter +from src.faceDetection.face_sdk.data_processor.train_dataset import ImageDataset +from src.faceDetection.face_sdk.head.head_def import HeadFactory logger.basicConfig(level=logger.INFO, format='%(levelname)s %(asctime)s %(filename)s: %(lineno)d] %(message)s', diff --git a/addition_module/face_lightning/KDF/training_mode/loss_conf.yaml b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/training_mode/loss_conf.yaml similarity index 100% rename from addition_module/face_lightning/KDF/training_mode/loss_conf.yaml rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/training_mode/loss_conf.yaml diff --git a/addition_module/face_lightning/KDF/training_mode/student_backbone_conf.yaml b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/training_mode/student_backbone_conf.yaml similarity index 100% rename from addition_module/face_lightning/KDF/training_mode/student_backbone_conf.yaml rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/training_mode/student_backbone_conf.yaml diff --git a/addition_module/face_lightning/KDF/training_mode/teacher_backbone_conf.yaml b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/training_mode/teacher_backbone_conf.yaml similarity index 100% rename from addition_module/face_lightning/KDF/training_mode/teacher_backbone_conf.yaml rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/training_mode/teacher_backbone_conf.yaml diff --git a/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/utils/__init__.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/addition_module/face_lightning/KDF/utils/net_util.py b/src/faceDetection/face_sdk/addition_module/face_lightning/KDF/utils/net_util.py similarity index 100% rename from addition_module/face_lightning/KDF/utils/net_util.py rename to src/faceDetection/face_sdk/addition_module/face_lightning/KDF/utils/net_util.py diff --git a/addition_module/face_lightning/README.md b/src/faceDetection/face_sdk/addition_module/face_lightning/README.md similarity index 100% rename from addition_module/face_lightning/README.md rename to src/faceDetection/face_sdk/addition_module/face_lightning/README.md diff --git a/src/faceDetection/face_sdk/addition_module/face_lightning/__init__.py b/src/faceDetection/face_sdk/addition_module/face_lightning/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/addition_module/face_mask_adding/FMA-3D/.gitignore b/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/.gitignore similarity index 100% rename from addition_module/face_mask_adding/FMA-3D/.gitignore rename to src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/.gitignore diff --git a/addition_module/face_mask_adding/FMA-3D/Data/images/FMA-3D.jpg b/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/images/FMA-3D.jpg similarity index 100% rename from addition_module/face_mask_adding/FMA-3D/Data/images/FMA-3D.jpg rename to src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/images/FMA-3D.jpg diff --git a/addition_module/face_mask_adding/FMA-3D/Data/images/mask-sample.jpg b/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/images/mask-sample.jpg similarity index 100% rename from addition_module/face_mask_adding/FMA-3D/Data/images/mask-sample.jpg rename to src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/images/mask-sample.jpg diff --git a/addition_module/face_mask_adding/FMA-3D/Data/mask-data/0.png b/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/mask-data/0.png similarity index 100% rename from addition_module/face_mask_adding/FMA-3D/Data/mask-data/0.png rename to src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/mask-data/0.png diff --git a/addition_module/face_mask_adding/FMA-3D/Data/mask-data/1.png b/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/mask-data/1.png similarity index 100% rename from addition_module/face_mask_adding/FMA-3D/Data/mask-data/1.png rename to src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/mask-data/1.png diff --git a/addition_module/face_mask_adding/FMA-3D/Data/mask-data/2.png b/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/mask-data/2.png similarity index 100% rename from addition_module/face_mask_adding/FMA-3D/Data/mask-data/2.png rename to src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/mask-data/2.png diff --git a/addition_module/face_mask_adding/FMA-3D/Data/mask-data/3.png b/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/mask-data/3.png similarity index 100% rename from addition_module/face_mask_adding/FMA-3D/Data/mask-data/3.png rename to src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/mask-data/3.png diff --git a/addition_module/face_mask_adding/FMA-3D/Data/mask-data/4.png b/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/mask-data/4.png similarity index 100% rename from addition_module/face_mask_adding/FMA-3D/Data/mask-data/4.png rename to src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/mask-data/4.png diff --git a/addition_module/face_mask_adding/FMA-3D/Data/mask-data/5.png b/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/mask-data/5.png similarity index 100% rename from addition_module/face_mask_adding/FMA-3D/Data/mask-data/5.png rename to src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/mask-data/5.png diff --git a/addition_module/face_mask_adding/FMA-3D/Data/mask-data/6.png b/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/mask-data/6.png similarity index 100% rename from addition_module/face_mask_adding/FMA-3D/Data/mask-data/6.png rename to src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/mask-data/6.png diff --git a/addition_module/face_mask_adding/FMA-3D/Data/mask-data/7.png b/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/mask-data/7.png similarity index 100% rename from addition_module/face_mask_adding/FMA-3D/Data/mask-data/7.png rename to src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/mask-data/7.png diff --git a/addition_module/face_mask_adding/FMA-3D/Data/test-data/test1.jpg b/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/test-data/test1.jpg similarity index 100% rename from addition_module/face_mask_adding/FMA-3D/Data/test-data/test1.jpg rename to src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/test-data/test1.jpg diff --git a/addition_module/face_mask_adding/FMA-3D/Data/test-data/test1_landmark.txt b/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/test-data/test1_landmark.txt similarity index 100% rename from addition_module/face_mask_adding/FMA-3D/Data/test-data/test1_landmark.txt rename to src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/test-data/test1_landmark.txt diff --git a/addition_module/face_mask_adding/FMA-3D/Data/uv-data/.ipynb_checkpoints/face_ind-checkpoint.txt b/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/uv-data/.ipynb_checkpoints/face_ind-checkpoint.txt similarity index 100% rename from addition_module/face_mask_adding/FMA-3D/Data/uv-data/.ipynb_checkpoints/face_ind-checkpoint.txt rename to src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/uv-data/.ipynb_checkpoints/face_ind-checkpoint.txt diff --git a/addition_module/face_mask_adding/FMA-3D/Data/uv-data/.ipynb_checkpoints/triangles-checkpoint.txt b/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/uv-data/.ipynb_checkpoints/triangles-checkpoint.txt similarity index 100% rename from addition_module/face_mask_adding/FMA-3D/Data/uv-data/.ipynb_checkpoints/triangles-checkpoint.txt rename to src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/uv-data/.ipynb_checkpoints/triangles-checkpoint.txt diff --git a/addition_module/face_mask_adding/FMA-3D/Data/uv-data/.ipynb_checkpoints/uv_face-checkpoint.png b/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/uv-data/.ipynb_checkpoints/uv_face-checkpoint.png similarity index 100% rename from addition_module/face_mask_adding/FMA-3D/Data/uv-data/.ipynb_checkpoints/uv_face-checkpoint.png rename to src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/uv-data/.ipynb_checkpoints/uv_face-checkpoint.png diff --git a/addition_module/face_mask_adding/FMA-3D/Data/uv-data/.ipynb_checkpoints/uv_face_eyes-checkpoint.png b/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/uv-data/.ipynb_checkpoints/uv_face_eyes-checkpoint.png similarity index 100% rename from addition_module/face_mask_adding/FMA-3D/Data/uv-data/.ipynb_checkpoints/uv_face_eyes-checkpoint.png rename to src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/uv-data/.ipynb_checkpoints/uv_face_eyes-checkpoint.png diff --git a/addition_module/face_mask_adding/FMA-3D/Data/uv-data/.ipynb_checkpoints/uv_face_mask-checkpoint.png b/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/uv-data/.ipynb_checkpoints/uv_face_mask-checkpoint.png similarity index 100% rename from addition_module/face_mask_adding/FMA-3D/Data/uv-data/.ipynb_checkpoints/uv_face_mask-checkpoint.png rename to src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/uv-data/.ipynb_checkpoints/uv_face_mask-checkpoint.png diff --git a/addition_module/face_mask_adding/FMA-3D/Data/uv-data/.ipynb_checkpoints/uv_kpt_ind-checkpoint.txt b/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/uv-data/.ipynb_checkpoints/uv_kpt_ind-checkpoint.txt similarity index 100% rename from addition_module/face_mask_adding/FMA-3D/Data/uv-data/.ipynb_checkpoints/uv_kpt_ind-checkpoint.txt rename to src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/uv-data/.ipynb_checkpoints/uv_kpt_ind-checkpoint.txt diff --git a/addition_module/face_mask_adding/FMA-3D/Data/uv-data/.ipynb_checkpoints/uv_kpt_mask-checkpoint.png b/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/uv-data/.ipynb_checkpoints/uv_kpt_mask-checkpoint.png similarity index 100% rename from addition_module/face_mask_adding/FMA-3D/Data/uv-data/.ipynb_checkpoints/uv_kpt_mask-checkpoint.png rename to src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/uv-data/.ipynb_checkpoints/uv_kpt_mask-checkpoint.png diff --git a/addition_module/face_mask_adding/FMA-3D/Data/uv-data/.ipynb_checkpoints/uv_weight_mask-checkpoint.png b/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/uv-data/.ipynb_checkpoints/uv_weight_mask-checkpoint.png similarity index 100% rename from addition_module/face_mask_adding/FMA-3D/Data/uv-data/.ipynb_checkpoints/uv_weight_mask-checkpoint.png rename to src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/uv-data/.ipynb_checkpoints/uv_weight_mask-checkpoint.png diff --git a/addition_module/face_mask_adding/FMA-3D/Data/uv-data/face_ind.txt b/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/uv-data/face_ind.txt similarity index 100% rename from addition_module/face_mask_adding/FMA-3D/Data/uv-data/face_ind.txt rename to src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/uv-data/face_ind.txt diff --git a/addition_module/face_mask_adding/FMA-3D/Data/uv-data/triangles.txt b/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/uv-data/triangles.txt similarity index 100% rename from addition_module/face_mask_adding/FMA-3D/Data/uv-data/triangles.txt rename to src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/uv-data/triangles.txt diff --git a/addition_module/face_mask_adding/FMA-3D/Data/uv-data/uv_face_mask.png b/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/uv-data/uv_face_mask.png similarity index 100% rename from addition_module/face_mask_adding/FMA-3D/Data/uv-data/uv_face_mask.png rename to src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/Data/uv-data/uv_face_mask.png diff --git a/addition_module/face_mask_adding/FMA-3D/README.md b/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/README.md similarity index 100% rename from addition_module/face_mask_adding/FMA-3D/README.md rename to src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/README.md diff --git a/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/__init__.py b/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/addition_module/face_mask_adding/FMA-3D/add_mask_all.py b/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/add_mask_all.py similarity index 100% rename from addition_module/face_mask_adding/FMA-3D/add_mask_all.py rename to src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/add_mask_all.py diff --git a/addition_module/face_mask_adding/FMA-3D/add_mask_one.py b/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/add_mask_one.py similarity index 100% rename from addition_module/face_mask_adding/FMA-3D/add_mask_one.py rename to src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/add_mask_one.py diff --git a/addition_module/face_mask_adding/FMA-3D/face_masker.py b/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/face_masker.py similarity index 99% rename from addition_module/face_mask_adding/FMA-3D/face_masker.py rename to src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/face_masker.py index 578d524..8065b2d 100644 --- a/addition_module/face_mask_adding/FMA-3D/face_masker.py +++ b/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/face_masker.py @@ -5,7 +5,6 @@ """ import os -from random import randint import warnings warnings.filterwarnings('ignore') import cv2 diff --git a/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/model/__init__.py b/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/model/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/addition_module/face_mask_adding/FMA-3D/model/prnet.pth b/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/model/prnet.pth similarity index 100% rename from addition_module/face_mask_adding/FMA-3D/model/prnet.pth rename to src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/model/prnet.pth diff --git a/addition_module/face_mask_adding/FMA-3D/model/prnet.py b/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/model/prnet.py similarity index 100% rename from addition_module/face_mask_adding/FMA-3D/model/prnet.py rename to src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/model/prnet.py diff --git a/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/utils/__init__.py b/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/addition_module/face_mask_adding/FMA-3D/utils/cython/README.md b/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/utils/cython/README.md similarity index 100% rename from addition_module/face_mask_adding/FMA-3D/utils/cython/README.md rename to src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/utils/cython/README.md diff --git a/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/utils/cython/__init__.py b/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/utils/cython/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/addition_module/face_mask_adding/FMA-3D/utils/cython/render.pyx b/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/utils/cython/render.pyx similarity index 100% rename from addition_module/face_mask_adding/FMA-3D/utils/cython/render.pyx rename to src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/utils/cython/render.pyx diff --git a/addition_module/face_mask_adding/FMA-3D/utils/cython/setup.py b/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/utils/cython/setup.py similarity index 100% rename from addition_module/face_mask_adding/FMA-3D/utils/cython/setup.py rename to src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/utils/cython/setup.py diff --git a/addition_module/face_mask_adding/FMA-3D/utils/read_info.py b/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/utils/read_info.py similarity index 100% rename from addition_module/face_mask_adding/FMA-3D/utils/read_info.py rename to src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/utils/read_info.py diff --git a/addition_module/face_mask_adding/FMA-3D/utils/render.py b/src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/utils/render.py similarity index 100% rename from addition_module/face_mask_adding/FMA-3D/utils/render.py rename to src/faceDetection/face_sdk/addition_module/face_mask_adding/FMA-3D/utils/render.py diff --git a/addition_module/face_mask_adding/README.md b/src/faceDetection/face_sdk/addition_module/face_mask_adding/README.md similarity index 100% rename from addition_module/face_mask_adding/README.md rename to src/faceDetection/face_sdk/addition_module/face_mask_adding/README.md diff --git a/src/faceDetection/face_sdk/addition_module/face_mask_adding/__init__.py b/src/faceDetection/face_sdk/addition_module/face_mask_adding/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/faceDetection/face_sdk/addition_module/model_convertor/__init__.py b/src/faceDetection/face_sdk/addition_module/model_convertor/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/faceDetection/face_sdk/addition_module/model_convertor/convert_to_onnx/__init__.py b/src/faceDetection/face_sdk/addition_module/model_convertor/convert_to_onnx/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/addition_module/model_convertor/convert_to_onnx/convert.sh b/src/faceDetection/face_sdk/addition_module/model_convertor/convert_to_onnx/convert.sh similarity index 100% rename from addition_module/model_convertor/convert_to_onnx/convert.sh rename to src/faceDetection/face_sdk/addition_module/model_convertor/convert_to_onnx/convert.sh diff --git a/addition_module/model_convertor/convert_to_onnx/model_convert.py b/src/faceDetection/face_sdk/addition_module/model_convertor/convert_to_onnx/model_convert.py similarity index 96% rename from addition_module/model_convertor/convert_to_onnx/model_convert.py rename to src/faceDetection/face_sdk/addition_module/model_convertor/convert_to_onnx/model_convert.py index 9a3728f..02f85ec 100644 --- a/addition_module/model_convertor/convert_to_onnx/model_convert.py +++ b/src/faceDetection/face_sdk/addition_module/model_convertor/convert_to_onnx/model_convert.py @@ -10,8 +10,8 @@ import argparse import onnxruntime import numpy as np -sys.path.append('../../../') -from backbone.MobileFaceNets import MobileFaceNet +sys.path.append('../../../../../../') +from src.faceDetection.face_sdk.backbone.MobileFaceNets import MobileFaceNet class OnnxConvertor: """Convert a pytorch model to a onnx model. diff --git a/addition_module/model_convertor/convert_to_onnx/test_images/11-FaceId-0_align.jpg b/src/faceDetection/face_sdk/addition_module/model_convertor/convert_to_onnx/test_images/11-FaceId-0_align.jpg similarity index 100% rename from addition_module/model_convertor/convert_to_onnx/test_images/11-FaceId-0_align.jpg rename to src/faceDetection/face_sdk/addition_module/model_convertor/convert_to_onnx/test_images/11-FaceId-0_align.jpg diff --git a/addition_module/model_convertor/convert_to_onnx/test_images/12-FaceId-0_align.jpg b/src/faceDetection/face_sdk/addition_module/model_convertor/convert_to_onnx/test_images/12-FaceId-0_align.jpg similarity index 100% rename from addition_module/model_convertor/convert_to_onnx/test_images/12-FaceId-0_align.jpg rename to src/faceDetection/face_sdk/addition_module/model_convertor/convert_to_onnx/test_images/12-FaceId-0_align.jpg diff --git a/src/faceDetection/face_sdk/addition_module/model_convertor/convert_to_sdk/__init__.py b/src/faceDetection/face_sdk/addition_module/model_convertor/convert_to_sdk/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/addition_module/model_convertor/convert_to_sdk/model_convert.py b/src/faceDetection/face_sdk/addition_module/model_convertor/convert_to_sdk/model_convert.py similarity index 97% rename from addition_module/model_convertor/convert_to_sdk/model_convert.py rename to src/faceDetection/face_sdk/addition_module/model_convertor/convert_to_sdk/model_convert.py index 4e8b406..7c41545 100644 --- a/addition_module/model_convertor/convert_to_sdk/model_convert.py +++ b/src/faceDetection/face_sdk/addition_module/model_convertor/convert_to_sdk/model_convert.py @@ -5,7 +5,7 @@ """ import sys import torch -sys.path.append('.') +sys.path.append('') from models.network_def.mobilefacenet_def import MobileFaceNet diff --git a/addition_module/model_convertor/readme.md b/src/faceDetection/face_sdk/addition_module/model_convertor/readme.md similarity index 100% rename from addition_module/model_convertor/readme.md rename to src/faceDetection/face_sdk/addition_module/model_convertor/readme.md diff --git a/backbone/AttentionNets.py b/src/faceDetection/face_sdk/backbone/AttentionNets.py similarity index 100% rename from backbone/AttentionNets.py rename to src/faceDetection/face_sdk/backbone/AttentionNets.py diff --git a/backbone/EfficientNets.py b/src/faceDetection/face_sdk/backbone/EfficientNets.py similarity index 100% rename from backbone/EfficientNets.py rename to src/faceDetection/face_sdk/backbone/EfficientNets.py diff --git a/backbone/GhostNet.py b/src/faceDetection/face_sdk/backbone/GhostNet.py similarity index 100% rename from backbone/GhostNet.py rename to src/faceDetection/face_sdk/backbone/GhostNet.py diff --git a/backbone/HRNet.py b/src/faceDetection/face_sdk/backbone/HRNet.py similarity index 100% rename from backbone/HRNet.py rename to src/faceDetection/face_sdk/backbone/HRNet.py diff --git a/backbone/LightCNN.py b/src/faceDetection/face_sdk/backbone/LightCNN.py similarity index 97% rename from backbone/LightCNN.py rename to src/faceDetection/face_sdk/backbone/LightCNN.py index 91277f7..c820216 100644 --- a/backbone/LightCNN.py +++ b/src/faceDetection/face_sdk/backbone/LightCNN.py @@ -1,200 +1,200 @@ -''' - implement Light CNN - @author: Alfred Xiang Wu - @date: 2017.07.04 -''' - -import math -import torch -import torch.nn as nn -import torch.nn.functional as F - -class Flatten(nn.Module): - def forward(self, input): - return input.view(input.size(0), -1) - -class mfm(nn.Module): - def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1, type=1): - super(mfm, self).__init__() - self.out_channels = out_channels - if type == 1: - self.filter = nn.Conv2d(in_channels, 2*out_channels, kernel_size=kernel_size, stride=stride, padding=padding) - else: - self.filter = nn.Linear(in_channels, 2*out_channels) - - def forward(self, x): - x = self.filter(x) - out = torch.split(x, self.out_channels, 1) - return torch.max(out[0], out[1]) - -class group(nn.Module): - def __init__(self, in_channels, out_channels, kernel_size, stride, padding): - super(group, self).__init__() - self.conv_a = mfm(in_channels, in_channels, 1, 1, 0) - self.conv = mfm(in_channels, out_channels, kernel_size, stride, padding) - - def forward(self, x): - x = self.conv_a(x) - x = self.conv(x) - return x - -class resblock(nn.Module): - def __init__(self, in_channels, out_channels): - super(resblock, self).__init__() - self.conv1 = mfm(in_channels, out_channels, kernel_size=3, stride=1, padding=1) - self.conv2 = mfm(in_channels, out_channels, kernel_size=3, stride=1, padding=1) - - def forward(self, x): - res = x - out = self.conv1(x) - out = self.conv2(out) - out = out + res - return out - -class network_9layers(nn.Module): - def __init__(self, num_classes=79077): - super(network_9layers, self).__init__() - self.features = nn.Sequential( - mfm(1, 48, 5, 1, 2), - nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True), - group(48, 96, 3, 1, 1), - nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True), - group(96, 192, 3, 1, 1), - nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True), - group(192, 128, 3, 1, 1), - group(128, 128, 3, 1, 1), - nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True), - ) - self.fc1 = mfm(8*8*128, 256, type=0) - self.fc2 = nn.Linear(256, num_classes) - - def forward(self, x): - x = self.features(x) - x = x.view(x.size(0), -1) - x = self.fc1(x) - x = F.dropout(x, training=self.training) - out = self.fc2(x) - return out, x - -class network_29layers(nn.Module): - def __init__(self, block, layers, num_classes=79077): - super(network_29layers, self).__init__() - #self.conv1 = mfm(1, 48, 5, 1, 2) - self.conv1 = mfm(3, 48, 5, 1, 2) - self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True) - self.block1 = self._make_layer(block, layers[0], 48, 48) - self.group1 = group(48, 96, 3, 1, 1) - self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True) - self.block2 = self._make_layer(block, layers[1], 96, 96) - self.group2 = group(96, 192, 3, 1, 1) - self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True) - self.block3 = self._make_layer(block, layers[2], 192, 192) - self.group3 = group(192, 128, 3, 1, 1) - self.block4 = self._make_layer(block, layers[3], 128, 128) - self.group4 = group(128, 128, 3, 1, 1) - self.pool4 = nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True) - self.fc = mfm(8*8*128, 256, type=0) - self.fc2 = nn.Linear(256, num_classes) - - def _make_layer(self, block, num_blocks, in_channels, out_channels): - layers = [] - for i in range(0, num_blocks): - layers.append(block(in_channels, out_channels)) - return nn.Sequential(*layers) - - def forward(self, x): - x = self.conv1(x) - x = self.pool1(x) - - x = self.block1(x) - x = self.group1(x) - x = self.pool2(x) - - x = self.block2(x) - x = self.group2(x) - x = self.pool3(x) - - x = self.block3(x) - x = self.group3(x) - x = self.block4(x) - x = self.group4(x) - x = self.pool4(x) - - x = x.view(x.size(0), -1) - fc = self.fc(x) - fc = F.dropout(fc, training=self.training) - out = self.fc2(fc) - return out, fc - - -class network_29layers_v2(nn.Module): - def __init__(self, block, layers, drop_ratio, out_h, out_w, feat_dim): - super(network_29layers_v2, self).__init__() - #self.conv1 = mfm(1, 48, 5, 1, 2) - self.conv1 = mfm(3, 48, 5, 1, 2) - self.block1 = self._make_layer(block, layers[0], 48, 48) - self.group1 = group(48, 96, 3, 1, 1) - self.block2 = self._make_layer(block, layers[1], 96, 96) - self.group2 = group(96, 192, 3, 1, 1) - self.block3 = self._make_layer(block, layers[2], 192, 192) - self.group3 = group(192, 128, 3, 1, 1) - self.block4 = self._make_layer(block, layers[3], 128, 128) - self.group4 = group(128, 128, 3, 1, 1) - #self.fc = nn.Linear(8*8*128, 256) - #self.fc2 = nn.Linear(256, num_classes, bias=False) - self.output_layer = nn.Sequential(nn.BatchNorm2d(128), - nn.Dropout(drop_ratio), - Flatten(), - nn.Linear(128 * out_h * out_w, feat_dim), - nn.BatchNorm1d(feat_dim)) - - def _make_layer(self, block, num_blocks, in_channels, out_channels): - layers = [] - for i in range(0, num_blocks): - layers.append(block(in_channels, out_channels)) - return nn.Sequential(*layers) - - def forward(self, x): - x = self.conv1(x) - x = F.max_pool2d(x, 2) + F.avg_pool2d(x, 2) - - x = self.block1(x) - x = self.group1(x) - x = F.max_pool2d(x, 2) + F.avg_pool2d(x, 2) - - x = self.block2(x) - x = self.group2(x) - x = F.max_pool2d(x, 2) + F.avg_pool2d(x, 2) - - x = self.block3(x) - x = self.group3(x) - x = self.block4(x) - x = self.group4(x) - x = F.max_pool2d(x, 2) + F.avg_pool2d(x, 2) # 7*7 - - #x = x.view(x.size(0), -1) - #fc = self.fc(x) - #x = F.dropout(fc, training=self.training) - #out = self.fc2(x) - #return out, fc - x = self.output_layer(x) - return x - - -def LightCNN_9Layers(drop_ratio, out_h, out_w, feat_dim): - model = network_9layers(drop_ratio, out_h, out_w, feat_dim) - return model - -def LightCNN_29Layers(drop_ratio, out_h, out_w, feat_dim): - model = network_29layers(resblock, [1, 2, 3, 4], drop_ratio, out_h, out_w, feat_dim) - return model - -def LightCNN_29Layers_v2(drop_ratio, out_h, out_w, feat_dim): - model = network_29layers_v2(resblock, [1, 2, 3, 4], drop_ratio, out_h, out_w, feat_dim) - return model - -def LightCNN(depth, drop_ratio, out_h, out_w, feat_dim): - if depth == 9: - return LightCNN_9Layers(drop_ratio, out_h, out_w, feat_dim) - elif depth == 29: - return LightCNN_29Layers_v2(drop_ratio, out_h, out_w, feat_dim) +''' + implement Light CNN + @author: Alfred Xiang Wu + @date: 2017.07.04 +''' + +import math +import torch +import torch.nn as nn +import torch.nn.functional as F + +class Flatten(nn.Module): + def forward(self, input): + return input.view(input.size(0), -1) + +class mfm(nn.Module): + def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1, type=1): + super(mfm, self).__init__() + self.out_channels = out_channels + if type == 1: + self.filter = nn.Conv2d(in_channels, 2*out_channels, kernel_size=kernel_size, stride=stride, padding=padding) + else: + self.filter = nn.Linear(in_channels, 2*out_channels) + + def forward(self, x): + x = self.filter(x) + out = torch.split(x, self.out_channels, 1) + return torch.max(out[0], out[1]) + +class group(nn.Module): + def __init__(self, in_channels, out_channels, kernel_size, stride, padding): + super(group, self).__init__() + self.conv_a = mfm(in_channels, in_channels, 1, 1, 0) + self.conv = mfm(in_channels, out_channels, kernel_size, stride, padding) + + def forward(self, x): + x = self.conv_a(x) + x = self.conv(x) + return x + +class resblock(nn.Module): + def __init__(self, in_channels, out_channels): + super(resblock, self).__init__() + self.conv1 = mfm(in_channels, out_channels, kernel_size=3, stride=1, padding=1) + self.conv2 = mfm(in_channels, out_channels, kernel_size=3, stride=1, padding=1) + + def forward(self, x): + res = x + out = self.conv1(x) + out = self.conv2(out) + out = out + res + return out + +class network_9layers(nn.Module): + def __init__(self, num_classes=79077): + super(network_9layers, self).__init__() + self.features = nn.Sequential( + mfm(1, 48, 5, 1, 2), + nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True), + group(48, 96, 3, 1, 1), + nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True), + group(96, 192, 3, 1, 1), + nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True), + group(192, 128, 3, 1, 1), + group(128, 128, 3, 1, 1), + nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True), + ) + self.fc1 = mfm(8*8*128, 256, type=0) + self.fc2 = nn.Linear(256, num_classes) + + def forward(self, x): + x = self.features(x) + x = x.view(x.size(0), -1) + x = self.fc1(x) + x = F.dropout(x, training=self.training) + out = self.fc2(x) + return out, x + +class network_29layers(nn.Module): + def __init__(self, block, layers, num_classes=79077): + super(network_29layers, self).__init__() + #self.conv1 = mfm(1, 48, 5, 1, 2) + self.conv1 = mfm(3, 48, 5, 1, 2) + self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True) + self.block1 = self._make_layer(block, layers[0], 48, 48) + self.group1 = group(48, 96, 3, 1, 1) + self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True) + self.block2 = self._make_layer(block, layers[1], 96, 96) + self.group2 = group(96, 192, 3, 1, 1) + self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True) + self.block3 = self._make_layer(block, layers[2], 192, 192) + self.group3 = group(192, 128, 3, 1, 1) + self.block4 = self._make_layer(block, layers[3], 128, 128) + self.group4 = group(128, 128, 3, 1, 1) + self.pool4 = nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True) + self.fc = mfm(8*8*128, 256, type=0) + self.fc2 = nn.Linear(256, num_classes) + + def _make_layer(self, block, num_blocks, in_channels, out_channels): + layers = [] + for i in range(0, num_blocks): + layers.append(block(in_channels, out_channels)) + return nn.Sequential(*layers) + + def forward(self, x): + x = self.conv1(x) + x = self.pool1(x) + + x = self.block1(x) + x = self.group1(x) + x = self.pool2(x) + + x = self.block2(x) + x = self.group2(x) + x = self.pool3(x) + + x = self.block3(x) + x = self.group3(x) + x = self.block4(x) + x = self.group4(x) + x = self.pool4(x) + + x = x.view(x.size(0), -1) + fc = self.fc(x) + fc = F.dropout(fc, training=self.training) + out = self.fc2(fc) + return out, fc + + +class network_29layers_v2(nn.Module): + def __init__(self, block, layers, drop_ratio, out_h, out_w, feat_dim): + super(network_29layers_v2, self).__init__() + #self.conv1 = mfm(1, 48, 5, 1, 2) + self.conv1 = mfm(3, 48, 5, 1, 2) + self.block1 = self._make_layer(block, layers[0], 48, 48) + self.group1 = group(48, 96, 3, 1, 1) + self.block2 = self._make_layer(block, layers[1], 96, 96) + self.group2 = group(96, 192, 3, 1, 1) + self.block3 = self._make_layer(block, layers[2], 192, 192) + self.group3 = group(192, 128, 3, 1, 1) + self.block4 = self._make_layer(block, layers[3], 128, 128) + self.group4 = group(128, 128, 3, 1, 1) + #self.fc = nn.Linear(8*8*128, 256) + #self.fc2 = nn.Linear(256, num_classes, bias=False) + self.output_layer = nn.Sequential(nn.BatchNorm2d(128), + nn.Dropout(drop_ratio), + Flatten(), + nn.Linear(128 * out_h * out_w, feat_dim), + nn.BatchNorm1d(feat_dim)) + + def _make_layer(self, block, num_blocks, in_channels, out_channels): + layers = [] + for i in range(0, num_blocks): + layers.append(block(in_channels, out_channels)) + return nn.Sequential(*layers) + + def forward(self, x): + x = self.conv1(x) + x = F.max_pool2d(x, 2) + F.avg_pool2d(x, 2) + + x = self.block1(x) + x = self.group1(x) + x = F.max_pool2d(x, 2) + F.avg_pool2d(x, 2) + + x = self.block2(x) + x = self.group2(x) + x = F.max_pool2d(x, 2) + F.avg_pool2d(x, 2) + + x = self.block3(x) + x = self.group3(x) + x = self.block4(x) + x = self.group4(x) + x = F.max_pool2d(x, 2) + F.avg_pool2d(x, 2) # 7*7 + + #x = x.view(x.size(0), -1) + #fc = self.fc(x) + #x = F.dropout(fc, training=self.training) + #out = self.fc2(x) + #return out, fc + x = self.output_layer(x) + return x + + +def LightCNN_9Layers(drop_ratio, out_h, out_w, feat_dim): + model = network_9layers(drop_ratio, out_h, out_w, feat_dim) + return model + +def LightCNN_29Layers(drop_ratio, out_h, out_w, feat_dim): + model = network_29layers(resblock, [1, 2, 3, 4], drop_ratio, out_h, out_w, feat_dim) + return model + +def LightCNN_29Layers_v2(drop_ratio, out_h, out_w, feat_dim): + model = network_29layers_v2(resblock, [1, 2, 3, 4], drop_ratio, out_h, out_w, feat_dim) + return model + +def LightCNN(depth, drop_ratio, out_h, out_w, feat_dim): + if depth == 9: + return LightCNN_9Layers(drop_ratio, out_h, out_w, feat_dim) + elif depth == 29: + return LightCNN_29Layers_v2(drop_ratio, out_h, out_w, feat_dim) diff --git a/backbone/MobileFaceNets.py b/src/faceDetection/face_sdk/backbone/MobileFaceNets.py similarity index 100% rename from backbone/MobileFaceNets.py rename to src/faceDetection/face_sdk/backbone/MobileFaceNets.py diff --git a/backbone/ReXNets.py b/src/faceDetection/face_sdk/backbone/ReXNets.py similarity index 100% rename from backbone/ReXNets.py rename to src/faceDetection/face_sdk/backbone/ReXNets.py diff --git a/backbone/RepVGG.py b/src/faceDetection/face_sdk/backbone/RepVGG.py similarity index 100% rename from backbone/RepVGG.py rename to src/faceDetection/face_sdk/backbone/RepVGG.py diff --git a/backbone/ResNets.py b/src/faceDetection/face_sdk/backbone/ResNets.py similarity index 100% rename from backbone/ResNets.py rename to src/faceDetection/face_sdk/backbone/ResNets.py diff --git a/backbone/Swin_Transformer.py b/src/faceDetection/face_sdk/backbone/Swin_Transformer.py similarity index 100% rename from backbone/Swin_Transformer.py rename to src/faceDetection/face_sdk/backbone/Swin_Transformer.py diff --git a/backbone/TF_NAS.py b/src/faceDetection/face_sdk/backbone/TF_NAS.py similarity index 100% rename from backbone/TF_NAS.py rename to src/faceDetection/face_sdk/backbone/TF_NAS.py diff --git a/src/faceDetection/face_sdk/backbone/__init__.py b/src/faceDetection/face_sdk/backbone/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/backbone/backbone_def.py b/src/faceDetection/face_sdk/backbone/backbone_def.py similarity index 90% rename from backbone/backbone_def.py rename to src/faceDetection/face_sdk/backbone/backbone_def.py index 2c0659d..222df13 100644 --- a/backbone/backbone_def.py +++ b/src/faceDetection/face_sdk/backbone/backbone_def.py @@ -6,20 +6,20 @@ import sys import yaml -sys.path.append('../../') -from backbone.ResNets import Resnet -from backbone.MobileFaceNets import MobileFaceNet -from backbone.EfficientNets import EfficientNet -from backbone.EfficientNets import efficientnet -from backbone.HRNet import HighResolutionNet -from backbone.GhostNet import GhostNet -from backbone.AttentionNets import ResidualAttentionNet -from backbone.TF_NAS import TF_NAS_A -from backbone.resnest.resnest import ResNeSt -from backbone.ReXNets import ReXNetV1 -from backbone.LightCNN import LightCNN -from backbone.RepVGG import RepVGG -from backbone.Swin_Transformer import SwinTransformer +sys.path.append('../../../../../') +from src.faceDetection.face_sdk.backbone.ResNets import Resnet +from src.faceDetection.face_sdk.backbone.MobileFaceNets import MobileFaceNet +from src.faceDetection.face_sdk.backbone.EfficientNets import EfficientNet +from src.faceDetection.face_sdk.backbone.EfficientNets import efficientnet +from src.faceDetection.face_sdk.backbone.HRNet import HighResolutionNet +from src.faceDetection.face_sdk.backbone.GhostNet import GhostNet +from src.faceDetection.face_sdk.backbone.AttentionNets import ResidualAttentionNet +from src.faceDetection.face_sdk.backbone.TF_NAS import TF_NAS_A +from src.faceDetection.face_sdk.backbone.resnest.resnest import ResNeSt +from src.faceDetection.face_sdk.backbone.ReXNets import ReXNetV1 +from src.faceDetection.face_sdk.backbone.LightCNN import LightCNN +from src.faceDetection.face_sdk.backbone.RepVGG import RepVGG +from src.faceDetection.face_sdk.backbone.Swin_Transformer import SwinTransformer class BackboneFactory: """Factory to produce backbone according the backbone_conf.yaml. diff --git a/backbone/resnest/__init__.py b/src/faceDetection/face_sdk/backbone/resnest/__init__.py similarity index 100% rename from backbone/resnest/__init__.py rename to src/faceDetection/face_sdk/backbone/resnest/__init__.py diff --git a/backbone/resnest/ablation.py b/src/faceDetection/face_sdk/backbone/resnest/ablation.py similarity index 100% rename from backbone/resnest/ablation.py rename to src/faceDetection/face_sdk/backbone/resnest/ablation.py diff --git a/backbone/resnest/resnest.py b/src/faceDetection/face_sdk/backbone/resnest/resnest.py similarity index 100% rename from backbone/resnest/resnest.py rename to src/faceDetection/face_sdk/backbone/resnest/resnest.py diff --git a/backbone/resnest/resnet.py b/src/faceDetection/face_sdk/backbone/resnest/resnet.py similarity index 100% rename from backbone/resnest/resnet.py rename to src/faceDetection/face_sdk/backbone/resnest/resnet.py diff --git a/backbone/resnest/splat.py b/src/faceDetection/face_sdk/backbone/resnest/splat.py similarity index 100% rename from backbone/resnest/splat.py rename to src/faceDetection/face_sdk/backbone/resnest/splat.py diff --git a/data/files/MS-Celeb-1M-v1c-r-shallow_train_list.txt b/src/faceDetection/face_sdk/data/files/MS-Celeb-1M-v1c-r-shallow_train_list.txt similarity index 100% rename from data/files/MS-Celeb-1M-v1c-r-shallow_train_list.txt rename to src/faceDetection/face_sdk/data/files/MS-Celeb-1M-v1c-r-shallow_train_list.txt diff --git a/data/files/MS-Celeb-1M-v1c-r_id_list.txt b/src/faceDetection/face_sdk/data/files/MS-Celeb-1M-v1c-r_id_list.txt similarity index 100% rename from data/files/MS-Celeb-1M-v1c-r_id_list.txt rename to src/faceDetection/face_sdk/data/files/MS-Celeb-1M-v1c-r_id_list.txt diff --git a/data/files/MS-Celeb-1M-v1c-r_train_list.txt b/src/faceDetection/face_sdk/data/files/MS-Celeb-1M-v1c-r_train_list.txt similarity index 100% rename from data/files/MS-Celeb-1M-v1c-r_train_list.txt rename to src/faceDetection/face_sdk/data/files/MS-Celeb-1M-v1c-r_train_list.txt diff --git a/data/files/facescrub2template_name.txt b/src/faceDetection/face_sdk/data/files/facescrub2template_name.txt similarity index 100% rename from data/files/facescrub2template_name.txt rename to src/faceDetection/face_sdk/data/files/facescrub2template_name.txt diff --git a/data/files/facescrub_face_info.txt b/src/faceDetection/face_sdk/data/files/facescrub_face_info.txt similarity index 100% rename from data/files/facescrub_face_info.txt rename to src/faceDetection/face_sdk/data/files/facescrub_face_info.txt diff --git a/data/files/lfw_face_info.txt b/src/faceDetection/face_sdk/data/files/lfw_face_info.txt similarity index 100% rename from data/files/lfw_face_info.txt rename to src/faceDetection/face_sdk/data/files/lfw_face_info.txt diff --git a/data/images/arch.jpg b/src/faceDetection/face_sdk/data/images/arch.jpg similarity index 100% rename from data/images/arch.jpg rename to src/faceDetection/face_sdk/data/images/arch.jpg diff --git a/data/images/face_parsing.jpg b/src/faceDetection/face_sdk/data/images/face_parsing.jpg similarity index 100% rename from data/images/face_parsing.jpg rename to src/faceDetection/face_sdk/data/images/face_parsing.jpg diff --git a/data/images/logo.png b/src/faceDetection/face_sdk/data/images/logo.png similarity index 100% rename from data/images/logo.png rename to src/faceDetection/face_sdk/data/images/logo.png diff --git a/data/images/sample_mask.jpg b/src/faceDetection/face_sdk/data/images/sample_mask.jpg similarity index 100% rename from data/images/sample_mask.jpg rename to src/faceDetection/face_sdk/data/images/sample_mask.jpg diff --git a/data/images/sample_non_mask.jpg b/src/faceDetection/face_sdk/data/images/sample_non_mask.jpg similarity index 100% rename from data/images/sample_non_mask.jpg rename to src/faceDetection/face_sdk/data/images/sample_non_mask.jpg diff --git a/src/faceDetection/face_sdk/data_processor/__init__.py b/src/faceDetection/face_sdk/data_processor/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/data_processor/test_dataset.py b/src/faceDetection/face_sdk/data_processor/test_dataset.py similarity index 97% rename from data_processor/test_dataset.py rename to src/faceDetection/face_sdk/data_processor/test_dataset.py index f95d1f6..a2fecfc 100644 --- a/data_processor/test_dataset.py +++ b/src/faceDetection/face_sdk/data_processor/test_dataset.py @@ -1,48 +1,48 @@ -""" -@author: Jun Wang -@date: 20201101 -@contact: jun21wangustc@gmail.com -""" - -import os -import logging as logger -import cv2 -import numpy as np -import torch -from torch.utils.data import Dataset - -logger.basicConfig(level=logger.INFO, - format='%(levelname)s %(asctime)s %(filename)s: %(lineno)d] %(message)s', - datefmt='%Y-%m-%d %H:%M:%S') - -class CommonTestDataset(Dataset): - """ Data processor for model evaluation. - - Attributes: - image_root(str): root directory of test set. - image_list_file(str): path of the image list file. - crop_eye(bool): crop eye(upper face) as input or not. - """ - def __init__(self, image_root, image_list_file, crop_eye=False): - self.image_root = image_root - self.image_list = [] - image_list_buf = open(image_list_file) - line = image_list_buf.readline().strip() - while line: - self.image_list.append(line) - line = image_list_buf.readline().strip() - self.mean = 127.5 - self.std = 128.0 - self.crop_eye = crop_eye - def __len__(self): - return len(self.image_list) - def __getitem__(self, index): - short_image_path = self.image_list[index] - image_path = os.path.join(self.image_root, short_image_path) - image = cv2.imdecode(np.fromfile(image_path, dtype=np.uint8), cv2.IMREAD_UNCHANGED) - #image = cv2.resize(image, (128, 128)) - if self.crop_eye: - image = image[:60, :] - image = (image.transpose((2, 0, 1)) - self.mean) / self.std - image = torch.from_numpy(image.astype(np.float32)) - return image, short_image_path +""" +@author: Jun Wang +@date: 20201101 +@contact: jun21wangustc@gmail.com +""" + +import os +import logging as logger +import cv2 +import numpy as np +import torch +from torch.utils.data import Dataset + +logger.basicConfig(level=logger.INFO, + format='%(levelname)s %(asctime)s %(filename)s: %(lineno)d] %(message)s', + datefmt='%Y-%m-%d %H:%M:%S') + +class CommonTestDataset(Dataset): + """ Data processor for model evaluation. + + Attributes: + image_root(str): root directory of test set. + image_list_file(str): path of the image list file. + crop_eye(bool): crop eye(upper face) as input or not. + """ + def __init__(self, image_root, image_list_file, crop_eye=False): + self.image_root = image_root + self.image_list = [] + image_list_buf = open(image_list_file) + line = image_list_buf.readline().strip() + while line: + self.image_list.append(line) + line = image_list_buf.readline().strip() + self.mean = 127.5 + self.std = 128.0 + self.crop_eye = crop_eye + def __len__(self): + return len(self.image_list) + def __getitem__(self, index): + short_image_path = self.image_list[index] + image_path = os.path.join(self.image_root, short_image_path) + image = cv2.imdecode(np.fromfile(image_path, dtype=np.uint8), cv2.IMREAD_UNCHANGED) + #image = cv2.resize(image, (128, 128)) + if self.crop_eye: + image = image[:60, :] + image = (image.transpose((2, 0, 1)) - self.mean) / self.std + image = torch.from_numpy(image.astype(np.float32)) + return image, short_image_path diff --git a/data_processor/train_dataset.py b/src/faceDetection/face_sdk/data_processor/train_dataset.py similarity index 97% rename from data_processor/train_dataset.py rename to src/faceDetection/face_sdk/data_processor/train_dataset.py index fb7ee0b..5b566be 100644 --- a/data_processor/train_dataset.py +++ b/src/faceDetection/face_sdk/data_processor/train_dataset.py @@ -1,121 +1,121 @@ -""" -@author: Hang Du, Jun Wang -@date: 20201101 -@contact: jun21wangustc@gmail.com -""" - -import os -import random -import cv2 -import torch -import numpy as np -from torch.utils.data import Dataset - -def transform(image): - """ Transform a image by cv2. - """ - img_size = image.shape[0] - # random crop - if random.random() > 0.5: - crop_size = 9 - x1_offset = np.random.randint(0, crop_size, size=1)[0] - y1_offset = np.random.randint(0, crop_size, size=1)[0] - x2_offset = np.random.randint(img_size-crop_size, img_size, size=1)[0] - y2_offset = np.random.randint(img_size-crop_size, img_size, size=1)[0] - image = image[x1_offset:x2_offset,y1_offset:y2_offset] - image = cv2.resize(image,(img_size,img_size)) - # horizontal flipping - if random.random() > 0.5: - image = cv2.flip(image, 1) - # grayscale conversion - if random.random() > 0.8: - image= cv2.cvtColor(image,cv2.COLOR_BGR2GRAY) - # rotation - if random.random() > 0.5: - theta = (random.randint(-10,10)) * np.pi / 180 - M_rotate = np.array([[np.cos(theta), -np.sin(theta), 0],[np.sin(theta), np.cos(theta), 0]], dtype=np.float32) - image = cv2.warpAffine(image, M_rotate, (img_size, img_size)) - # normalizing - if image.ndim == 2: - image = (image - 127.5) * 0.0078125 - new_image = np.zeros([3,img_size,img_size], np.float32) - new_image[0,:,:] = image - image = torch.from_numpy(new_image.astype(np.float32)) - else: - image = (image.transpose((2, 0, 1)) - 127.5) * 0.0078125 - image = torch.from_numpy(image.astype(np.float32)) - return image - - -class ImageDataset(Dataset): - def __init__(self, data_root, train_file, crop_eye=False): - self.data_root = data_root - self.train_list = [] - train_file_buf = open(train_file) - line = train_file_buf.readline().strip() - while line: - image_path, image_label = line.split(' ') - self.train_list.append((image_path, int(image_label))) - line = train_file_buf.readline().strip() - self.crop_eye = crop_eye - def __len__(self): - return len(self.train_list) - def __getitem__(self, index): - image_path, image_label = self.train_list[index] - image_path = os.path.join(self.data_root, image_path) - image = cv2.imread(image_path) - if self.crop_eye: - image = image[:60, :] - #image = cv2.resize(image, (128, 128)) #128 * 128 - if random.random() > 0.5: - image = cv2.flip(image, 1) - if image.ndim == 2: - image = image[:, :, np.newaxis] - image = (image.transpose((2, 0, 1)) - 127.5) * 0.0078125 - image = torch.from_numpy(image.astype(np.float32)) - return image, image_label - -class ImageDataset_SST(Dataset): - def __init__(self, data_root, train_file, exclude_id_set): - self.data_root = data_root - label_set = set() - # get id2image_path_list - self.id2image_path_list = {} - train_file_buf = open(train_file) - line = train_file_buf.readline().strip() - while line: - image_path, label = line.split(' ') - label = int(label) - if label in exclude_id_set: - line = train_file_buf.readline().strip() - continue - label_set.add(label) - if not label in self.id2image_path_list: - self.id2image_path_list[label] = [] - self.id2image_path_list[label].append(image_path) - line = train_file_buf.readline().strip() - self.train_list = list(label_set) - print('Valid ids: %d.' % len(self.train_list)) - - def __len__(self): - return len(self.train_list) - - def __getitem__(self, index): - cur_id = self.train_list[index] - cur_image_path_list = self.id2image_path_list[cur_id] - if len(cur_image_path_list) == 1: - image_path1 = cur_image_path_list[0] - image_path2 = cur_image_path_list[0] - else: - training_samples = random.sample(cur_image_path_list, 2) - image_path1 = training_samples[0] - image_path2 = training_samples[1] - image_path1 = os.path.join(self.data_root, image_path1) - image_path2 = os.path.join(self.data_root, image_path2) - image1 = cv2.imread(image_path1) - image2 = cv2.imread(image_path2) - image1 = transform(image1) - image2 = transform(image2) - if random.random() > 0.5: - return image2, image1, cur_id - return image1, image2, cur_id +""" +@author: Hang Du, Jun Wang +@date: 20201101 +@contact: jun21wangustc@gmail.com +""" + +import os +import random +import cv2 +import torch +import numpy as np +from torch.utils.data import Dataset + +def transform(image): + """ Transform a image by cv2. + """ + img_size = image.shape[0] + # random crop + if random.random() > 0.5: + crop_size = 9 + x1_offset = np.random.randint(0, crop_size, size=1)[0] + y1_offset = np.random.randint(0, crop_size, size=1)[0] + x2_offset = np.random.randint(img_size-crop_size, img_size, size=1)[0] + y2_offset = np.random.randint(img_size-crop_size, img_size, size=1)[0] + image = image[x1_offset:x2_offset,y1_offset:y2_offset] + image = cv2.resize(image,(img_size,img_size)) + # horizontal flipping + if random.random() > 0.5: + image = cv2.flip(image, 1) + # grayscale conversion + if random.random() > 0.8: + image= cv2.cvtColor(image,cv2.COLOR_BGR2GRAY) + # rotation + if random.random() > 0.5: + theta = (random.randint(-10,10)) * np.pi / 180 + M_rotate = np.array([[np.cos(theta), -np.sin(theta), 0],[np.sin(theta), np.cos(theta), 0]], dtype=np.float32) + image = cv2.warpAffine(image, M_rotate, (img_size, img_size)) + # normalizing + if image.ndim == 2: + image = (image - 127.5) * 0.0078125 + new_image = np.zeros([3,img_size,img_size], np.float32) + new_image[0,:,:] = image + image = torch.from_numpy(new_image.astype(np.float32)) + else: + image = (image.transpose((2, 0, 1)) - 127.5) * 0.0078125 + image = torch.from_numpy(image.astype(np.float32)) + return image + + +class ImageDataset(Dataset): + def __init__(self, data_root, train_file, crop_eye=False): + self.data_root = data_root + self.train_list = [] + train_file_buf = open(train_file) + line = train_file_buf.readline().strip() + while line: + image_path, image_label = line.split(' ') + self.train_list.append((image_path, int(image_label))) + line = train_file_buf.readline().strip() + self.crop_eye = crop_eye + def __len__(self): + return len(self.train_list) + def __getitem__(self, index): + image_path, image_label = self.train_list[index] + image_path = os.path.join(self.data_root, image_path) + image = cv2.imread(image_path) + if self.crop_eye: + image = image[:60, :] + #image = cv2.resize(image, (128, 128)) #128 * 128 + if random.random() > 0.5: + image = cv2.flip(image, 1) + if image.ndim == 2: + image = image[:, :, np.newaxis] + image = (image.transpose((2, 0, 1)) - 127.5) * 0.0078125 + image = torch.from_numpy(image.astype(np.float32)) + return image, image_label + +class ImageDataset_SST(Dataset): + def __init__(self, data_root, train_file, exclude_id_set): + self.data_root = data_root + label_set = set() + # get id2image_path_list + self.id2image_path_list = {} + train_file_buf = open(train_file) + line = train_file_buf.readline().strip() + while line: + image_path, label = line.split(' ') + label = int(label) + if label in exclude_id_set: + line = train_file_buf.readline().strip() + continue + label_set.add(label) + if not label in self.id2image_path_list: + self.id2image_path_list[label] = [] + self.id2image_path_list[label].append(image_path) + line = train_file_buf.readline().strip() + self.train_list = list(label_set) + print('Valid ids: %d.' % len(self.train_list)) + + def __len__(self): + return len(self.train_list) + + def __getitem__(self, index): + cur_id = self.train_list[index] + cur_image_path_list = self.id2image_path_list[cur_id] + if len(cur_image_path_list) == 1: + image_path1 = cur_image_path_list[0] + image_path2 = cur_image_path_list[0] + else: + training_samples = random.sample(cur_image_path_list, 2) + image_path1 = training_samples[0] + image_path2 = training_samples[1] + image_path1 = os.path.join(self.data_root, image_path1) + image_path2 = os.path.join(self.data_root, image_path2) + image1 = cv2.imread(image_path1) + image2 = cv2.imread(image_path2) + image1 = transform(image1) + image2 = transform(image2) + if random.random() > 0.5: + return image2, image1, cur_id + return image1, image2, cur_id diff --git a/docker/Dockerfile b/src/faceDetection/face_sdk/docker/Dockerfile similarity index 100% rename from docker/Dockerfile rename to src/faceDetection/face_sdk/docker/Dockerfile diff --git a/docker/README.MD b/src/faceDetection/face_sdk/docker/README.MD similarity index 100% rename from docker/README.MD rename to src/faceDetection/face_sdk/docker/README.MD diff --git a/face_sdk/README.md b/src/faceDetection/face_sdk/face_sdk/README.md similarity index 100% rename from face_sdk/README.md rename to src/faceDetection/face_sdk/face_sdk/README.md diff --git a/src/faceDetection/face_sdk/face_sdk/__init__.py b/src/faceDetection/face_sdk/face_sdk/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/faceDetection/face_sdk/face_sdk/api_usage/__init__.py b/src/faceDetection/face_sdk/face_sdk/api_usage/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/face_sdk/api_usage/face_alignment.py b/src/faceDetection/face_sdk/face_sdk/api_usage/face_alignment.py similarity index 96% rename from face_sdk/api_usage/face_alignment.py rename to src/faceDetection/face_sdk/face_sdk/api_usage/face_alignment.py index 3d1fb92..4674df4 100644 --- a/face_sdk/api_usage/face_alignment.py +++ b/src/faceDetection/face_sdk/face_sdk/api_usage/face_alignment.py @@ -1,78 +1,78 @@ -""" -@author: JiXuan Xu, Jun Wang -@date: 20201023 -@contact: jun21wangustc@gmail.com -""" -import sys -sys.path.append('.') -import logging.config -logging.config.fileConfig("config/logging.conf") -logger = logging.getLogger('api') - -import yaml -import cv2 -import numpy as np -from core.model_loader.face_alignment.FaceAlignModelLoader import FaceAlignModelLoader -from core.model_handler.face_alignment.FaceAlignModelHandler import FaceAlignModelHandler - -with open('config/model_conf.yaml') as f: - model_conf = yaml.load(f) - -if __name__ == '__main__': - # common setting for all model, need not modify. - model_path = 'models' - - # model setting, modified along with model - scene = 'non-mask' - model_category = 'face_alignment' - model_name = model_conf[scene][model_category] - - logger.info('Start to load the face landmark model...') - # load model - try: - faceAlignModelLoader = FaceAlignModelLoader(model_path, model_category, model_name) - except Exception as e: - logger.error('Failed to parse model configuration file!') - logger.error(e) - sys.exit(-1) - else: - logger.info('Successfully parsed the model configuration file model_meta.json!') - - try: - model, cfg = faceAlignModelLoader.load_model() - except Exception as e: - logger.error('Model loading failed!') - logger.error(e) - sys.exit(-1) - else: - logger.info('Successfully loaded the face landmark model!') - - faceAlignModelHandler = FaceAlignModelHandler(model, 'cuda:0', cfg) - - # read image - image_path = 'api_usage/test_images/test1.jpg' - image_det_txt_path = 'api_usage/test_images/test1_detect_res.txt' - image = cv2.imread(image_path, cv2.IMREAD_COLOR) - with open(image_det_txt_path, 'r') as f: - lines = f.readlines() - try: - for i, line in enumerate(lines): - line = line.strip().split() - det = np.asarray(list(map(int, line[0:4])), dtype=np.int32) - landmarks = faceAlignModelHandler.inference_on_image(image, det) - - save_path_img = 'api_usage/temp/test1_' + 'landmark_res' + str(i) + '.jpg' - save_path_txt = 'api_usage/temp/test1_' + 'landmark_res' + str(i) + '.txt' - image_show = image.copy() - with open(save_path_txt, "w") as fd: - for (x, y) in landmarks.astype(np.int32): - cv2.circle(image_show, (x, y), 2, (255, 0, 0),-1) - line = str(x) + ' ' + str(y) + ' ' - fd.write(line) - cv2.imwrite(save_path_img, image_show) - except Exception as e: - logger.error('Face landmark failed!') - logger.error(e) - sys.exit(-1) - else: - logger.info('Successful face landmark!') +""" +@author: JiXuan Xu, Jun Wang +@date: 20201023 +@contact: jun21wangustc@gmail.com +""" +import sys +sys.path.append('') +import logging.config +logging.config.fileConfig("config/logging.conf") +logger = logging.getLogger('api') + +import yaml +import cv2 +import numpy as np +from core.model_loader.face_alignment.FaceAlignModelLoader import FaceAlignModelLoader +from core.model_handler.face_alignment.FaceAlignModelHandler import FaceAlignModelHandler + +with open('config/model_conf.yaml') as f: + model_conf = yaml.load(f) + +if __name__ == '__main__': + # common setting for all model, need not modify. + model_path = 'models' + + # model setting, modified along with model + scene = 'non-mask' + model_category = 'face_alignment' + model_name = model_conf[scene][model_category] + + logger.info('Start to load the face landmark model...') + # load model + try: + faceAlignModelLoader = FaceAlignModelLoader(model_path, model_category, model_name) + except Exception as e: + logger.error('Failed to parse model configuration file!') + logger.error(e) + sys.exit(-1) + else: + logger.info('Successfully parsed the model configuration file model_meta.json!') + + try: + model, cfg = faceAlignModelLoader.load_model() + except Exception as e: + logger.error('Model loading failed!') + logger.error(e) + sys.exit(-1) + else: + logger.info('Successfully loaded the face landmark model!') + + faceAlignModelHandler = FaceAlignModelHandler(model, 'cuda:0', cfg) + + # read image + image_path = 'api_usage/test_images/test1.jpg' + image_det_txt_path = 'api_usage/test_images/test1_detect_res.txt' + image = cv2.imread(image_path, cv2.IMREAD_COLOR) + with open(image_det_txt_path, 'r') as f: + lines = f.readlines() + try: + for i, line in enumerate(lines): + line = line.strip().split() + det = np.asarray(list(map(int, line[0:4])), dtype=np.int32) + landmarks = faceAlignModelHandler.inference_on_image(image, det) + + save_path_img = 'api_usage/temp/test1_' + 'landmark_res' + str(i) + '.jpg' + save_path_txt = 'api_usage/temp/test1_' + 'landmark_res' + str(i) + '.txt' + image_show = image.copy() + with open(save_path_txt, "w") as fd: + for (x, y) in landmarks.astype(np.int32): + cv2.circle(image_show, (x, y), 2, (255, 0, 0),-1) + line = str(x) + ' ' + str(y) + ' ' + fd.write(line) + cv2.imwrite(save_path_img, image_show) + except Exception as e: + logger.error('Face landmark failed!') + logger.error(e) + sys.exit(-1) + else: + logger.info('Successful face landmark!') diff --git a/face_sdk/api_usage/face_crop.py b/src/faceDetection/face_sdk/face_sdk/api_usage/face_crop.py similarity index 97% rename from face_sdk/api_usage/face_crop.py rename to src/faceDetection/face_sdk/face_sdk/api_usage/face_crop.py index 717fc92..3051bd4 100644 --- a/face_sdk/api_usage/face_crop.py +++ b/src/faceDetection/face_sdk/face_sdk/api_usage/face_crop.py @@ -4,7 +4,7 @@ @contact: jun21wangustc@gmail.com """ import sys -sys.path.append('.') +sys.path.append('') import logging mpl_logger = logging.getLogger('matplotlib') mpl_logger.setLevel(logging.WARNING) diff --git a/face_sdk/api_usage/face_detect.py b/src/faceDetection/face_sdk/face_sdk/api_usage/face_detect.py similarity index 96% rename from face_sdk/api_usage/face_detect.py rename to src/faceDetection/face_sdk/face_sdk/api_usage/face_detect.py index 3403c0c..c7a7e14 100644 --- a/face_sdk/api_usage/face_detect.py +++ b/src/faceDetection/face_sdk/face_sdk/api_usage/face_detect.py @@ -1,80 +1,80 @@ -""" -@author: JiXuan Xu, Jun Wang -@date: 20201019 -@contact: jun21wangustc@gmail.com -""" -import sys -sys.path.append('.') -import logging.config -logging.config.fileConfig("config/logging.conf") -logger = logging.getLogger('api') - -import yaml -import cv2 -import numpy as np -from core.model_loader.face_detection.FaceDetModelLoader import FaceDetModelLoader -from core.model_handler.face_detection.FaceDetModelHandler import FaceDetModelHandler - -with open('config/model_conf.yaml') as f: - model_conf = yaml.load(f) - -if __name__ == '__main__': - # common setting for all model, need not modify. - model_path = 'models' - - # model setting, modified along with model - scene = 'non-mask' - model_category = 'face_detection' - model_name = model_conf[scene][model_category] - - logger.info('Start to load the face detection model...') - # load model - try: - faceDetModelLoader = FaceDetModelLoader(model_path, model_category, model_name) - except Exception as e: - logger.error('Failed to parse model configuration file!') - logger.error(e) - sys.exit(-1) - else: - logger.info('Successfully parsed the model configuration file model_meta.json!') - - try: - model, cfg = faceDetModelLoader.load_model() - except Exception as e: - logger.error('Model loading failed!') - logger.error(e) - sys.exit(-1) - else: - logger.info('Successfully loaded the face detection model!') - - # read image - image_path = 'api_usage/test_images/test1.jpg' - image = cv2.imread(image_path, cv2.IMREAD_COLOR) - faceDetModelHandler = FaceDetModelHandler(model, 'cuda:0', cfg) - - try: - dets = faceDetModelHandler.inference_on_image(image) - except Exception as e: - logger.error('Face detection failed!') - logger.error(e) - sys.exit(-1) - else: - logger.info('Successful face detection!') - - # gen result - save_path_img = 'api_usage/temp/test1_detect_res.jpg' - save_path_txt = 'api_usage/temp/test1_detect_res.txt' - - bboxs = dets - with open(save_path_txt, "w") as fd: - for box in bboxs: - line = str(int(box[0])) + " " + str(int(box[1])) + " " + \ - str(int(box[2])) + " " + str(int(box[3])) + " " + \ - str(box[4]) + " \n" - fd.write(line) - - for box in bboxs: - box = list(map(int, box)) - cv2.rectangle(image, (box[0], box[1]), (box[2], box[3]), (0, 0, 255), 2) - cv2.imwrite(save_path_img, image) - logger.info('Successfully generate face detection results!') +""" +@author: JiXuan Xu, Jun Wang +@date: 20201019 +@contact: jun21wangustc@gmail.com +""" +import sys +sys.path.append('') +import logging.config +logging.config.fileConfig("config/logging.conf") +logger = logging.getLogger('api') + +import yaml +import cv2 +import numpy as np +from core.model_loader.face_detection.FaceDetModelLoader import FaceDetModelLoader +from core.model_handler.face_detection.FaceDetModelHandler import FaceDetModelHandler + +with open('config/model_conf.yaml') as f: + model_conf = yaml.load(f) + +if __name__ == '__main__': + # common setting for all model, need not modify. + model_path = 'models' + + # model setting, modified along with model + scene = 'non-mask' + model_category = 'face_detection' + model_name = model_conf[scene][model_category] + + logger.info('Start to load the face detection model...') + # load model + try: + faceDetModelLoader = FaceDetModelLoader(model_path, model_category, model_name) + except Exception as e: + logger.error('Failed to parse model configuration file!') + logger.error(e) + sys.exit(-1) + else: + logger.info('Successfully parsed the model configuration file model_meta.json!') + + try: + model, cfg = faceDetModelLoader.load_model() + except Exception as e: + logger.error('Model loading failed!') + logger.error(e) + sys.exit(-1) + else: + logger.info('Successfully loaded the face detection model!') + + # read image + image_path = 'api_usage/test_images/test1.jpg' + image = cv2.imread(image_path, cv2.IMREAD_COLOR) + faceDetModelHandler = FaceDetModelHandler(model, 'cuda:0', cfg) + + try: + dets = faceDetModelHandler.inference_on_image(image) + except Exception as e: + logger.error('Face detection failed!') + logger.error(e) + sys.exit(-1) + else: + logger.info('Successful face detection!') + + # gen result + save_path_img = 'api_usage/temp/test1_detect_res.jpg' + save_path_txt = 'api_usage/temp/test1_detect_res.txt' + + bboxs = dets + with open(save_path_txt, "w") as fd: + for box in bboxs: + line = str(int(box[0])) + " " + str(int(box[1])) + " " + \ + str(int(box[2])) + " " + str(int(box[3])) + " " + \ + str(box[4]) + " \n" + fd.write(line) + + for box in bboxs: + box = list(map(int, box)) + cv2.rectangle(image, (box[0], box[1]), (box[2], box[3]), (0, 0, 255), 2) + cv2.imwrite(save_path_img, image) + logger.info('Successfully generate face detection results!') diff --git a/face_sdk/api_usage/face_feature.py b/src/faceDetection/face_sdk/face_sdk/api_usage/face_feature.py similarity index 98% rename from face_sdk/api_usage/face_feature.py rename to src/faceDetection/face_sdk/face_sdk/api_usage/face_feature.py index 056feed..3358f25 100644 --- a/face_sdk/api_usage/face_feature.py +++ b/src/faceDetection/face_sdk/face_sdk/api_usage/face_feature.py @@ -4,7 +4,7 @@ @contact: jun21wangustc@gmail.com """ import sys -sys.path.append('.') +sys.path.append('') import logging.config logging.config.fileConfig("config/logging.conf") logger = logging.getLogger('api') diff --git a/face_sdk/api_usage/face_parsing.py b/src/faceDetection/face_sdk/face_sdk/api_usage/face_parsing.py similarity index 99% rename from face_sdk/api_usage/face_parsing.py rename to src/faceDetection/face_sdk/face_sdk/api_usage/face_parsing.py index b7e681a..3de5bae 100644 --- a/face_sdk/api_usage/face_parsing.py +++ b/src/faceDetection/face_sdk/face_sdk/api_usage/face_parsing.py @@ -5,7 +5,7 @@ """ import sys -sys.path.append('.') +sys.path.append('') import logging mpl_logger = logging.getLogger('matplotlib') mpl_logger.setLevel(logging.WARNING) diff --git a/face_sdk/api_usage/face_pipline.py b/src/faceDetection/face_sdk/face_sdk/api_usage/face_pipline.py similarity index 99% rename from face_sdk/api_usage/face_pipline.py rename to src/faceDetection/face_sdk/face_sdk/api_usage/face_pipline.py index 37d1d01..4130e5c 100644 --- a/face_sdk/api_usage/face_pipline.py +++ b/src/faceDetection/face_sdk/face_sdk/api_usage/face_pipline.py @@ -4,7 +4,7 @@ @contact: jun21wangustc@gmail.com """ import sys -sys.path.append('.') +sys.path.append('') import logging mpl_logger = logging.getLogger('matplotlib') mpl_logger.setLevel(logging.WARNING) diff --git a/face_sdk/api_usage/temp/readme.md b/src/faceDetection/face_sdk/face_sdk/api_usage/temp/readme.md similarity index 100% rename from face_sdk/api_usage/temp/readme.md rename to src/faceDetection/face_sdk/face_sdk/api_usage/temp/readme.md diff --git a/face_sdk/api_usage/test_images/test1.jpg b/src/faceDetection/face_sdk/face_sdk/api_usage/test_images/test1.jpg similarity index 100% rename from face_sdk/api_usage/test_images/test1.jpg rename to src/faceDetection/face_sdk/face_sdk/api_usage/test_images/test1.jpg diff --git a/face_sdk/api_usage/test_images/test1_cropped.jpg b/src/faceDetection/face_sdk/face_sdk/api_usage/test_images/test1_cropped.jpg similarity index 100% rename from face_sdk/api_usage/test_images/test1_cropped.jpg rename to src/faceDetection/face_sdk/face_sdk/api_usage/test_images/test1_cropped.jpg diff --git a/face_sdk/api_usage/test_images/test1_detect_res.txt b/src/faceDetection/face_sdk/face_sdk/api_usage/test_images/test1_detect_res.txt similarity index 100% rename from face_sdk/api_usage/test_images/test1_detect_res.txt rename to src/faceDetection/face_sdk/face_sdk/api_usage/test_images/test1_detect_res.txt diff --git a/face_sdk/api_usage/test_images/test1_landmark_res0.txt b/src/faceDetection/face_sdk/face_sdk/api_usage/test_images/test1_landmark_res0.txt similarity index 100% rename from face_sdk/api_usage/test_images/test1_landmark_res0.txt rename to src/faceDetection/face_sdk/face_sdk/api_usage/test_images/test1_landmark_res0.txt diff --git a/face_sdk/api_usage/test_images/test2.jpg b/src/faceDetection/face_sdk/face_sdk/api_usage/test_images/test2.jpg similarity index 100% rename from face_sdk/api_usage/test_images/test2.jpg rename to src/faceDetection/face_sdk/face_sdk/api_usage/test_images/test2.jpg diff --git a/face_sdk/api_usage/test_images/test3_mask.jpg b/src/faceDetection/face_sdk/face_sdk/api_usage/test_images/test3_mask.jpg similarity index 100% rename from face_sdk/api_usage/test_images/test3_mask.jpg rename to src/faceDetection/face_sdk/face_sdk/api_usage/test_images/test3_mask.jpg diff --git a/face_sdk/api_usage/test_images/test4_mask.jpg b/src/faceDetection/face_sdk/face_sdk/api_usage/test_images/test4_mask.jpg similarity index 100% rename from face_sdk/api_usage/test_images/test4_mask.jpg rename to src/faceDetection/face_sdk/face_sdk/api_usage/test_images/test4_mask.jpg diff --git a/face_sdk/config/logging.conf b/src/faceDetection/face_sdk/face_sdk/config/logging.conf similarity index 100% rename from face_sdk/config/logging.conf rename to src/faceDetection/face_sdk/face_sdk/config/logging.conf diff --git a/face_sdk/config/model_conf.yaml b/src/faceDetection/face_sdk/face_sdk/config/model_conf.yaml similarity index 100% rename from face_sdk/config/model_conf.yaml rename to src/faceDetection/face_sdk/face_sdk/config/model_conf.yaml diff --git a/src/faceDetection/face_sdk/face_sdk/core/__init__.py b/src/faceDetection/face_sdk/face_sdk/core/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/face_sdk/core/image_cropper/BaseImageCropper.py b/src/faceDetection/face_sdk/face_sdk/core/image_cropper/BaseImageCropper.py similarity index 100% rename from face_sdk/core/image_cropper/BaseImageCropper.py rename to src/faceDetection/face_sdk/face_sdk/core/image_cropper/BaseImageCropper.py diff --git a/src/faceDetection/face_sdk/face_sdk/core/image_cropper/__init__.py b/src/faceDetection/face_sdk/face_sdk/core/image_cropper/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/face_sdk/core/image_cropper/arcface_cropper/FaceRecImageCropper.py b/src/faceDetection/face_sdk/face_sdk/core/image_cropper/arcface_cropper/FaceRecImageCropper.py similarity index 94% rename from face_sdk/core/image_cropper/arcface_cropper/FaceRecImageCropper.py rename to src/faceDetection/face_sdk/face_sdk/core/image_cropper/arcface_cropper/FaceRecImageCropper.py index 7ea1a85..f1d0242 100644 --- a/face_sdk/core/image_cropper/arcface_cropper/FaceRecImageCropper.py +++ b/src/faceDetection/face_sdk/face_sdk/core/image_cropper/arcface_cropper/FaceRecImageCropper.py @@ -11,8 +11,8 @@ import numpy as np from skimage import transform as trans -from core.image_cropper.BaseImageCropper import BaseImageCropper -from utils.lms_trans import lms106_2_lms5, lms25_2_lms5 +from src.faceDetection.face_sdk.face_sdk.core.image_cropper.BaseImageCropper import BaseImageCropper +from src.faceDetection.face_sdk.face_sdk.utils.lms_trans import lms106_2_lms5, lms25_2_lms5 src1 = np.array([ [51.642,50.115], diff --git a/src/faceDetection/face_sdk/face_sdk/core/image_cropper/arcface_cropper/__init__.py b/src/faceDetection/face_sdk/face_sdk/core/image_cropper/arcface_cropper/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/face_sdk/core/model_handler/BaseModelHandler.py b/src/faceDetection/face_sdk/face_sdk/core/model_handler/BaseModelHandler.py similarity index 100% rename from face_sdk/core/model_handler/BaseModelHandler.py rename to src/faceDetection/face_sdk/face_sdk/core/model_handler/BaseModelHandler.py diff --git a/src/faceDetection/face_sdk/face_sdk/core/model_handler/__init__.py b/src/faceDetection/face_sdk/face_sdk/core/model_handler/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/face_sdk/core/model_handler/face_alignment/FaceAlignModelHandler.py b/src/faceDetection/face_sdk/face_sdk/core/model_handler/face_alignment/FaceAlignModelHandler.py similarity index 94% rename from face_sdk/core/model_handler/face_alignment/FaceAlignModelHandler.py rename to src/faceDetection/face_sdk/face_sdk/core/model_handler/face_alignment/FaceAlignModelHandler.py index b8117a3..6b8624a 100644 --- a/face_sdk/core/model_handler/face_alignment/FaceAlignModelHandler.py +++ b/src/faceDetection/face_sdk/face_sdk/core/model_handler/face_alignment/FaceAlignModelHandler.py @@ -4,7 +4,7 @@ @contact: jun21wangustc@gmail.com """ import logging.config -logging.config.fileConfig("config/logging.conf") +#logging.config.fileConfig("config/logging.conf") logger = logging.getLogger('sdk') import cv2 @@ -12,8 +12,8 @@ import numpy as np import torch.backends.cudnn as cudnn -from core.model_handler.BaseModelHandler import BaseModelHandler -from utils.BuzException import * +from src.faceDetection.face_sdk.face_sdk.core.model_handler.BaseModelHandler import BaseModelHandler +from src.faceDetection.face_sdk.face_sdk.utils.BuzException import * from torchvision import transforms class FaceAlignModelHandler(BaseModelHandler): diff --git a/src/faceDetection/face_sdk/face_sdk/core/model_handler/face_alignment/__init__.py b/src/faceDetection/face_sdk/face_sdk/core/model_handler/face_alignment/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/face_sdk/core/model_handler/face_detection/FaceDetModelHandler.py b/src/faceDetection/face_sdk/face_sdk/core/model_handler/face_detection/FaceDetModelHandler.py similarity index 96% rename from face_sdk/core/model_handler/face_detection/FaceDetModelHandler.py rename to src/faceDetection/face_sdk/face_sdk/core/model_handler/face_detection/FaceDetModelHandler.py index b458de2..73d0cc8 100644 --- a/face_sdk/core/model_handler/face_detection/FaceDetModelHandler.py +++ b/src/faceDetection/face_sdk/face_sdk/core/model_handler/face_detection/FaceDetModelHandler.py @@ -5,7 +5,7 @@ """ import logging.config -logging.config.fileConfig("config/logging.conf") +# logging.config.fileConfig("config/logging.conf") logger = logging.getLogger('sdk') import torch @@ -14,9 +14,8 @@ from itertools import product as product import torch.backends.cudnn as cudnn -from core.model_handler.BaseModelHandler import BaseModelHandler -from utils.BuzException import * - +from src.faceDetection.face_sdk.face_sdk.core.model_handler.BaseModelHandler import BaseModelHandler +from src.faceDetection.face_sdk.face_sdk.utils.BuzException import * class FaceDetModelHandler(BaseModelHandler): """Implementation of face detection model handler diff --git a/src/faceDetection/face_sdk/face_sdk/core/model_handler/face_detection/__init__.py b/src/faceDetection/face_sdk/face_sdk/core/model_handler/face_detection/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/face_sdk/core/model_handler/face_parsing/FaceParsingModelHandler.py b/src/faceDetection/face_sdk/face_sdk/core/model_handler/face_parsing/FaceParsingModelHandler.py similarity index 100% rename from face_sdk/core/model_handler/face_parsing/FaceParsingModelHandler.py rename to src/faceDetection/face_sdk/face_sdk/core/model_handler/face_parsing/FaceParsingModelHandler.py diff --git a/src/faceDetection/face_sdk/face_sdk/core/model_handler/face_parsing/__init__.py b/src/faceDetection/face_sdk/face_sdk/core/model_handler/face_parsing/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/face_sdk/core/model_handler/face_recognition/FaceRecModelHandler.py b/src/faceDetection/face_sdk/face_sdk/core/model_handler/face_recognition/FaceRecModelHandler.py similarity index 90% rename from face_sdk/core/model_handler/face_recognition/FaceRecModelHandler.py rename to src/faceDetection/face_sdk/face_sdk/core/model_handler/face_recognition/FaceRecModelHandler.py index a9d8978..77733c9 100644 --- a/face_sdk/core/model_handler/face_recognition/FaceRecModelHandler.py +++ b/src/faceDetection/face_sdk/face_sdk/core/model_handler/face_recognition/FaceRecModelHandler.py @@ -4,14 +4,14 @@ @contact: jun21wangustc@gmail.com """ import logging.config -logging.config.fileConfig("config/logging.conf") +#logging.config.fileConfig("config/logging.conf") logger = logging.getLogger('sdk') import numpy as np import torch -from core.model_handler.BaseModelHandler import BaseModelHandler -from utils.BuzException import * +from src.faceDetection.face_sdk.face_sdk.core.model_handler.BaseModelHandler import BaseModelHandler +from src.faceDetection.face_sdk.face_sdk.utils.BuzException import * class FaceRecModelHandler(BaseModelHandler): """Implementation of face recognition model handler diff --git a/src/faceDetection/face_sdk/face_sdk/core/model_handler/face_recognition/__init__.py b/src/faceDetection/face_sdk/face_sdk/core/model_handler/face_recognition/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/face_sdk/core/model_loader/BaseModelLoader.py b/src/faceDetection/face_sdk/face_sdk/core/model_loader/BaseModelLoader.py similarity index 97% rename from face_sdk/core/model_loader/BaseModelLoader.py rename to src/faceDetection/face_sdk/face_sdk/core/model_loader/BaseModelLoader.py index 2577851..c7a1cb5 100644 --- a/face_sdk/core/model_loader/BaseModelLoader.py +++ b/src/faceDetection/face_sdk/face_sdk/core/model_loader/BaseModelLoader.py @@ -7,7 +7,7 @@ import sys sys.path.append('models/network_def') import logging.config -logging.config.fileConfig("config/logging.conf") +# logging.config.fileConfig("config/logging.conf") logger = logging.getLogger('sdk') from abc import ABCMeta, abstractmethod diff --git a/src/faceDetection/face_sdk/face_sdk/core/model_loader/__init__.py b/src/faceDetection/face_sdk/face_sdk/core/model_loader/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/face_sdk/core/model_loader/face_alignment/FaceAlignModelLoader.py b/src/faceDetection/face_sdk/face_sdk/core/model_loader/face_alignment/FaceAlignModelLoader.py similarity index 64% rename from face_sdk/core/model_loader/face_alignment/FaceAlignModelLoader.py rename to src/faceDetection/face_sdk/face_sdk/core/model_loader/face_alignment/FaceAlignModelLoader.py index efb9276..66530c5 100644 --- a/face_sdk/core/model_loader/face_alignment/FaceAlignModelLoader.py +++ b/src/faceDetection/face_sdk/face_sdk/core/model_loader/face_alignment/FaceAlignModelLoader.py @@ -1,30 +1,34 @@ -""" -@author: JiXuan Xu, Jun Wang -@date: 20201023 -@contact: jun21wangustc@gmail.com -""" -import logging.config -logging.config.fileConfig("config/logging.conf") -logger = logging.getLogger('sdk') - -import torch - -from core.model_loader.BaseModelLoader import BaseModelLoader - -class FaceAlignModelLoader(BaseModelLoader): - def __init__(self, model_path, model_category, model_name, meta_file='model_meta.json'): - logger.info('Start to analyze the face landmark model, model path: %s, model category: %s，model name: %s' % - (model_path, model_category, model_name)) - super().__init__(model_path, model_category, model_name, meta_file) - self.cfg['img_size'] = self.meta_conf['input_width'] - - def load_model(self): - try: - model = torch.load(self.cfg['model_file_path']) - except Exception as e: - logger.error('The model failed to load, please check the model path: %s!' - % self.cfg['model_file_path']) - raise e - else: - logger.info('Successfully loaded the face landmark model!') - return model, self.cfg +""" +@author: JiXuan Xu, Jun Wang +@date: 20201023 +@contact: jun21wangustc@gmail.com +""" +import logging.config +#logging.config.fileConfig("config/logging.conf") +logger = logging.getLogger('sdk') + +import torch + +from src.faceDetection.face_sdk.face_sdk.core.model_loader.BaseModelLoader import BaseModelLoader + +class FaceAlignModelLoader(BaseModelLoader): + def __init__(self, model_path, model_category, model_name, meta_file='model_meta.json'): + logger.info('Start to analyze the face landmark model, model path: %s, model category: %s，model name: %s' % + (model_path, model_category, model_name)) + super().__init__(model_path, model_category, model_name, meta_file) + self.cfg['img_size'] = self.meta_conf['input_width'] + + def load_model(self): + try: + model = torch.load(self.cfg['model_file_path'], map_location=torch.device('cpu'), weights_only=False) + if hasattr(model, 'module'): + model = model.module # Extract the actual model from DataParallel wrapper + model = model.to('cpu') # Ensure the model is on CPU + model.eval() # Set the model to evaluation mode + except Exception as e: + logger.error('The model failed to load, please check the model path: %s!' + % self.cfg['model_file_path']) + raise e + else: + logger.info('Successfully loaded the face landmark model!') + return model, self.cfg diff --git a/src/faceDetection/face_sdk/face_sdk/core/model_loader/face_alignment/__init__.py b/src/faceDetection/face_sdk/face_sdk/core/model_loader/face_alignment/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/face_sdk/core/model_loader/face_detection/FaceDetModelLoader.py b/src/faceDetection/face_sdk/face_sdk/core/model_loader/face_detection/FaceDetModelLoader.py similarity index 71% rename from face_sdk/core/model_loader/face_detection/FaceDetModelLoader.py rename to src/faceDetection/face_sdk/face_sdk/core/model_loader/face_detection/FaceDetModelLoader.py index a441263..d4c3176 100644 --- a/face_sdk/core/model_loader/face_detection/FaceDetModelLoader.py +++ b/src/faceDetection/face_sdk/face_sdk/core/model_loader/face_detection/FaceDetModelLoader.py @@ -1,35 +1,40 @@ -""" -@author: JiXuan Xu, Jun Wang -@date: 20201019 -@contact: jun21wangustc@gmail.com -""" -import logging.config -logging.config.fileConfig("config/logging.conf") -logger = logging.getLogger('sdk') - -import torch - -from core.model_loader.BaseModelLoader import BaseModelLoader - -class FaceDetModelLoader(BaseModelLoader): - def __init__(self, model_path, model_category, model_name, meta_file='model_meta.json'): - logger.info('Start to analyze the face detection model, model path: %s, model category: %s，model name: %s' % - (model_path, model_category, model_name)) - super().__init__(model_path, model_category, model_name, meta_file) - self.cfg['min_sizes'] = self.meta_conf['min_sizes'] - self.cfg['steps'] = self.meta_conf['steps'] - self.cfg['variance'] = self.meta_conf['variance'] - self.cfg['in_channel'] = self.meta_conf['in_channel'] - self.cfg['out_channel'] = self.meta_conf['out_channel'] - self.cfg['confidence_threshold'] = self.meta_conf['confidence_threshold'] - - def load_model(self): - try: - model = torch.load(self.cfg['model_file_path']) - except Exception as e: - logger.error('The model failed to load, please check the model path: %s!' - % self.cfg['model_file_path']) - raise e - else: - logger.info('Successfully loaded the face detection model!') - return model, self.cfg +""" +@author: JiXuan Xu, Jun Wang +@date: 20201019 +@contact: jun21wangustc@gmail.com +""" +import logging.config + +from src.faceDetection.face_sdk.face_sdk.core.model_loader.BaseModelLoader import BaseModelLoader + +# logging.config.fileConfig("config/logging.conf") +logger = logging.getLogger('sdk') + +import torch + +class FaceDetModelLoader(BaseModelLoader): + def __init__(self, model_path, model_category, model_name, meta_file='model_meta.json'): + logger.info('Start to analyze the face detection model, model path: %s, model category: %s，model name: %s' % + (model_path, model_category, model_name)) + super().__init__(model_path, model_category, model_name, meta_file) + self.cfg['min_sizes'] = self.meta_conf['min_sizes'] + self.cfg['steps'] = self.meta_conf['steps'] + self.cfg['variance'] = self.meta_conf['variance'] + self.cfg['in_channel'] = self.meta_conf['in_channel'] + self.cfg['out_channel'] = self.meta_conf['out_channel'] + self.cfg['confidence_threshold'] = self.meta_conf['confidence_threshold'] + + def load_model(self): + try: + model = torch.load(self.cfg['model_file_path'], map_location=torch.device('cpu'), weights_only=False) + if hasattr(model, 'module'): + model = model.module # Extract the actual model from DataParallel wrapper + model = model.to('cpu') # Ensure the model is on CPU + model.eval() # Set the model to evaluation mode + except Exception as e: + logger.error('The model failed to load, please check the model path: %s!' + % self.cfg['model_file_path']) + raise e + else: + logger.info('Successfully loaded the face detection model!') + return model, self.cfg diff --git a/src/faceDetection/face_sdk/face_sdk/core/model_loader/face_detection/__init__.py b/src/faceDetection/face_sdk/face_sdk/core/model_loader/face_detection/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/face_sdk/core/model_loader/face_parsing/FaceParsingModelLoader.py b/src/faceDetection/face_sdk/face_sdk/core/model_loader/face_parsing/FaceParsingModelLoader.py similarity index 100% rename from face_sdk/core/model_loader/face_parsing/FaceParsingModelLoader.py rename to src/faceDetection/face_sdk/face_sdk/core/model_loader/face_parsing/FaceParsingModelLoader.py diff --git a/src/faceDetection/face_sdk/face_sdk/core/model_loader/face_parsing/__init__.py b/src/faceDetection/face_sdk/face_sdk/core/model_loader/face_parsing/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/face_sdk/core/model_loader/face_recognition/FaceRecModelLoader.py b/src/faceDetection/face_sdk/face_sdk/core/model_loader/face_recognition/FaceRecModelLoader.py similarity index 65% rename from face_sdk/core/model_loader/face_recognition/FaceRecModelLoader.py rename to src/faceDetection/face_sdk/face_sdk/core/model_loader/face_recognition/FaceRecModelLoader.py index ce880ce..5060272 100644 --- a/face_sdk/core/model_loader/face_recognition/FaceRecModelLoader.py +++ b/src/faceDetection/face_sdk/face_sdk/core/model_loader/face_recognition/FaceRecModelLoader.py @@ -4,12 +4,12 @@ @contact: jun21wangustc@gmail.com """ import logging.config -logging.config.fileConfig("config/logging.conf") +#logging.config.fileConfig("config/logging.conf") logger = logging.getLogger('sdk') import torch -from core.model_loader.BaseModelLoader import BaseModelLoader +from src.faceDetection.face_sdk.face_sdk.core.model_loader.BaseModelLoader import BaseModelLoader class FaceRecModelLoader(BaseModelLoader): def __init__(self, model_path, model_category, model_name, meta_file='model_meta.json'): @@ -18,10 +18,14 @@ def __init__(self, model_path, model_category, model_name, meta_file='model_meta super().__init__(model_path, model_category, model_name, meta_file) self.cfg['mean'] = self.meta_conf['mean'] self.cfg['std'] = self.meta_conf['std'] - + def load_model(self): try: - model = torch.load(self.cfg['model_file_path']) + model = torch.load(self.cfg['model_file_path'], map_location=torch.device('cpu'), weights_only=False) + if hasattr(model, 'module'): + model = model.module # Extract the actual model from DataParallel wrapper + model = model.to('cpu') # Ensure the model is on CPU + model.eval() # Set the model to evaluation mode except Exception as e: logger.error('The model failed to load, please check the model path: %s!' % self.cfg['model_file_path']) diff --git a/src/faceDetection/face_sdk/face_sdk/core/model_loader/face_recognition/__init__.py b/src/faceDetection/face_sdk/face_sdk/core/model_loader/face_recognition/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/face_sdk/logs/readme.md b/src/faceDetection/face_sdk/face_sdk/logs/readme.md similarity index 100% rename from face_sdk/logs/readme.md rename to src/faceDetection/face_sdk/face_sdk/logs/readme.md diff --git a/face_sdk/logs/sdk.log b/src/faceDetection/face_sdk/face_sdk/logs/sdk.log similarity index 100% rename from face_sdk/logs/sdk.log rename to src/faceDetection/face_sdk/face_sdk/logs/sdk.log diff --git a/src/faceDetection/face_sdk/face_sdk/models/__init__.py b/src/faceDetection/face_sdk/face_sdk/models/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/face_sdk/models/face_alignment/face_alignment_1.0/face_landmark_pfld.pkl b/src/faceDetection/face_sdk/face_sdk/models/face_alignment/face_alignment_1.0/face_landmark_pfld.pkl similarity index 100% rename from face_sdk/models/face_alignment/face_alignment_1.0/face_landmark_pfld.pkl rename to src/faceDetection/face_sdk/face_sdk/models/face_alignment/face_alignment_1.0/face_landmark_pfld.pkl diff --git a/face_sdk/models/face_alignment/face_alignment_1.0/model_meta.json b/src/faceDetection/face_sdk/face_sdk/models/face_alignment/face_alignment_1.0/model_meta.json similarity index 100% rename from face_sdk/models/face_alignment/face_alignment_1.0/model_meta.json rename to src/faceDetection/face_sdk/face_sdk/models/face_alignment/face_alignment_1.0/model_meta.json diff --git a/face_sdk/models/face_alignment/face_alignment_2.0/face_landmark_pfld.pkl b/src/faceDetection/face_sdk/face_sdk/models/face_alignment/face_alignment_2.0/face_landmark_pfld.pkl similarity index 100% rename from face_sdk/models/face_alignment/face_alignment_2.0/face_landmark_pfld.pkl rename to src/faceDetection/face_sdk/face_sdk/models/face_alignment/face_alignment_2.0/face_landmark_pfld.pkl diff --git a/face_sdk/models/face_alignment/face_alignment_2.0/model_meta.json b/src/faceDetection/face_sdk/face_sdk/models/face_alignment/face_alignment_2.0/model_meta.json similarity index 100% rename from face_sdk/models/face_alignment/face_alignment_2.0/model_meta.json rename to src/faceDetection/face_sdk/face_sdk/models/face_alignment/face_alignment_2.0/model_meta.json diff --git a/face_sdk/models/face_detection/face_detection_1.0/face_detection_retina.pkl b/src/faceDetection/face_sdk/face_sdk/models/face_detection/face_detection_1.0/face_detection_retina.pkl similarity index 100% rename from face_sdk/models/face_detection/face_detection_1.0/face_detection_retina.pkl rename to src/faceDetection/face_sdk/face_sdk/models/face_detection/face_detection_1.0/face_detection_retina.pkl diff --git a/face_sdk/models/face_detection/face_detection_1.0/model_meta.json b/src/faceDetection/face_sdk/face_sdk/models/face_detection/face_detection_1.0/model_meta.json similarity index 100% rename from face_sdk/models/face_detection/face_detection_1.0/model_meta.json rename to src/faceDetection/face_sdk/face_sdk/models/face_detection/face_detection_1.0/model_meta.json diff --git a/face_sdk/models/face_detection/face_detection_2.0/face_detection_retina.pkl b/src/faceDetection/face_sdk/face_sdk/models/face_detection/face_detection_2.0/face_detection_retina.pkl similarity index 100% rename from face_sdk/models/face_detection/face_detection_2.0/face_detection_retina.pkl rename to src/faceDetection/face_sdk/face_sdk/models/face_detection/face_detection_2.0/face_detection_retina.pkl diff --git a/face_sdk/models/face_detection/face_detection_2.0/model_meta.json b/src/faceDetection/face_sdk/face_sdk/models/face_detection/face_detection_2.0/model_meta.json similarity index 100% rename from face_sdk/models/face_detection/face_detection_2.0/model_meta.json rename to src/faceDetection/face_sdk/face_sdk/models/face_detection/face_detection_2.0/model_meta.json diff --git a/face_sdk/models/face_parsing/README.md b/src/faceDetection/face_sdk/face_sdk/models/face_parsing/README.md similarity index 100% rename from face_sdk/models/face_parsing/README.md rename to src/faceDetection/face_sdk/face_sdk/models/face_parsing/README.md diff --git a/face_sdk/models/face_parsing/face_parsing_1.0/model_meta.json b/src/faceDetection/face_sdk/face_sdk/models/face_parsing/face_parsing_1.0/model_meta.json similarity index 100% rename from face_sdk/models/face_parsing/face_parsing_1.0/model_meta.json rename to src/faceDetection/face_sdk/face_sdk/models/face_parsing/face_parsing_1.0/model_meta.json diff --git a/face_sdk/models/face_recognition/face_recognition_1.0/face_recognition_mv.pkl b/src/faceDetection/face_sdk/face_sdk/models/face_recognition/face_recognition_1.0/face_recognition_mv.pkl similarity index 100% rename from face_sdk/models/face_recognition/face_recognition_1.0/face_recognition_mv.pkl rename to src/faceDetection/face_sdk/face_sdk/models/face_recognition/face_recognition_1.0/face_recognition_mv.pkl diff --git a/face_sdk/models/face_recognition/face_recognition_1.0/model_meta.json b/src/faceDetection/face_sdk/face_sdk/models/face_recognition/face_recognition_1.0/model_meta.json similarity index 100% rename from face_sdk/models/face_recognition/face_recognition_1.0/model_meta.json rename to src/faceDetection/face_sdk/face_sdk/models/face_recognition/face_recognition_1.0/model_meta.json diff --git a/face_sdk/models/face_recognition/face_recognition_2.0/face_recognition_mv.pkl b/src/faceDetection/face_sdk/face_sdk/models/face_recognition/face_recognition_2.0/face_recognition_mv.pkl similarity index 100% rename from face_sdk/models/face_recognition/face_recognition_2.0/face_recognition_mv.pkl rename to src/faceDetection/face_sdk/face_sdk/models/face_recognition/face_recognition_2.0/face_recognition_mv.pkl diff --git a/face_sdk/models/face_recognition/face_recognition_2.0/model_meta.json b/src/faceDetection/face_sdk/face_sdk/models/face_recognition/face_recognition_2.0/model_meta.json similarity index 100% rename from face_sdk/models/face_recognition/face_recognition_2.0/model_meta.json rename to src/faceDetection/face_sdk/face_sdk/models/face_recognition/face_recognition_2.0/model_meta.json diff --git a/src/faceDetection/face_sdk/face_sdk/models/network_def/__init__.py b/src/faceDetection/face_sdk/face_sdk/models/network_def/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/face_sdk/models/network_def/mobilefacenet_def.py b/src/faceDetection/face_sdk/face_sdk/models/network_def/mobilefacenet_def.py similarity index 100% rename from face_sdk/models/network_def/mobilefacenet_def.py rename to src/faceDetection/face_sdk/face_sdk/models/network_def/mobilefacenet_def.py diff --git a/face_sdk/models/network_def/mobilev3_pfld.py b/src/faceDetection/face_sdk/face_sdk/models/network_def/mobilev3_pfld.py similarity index 96% rename from face_sdk/models/network_def/mobilev3_pfld.py rename to src/faceDetection/face_sdk/face_sdk/models/network_def/mobilev3_pfld.py index c53e359..c0c1e29 100644 --- a/face_sdk/models/network_def/mobilev3_pfld.py +++ b/src/faceDetection/face_sdk/face_sdk/models/network_def/mobilev3_pfld.py @@ -1,195 +1,195 @@ -# derive from: -# https://github.com/Hsintao/pfld_106_face_landmarks/blob/master/models/mobilev3_pfld.py - -import torch -import torch.nn as nn -import torch.nn.functional as F - - -def conv_bn(inp, oup, kernel_size, stride, padding=1, conv_layer=nn.Conv2d, norm_layer=nn.BatchNorm2d, nlin_layer=nn.ReLU): - return nn.Sequential( - conv_layer(inp, oup, kernel_size, stride, padding, bias=False), - norm_layer(oup), - nlin_layer(inplace=True) - ) - - -def conv_1x1_bn(inp, oup, conv_layer=nn.Conv2d, norm_layer=nn.BatchNorm2d, nlin_layer=nn.ReLU): - return nn.Sequential( - conv_layer(inp, oup, 1, 1, 0, bias=False), - norm_layer(oup), - nlin_layer(inplace=True) - ) - - -class Hswish(nn.Module): - def __init__(self, inplace=True): - super(Hswish, self).__init__() - self.inplace = inplace - - def forward(self, x): - return x * F.relu6(x + 3., inplace=self.inplace) / 6. - - -class Hsigmoid(nn.Module): - def __init__(self, inplace=True): - super(Hsigmoid, self).__init__() - self.inplace = inplace - - def forward(self, x): - return F.relu6(x + 3., inplace=self.inplace) / 6. - - -class SEModule(nn.Module): - def __init__(self, channel, reduction=4): - super(SEModule, self).__init__() - self.avg_pool = nn.AdaptiveAvgPool2d(1) - self.fc = nn.Sequential( - nn.Linear(channel, channel // reduction, bias=False), - nn.ReLU(inplace=True), - nn.Linear(channel // reduction, channel, bias=False), - Hsigmoid() - ) - - def forward(self, x): - b, c, h, w = x.size() - # F.avg_pool2d() - y = self.avg_pool(x).view(b, c) - y = self.fc(y).view(b, c, 1, 1) - return x * y - - -class Identity(nn.Module): - def __init__(self, channel): - super(Identity, self).__init__() - - def forward(self, x): - return x - - -class MobileBottleneck(nn.Module): - def __init__(self, inp, oup, kernel, stride, exp, se=False, nl='RE'): - super(MobileBottleneck, self).__init__() - assert stride in [1, 2] - assert kernel in [3, 5] - padding = (kernel - 1) // 2 - self.use_res_connect = stride == 1 and inp == oup - - conv_layer = nn.Conv2d - norm_layer = nn.BatchNorm2d - if nl == 'RE': - nlin_layer = nn.ReLU # or ReLU6 - elif nl == 'HS': - nlin_layer = Hswish - else: - raise NotImplementedError - if se: - SELayer = SEModule - else: - SELayer = Identity - - self.conv = nn.Sequential( - # pw - conv_layer(inp, exp, 1, 1, 0, bias=False), - norm_layer(exp), - nlin_layer(inplace=True), - # dw - conv_layer(exp, exp, kernel, stride, padding, groups=exp, bias=False), - norm_layer(exp), - SELayer(exp), - nlin_layer(inplace=True), - # pw-linear - conv_layer(exp, oup, 1, 1, 0, bias=False), - norm_layer(oup), - ) - - def forward(self, x): - if self.use_res_connect: - return x + self.conv(x) - else: - return self.conv(x) - - -class PFLDInference(nn.Module): - def __init__(self): - super(PFLDInference, self).__init__() - self.use_attention = True - self.conv_bn1 = conv_bn(3, 16, 3, stride=1, nlin_layer=Hswish) - self.conv_bn2 = MobileBottleneck(16, 16, 3, 1, 16, False, 'RE') - - self.conv3_1 = MobileBottleneck(16, 24, 3, 2, 64, False, 'RE') - - self.block3_2 = MobileBottleneck(24, 24, 3, 1, 72, False, "RE") - self.block3_3 = MobileBottleneck(24, 40, 5, 2, 72, self.use_attention, "RE") - self.block3_4 = MobileBottleneck(40, 40, 5, 1, 120, self.use_attention, "RE") - self.block3_5 = MobileBottleneck(40, 40, 5, 1, 120, self.use_attention, "RE") - - self.conv4_1 = MobileBottleneck(40, 80, 3, 2, 240, False, "RE") - - self.conv5_1 = MobileBottleneck(80, 80, 3, 1, 200, False, "HS") - self.block5_2 = MobileBottleneck(80, 112, 3, 1, 480, self.use_attention, "HS") - self.block5_3 = MobileBottleneck(112, 112, 3, 1, 672, self.use_attention, "HS") - self.block5_4 = MobileBottleneck(112, 160, 3, 1, 672, self.use_attention, "HS") - - self.conv6_1 = MobileBottleneck(160, 16, 3, 1, 320, False, "HS") # [16, 14, 14] - - self.conv7 = nn.Conv2d(16, 32, 3, 2, padding=1) - self.conv8 = nn.Conv2d(32, 128, 7, 1, 0) - self.avg_pool1 = nn.AvgPool2d(14) - self.avg_pool2 = nn.AvgPool2d(7) - self.fc = nn.Linear(176, 106 * 2) - - def forward(self, x): # x: 3, 112, 112 - x = self.conv_bn1(x) # [64, 56, 56] - x = self.conv_bn2(x) # [64, 56, 56] - x = self.conv3_1(x) - x = self.block3_2(x) - x = self.block3_3(x) - x = self.block3_4(x) - out1 = self.block3_5(x) - - x = self.conv4_1(out1) - - x = self.conv5_1(x) - x = self.block5_2(x) - x = self.block5_3(x) - x = self.block5_4(x) - x = self.conv6_1(x) - x1 = self.avg_pool1(x) - x1 = x1.view(x1.size(0), -1) - - x = self.conv7(x) - x2 = self.avg_pool2(x) - x2 = x2.view(x2.size(0), -1) - - x3 = self.conv8(x) - x3 = x3.view(x1.size(0), -1) - - multi_scale = torch.cat([x1, x2, x3], 1) - landmarks = self.fc(multi_scale) - - return out1, landmarks - - -class AuxiliaryNet(nn.Module): - def __init__(self): - super(AuxiliaryNet, self).__init__() - self.conv1 = conv_bn(40, 128, 3, 2) - self.conv2 = conv_bn(128, 128, 3, 1) - self.conv3 = conv_bn(128, 32, 3, 2) - self.conv4 = conv_bn(32, 128, 3, 1, padding=0) - self.max_pool1 = nn.MaxPool2d(5) - self.fc1 = nn.Linear(128, 32) - self.fc2 = nn.Linear(32, 3) - - def forward(self, x): - x = self.conv1(x) - x = self.conv2(x) - x = self.conv3(x) - x = self.conv4(x) - x = self.max_pool1(x) - x = x.view(x.size(0), -1) - x = self.fc1(x) - x = self.fc2(x) - - return x +# derive from: +# https://github.com/Hsintao/pfld_106_face_landmarks/blob/master/models/mobilev3_pfld.py + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +def conv_bn(inp, oup, kernel_size, stride, padding=1, conv_layer=nn.Conv2d, norm_layer=nn.BatchNorm2d, nlin_layer=nn.ReLU): + return nn.Sequential( + conv_layer(inp, oup, kernel_size, stride, padding, bias=False), + norm_layer(oup), + nlin_layer(inplace=True) + ) + + +def conv_1x1_bn(inp, oup, conv_layer=nn.Conv2d, norm_layer=nn.BatchNorm2d, nlin_layer=nn.ReLU): + return nn.Sequential( + conv_layer(inp, oup, 1, 1, 0, bias=False), + norm_layer(oup), + nlin_layer(inplace=True) + ) + + +class Hswish(nn.Module): + def __init__(self, inplace=True): + super(Hswish, self).__init__() + self.inplace = inplace + + def forward(self, x): + return x * F.relu6(x + 3., inplace=self.inplace) / 6. + + +class Hsigmoid(nn.Module): + def __init__(self, inplace=True): + super(Hsigmoid, self).__init__() + self.inplace = inplace + + def forward(self, x): + return F.relu6(x + 3., inplace=self.inplace) / 6. + + +class SEModule(nn.Module): + def __init__(self, channel, reduction=4): + super(SEModule, self).__init__() + self.avg_pool = nn.AdaptiveAvgPool2d(1) + self.fc = nn.Sequential( + nn.Linear(channel, channel // reduction, bias=False), + nn.ReLU(inplace=True), + nn.Linear(channel // reduction, channel, bias=False), + Hsigmoid() + ) + + def forward(self, x): + b, c, h, w = x.size() + # F.avg_pool2d() + y = self.avg_pool(x).view(b, c) + y = self.fc(y).view(b, c, 1, 1) + return x * y + + +class Identity(nn.Module): + def __init__(self, channel): + super(Identity, self).__init__() + + def forward(self, x): + return x + + +class MobileBottleneck(nn.Module): + def __init__(self, inp, oup, kernel, stride, exp, se=False, nl='RE'): + super(MobileBottleneck, self).__init__() + assert stride in [1, 2] + assert kernel in [3, 5] + padding = (kernel - 1) // 2 + self.use_res_connect = stride == 1 and inp == oup + + conv_layer = nn.Conv2d + norm_layer = nn.BatchNorm2d + if nl == 'RE': + nlin_layer = nn.ReLU # or ReLU6 + elif nl == 'HS': + nlin_layer = Hswish + else: + raise NotImplementedError + if se: + SELayer = SEModule + else: + SELayer = Identity + + self.conv = nn.Sequential( + # pw + conv_layer(inp, exp, 1, 1, 0, bias=False), + norm_layer(exp), + nlin_layer(inplace=True), + # dw + conv_layer(exp, exp, kernel, stride, padding, groups=exp, bias=False), + norm_layer(exp), + SELayer(exp), + nlin_layer(inplace=True), + # pw-linear + conv_layer(exp, oup, 1, 1, 0, bias=False), + norm_layer(oup), + ) + + def forward(self, x): + if self.use_res_connect: + return x + self.conv(x) + else: + return self.conv(x) + + +class PFLDInference(nn.Module): + def __init__(self): + super(PFLDInference, self).__init__() + self.use_attention = True + self.conv_bn1 = conv_bn(3, 16, 3, stride=1, nlin_layer=Hswish) + self.conv_bn2 = MobileBottleneck(16, 16, 3, 1, 16, False, 'RE') + + self.conv3_1 = MobileBottleneck(16, 24, 3, 2, 64, False, 'RE') + + self.block3_2 = MobileBottleneck(24, 24, 3, 1, 72, False, "RE") + self.block3_3 = MobileBottleneck(24, 40, 5, 2, 72, self.use_attention, "RE") + self.block3_4 = MobileBottleneck(40, 40, 5, 1, 120, self.use_attention, "RE") + self.block3_5 = MobileBottleneck(40, 40, 5, 1, 120, self.use_attention, "RE") + + self.conv4_1 = MobileBottleneck(40, 80, 3, 2, 240, False, "RE") + + self.conv5_1 = MobileBottleneck(80, 80, 3, 1, 200, False, "HS") + self.block5_2 = MobileBottleneck(80, 112, 3, 1, 480, self.use_attention, "HS") + self.block5_3 = MobileBottleneck(112, 112, 3, 1, 672, self.use_attention, "HS") + self.block5_4 = MobileBottleneck(112, 160, 3, 1, 672, self.use_attention, "HS") + + self.conv6_1 = MobileBottleneck(160, 16, 3, 1, 320, False, "HS") # [16, 14, 14] + + self.conv7 = nn.Conv2d(16, 32, 3, 2, padding=1) + self.conv8 = nn.Conv2d(32, 128, 7, 1, 0) + self.avg_pool1 = nn.AvgPool2d(14) + self.avg_pool2 = nn.AvgPool2d(7) + self.fc = nn.Linear(176, 106 * 2) + + def forward(self, x): # x: 3, 112, 112 + x = self.conv_bn1(x) # [64, 56, 56] + x = self.conv_bn2(x) # [64, 56, 56] + x = self.conv3_1(x) + x = self.block3_2(x) + x = self.block3_3(x) + x = self.block3_4(x) + out1 = self.block3_5(x) + + x = self.conv4_1(out1) + + x = self.conv5_1(x) + x = self.block5_2(x) + x = self.block5_3(x) + x = self.block5_4(x) + x = self.conv6_1(x) + x1 = self.avg_pool1(x) + x1 = x1.view(x1.size(0), -1) + + x = self.conv7(x) + x2 = self.avg_pool2(x) + x2 = x2.view(x2.size(0), -1) + + x3 = self.conv8(x) + x3 = x3.view(x1.size(0), -1) + + multi_scale = torch.cat([x1, x2, x3], 1) + landmarks = self.fc(multi_scale) + + return out1, landmarks + + +class AuxiliaryNet(nn.Module): + def __init__(self): + super(AuxiliaryNet, self).__init__() + self.conv1 = conv_bn(40, 128, 3, 2) + self.conv2 = conv_bn(128, 128, 3, 1) + self.conv3 = conv_bn(128, 32, 3, 2) + self.conv4 = conv_bn(32, 128, 3, 1, padding=0) + self.max_pool1 = nn.MaxPool2d(5) + self.fc1 = nn.Linear(128, 32) + self.fc2 = nn.Linear(32, 3) + + def forward(self, x): + x = self.conv1(x) + x = self.conv2(x) + x = self.conv3(x) + x = self.conv4(x) + x = self.max_pool1(x) + x = x.view(x.size(0), -1) + x = self.fc1(x) + x = self.fc2(x) + + return x diff --git a/face_sdk/models/network_def/retinaface_def.py b/src/faceDetection/face_sdk/face_sdk/models/network_def/retinaface_def.py similarity index 100% rename from face_sdk/models/network_def/retinaface_def.py rename to src/faceDetection/face_sdk/face_sdk/models/network_def/retinaface_def.py diff --git a/face_sdk/utils/BuzException.py b/src/faceDetection/face_sdk/face_sdk/utils/BuzException.py similarity index 100% rename from face_sdk/utils/BuzException.py rename to src/faceDetection/face_sdk/face_sdk/utils/BuzException.py diff --git a/src/faceDetection/face_sdk/face_sdk/utils/__init__.py b/src/faceDetection/face_sdk/face_sdk/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/face_sdk/utils/draw.py b/src/faceDetection/face_sdk/face_sdk/utils/draw.py similarity index 100% rename from face_sdk/utils/draw.py rename to src/faceDetection/face_sdk/face_sdk/utils/draw.py diff --git a/face_sdk/utils/lms_trans.py b/src/faceDetection/face_sdk/face_sdk/utils/lms_trans.py similarity index 100% rename from face_sdk/utils/lms_trans.py rename to src/faceDetection/face_sdk/face_sdk/utils/lms_trans.py diff --git a/face_sdk/utils/show.py b/src/faceDetection/face_sdk/face_sdk/utils/show.py similarity index 100% rename from face_sdk/utils/show.py rename to src/faceDetection/face_sdk/face_sdk/utils/show.py diff --git a/face_sdk/utils/transform.py b/src/faceDetection/face_sdk/face_sdk/utils/transform.py similarity index 100% rename from face_sdk/utils/transform.py rename to src/faceDetection/face_sdk/face_sdk/utils/transform.py diff --git a/head/AM_Softmax.py b/src/faceDetection/face_sdk/head/AM_Softmax.py similarity index 100% rename from head/AM_Softmax.py rename to src/faceDetection/face_sdk/head/AM_Softmax.py diff --git a/head/AdaCos.py b/src/faceDetection/face_sdk/head/AdaCos.py similarity index 100% rename from head/AdaCos.py rename to src/faceDetection/face_sdk/head/AdaCos.py diff --git a/head/AdaM_Softmax.py b/src/faceDetection/face_sdk/head/AdaM_Softmax.py similarity index 97% rename from head/AdaM_Softmax.py rename to src/faceDetection/face_sdk/head/AdaM_Softmax.py index ad7eff7..f11f033 100644 --- a/head/AdaM_Softmax.py +++ b/src/faceDetection/face_sdk/head/AdaM_Softmax.py @@ -1,41 +1,41 @@ -""" -@author: Hang Du, Jun Wang -@date: 20201128 -@contact: jun21wangustc@gmail.com -""" -import torch -import torch.nn.functional as F -from torch.nn import Module, Parameter - -class Adam_Softmax(Module): - """Implementation for "AdaptiveFace: Adaptive Margin and Sampling for Face Recognition". - """ - def __init__(self, feat_dim, num_class, scale=30.0, lamda=70.0): - super(Adam_Softmax, self).__init__() - self.num_class = num_class - self.scale = scale - self.lamda = lamda - self.kernel = Parameter(torch.Tensor(feat_dim, num_class)) - self.kernel.data.uniform_(-1, 1).renorm_(2, 1, 1e-5).mul_(1e5) - self.adam = Parameter(torch.Tensor(1, num_class)) - self.adam.data.uniform_(0.3,0.4) - def forward(self, feats, labels): - kernel_norm = F.normalize(self.kernel, dim=0) - feats = F.normalize(feats) - cos_theta = torch.mm(feats, kernel_norm) - cos_theta = cos_theta.clamp(-1,1) - # margin in [0,1] for cosface. - self.adam.data.clamp_(0,1) - - margin = self.adam[:, labels].view(-1, 1) - cos_theta_m = cos_theta - margin - index = torch.zeros_like(cos_theta) - index.scatter_(1,labels.data.view(-1,1),1) - index = index.byte().bool() - output = cos_theta * 1.0 - output[index] = cos_theta_m[index] - output *= self.scale - - #ensure the loss > 0 - Lm = -1* torch.sum(self.adam, dim=1)/self.num_class + 1 - return output, self.lamda*Lm +""" +@author: Hang Du, Jun Wang +@date: 20201128 +@contact: jun21wangustc@gmail.com +""" +import torch +import torch.nn.functional as F +from torch.nn import Module, Parameter + +class Adam_Softmax(Module): + """Implementation for "AdaptiveFace: Adaptive Margin and Sampling for Face Recognition". + """ + def __init__(self, feat_dim, num_class, scale=30.0, lamda=70.0): + super(Adam_Softmax, self).__init__() + self.num_class = num_class + self.scale = scale + self.lamda = lamda + self.kernel = Parameter(torch.Tensor(feat_dim, num_class)) + self.kernel.data.uniform_(-1, 1).renorm_(2, 1, 1e-5).mul_(1e5) + self.adam = Parameter(torch.Tensor(1, num_class)) + self.adam.data.uniform_(0.3,0.4) + def forward(self, feats, labels): + kernel_norm = F.normalize(self.kernel, dim=0) + feats = F.normalize(feats) + cos_theta = torch.mm(feats, kernel_norm) + cos_theta = cos_theta.clamp(-1,1) + # margin in [0,1] for cosface. + self.adam.data.clamp_(0,1) + + margin = self.adam[:, labels].view(-1, 1) + cos_theta_m = cos_theta - margin + index = torch.zeros_like(cos_theta) + index.scatter_(1,labels.data.view(-1,1),1) + index = index.byte().bool() + output = cos_theta * 1.0 + output[index] = cos_theta_m[index] + output *= self.scale + + #ensure the loss > 0 + Lm = -1* torch.sum(self.adam, dim=1)/self.num_class + 1 + return output, self.lamda*Lm diff --git a/head/ArcFace.py b/src/faceDetection/face_sdk/head/ArcFace.py similarity index 100% rename from head/ArcFace.py rename to src/faceDetection/face_sdk/head/ArcFace.py diff --git a/head/ArcNegFace.py b/src/faceDetection/face_sdk/head/ArcNegFace.py similarity index 100% rename from head/ArcNegFace.py rename to src/faceDetection/face_sdk/head/ArcNegFace.py diff --git a/head/CircleLoss.py b/src/faceDetection/face_sdk/head/CircleLoss.py similarity index 100% rename from head/CircleLoss.py rename to src/faceDetection/face_sdk/head/CircleLoss.py diff --git a/head/CurricularFace.py b/src/faceDetection/face_sdk/head/CurricularFace.py similarity index 100% rename from head/CurricularFace.py rename to src/faceDetection/face_sdk/head/CurricularFace.py diff --git a/head/MV_Softmax.py b/src/faceDetection/face_sdk/head/MV_Softmax.py similarity index 100% rename from head/MV_Softmax.py rename to src/faceDetection/face_sdk/head/MV_Softmax.py diff --git a/head/MagFace.py b/src/faceDetection/face_sdk/head/MagFace.py similarity index 100% rename from head/MagFace.py rename to src/faceDetection/face_sdk/head/MagFace.py diff --git a/head/NPCFace.py b/src/faceDetection/face_sdk/head/NPCFace.py similarity index 97% rename from head/NPCFace.py rename to src/faceDetection/face_sdk/head/NPCFace.py index 8a55400..98074ec 100644 --- a/head/NPCFace.py +++ b/src/faceDetection/face_sdk/head/NPCFace.py @@ -1,56 +1,56 @@ -""" -@author: Hand Du, Jun Wang -@date: 20201019 -@contact: jun21wangustc@gmail.com -""" - -import math -import torch -import torch.nn.functional as F -from torch.nn import Module, Parameter - -class NPCFace(Module): - """Implementation for "NPCFace: A Negative-Positive Cooperation - Supervision for Training Large-scale Face Recognition" - """ - def __init__(self, feat_dim=512, num_class=86876, margin=0.5, scale=64): - super(NPCFace, self).__init__() - self.kernel = Parameter(torch.Tensor(feat_dim, num_class)) - self.kernel.data.uniform_(-1, 1).renorm_(2, 1, 1e-5).mul_(1e5) - self.margin = margin - self.scale = scale - self.cos_m = math.cos(margin) - self.sin_m = math.sin(margin) - self.m0 = 0.40 - self.m1 = 0.20 - self.t = 1.10 - self.a = 0.20 - self.cos_m0 = math.cos(self.m0) - self.sin_m0 = math.sin(self.m0) - self.num_class = num_class - - def forward(self, x, label): - kernel_norm = F.normalize(self.kernel, dim=0) - x = F.normalize(x) - cos_theta = torch.mm(x, kernel_norm) - cos_theta = cos_theta.clamp(-1, 1) - batch_size = label.size(0) - gt = cos_theta[torch.arange(0, batch_size), label].view(-1, 1) - sin_theta = torch.sqrt(1.0 - torch.pow(gt, 2)) - cos_theta_m = gt * self.cos_m - sin_theta * self.sin_m - with torch.no_grad(): - hard_mask = (cos_theta > cos_theta_m).type(torch.FloatTensor).cuda() - hard_mask.scatter_(1, label.data.view(-1, 1), 0) - hard_cos = torch.where(hard_mask > 0, cos_theta, torch.zeros_like(cos_theta)) - sum_hard_cos = torch.sum(hard_cos,dim=1).view(-1, 1) - sum_hard_mask = torch.sum(hard_mask, dim=1).view(-1,1) - sum_hard_mask = sum_hard_mask.clamp(1, self.num_class) - avg_hard_cos = sum_hard_cos / sum_hard_mask - newm = self.m0 + self.m1 * avg_hard_cos - cos_newm = torch.cos(newm) - sin_newm = torch.sin(newm) - final_gt = torch.where(gt > 0, gt * cos_newm - sin_theta * sin_newm , gt) - cos_theta = torch.where(cos_theta > cos_theta_m, self.t * cos_theta + self.a , cos_theta) - cos_theta.scatter_(1, label.data.view(-1, 1), final_gt) - cos_theta *= self.scale - return cos_theta +""" +@author: Hand Du, Jun Wang +@date: 20201019 +@contact: jun21wangustc@gmail.com +""" + +import math +import torch +import torch.nn.functional as F +from torch.nn import Module, Parameter + +class NPCFace(Module): + """Implementation for "NPCFace: A Negative-Positive Cooperation + Supervision for Training Large-scale Face Recognition" + """ + def __init__(self, feat_dim=512, num_class=86876, margin=0.5, scale=64): + super(NPCFace, self).__init__() + self.kernel = Parameter(torch.Tensor(feat_dim, num_class)) + self.kernel.data.uniform_(-1, 1).renorm_(2, 1, 1e-5).mul_(1e5) + self.margin = margin + self.scale = scale + self.cos_m = math.cos(margin) + self.sin_m = math.sin(margin) + self.m0 = 0.40 + self.m1 = 0.20 + self.t = 1.10 + self.a = 0.20 + self.cos_m0 = math.cos(self.m0) + self.sin_m0 = math.sin(self.m0) + self.num_class = num_class + + def forward(self, x, label): + kernel_norm = F.normalize(self.kernel, dim=0) + x = F.normalize(x) + cos_theta = torch.mm(x, kernel_norm) + cos_theta = cos_theta.clamp(-1, 1) + batch_size = label.size(0) + gt = cos_theta[torch.arange(0, batch_size), label].view(-1, 1) + sin_theta = torch.sqrt(1.0 - torch.pow(gt, 2)) + cos_theta_m = gt * self.cos_m - sin_theta * self.sin_m + with torch.no_grad(): + hard_mask = (cos_theta > cos_theta_m).type(torch.FloatTensor).cuda() + hard_mask.scatter_(1, label.data.view(-1, 1), 0) + hard_cos = torch.where(hard_mask > 0, cos_theta, torch.zeros_like(cos_theta)) + sum_hard_cos = torch.sum(hard_cos,dim=1).view(-1, 1) + sum_hard_mask = torch.sum(hard_mask, dim=1).view(-1,1) + sum_hard_mask = sum_hard_mask.clamp(1, self.num_class) + avg_hard_cos = sum_hard_cos / sum_hard_mask + newm = self.m0 + self.m1 * avg_hard_cos + cos_newm = torch.cos(newm) + sin_newm = torch.sin(newm) + final_gt = torch.where(gt > 0, gt * cos_newm - sin_theta * sin_newm , gt) + cos_theta = torch.where(cos_theta > cos_theta_m, self.t * cos_theta + self.a , cos_theta) + cos_theta.scatter_(1, label.data.view(-1, 1), final_gt) + cos_theta *= self.scale + return cos_theta diff --git a/head/SST_Prototype.py b/src/faceDetection/face_sdk/head/SST_Prototype.py similarity index 100% rename from head/SST_Prototype.py rename to src/faceDetection/face_sdk/head/SST_Prototype.py diff --git a/src/faceDetection/face_sdk/head/__init__.py b/src/faceDetection/face_sdk/head/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/head/head_def.py b/src/faceDetection/face_sdk/head/head_def.py similarity index 89% rename from head/head_def.py rename to src/faceDetection/face_sdk/head/head_def.py index df5c11b..07830bf 100644 --- a/head/head_def.py +++ b/src/faceDetection/face_sdk/head/head_def.py @@ -6,18 +6,18 @@ import sys import yaml -sys.path.append('../../') -from head.AdaCos import AdaCos -from head.AdaM_Softmax import Adam_Softmax -from head.AM_Softmax import AM_Softmax -from head.ArcFace import ArcFace -from head.CircleLoss import CircleLoss -from head.CurricularFace import CurricularFace -from head.MV_Softmax import MV_Softmax -from head.NPCFace import NPCFace -from head.SST_Prototype import SST_Prototype -from head.ArcNegFace import ArcNegFace -from head.MagFace import MagFace +sys.path.append('../../../../../') +from src.faceDetection.face_sdk.head.AdaCos import AdaCos +from src.faceDetection.face_sdk.head.AdaM_Softmax import Adam_Softmax +from src.faceDetection.face_sdk.head.AM_Softmax import AM_Softmax +from src.faceDetection.face_sdk.head.ArcFace import ArcFace +from src.faceDetection.face_sdk.head.CircleLoss import CircleLoss +from src.faceDetection.face_sdk.head.CurricularFace import CurricularFace +from src.faceDetection.face_sdk.head.MV_Softmax import MV_Softmax +from src.faceDetection.face_sdk.head.NPCFace import NPCFace +from src.faceDetection.face_sdk.head.SST_Prototype import SST_Prototype +from src.faceDetection.face_sdk.head.ArcNegFace import ArcNegFace +from src.faceDetection.face_sdk.head.MagFace import MagFace class HeadFactory: """Factory to produce head according to the head_conf.yaml diff --git a/reference/readme.md b/src/faceDetection/face_sdk/reference/readme.md similarity index 100% rename from reference/readme.md rename to src/faceDetection/face_sdk/reference/readme.md diff --git a/test_protocol/README.md b/src/faceDetection/face_sdk/test_protocol/README.md similarity index 100% rename from test_protocol/README.md rename to src/faceDetection/face_sdk/test_protocol/README.md diff --git a/src/faceDetection/face_sdk/test_protocol/__init__.py b/src/faceDetection/face_sdk/test_protocol/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/test_protocol/backbone_conf.yaml b/src/faceDetection/face_sdk/test_protocol/backbone_conf.yaml similarity index 100% rename from test_protocol/backbone_conf.yaml rename to src/faceDetection/face_sdk/test_protocol/backbone_conf.yaml diff --git a/test_protocol/data_conf.yaml b/src/faceDetection/face_sdk/test_protocol/data_conf.yaml similarity index 100% rename from test_protocol/data_conf.yaml rename to src/faceDetection/face_sdk/test_protocol/data_conf.yaml diff --git a/test_protocol/extract_feature.py b/src/faceDetection/face_sdk/test_protocol/extract_feature.py similarity index 90% rename from test_protocol/extract_feature.py rename to src/faceDetection/face_sdk/test_protocol/extract_feature.py index 1276bd6..e2304d7 100644 --- a/test_protocol/extract_feature.py +++ b/src/faceDetection/face_sdk/test_protocol/extract_feature.py @@ -7,13 +7,12 @@ import sys import yaml import argparse -import torch -from torch.utils.data import Dataset, DataLoader +from torch.utils.data import DataLoader from utils.model_loader import ModelLoader from utils.extractor.feature_extractor import CommonExtractor -sys.path.append('..') -from data_processor.test_dataset import CommonTestDataset -from backbone.backbone_def import BackboneFactory +sys.path.append('../../../..') +from src.faceDetection.face_sdk.data_processor.test_dataset import CommonTestDataset +from src.faceDetection.face_sdk.backbone.backbone_def import BackboneFactory if __name__ == '__main__': conf = argparse.ArgumentParser(description='extract features for megaface.') diff --git a/test_protocol/extract_feature.sh b/src/faceDetection/face_sdk/test_protocol/extract_feature.sh similarity index 100% rename from test_protocol/extract_feature.sh rename to src/faceDetection/face_sdk/test_protocol/extract_feature.sh diff --git a/src/faceDetection/face_sdk/test_protocol/ijbc/__init__.py b/src/faceDetection/face_sdk/test_protocol/ijbc/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/faceDetection/face_sdk/test_protocol/ijbc/face_cropper/__init__.py b/src/faceDetection/face_sdk/test_protocol/ijbc/face_cropper/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/test_protocol/ijbc/face_cropper/crop_ijbc_by_arcface.py b/src/faceDetection/face_sdk/test_protocol/ijbc/face_cropper/crop_ijbc_by_arcface.py similarity index 100% rename from test_protocol/ijbc/face_cropper/crop_ijbc_by_arcface.py rename to src/faceDetection/face_sdk/test_protocol/ijbc/face_cropper/crop_ijbc_by_arcface.py diff --git a/test_protocol/ijbc/ijbc_evaluator.py b/src/faceDetection/face_sdk/test_protocol/ijbc/ijbc_evaluator.py similarity index 100% rename from test_protocol/ijbc/ijbc_evaluator.py rename to src/faceDetection/face_sdk/test_protocol/ijbc/ijbc_evaluator.py diff --git a/src/faceDetection/face_sdk/test_protocol/lfw/__init__.py b/src/faceDetection/face_sdk/test_protocol/lfw/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/faceDetection/face_sdk/test_protocol/lfw/face_cropper/__init__.py b/src/faceDetection/face_sdk/test_protocol/lfw/face_cropper/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/test_protocol/lfw/face_cropper/crop_agedb_by_arcface.py b/src/faceDetection/face_sdk/test_protocol/lfw/face_cropper/crop_agedb_by_arcface.py similarity index 100% rename from test_protocol/lfw/face_cropper/crop_agedb_by_arcface.py rename to src/faceDetection/face_sdk/test_protocol/lfw/face_cropper/crop_agedb_by_arcface.py diff --git a/test_protocol/lfw/face_cropper/crop_calfw_by_arcface.py b/src/faceDetection/face_sdk/test_protocol/lfw/face_cropper/crop_calfw_by_arcface.py similarity index 100% rename from test_protocol/lfw/face_cropper/crop_calfw_by_arcface.py rename to src/faceDetection/face_sdk/test_protocol/lfw/face_cropper/crop_calfw_by_arcface.py diff --git a/test_protocol/lfw/face_cropper/crop_cplfw_by_arcface.py b/src/faceDetection/face_sdk/test_protocol/lfw/face_cropper/crop_cplfw_by_arcface.py similarity index 100% rename from test_protocol/lfw/face_cropper/crop_cplfw_by_arcface.py rename to src/faceDetection/face_sdk/test_protocol/lfw/face_cropper/crop_cplfw_by_arcface.py diff --git a/test_protocol/lfw/face_cropper/crop_lfw_by_arcface.py b/src/faceDetection/face_sdk/test_protocol/lfw/face_cropper/crop_lfw_by_arcface.py similarity index 100% rename from test_protocol/lfw/face_cropper/crop_lfw_by_arcface.py rename to src/faceDetection/face_sdk/test_protocol/lfw/face_cropper/crop_lfw_by_arcface.py diff --git a/test_protocol/lfw/face_cropper/crop_rfw_by_arcface.py b/src/faceDetection/face_sdk/test_protocol/lfw/face_cropper/crop_rfw_by_arcface.py similarity index 100% rename from test_protocol/lfw/face_cropper/crop_rfw_by_arcface.py rename to src/faceDetection/face_sdk/test_protocol/lfw/face_cropper/crop_rfw_by_arcface.py diff --git a/test_protocol/lfw/lfw_evaluator.py b/src/faceDetection/face_sdk/test_protocol/lfw/lfw_evaluator.py similarity index 100% rename from test_protocol/lfw/lfw_evaluator.py rename to src/faceDetection/face_sdk/test_protocol/lfw/lfw_evaluator.py diff --git a/test_protocol/lfw/pairs_parser.py b/src/faceDetection/face_sdk/test_protocol/lfw/pairs_parser.py similarity index 100% rename from test_protocol/lfw/pairs_parser.py rename to src/faceDetection/face_sdk/test_protocol/lfw/pairs_parser.py diff --git a/src/faceDetection/face_sdk/test_protocol/megaface/__init__.py b/src/faceDetection/face_sdk/test_protocol/megaface/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/faceDetection/face_sdk/test_protocol/megaface/face_cropper/__init__.py b/src/faceDetection/face_sdk/test_protocol/megaface/face_cropper/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/test_protocol/megaface/face_cropper/crop_eye.py b/src/faceDetection/face_sdk/test_protocol/megaface/face_cropper/crop_eye.py similarity index 100% rename from test_protocol/megaface/face_cropper/crop_eye.py rename to src/faceDetection/face_sdk/test_protocol/megaface/face_cropper/crop_eye.py diff --git a/test_protocol/megaface/face_cropper/crop_facescrub_by_arcface.py b/src/faceDetection/face_sdk/test_protocol/megaface/face_cropper/crop_facescrub_by_arcface.py similarity index 100% rename from test_protocol/megaface/face_cropper/crop_facescrub_by_arcface.py rename to src/faceDetection/face_sdk/test_protocol/megaface/face_cropper/crop_facescrub_by_arcface.py diff --git a/test_protocol/megaface/face_cropper/crop_megaface_by_arcface.py b/src/faceDetection/face_sdk/test_protocol/megaface/face_cropper/crop_megaface_by_arcface.py similarity index 100% rename from test_protocol/megaface/face_cropper/crop_megaface_by_arcface.py rename to src/faceDetection/face_sdk/test_protocol/megaface/face_cropper/crop_megaface_by_arcface.py diff --git a/test_protocol/megaface/megaface_evaluator.py b/src/faceDetection/face_sdk/test_protocol/megaface/megaface_evaluator.py similarity index 100% rename from test_protocol/megaface/megaface_evaluator.py rename to src/faceDetection/face_sdk/test_protocol/megaface/megaface_evaluator.py diff --git a/test_protocol/remove_noises.py b/src/faceDetection/face_sdk/test_protocol/remove_noises.py similarity index 100% rename from test_protocol/remove_noises.py rename to src/faceDetection/face_sdk/test_protocol/remove_noises.py diff --git a/test_protocol/remove_noises.sh b/src/faceDetection/face_sdk/test_protocol/remove_noises.sh similarity index 100% rename from test_protocol/remove_noises.sh rename to src/faceDetection/face_sdk/test_protocol/remove_noises.sh diff --git a/addition_module/face_lightning/KDF/test_protocol/test_ijbc.py b/src/faceDetection/face_sdk/test_protocol/test_ijbc.py similarity index 94% rename from addition_module/face_lightning/KDF/test_protocol/test_ijbc.py rename to src/faceDetection/face_sdk/test_protocol/test_ijbc.py index 828b0d6..aef4604 100644 --- a/addition_module/face_lightning/KDF/test_protocol/test_ijbc.py +++ b/src/faceDetection/face_sdk/test_protocol/test_ijbc.py @@ -13,9 +13,9 @@ from ijbc.ijbc_evaluator import IJBCEvaluator from utils.model_loader import ModelLoader from utils.extractor.feature_extractor import CommonExtractor -sys.path.append('..') -from data_processor.test_dataset import CommonTestDataset -from backbone.backbone_def import BackboneFactory +sys.path.append('../../../..') +from src.faceDetection.face_sdk.data_processor.test_dataset import CommonTestDataset +from src.faceDetection.face_sdk.backbone.backbone_def import BackboneFactory def accu_key(elem): return elem[1] diff --git a/test_protocol/test_ijbc.sh b/src/faceDetection/face_sdk/test_protocol/test_ijbc.sh similarity index 100% rename from test_protocol/test_ijbc.sh rename to src/faceDetection/face_sdk/test_protocol/test_ijbc.sh diff --git a/addition_module/face_lightning/KDF/test_protocol/test_lfw.py b/src/faceDetection/face_sdk/test_protocol/test_lfw.py similarity index 94% rename from addition_module/face_lightning/KDF/test_protocol/test_lfw.py rename to src/faceDetection/face_sdk/test_protocol/test_lfw.py index 3485e5d..c20806d 100644 --- a/addition_module/face_lightning/KDF/test_protocol/test_lfw.py +++ b/src/faceDetection/face_sdk/test_protocol/test_lfw.py @@ -14,10 +14,9 @@ from lfw.lfw_evaluator import LFWEvaluator from utils.model_loader import ModelLoader from utils.extractor.feature_extractor import CommonExtractor -sys.path.append('..') -from backbone.backbone_def import BackboneFactory -sys.path.append('../../../../') -from data_processor.test_dataset import CommonTestDataset +sys.path.append('../../../..') +from src.faceDetection.face_sdk.data_processor.test_dataset import CommonTestDataset +from src.faceDetection.face_sdk.backbone.backbone_def import BackboneFactory def accu_key(elem): return elem[1] diff --git a/test_protocol/test_lfw.sh b/src/faceDetection/face_sdk/test_protocol/test_lfw.sh similarity index 100% rename from test_protocol/test_lfw.sh rename to src/faceDetection/face_sdk/test_protocol/test_lfw.sh diff --git a/test_protocol/test_megaface.py b/src/faceDetection/face_sdk/test_protocol/test_megaface.py similarity index 100% rename from test_protocol/test_megaface.py rename to src/faceDetection/face_sdk/test_protocol/test_megaface.py diff --git a/test_protocol/test_megaface.sh b/src/faceDetection/face_sdk/test_protocol/test_megaface.sh similarity index 100% rename from test_protocol/test_megaface.sh rename to src/faceDetection/face_sdk/test_protocol/test_megaface.sh diff --git a/src/faceDetection/face_sdk/test_protocol/utils/__init__.py b/src/faceDetection/face_sdk/test_protocol/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/faceDetection/face_sdk/test_protocol/utils/extractor/__init__.py b/src/faceDetection/face_sdk/test_protocol/utils/extractor/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/test_protocol/utils/extractor/feature_extract_save.py b/src/faceDetection/face_sdk/test_protocol/utils/extractor/feature_extract_save.py similarity index 100% rename from test_protocol/utils/extractor/feature_extract_save.py rename to src/faceDetection/face_sdk/test_protocol/utils/extractor/feature_extract_save.py diff --git a/test_protocol/utils/extractor/feature_extractor.py b/src/faceDetection/face_sdk/test_protocol/utils/extractor/feature_extractor.py similarity index 97% rename from test_protocol/utils/extractor/feature_extractor.py rename to src/faceDetection/face_sdk/test_protocol/utils/extractor/feature_extractor.py index 8bbfcdf..0906537 100644 --- a/test_protocol/utils/extractor/feature_extractor.py +++ b/src/faceDetection/face_sdk/test_protocol/utils/extractor/feature_extractor.py @@ -1,99 +1,99 @@ -""" -@author: Jun Wang -@date: 20201016 -@contact: jun21wangustc@gmail.com -""" - -import os -import logging as logger -import numpy as np -import torch -import torch.nn.functional as F -logger.basicConfig(level=logger.INFO, - format='%(levelname)s %(asctime)s %(filename)s: %(lineno)d] %(message)s', - datefmt='%Y-%m-%d %H:%M:%S') - -class CommonExtractor: - """Common feature extractor. - - Attributes: - device(object): device to init model. - """ - def __init__(self, device): - self.device = torch.device(device) - - def extract_online(self, model, data_loader): - """Extract and return features. - - Args: - model(object): initialized model. - data_loader(object): load data to be extracted. - - Returns: - image_name2feature(dict): key is the name of image, value is feature of image. - """ - model.eval() - image_name2feature = {} - with torch.no_grad(): - for batch_idx, (images, filenames) in enumerate(data_loader): - images = images.to(self.device) - features = model(images) - features = F.normalize(features) - features = features.cpu().numpy() - for filename, feature in zip(filenames, features): - image_name2feature[filename] = feature - return image_name2feature - - def extract_offline(self, feats_root, model, data_loader): - """Extract and save features. - - Args: - feats_root(str): the path to save features. - model(object): initialized model. - data_loader(object): load data to be extracted. - """ - model.eval() - with torch.no_grad(): - for batch_idx, (images, filenames) in enumerate(data_loader): - images = images.to(self.device) - features = model(images) - features = F.normalize(features) - features = features.cpu().numpy() - for filename, feature in zip(filenames, features): - feature_name = os.path.splitext(filename)[0] - feature_path = os.path.join(feats_root, feature_name + '.npy') - feature_dir = os.path.dirname(feature_path) - if not os.path.exists(feature_dir): - os.makedirs(feature_dir) - np.save(feature_path, feature) - if (batch_idx + 1) % 10 == 0: - logger.info('Finished batches: %d/%d.' % (batch_idx+1, len(data_loader))) - -class FeatureHandler: - """Some method to deal with features. - - Atributes: - feats_root(str): the directory which the fetures in. - """ - def __init__(self, feats_root): - self.feats_root = feats_root - - def load_feature(self): - """Load features to memory. - - Returns: - image_name2feature(dict): key is the name of image, value is feature of image. - """ - image_name2feature = {} - for root, dirs, files in os.walk(self.feats_root): - for cur_file in files: - if cur_file.endswith('.npy'): - cur_file_path = os.path.join(root, cur_file) - cur_feats = np.load(cur_file_path) - if self.feats_root.endswith('/'): - cur_short_path = cur_file_path[len(self.feats_root) : ] - else: - cur_short_path = cur_file_path[len(self.feats_root) + 1 : ] - cur_key = cur_short_path.replace('.npy', '.jpg') - image_name2feature[cur_key] = cur_feats - return image_name2feature +""" +@author: Jun Wang +@date: 20201016 +@contact: jun21wangustc@gmail.com +""" + +import os +import logging as logger +import numpy as np +import torch +import torch.nn.functional as F +logger.basicConfig(level=logger.INFO, + format='%(levelname)s %(asctime)s %(filename)s: %(lineno)d] %(message)s', + datefmt='%Y-%m-%d %H:%M:%S') + +class CommonExtractor: + """Common feature extractor. + + Attributes: + device(object): device to init model. + """ + def __init__(self, device): + self.device = torch.device(device) + + def extract_online(self, model, data_loader): + """Extract and return features. + + Args: + model(object): initialized model. + data_loader(object): load data to be extracted. + + Returns: + image_name2feature(dict): key is the name of image, value is feature of image. + """ + model.eval() + image_name2feature = {} + with torch.no_grad(): + for batch_idx, (images, filenames) in enumerate(data_loader): + images = images.to(self.device) + features = model(images) + features = F.normalize(features) + features = features.cpu().numpy() + for filename, feature in zip(filenames, features): + image_name2feature[filename] = feature + return image_name2feature + + def extract_offline(self, feats_root, model, data_loader): + """Extract and save features. + + Args: + feats_root(str): the path to save features. + model(object): initialized model. + data_loader(object): load data to be extracted. + """ + model.eval() + with torch.no_grad(): + for batch_idx, (images, filenames) in enumerate(data_loader): + images = images.to(self.device) + features = model(images) + features = F.normalize(features) + features = features.cpu().numpy() + for filename, feature in zip(filenames, features): + feature_name = os.path.splitext(filename)[0] + feature_path = os.path.join(feats_root, feature_name + '.npy') + feature_dir = os.path.dirname(feature_path) + if not os.path.exists(feature_dir): + os.makedirs(feature_dir) + np.save(feature_path, feature) + if (batch_idx + 1) % 10 == 0: + logger.info('Finished batches: %d/%d.' % (batch_idx+1, len(data_loader))) + +class FeatureHandler: + """Some method to deal with features. + + Atributes: + feats_root(str): the directory which the fetures in. + """ + def __init__(self, feats_root): + self.feats_root = feats_root + + def load_feature(self): + """Load features to memory. + + Returns: + image_name2feature(dict): key is the name of image, value is feature of image. + """ + image_name2feature = {} + for root, dirs, files in os.walk(self.feats_root): + for cur_file in files: + if cur_file.endswith('.npy'): + cur_file_path = os.path.join(root, cur_file) + cur_feats = np.load(cur_file_path) + if self.feats_root.endswith('/'): + cur_short_path = cur_file_path[len(self.feats_root) : ] + else: + cur_short_path = cur_file_path[len(self.feats_root) + 1 : ] + cur_key = cur_short_path.replace('.npy', '.jpg') + image_name2feature[cur_key] = cur_feats + return image_name2feature diff --git a/test_protocol/utils/feat_concat.py b/src/faceDetection/face_sdk/test_protocol/utils/feat_concat.py similarity index 100% rename from test_protocol/utils/feat_concat.py rename to src/faceDetection/face_sdk/test_protocol/utils/feat_concat.py diff --git a/test_protocol/utils/model_loader.py b/src/faceDetection/face_sdk/test_protocol/utils/model_loader.py similarity index 100% rename from test_protocol/utils/model_loader.py rename to src/faceDetection/face_sdk/test_protocol/utils/model_loader.py diff --git a/test_protocol/utils/online_val.py b/src/faceDetection/face_sdk/test_protocol/utils/online_val.py similarity index 86% rename from test_protocol/utils/online_val.py rename to src/faceDetection/face_sdk/test_protocol/utils/online_val.py index d7a43b9..1c2dc96 100644 --- a/test_protocol/utils/online_val.py +++ b/src/faceDetection/face_sdk/test_protocol/utils/online_val.py @@ -9,12 +9,12 @@ import logging as logger from torch.utils.data import DataLoader -sys.path.append('../../') -from data_processor.test_dataset import CommonTestDataset -from test_protocol.utils.extractor.feature_extractor import CommonExtractor -from test_protocol.lfw.lfw_evaluator import LFWEvaluator -from test_protocol.lfw.pairs_parser import PairsParserFactory -from test_protocol.ijbc.ijbc_evaluator import IJBCEvaluator +sys.path.append('../../../../../') +from src.faceDetection.face_sdk.data_processor.test_dataset import CommonTestDataset +from src.faceDetection.face_sdk.test_protocol.utils.extractor.feature_extractor import CommonExtractor +from src.faceDetection.face_sdk.test_protocol.lfw.lfw_evaluator import LFWEvaluator +from src.faceDetection.face_sdk.test_protocol.lfw.pairs_parser import PairsParserFactory +from src.faceDetection.face_sdk.test_protocol.ijbc.ijbc_evaluator import IJBCEvaluator logger.basicConfig(level=logger.INFO, format='%(levelname)s %(asctime)s %(filename)s: %(lineno)d] %(message)s', diff --git a/test_protocol/utils/pair_analysis.py b/src/faceDetection/face_sdk/test_protocol/utils/pair_analysis.py similarity index 100% rename from test_protocol/utils/pair_analysis.py rename to src/faceDetection/face_sdk/test_protocol/utils/pair_analysis.py diff --git a/training_mode/README.md b/src/faceDetection/face_sdk/training_mode/README.md similarity index 100% rename from training_mode/README.md rename to src/faceDetection/face_sdk/training_mode/README.md diff --git a/src/faceDetection/face_sdk/training_mode/__init__.py b/src/faceDetection/face_sdk/training_mode/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/training_mode/backbone_conf.yaml b/src/faceDetection/face_sdk/training_mode/backbone_conf.yaml similarity index 100% rename from training_mode/backbone_conf.yaml rename to src/faceDetection/face_sdk/training_mode/backbone_conf.yaml diff --git a/training_mode/conventional_training/README.md b/src/faceDetection/face_sdk/training_mode/conventional_training/README.md similarity index 100% rename from training_mode/conventional_training/README.md rename to src/faceDetection/face_sdk/training_mode/conventional_training/README.md diff --git a/src/faceDetection/face_sdk/training_mode/conventional_training/__init__.py b/src/faceDetection/face_sdk/training_mode/conventional_training/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/training_mode/conventional_training/train.py b/src/faceDetection/face_sdk/training_mode/conventional_training/train.py similarity index 96% rename from training_mode/conventional_training/train.py rename to src/faceDetection/face_sdk/training_mode/conventional_training/train.py index 5d3395c..b41bbc0 100644 --- a/training_mode/conventional_training/train.py +++ b/src/faceDetection/face_sdk/training_mode/conventional_training/train.py @@ -14,11 +14,11 @@ from torch.utils.data import DataLoader from tensorboardX import SummaryWriter -sys.path.append('../../') -from utils.AverageMeter import AverageMeter -from data_processor.train_dataset import ImageDataset -from backbone.backbone_def import BackboneFactory -from head.head_def import HeadFactory +sys.path.append('../../../../../') +from src.faceDetection.face_sdk.utils.AverageMeter import AverageMeter +from src.faceDetection.face_sdk.data_processor.train_dataset import ImageDataset +from src.faceDetection.face_sdk.backbone.backbone_def import BackboneFactory +from src.faceDetection.face_sdk.head.head_def import HeadFactory logger.basicConfig(level=logger.INFO, format='%(levelname)s %(asctime)s %(filename)s: %(lineno)d] %(message)s', diff --git a/training_mode/conventional_training/train.sh b/src/faceDetection/face_sdk/training_mode/conventional_training/train.sh similarity index 100% rename from training_mode/conventional_training/train.sh rename to src/faceDetection/face_sdk/training_mode/conventional_training/train.sh diff --git a/training_mode/conventional_training/train_amp.py b/src/faceDetection/face_sdk/training_mode/conventional_training/train_amp.py similarity index 96% rename from training_mode/conventional_training/train_amp.py rename to src/faceDetection/face_sdk/training_mode/conventional_training/train_amp.py index bdfbd65..45f8e49 100644 --- a/training_mode/conventional_training/train_amp.py +++ b/src/faceDetection/face_sdk/training_mode/conventional_training/train_amp.py @@ -15,11 +15,11 @@ from tensorboardX import SummaryWriter from apex import amp -sys.path.append('../../') -from utils.AverageMeter import AverageMeter -from data_processor.train_dataset import ImageDataset -from backbone.backbone_def import BackboneFactory -from head.head_def import HeadFactory +sys.path.append('../../../../../') +from src.faceDetection.face_sdk.utils.AverageMeter import AverageMeter +from src.faceDetection.face_sdk.data_processor.train_dataset import ImageDataset +from src.faceDetection.face_sdk.backbone.backbone_def import BackboneFactory +from src.faceDetection.face_sdk.head.head_def import HeadFactory logger.basicConfig(level=logger.INFO, format='%(levelname)s %(asctime)s %(filename)s: %(lineno)d] %(message)s', diff --git a/src/faceDetection/face_sdk/training_mode/distributed_training/__init__.py b/src/faceDetection/face_sdk/training_mode/distributed_training/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/training_mode/distributed_training/train.py b/src/faceDetection/face_sdk/training_mode/distributed_training/train.py similarity index 96% rename from training_mode/distributed_training/train.py rename to src/faceDetection/face_sdk/training_mode/distributed_training/train.py index e58a48e..1b857fc 100644 --- a/training_mode/distributed_training/train.py +++ b/src/faceDetection/face_sdk/training_mode/distributed_training/train.py @@ -5,7 +5,6 @@ """ import os import sys -import shutil import argparse import logging as logger @@ -16,11 +15,11 @@ from torch.utils.data import DataLoader from tensorboardX import SummaryWriter -sys.path.append('../../') -from utils.AverageMeter import AverageMeter -from data_processor.train_dataset import ImageDataset -from backbone.backbone_def import BackboneFactory -from head.head_def import HeadFactory +sys.path.append('../../../../../') +from src.faceDetection.face_sdk.utils.AverageMeter import AverageMeter +from src.faceDetection.face_sdk.data_processor.train_dataset import ImageDataset +from src.faceDetection.face_sdk.backbone.backbone_def import BackboneFactory +from src.faceDetection.face_sdk.head.head_def import HeadFactory logger.basicConfig(level=logger.INFO, format='%(levelname)s %(asctime)s %(filename)s: %(lineno)d] %(message)s', diff --git a/training_mode/distributed_training/train.sh b/src/faceDetection/face_sdk/training_mode/distributed_training/train.sh similarity index 99% rename from training_mode/distributed_training/train.sh rename to src/faceDetection/face_sdk/training_mode/distributed_training/train.sh index ea08562..f14aca4 100644 --- a/training_mode/distributed_training/train.sh +++ b/src/faceDetection/face_sdk/training_mode/distributed_training/train.sh @@ -1,2 +1,2 @@ -export OMP_NUM_THREADS=4 -python3 -m torch.distributed.launch --nproc_per_node=4 --nnodes=1 --node_rank=0 --master_addr="127.0.0.1" --master_port=1324 train.py --data_root 'msra_crop' --train_file 'msceleb_deepglint_train_file.txt' --backbone_type 'ResNet' --backbone_conf_file '../backbone_conf.yaml' --head_type 'MV-Softmax' --head_conf_file '../head_conf.yaml' --lr 0.1 --out_dir 'out_dir' --epoches 18 --step '10,13,16' --print_freq 200 --save_freq 3000 --batch_size 128 --momentum 0.9 --log_dir 'log' --tensorboardx_logdir 'mv-resnet' 2>&1 | tee log.log +export OMP_NUM_THREADS=4 +python3 -m torch.distributed.launch --nproc_per_node=4 --nnodes=1 --node_rank=0 --master_addr="127.0.0.1" --master_port=1324 train.py --data_root 'msra_crop' --train_file 'msceleb_deepglint_train_file.txt' --backbone_type 'ResNet' --backbone_conf_file '../backbone_conf.yaml' --head_type 'MV-Softmax' --head_conf_file '../head_conf.yaml' --lr 0.1 --out_dir 'out_dir' --epoches 18 --step '10,13,16' --print_freq 200 --save_freq 3000 --batch_size 128 --momentum 0.9 --log_dir 'log' --tensorboardx_logdir 'mv-resnet' 2>&1 | tee log.log diff --git a/training_mode/head_conf.yaml b/src/faceDetection/face_sdk/training_mode/head_conf.yaml similarity index 100% rename from training_mode/head_conf.yaml rename to src/faceDetection/face_sdk/training_mode/head_conf.yaml diff --git a/src/faceDetection/face_sdk/training_mode/semi-siamese_training/__init__.py b/src/faceDetection/face_sdk/training_mode/semi-siamese_training/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/training_mode/semi-siamese_training/train.py b/src/faceDetection/face_sdk/training_mode/semi-siamese_training/train.py similarity index 94% rename from training_mode/semi-siamese_training/train.py rename to src/faceDetection/face_sdk/training_mode/semi-siamese_training/train.py index d7f84ea..5f15a251 100644 --- a/training_mode/semi-siamese_training/train.py +++ b/src/faceDetection/face_sdk/training_mode/semi-siamese_training/train.py @@ -1,192 +1,192 @@ -""" -@author: Jun Wang -@date: 20201019 -@contact: jun21wangustc@gmail.com -""" -import os -import sys -import shutil -import argparse -import logging as logger - -import torch -from torch import optim -from torch.utils.data import DataLoader -from tensorboardX import SummaryWriter - -sys.path.append('../../') -from utils.AverageMeter import AverageMeter -from data_processor.train_dataset import ImageDataset_SST -from backbone.backbone_def import BackboneFactory -from head.head_def import HeadFactory -from test_protocol.utils.online_val import Evaluator - -logger.basicConfig(level=logger.INFO, - format='%(levelname)s %(asctime)s %(filename)s: %(lineno)d] %(message)s', - datefmt='%Y-%m-%d %H:%M:%S') - -def get_lr(optimizer): - """Get the current learning rate from optimizer. - """ - for param_group in optimizer.param_groups: - return param_group['lr'] - -def moving_average(probe, gallery, alpha): - """Update the gallery-set network in the momentum way.(MoCo) - """ - for param_probe, param_gallery in zip(probe.parameters(), gallery.parameters()): - param_gallery.data = \ - alpha* param_gallery.data + (1 - alpha) * param_probe.detach().data - -def train_BN(m): - classname = m.__class__.__name__ - if classname.find('BatchNorm') != -1: - m.train() - -def shuffle_BN(batch_size): - """ShuffleBN for batch, the same as MoCo https://arxiv.org/abs/1911.05722 ####### - """ - shuffle_ids = torch.randperm(batch_size).long().cuda() - reshuffle_ids = torch.zeros(batch_size).long().cuda() - reshuffle_ids.index_copy_(0, shuffle_ids, torch.arange(batch_size).long().cuda()) - return shuffle_ids, reshuffle_ids - -def train_one_epoch(data_loader, probe_net, gallery_net, prototype, optimizer, - criterion, cur_epoch, conf, loss_meter): - """Tain one epoch by semi-siamese training. - """ - for batch_idx, (images1, images2, id_indexes) in enumerate(data_loader): - batch_size = images1.size(0) - global_batch_idx = cur_epoch * len(data_loader) + batch_idx - images1 = images1.cuda() - images2 = images2.cuda() - # set inputs as probe or gallery - shuffle_ids, reshuffle_ids = shuffle_BN(batch_size) - images1_probe = probe_net(images1) - with torch.no_grad(): - images2 = images2[shuffle_ids] - images2_gallery = gallery_net(images2)[reshuffle_ids] - images2 = images2[reshuffle_ids] - shuffle_ids, reshuffle_ids = shuffle_BN(batch_size) - images2_probe = probe_net(images2) - with torch.no_grad(): - images1 = images1[shuffle_ids] - images1_gallery = gallery_net(images1)[reshuffle_ids] - images1 = images1[reshuffle_ids] - output1, output2, label, id_set = prototype( - images1_probe, images2_gallery, images2_probe, images1_gallery, id_indexes) - loss = (criterion(output1, label) + criterion(output2, label))/2 - # update - optimizer.zero_grad() - loss.backward() - optimizer.step() - moving_average(probe_net, gallery_net, conf.alpha) - loss_meter.update(loss.item(), batch_size) - if batch_idx % conf.print_freq == 0: - loss_val = loss_meter.avg - lr = get_lr(optimizer) - logger.info('Epoch %d, iter %d, lr %f, loss %f' % - (cur_epoch, batch_idx, lr, loss_val)) - conf.writer.add_scalar('Train_loss', loss_val, global_batch_idx) - conf.writer.add_scalar('Train_lr', lr, global_batch_idx) - if cur_epoch % conf.save_freq == 0 or cur_epoch == conf.epoches - 1: - saved_name = ('Epoch_{}.pt'.format(cur_epoch)) - torch.save(probe_net.state_dict(), os.path.join(conf.out_dir, saved_name)) - logger.info('save checkpoint %s to disk...' % saved_name) - return id_set - -def train(conf): - """Total training procedure. - """ - conf.device = torch.device('cuda:0') - criterion = torch.nn.CrossEntropyLoss().cuda(conf.device) - backbone_factory = BackboneFactory(conf.backbone_type, conf.backbone_conf_file) - probe_net = backbone_factory.get_backbone() - gallery_net = backbone_factory.get_backbone() - head_factory = HeadFactory(conf.head_type, conf.head_conf_file) - prototype = head_factory.get_head().cuda(conf.device) - probe_net = torch.nn.DataParallel(probe_net).cuda() - gallery_net = torch.nn.DataParallel(gallery_net).cuda() - optimizer = optim.SGD(probe_net.parameters(), lr=conf.lr, momentum=conf.momentum, weight_decay=5e-4) - lr_schedule = optim.lr_scheduler.MultiStepLR(optimizer, milestones=conf.milestones, gamma=0.1) - if conf.resume: - probe_net.load_state_dict(torch.load(args.pretrain_model)) - moving_average(probe_net, gallery_net, 0) - probe_net.train() - gallery_net.eval().apply(train_BN) - - exclude_id_set = set() - loss_meter = AverageMeter() - for epoch in range(conf.epoches): - data_loader = DataLoader( - ImageDataset_SST(conf.data_root, conf.train_file, exclude_id_set), - conf.batch_size, True, num_workers = 4, drop_last = True) - exclude_id_set = train_one_epoch(data_loader, probe_net, gallery_net, - prototype, optimizer, criterion, epoch, conf, loss_meter) - lr_schedule.step() - if conf.evaluate: - conf.evaluator.evaluate(probe_net) - -if __name__ == '__main__': - conf = argparse.ArgumentParser(description='semi-siamese_training for face recognition.') - conf.add_argument("--data_root", type = str, - help = "The root folder of training set.") - conf.add_argument("--train_file", type = str, - help = "The train file path.") - conf.add_argument('--backbone_type', type=str, default='Mobilefacenets', - help='Mobilefacenets, Resnet.') - conf.add_argument('--backbone_conf_file', type=str, - help='the path of backbone_conf.yaml.') - conf.add_argument("--head_type", type = str, - help = "mv-softmax, arcface, npc-face ...") - conf.add_argument("--head_conf_file", type = str, - help = "the path of head_conf.yaml..") - conf.add_argument('--lr', type = float, default = 0.1, - help='The initial learning rate.') - conf.add_argument("--out_dir", type=str, default='out_dir', - help=" The folder to save models.") - conf.add_argument('--epoches', type = int, default = 130, - help = 'The training epoches.') - conf.add_argument('--step', type = str, default = '60,100,120', - help = 'Step for lr.') - conf.add_argument('--print_freq', type = int, default = 10, - help = 'The print frequency for training state.') - conf.add_argument('--save_freq', type=int, default=1, - help='The save frequency for training state.') - conf.add_argument('--batch_size', type=int, default=128, - help='batch size over all gpus.') - conf.add_argument('--momentum', type=float, default=0.9, - help='The momentum for sgd.') - conf.add_argument('--alpha', type=float, default=0.999, - help='weight of moving_average') - conf.add_argument('--log_dir', type = str, default = 'log', - help = 'The directory to save log.log') - conf.add_argument('--tensorboardx_logdir', type = str, - help = 'The directory to save tensorboardx logs') - conf.add_argument('--pretrain_model', type = str, default = 'mv_epoch_8.pt', - help = 'The path of pretrained model') - conf.add_argument('--resume', '-r', action = 'store_true', default = False, - help = 'Resume from checkpoint or not.') - conf.add_argument('--evaluate', '-e', action = 'store_true', default = False, - help = 'Evaluate the training model.') - conf.add_argument('--test_set', type = str, default = 'LFW', - help = 'Test set to evaluate the model.') - conf.add_argument('--test_data_conf_file', type = str, - help = 'The path of test data conf file.') - args = conf.parse_args() - args.milestones = [int(num) for num in args.step.split(',')] - if not os.path.exists(args.out_dir): - os.makedirs(args.out_dir) - if not os.path.exists(args.log_dir): - os.makedirs(args.log_dir) - tensorboardx_logdir = os.path.join(args.log_dir, args.tensorboardx_logdir) - if os.path.exists(tensorboardx_logdir): - shutil.rmtree(tensorboardx_logdir) - writer = SummaryWriter(log_dir=tensorboardx_logdir) - args.writer = writer - if args.evaluate: - args.evaluator = Evaluator(args.test_set, args.test_data_conf_file) - logger.info('Start optimization.') - logger.info(args) - train(args) - logger.info('Optimization done!') +""" +@author: Jun Wang +@date: 20201019 +@contact: jun21wangustc@gmail.com +""" +import os +import sys +import shutil +import argparse +import logging as logger + +import torch +from torch import optim +from torch.utils.data import DataLoader +from tensorboardX import SummaryWriter + +sys.path.append('../../../../../') +from src.faceDetection.face_sdk.utils.AverageMeter import AverageMeter +from src.faceDetection.face_sdk.data_processor.train_dataset import ImageDataset_SST +from src.faceDetection.face_sdk.backbone.backbone_def import BackboneFactory +from src.faceDetection.face_sdk.head.head_def import HeadFactory +from src.faceDetection.face_sdk.test_protocol.utils.online_val import Evaluator + +logger.basicConfig(level=logger.INFO, + format='%(levelname)s %(asctime)s %(filename)s: %(lineno)d] %(message)s', + datefmt='%Y-%m-%d %H:%M:%S') + +def get_lr(optimizer): + """Get the current learning rate from optimizer. + """ + for param_group in optimizer.param_groups: + return param_group['lr'] + +def moving_average(probe, gallery, alpha): + """Update the gallery-set network in the momentum way.(MoCo) + """ + for param_probe, param_gallery in zip(probe.parameters(), gallery.parameters()): + param_gallery.data = \ + alpha* param_gallery.data + (1 - alpha) * param_probe.detach().data + +def train_BN(m): + classname = m.__class__.__name__ + if classname.find('BatchNorm') != -1: + m.train() + +def shuffle_BN(batch_size): + """ShuffleBN for batch, the same as MoCo https://arxiv.org/abs/1911.05722 ####### + """ + shuffle_ids = torch.randperm(batch_size).long().cuda() + reshuffle_ids = torch.zeros(batch_size).long().cuda() + reshuffle_ids.index_copy_(0, shuffle_ids, torch.arange(batch_size).long().cuda()) + return shuffle_ids, reshuffle_ids + +def train_one_epoch(data_loader, probe_net, gallery_net, prototype, optimizer, + criterion, cur_epoch, conf, loss_meter): + """Tain one epoch by semi-siamese training. + """ + for batch_idx, (images1, images2, id_indexes) in enumerate(data_loader): + batch_size = images1.size(0) + global_batch_idx = cur_epoch * len(data_loader) + batch_idx + images1 = images1.cuda() + images2 = images2.cuda() + # set inputs as probe or gallery + shuffle_ids, reshuffle_ids = shuffle_BN(batch_size) + images1_probe = probe_net(images1) + with torch.no_grad(): + images2 = images2[shuffle_ids] + images2_gallery = gallery_net(images2)[reshuffle_ids] + images2 = images2[reshuffle_ids] + shuffle_ids, reshuffle_ids = shuffle_BN(batch_size) + images2_probe = probe_net(images2) + with torch.no_grad(): + images1 = images1[shuffle_ids] + images1_gallery = gallery_net(images1)[reshuffle_ids] + images1 = images1[reshuffle_ids] + output1, output2, label, id_set = prototype( + images1_probe, images2_gallery, images2_probe, images1_gallery, id_indexes) + loss = (criterion(output1, label) + criterion(output2, label))/2 + # update + optimizer.zero_grad() + loss.backward() + optimizer.step() + moving_average(probe_net, gallery_net, conf.alpha) + loss_meter.update(loss.item(), batch_size) + if batch_idx % conf.print_freq == 0: + loss_val = loss_meter.avg + lr = get_lr(optimizer) + logger.info('Epoch %d, iter %d, lr %f, loss %f' % + (cur_epoch, batch_idx, lr, loss_val)) + conf.writer.add_scalar('Train_loss', loss_val, global_batch_idx) + conf.writer.add_scalar('Train_lr', lr, global_batch_idx) + if cur_epoch % conf.save_freq == 0 or cur_epoch == conf.epoches - 1: + saved_name = ('Epoch_{}.pt'.format(cur_epoch)) + torch.save(probe_net.state_dict(), os.path.join(conf.out_dir, saved_name)) + logger.info('save checkpoint %s to disk...' % saved_name) + return id_set + +def train(conf): + """Total training procedure. + """ + conf.device = torch.device('cuda:0') + criterion = torch.nn.CrossEntropyLoss().cuda(conf.device) + backbone_factory = BackboneFactory(conf.backbone_type, conf.backbone_conf_file) + probe_net = backbone_factory.get_backbone() + gallery_net = backbone_factory.get_backbone() + head_factory = HeadFactory(conf.head_type, conf.head_conf_file) + prototype = head_factory.get_head().cuda(conf.device) + probe_net = torch.nn.DataParallel(probe_net).cuda() + gallery_net = torch.nn.DataParallel(gallery_net).cuda() + optimizer = optim.SGD(probe_net.parameters(), lr=conf.lr, momentum=conf.momentum, weight_decay=5e-4) + lr_schedule = optim.lr_scheduler.MultiStepLR(optimizer, milestones=conf.milestones, gamma=0.1) + if conf.resume: + probe_net.load_state_dict(torch.load(args.pretrain_model)) + moving_average(probe_net, gallery_net, 0) + probe_net.train() + gallery_net.eval().apply(train_BN) + + exclude_id_set = set() + loss_meter = AverageMeter() + for epoch in range(conf.epoches): + data_loader = DataLoader( + ImageDataset_SST(conf.data_root, conf.train_file, exclude_id_set), + conf.batch_size, True, num_workers = 4, drop_last = True) + exclude_id_set = train_one_epoch(data_loader, probe_net, gallery_net, + prototype, optimizer, criterion, epoch, conf, loss_meter) + lr_schedule.step() + if conf.evaluate: + conf.evaluator.evaluate(probe_net) + +if __name__ == '__main__': + conf = argparse.ArgumentParser(description='semi-siamese_training for face recognition.') + conf.add_argument("--data_root", type = str, + help = "The root folder of training set.") + conf.add_argument("--train_file", type = str, + help = "The train file path.") + conf.add_argument('--backbone_type', type=str, default='Mobilefacenets', + help='Mobilefacenets, Resnet.') + conf.add_argument('--backbone_conf_file', type=str, + help='the path of backbone_conf.yaml.') + conf.add_argument("--head_type", type = str, + help = "mv-softmax, arcface, npc-face ...") + conf.add_argument("--head_conf_file", type = str, + help = "the path of head_conf.yaml..") + conf.add_argument('--lr', type = float, default = 0.1, + help='The initial learning rate.') + conf.add_argument("--out_dir", type=str, default='out_dir', + help=" The folder to save models.") + conf.add_argument('--epoches', type = int, default = 130, + help = 'The training epoches.') + conf.add_argument('--step', type = str, default = '60,100,120', + help = 'Step for lr.') + conf.add_argument('--print_freq', type = int, default = 10, + help = 'The print frequency for training state.') + conf.add_argument('--save_freq', type=int, default=1, + help='The save frequency for training state.') + conf.add_argument('--batch_size', type=int, default=128, + help='batch size over all gpus.') + conf.add_argument('--momentum', type=float, default=0.9, + help='The momentum for sgd.') + conf.add_argument('--alpha', type=float, default=0.999, + help='weight of moving_average') + conf.add_argument('--log_dir', type = str, default = 'log', + help = 'The directory to save log.log') + conf.add_argument('--tensorboardx_logdir', type = str, + help = 'The directory to save tensorboardx logs') + conf.add_argument('--pretrain_model', type = str, default = 'mv_epoch_8.pt', + help = 'The path of pretrained model') + conf.add_argument('--resume', '-r', action = 'store_true', default = False, + help = 'Resume from checkpoint or not.') + conf.add_argument('--evaluate', '-e', action = 'store_true', default = False, + help = 'Evaluate the training model.') + conf.add_argument('--test_set', type = str, default = 'LFW', + help = 'Test set to evaluate the model.') + conf.add_argument('--test_data_conf_file', type = str, + help = 'The path of test data conf file.') + args = conf.parse_args() + args.milestones = [int(num) for num in args.step.split(',')] + if not os.path.exists(args.out_dir): + os.makedirs(args.out_dir) + if not os.path.exists(args.log_dir): + os.makedirs(args.log_dir) + tensorboardx_logdir = os.path.join(args.log_dir, args.tensorboardx_logdir) + if os.path.exists(tensorboardx_logdir): + shutil.rmtree(tensorboardx_logdir) + writer = SummaryWriter(log_dir=tensorboardx_logdir) + args.writer = writer + if args.evaluate: + args.evaluator = Evaluator(args.test_set, args.test_data_conf_file) + logger.info('Start optimization.') + logger.info(args) + train(args) + logger.info('Optimization done!') diff --git a/training_mode/semi-siamese_training/train.sh b/src/faceDetection/face_sdk/training_mode/semi-siamese_training/train.sh similarity index 100% rename from training_mode/semi-siamese_training/train.sh rename to src/faceDetection/face_sdk/training_mode/semi-siamese_training/train.sh diff --git a/src/faceDetection/face_sdk/training_mode/siamese-triplet_training/__init__.py b/src/faceDetection/face_sdk/training_mode/siamese-triplet_training/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/training_mode/siamese-triplet_training/losses.py b/src/faceDetection/face_sdk/training_mode/siamese-triplet_training/losses.py similarity index 100% rename from training_mode/siamese-triplet_training/losses.py rename to src/faceDetection/face_sdk/training_mode/siamese-triplet_training/losses.py diff --git a/training_mode/siamese-triplet_training/pair_selector.py b/src/faceDetection/face_sdk/training_mode/siamese-triplet_training/pair_selector.py similarity index 100% rename from training_mode/siamese-triplet_training/pair_selector.py rename to src/faceDetection/face_sdk/training_mode/siamese-triplet_training/pair_selector.py diff --git a/training_mode/siamese-triplet_training/readme.md b/src/faceDetection/face_sdk/training_mode/siamese-triplet_training/readme.md similarity index 100% rename from training_mode/siamese-triplet_training/readme.md rename to src/faceDetection/face_sdk/training_mode/siamese-triplet_training/readme.md diff --git a/training_mode/siamese-triplet_training/train.py b/src/faceDetection/face_sdk/training_mode/siamese-triplet_training/train.py similarity index 94% rename from training_mode/siamese-triplet_training/train.py rename to src/faceDetection/face_sdk/training_mode/siamese-triplet_training/train.py index 31fb92a..26ae27a 100644 --- a/training_mode/siamese-triplet_training/train.py +++ b/src/faceDetection/face_sdk/training_mode/siamese-triplet_training/train.py @@ -14,15 +14,14 @@ from torch.utils.data import DataLoader from tensorboardX import SummaryWriter -from losses import OnlineContrastiveLoss, OnlineTripletLoss -from pair_selector import HardNegativePairSelector, FunctionNegativeTripletSelector, random_hard_negative +from losses import OnlineTripletLoss +from pair_selector import FunctionNegativeTripletSelector, random_hard_negative -sys.path.append('../../') -from utils.AverageMeter import AverageMeter -from data_processor.train_dataset import ImageDataset_SST -from backbone.backbone_def import BackboneFactory -from head.head_def import HeadFactory -from test_protocol.utils.online_val import Evaluator +sys.path.append('../../../../../') +from src.faceDetection.face_sdk.utils.AverageMeter import AverageMeter +from src.faceDetection.face_sdk.data_processor.train_dataset import ImageDataset_SST +from src.faceDetection.face_sdk.backbone.backbone_def import BackboneFactory +from src.faceDetection.face_sdk.test_protocol.utils.online_val import Evaluator logger.basicConfig(level=logger.INFO, format='%(levelname)s %(asctime)s %(filename)s: %(lineno)d] %(message)s', diff --git a/training_mode/siamese-triplet_training/train.sh b/src/faceDetection/face_sdk/training_mode/siamese-triplet_training/train.sh similarity index 100% rename from training_mode/siamese-triplet_training/train.sh rename to src/faceDetection/face_sdk/training_mode/siamese-triplet_training/train.sh diff --git a/training_mode/swin_training/README.md b/src/faceDetection/face_sdk/training_mode/swin_training/README.md similarity index 100% rename from training_mode/swin_training/README.md rename to src/faceDetection/face_sdk/training_mode/swin_training/README.md diff --git a/src/faceDetection/face_sdk/training_mode/swin_training/__init__.py b/src/faceDetection/face_sdk/training_mode/swin_training/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/training_mode/swin_training/lr_scheduler.py b/src/faceDetection/face_sdk/training_mode/swin_training/lr_scheduler.py similarity index 100% rename from training_mode/swin_training/lr_scheduler.py rename to src/faceDetection/face_sdk/training_mode/swin_training/lr_scheduler.py diff --git a/training_mode/swin_training/optimizer.py b/src/faceDetection/face_sdk/training_mode/swin_training/optimizer.py similarity index 100% rename from training_mode/swin_training/optimizer.py rename to src/faceDetection/face_sdk/training_mode/swin_training/optimizer.py diff --git a/training_mode/swin_training/train.py b/src/faceDetection/face_sdk/training_mode/swin_training/train.py similarity index 96% rename from training_mode/swin_training/train.py rename to src/faceDetection/face_sdk/training_mode/swin_training/train.py index b9b5f81..44a93c9 100644 --- a/training_mode/swin_training/train.py +++ b/src/faceDetection/face_sdk/training_mode/swin_training/train.py @@ -5,14 +5,12 @@ """ import os import sys -import shutil import argparse import logging as logger import torch import torch.distributed as dist import torch.utils.data.distributed -from torch import optim from torch.utils.data import DataLoader from tensorboardX import SummaryWriter from apex import amp @@ -20,11 +18,11 @@ from optimizer import build_optimizer from lr_scheduler import build_scheduler -sys.path.append('../../') -from utils.AverageMeter import AverageMeter -from data_processor.train_dataset import ImageDataset -from backbone.backbone_def import BackboneFactory -from head.head_def import HeadFactory +sys.path.append('../../../../../') +from src.faceDetection.face_sdk.utils.AverageMeter import AverageMeter +from src.faceDetection.face_sdk.data_processor.train_dataset import ImageDataset +from src.faceDetection.face_sdk.backbone.backbone_def import BackboneFactory +from src.faceDetection.face_sdk.head.head_def import HeadFactory logger.basicConfig(level=logger.INFO, format='%(levelname)s %(asctime)s %(filename)s: %(lineno)d] %(message)s', diff --git a/training_mode/swin_training/train.sh b/src/faceDetection/face_sdk/training_mode/swin_training/train.sh similarity index 100% rename from training_mode/swin_training/train.sh rename to src/faceDetection/face_sdk/training_mode/swin_training/train.sh diff --git a/utils/AverageMeter.py b/src/faceDetection/face_sdk/utils/AverageMeter.py similarity index 100% rename from utils/AverageMeter.py rename to src/faceDetection/face_sdk/utils/AverageMeter.py diff --git a/src/faceDetection/face_sdk/utils/__init__.py b/src/faceDetection/face_sdk/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/utils/compute_cost.py b/src/faceDetection/face_sdk/utils/compute_cost.py similarity index 87% rename from utils/compute_cost.py rename to src/faceDetection/face_sdk/utils/compute_cost.py index 44fe475..555c932 100644 --- a/utils/compute_cost.py +++ b/src/faceDetection/face_sdk/utils/compute_cost.py @@ -9,8 +9,8 @@ from thop import profile from thop import clever_format -sys.path.append('..') -from backbone.backbone_def import BackboneFactory +sys.path.append('../../../..') +from src.faceDetection.face_sdk.backbone.backbone_def import BackboneFactory #backbone_type = 'MobileFaceNet' #backbone_type = 'ResNet' diff --git a/utils/gen_train_file.py b/src/faceDetection/face_sdk/utils/gen_train_file.py similarity index 100% rename from utils/gen_train_file.py rename to src/faceDetection/face_sdk/utils/gen_train_file.py diff --git a/src/imageCapture/camera.py b/src/imageCapture/camera.py new file mode 100644 index 0000000..4713462 --- /dev/null +++ b/src/imageCapture/camera.py @@ -0,0 +1,115 @@ +import cv2 +import os +from datetime import datetime +from src.faceDetection.detection import FaceDetector +import subprocess +import platform + +class CameraHandler: + _SAVE_DIR = os.path.join(os.path.dirname(__file__), '..', 'data', 'model') + _HELP_TEXT = [ + "Make photo - 's'", + "Quit - 'q'", + "Begin analysis - 'a'", + "Open configurator - 'c'" + ] + WINDOW_NAME = "Preview" + FRAME_WIDTH = 1280 + FRAME_HEIGHT = 720 + + def __init__(self, camera_index=0): + os.makedirs(self._SAVE_DIR, exist_ok=True) + self._analyze_mode = False + self._face_detector = None + self.cap = cv2.VideoCapture(camera_index) + self._running = True + self.current_frame = None + + if not self.cap.isOpened(): + raise RuntimeError("Cannot open camera.") + + def start(self): + while self._running: + success, frame = self.cap.read() + self.current_frame = frame + + if not success: + print("Error while fetching frame.") + break + + if self._analyze_mode and self._face_detector: + frame = self._face_detector.analyze(frame) + + self._show_helper_texts(frame) + cv2.imshow(self.WINDOW_NAME, frame) + self._handle_key() + + self._cleanup() + + def _handle_key(self): + key = cv2.waitKey(1) & 0xFF + if key == ord('q'): + self._running = False + elif key == ord('s'): + self._save_frame() + elif key == ord('a'): + self._begin_analysis() + elif key == ord('c'): + self._launch_configurator() + + def _cleanup(self): + if hasattr(self, 'cap') and self.cap.isOpened(): + self.cap.release() + cv2.destroyAllWindows() + + def _show_helper_texts(self, frame): + for i, line in enumerate(self._HELP_TEXT): + cv2.putText(frame, line, (10, 25 + i * 50), + cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2) + + def _is_opened(self): + return self.cap.isOpened() + + def _save_frame(self): + # Pobierz nową, czystą klatkę z kamery (bez helperów) + ret, raw_frame = self.cap.read() + if not ret: + print("No frame to save.") + return + + filename = os.path.join( + self._SAVE_DIR, + f"wzorzec_{datetime.now().strftime('%Y%m%d_%H%M%S')}.jpg" + ) + + try: + is_saved = cv2.imwrite(filename, raw_frame) + if not is_saved: + raise ValueError(f"Cannot save file {filename}") + print(f"Saved file to {filename}") + except Exception as e: + print(f"Saving error: {e}") + + def _begin_analysis(self): + if self._analyze_mode: + self._analyze_mode = False + print("Analysis mode: OFF") + else: + self._face_detector = FaceDetector() + self._analyze_mode = True + print("Analysis mode: ON") + + def _launch_configurator(self): + config_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'configurator', 'configurator.py')) + print(f"Opening configurator: {config_path}") + try: + if platform.system() == "Windows": + subprocess.Popen(['python', config_path], shell=True) + else: + subprocess.Popen(['python3', config_path]) + except Exception as e: + print(f"Failed to launch configurator: {e}") + +if __name__ == "__main__": + camera = CameraHandler() + camera.start() \ No newline at end of file diff --git a/src/main.py b/src/main.py new file mode 100644 index 0000000..d84a385 --- /dev/null +++ b/src/main.py @@ -0,0 +1,10 @@ +import torch +from imageCapture.camera import CameraHandler +import os + +os.environ['CUDA_VISIBLE_DEVICES'] = '-1' +torch.backends.cudnn.enabled = False + +if __name__ == "__main__": + camera = CameraHandler() + camera.start()