This repository contains lib for infer SAN model. WARNING!!! this lib doesnt work whithout two additional element:
- Minimal lib for create SAN model
- Weights for SAN in state_dict format example you can finde there.
This repo is closely related with SAN lib repo.
- first you need to clone repo (git clone https://github.com/PososikTeam/landmark_detection)
- install landmark_detection (python install setup.py) this command must be ran in cloned repo dir
- download minimal lib for create SAN model, dir name for this library must be lib
- download weights or convert from checkpoint by using this notebook
- cut weights into minimal lib for create SAN model, weights name file must be san.pth
- now you can use this lib with following directory structur:
├── lib
│ ...
│ └── san.pth
├── code.py
where code.py is the place there you want to use landmark_detection.
from landmark_detector import LandmarkDetector
import cv2
def main():
img = cv2.cvtColor(cv2.imread('img.png'), cv2.COLOR_BGR2RGB)
box = [200, 10, 580, 500]
landmark_detector = LandmarkDetector()
input_dict = {'image' : img, 'box' : box}
answ_dict = landmark_detector.predict(input_dict)
print(answ_dict['landmarks'])
print(answ_dict['error_message'])
if '__name__' == 'main':
main()
answ_dict['landmarks']
contains x, y coordinates point and it probability, so shape = (num_points, 3),
answ_dict['error_message']
contains error message.
Class LandmarkDetector contains path as input parameter, this path is weights directory path, so you can add weights file in any place with add this place directory as input for LandmarkDetector class. Also that class have num_points parameter- counts of model nedd to detect point.
class LandmarkDetector():
def __init__(self, num_points = 68, path = None):