File size: 3,620 Bytes
5751632 edba0cc 5751632 edba0cc 5751632 9040d50 5751632 820b49d 5751632 820b49d 5751632 820b49d 5751632 820b49d 5751632 820b49d 5751632 820b49d 5751632 820b49d 9040d50 5751632 820b49d 9040d50 820b49d 9040d50 5751632 9040d50 5751632 820b49d 5751632 820b49d 9040d50 5751632 9040d50 5751632 9040d50 820b49d 9040d50 820b49d 9040d50 5751632 9040d50 5751632 9040d50 820b49d 9040d50 820b49d 5751632 820b49d 9040d50 5751632 9040d50 820b49d 5751632 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 | import os
import numpy as np
import cv2
import torch
from PIL import Image
# Make sure these modules exist in the same package folder
from .exp_recognition_model import trnscm, load_model, classes
#############################################################################################################################
# Caution: Don't change any of the filenames, function names and definitions #
# Always use the current_path + file_name for refering any files, without it we cannot access files on the server #
#############################################################################################################################
current_path = os.path.dirname(os.path.abspath(__file__))
_MODEL = None
_DEVICE = None
def _to_bgr_np(img):
"""
Convert various inputs -> OpenCV BGR numpy array.
Supports: filepath str, PIL.Image, numpy RGB/gray/RGBA.
"""
if img is None:
return None
# filepath string
if isinstance(img, str):
return cv2.imread(img)
# PIL image
if isinstance(img, Image.Image):
rgb = np.array(img.convert("RGB"))
return cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)
# numpy array
if isinstance(img, np.ndarray):
arr = img
# grayscale (H,W)
if arr.ndim == 2:
return cv2.cvtColor(arr, cv2.COLOR_GRAY2BGR)
# RGBA -> RGB
if arr.ndim == 3 and arr.shape[2] == 4:
arr = arr[:, :, :3]
# assume RGB (common from Gradio type="numpy")
if arr.ndim == 3 and arr.shape[2] == 3:
return cv2.cvtColor(arr, cv2.COLOR_RGB2BGR)
return arr
return None
def detected_face(image_bgr):
"""
Returns PIL grayscale cropped face with maximum area.
Returns 0 if not detected or cascade not available.
"""
face_haar = os.path.join(current_path, "haarcascade_frontalface_default.xml")
face_cascade = cv2.CascadeClassifier(face_haar)
# if cascade missing or failed to load
if face_cascade.empty():
return 0
gray = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2GRAY)
faces = face_cascade.detectMultiScale(gray, 1.3, 5)
if faces is None or len(faces) == 0:
return 0
# pick max area
x, y, w, h = max(faces, key=lambda f: f[2] * f[3])
crop = gray[y:y + h, x:x + w]
return Image.fromarray(crop)
def _load_cached_model():
"""
Load the expression model once and cache it.
"""
global _MODEL, _DEVICE
if _MODEL is not None:
return _MODEL, _DEVICE
_DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
ckpt_path = os.path.join(current_path, "best_resnet18_expression.pt")
if not os.path.exists(ckpt_path):
raise FileNotFoundError(f"Missing checkpoint: {ckpt_path}")
_MODEL = load_model(ckpt_path, device=str(_DEVICE), num_classes=len(classes))
return _MODEL, _DEVICE
def get_expression(img):
"""
img: can be numpy array (RGB), PIL.Image, or filepath string
returns: expression string
"""
model, device = _load_cached_model()
img_bgr = _to_bgr_np(img)
if img_bgr is None:
return "UNKNOWN"
face = detected_face(img_bgr)
if face == 0:
# fallback: use full frame grayscale
face = Image.fromarray(cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY))
x = trnscm(face).unsqueeze(0).to(device) # [1,3,H,W] if trnscm ensures RGB
with torch.no_grad():
logits = model(x)
pred = int(torch.argmax(logits, dim=1).item())
return classes.get(pred, "UNKNOWN").capitalize() |