kataria_opticals_api / landmarks.py
codernotme's picture
update
7f3db4a verified
import cv2
import numpy as np
import os
import urllib.request
from huggingface_hub import hf_hub_download
REPO_ID = "codernotme/kataria_optical"
MODEL_FILE = "lbfmodel.yaml"
_facemark = None
_face_cascade = None
def ensure_model():
global MODEL_FILE
if not os.path.exists(MODEL_FILE):
print(f"Attempting to download {MODEL_FILE} from HF Hub...")
try:
downloaded_path = hf_hub_download(repo_id=REPO_ID, filename=MODEL_FILE)
MODEL_FILE = downloaded_path
print("Download from HF Hub complete.")
except Exception as e:
print(f"HF Hub download failed ({e}), falling back to GitHub...")
MODEL_URL = "https://raw.githubusercontent.com/kurnianggoro/GSOC2017/master/data/lbfmodel.yaml"
try:
urllib.request.urlretrieve(MODEL_URL, MODEL_FILE)
print("Download from GitHub complete.")
except Exception as e2:
raise RuntimeError(f"Failed to download LBF model from all sources: {e2}")
def _get_facemark():
global _facemark
if _facemark is None:
ensure_model()
_facemark = cv2.face.createFacemarkLBF()
try:
_facemark.loadModel(MODEL_FILE)
except Exception as e:
raise RuntimeError(f"Could not load LBF model: {e}")
return _facemark
def _get_face_cascade():
global _face_cascade
if _face_cascade is None:
face_cascade_path = cv2.data.haarcascades + "haarcascade_frontalface_alt2.xml"
_face_cascade = cv2.CascadeClassifier(face_cascade_path)
if _face_cascade.empty():
_face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + "haarcascade_frontalface_default.xml")
return _face_cascade
def get_landmarks(image_input):
"""
Extracts 68 face landmarks using OpenCV FacemarkLBF.
Returns a list/array of (x, y) coordinates.
"""
facemark = _get_facemark()
if isinstance(image_input, str):
image = cv2.imread(image_input)
if image is None:
raise ValueError(f"Could not read image: {image_input}")
else:
if hasattr(image_input, "mode"):
image = cv2.cvtColor(np.array(image_input), cv2.COLOR_RGB2BGR)
else:
image = np.array(image_input)
if image.ndim == 3 and image.shape[2] == 3:
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
if image is None or image.size == 0:
raise ValueError("Could not read image data.")
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
# Detect faces using Haar Cascade (included in cv2)
face_cascade = _get_face_cascade()
faces = face_cascade.detectMultiScale(gray, 1.3, 5)
if len(faces) == 0:
raise ValueError("No face detected.")
# Find landmarks
ok, landmarks = facemark.fit(gray, faces)
if not ok:
raise ValueError("Could not detect landmarks.")
# landmarks comes as a list of numpy arrays, we extract the first face's landmarks
# landmarks[0] is shape (1, 68, 2)
return landmarks[0][0]