File size: 9,373 Bytes
7fa4c7a 33f434e 7fa4c7a 33f434e 7fa4c7a 33f434e 7fa4c7a 33f434e 7fa4c7a 33f434e 7fa4c7a 33f434e 4845d75 33f434e 7fa4c7a 33f434e 7fa4c7a 4845d75 7fa4c7a |
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 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 |
"""
Custom Inference Handler for Hugging Face Inference Endpoints
Combines Qwen2.5-VL embedding extraction + MLP classifiers
"""
import torch
import torch.nn as nn
from transformers import AutoProcessor, AutoModelForVision2Seq
from pathlib import Path
import numpy as np
from typing import Dict, Any
import av
import tempfile
class MLPClassifier(nn.Module):
"""MLP classifier matching training architecture"""
def __init__(self, input_dim, hidden_dim=512, num_classes=4, dropout=0.3):
super(MLPClassifier, self).__init__()
self.classifier = nn.Sequential(
nn.Linear(input_dim, hidden_dim),
nn.ReLU(),
nn.Dropout(dropout),
nn.Linear(hidden_dim, hidden_dim // 2),
nn.ReLU(),
nn.Dropout(dropout),
nn.Linear(hidden_dim // 2, num_classes)
)
def forward(self, x):
return self.classifier(x)
class EndpointHandler:
"""
Custom handler for HF Inference Endpoints
"""
def __init__(self, path: str):
"""
Initialize the handler
Args:
path: Path to the model directory on HF Hub
"""
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# Load Qwen2.5-VL model for embeddings
print("Loading Qwen2.5-VL model...")
self.processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct")
self.vision_model = AutoModelForVision2Seq.from_pretrained(
"Qwen/Qwen2.5-VL-7B-Instruct",
torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
device_map="auto" if torch.cuda.is_available() else None
)
self.vision_model.eval()
# Load MLP classifiers for each emotion
self.categories = ["Boredom", "Engagement", "Confusion", "Frustration"]
self.classifiers = {}
path = Path(path)
classifiers_dir = path / "classifiers"
print("Loading MLP classifiers...")
for category in self.categories:
checkpoint_path = classifiers_dir / f"mlp_{category}_best.pth"
if checkpoint_path.exists():
# Determine embedding dimension from checkpoint
checkpoint = torch.load(checkpoint_path, map_location='cpu', weights_only=False)
# Get input dimension from first layer
first_layer_weight = checkpoint['model_state_dict']['classifier.0.weight']
input_dim = first_layer_weight.shape[1]
# Initialize model
model = MLPClassifier(input_dim=input_dim, num_classes=4)
model.load_state_dict(checkpoint['model_state_dict'])
model.to(self.device)
model.eval()
self.classifiers[category] = model
print(f" ✓ Loaded {category} classifier")
else:
print(f" ✗ Missing {category} classifier at {checkpoint_path}")
self.fps = 1 # Frame sampling rate
def extract_image_embeddings(self, image_path: str) -> np.ndarray:
"""Extract embeddings from a single image using Qwen model"""
from PIL import Image
# Load image
image = Image.open(image_path).convert('RGB')
messages = [
{
"role": "user",
"content": [
{"type": "image", "image": image},
{"type": "text", "text": "Analyze this image."}
]
}
]
with torch.no_grad():
text = self.processor.apply_chat_template(
messages,
add_generation_prompt=True,
tokenize=False,
)
inputs = self.processor(
text=[text],
images=[image],
return_tensors="pt",
padding=True,
)
inputs = {k: v.to(self.vision_model.device) for k, v in inputs.items()}
outputs = self.vision_model(**inputs, output_hidden_states=True)
hidden_states = outputs.hidden_states[-1]
# Average pooling over sequence dimension
embeddings = hidden_states.mean(dim=1).squeeze(0)
embeddings = embeddings.cpu().numpy()
return embeddings
def extract_video_embeddings(self, video_path: str) -> np.ndarray:
"""Extract embeddings from video using Qwen model"""
messages = [
{
"role": "user",
"content": [
{"type": "video", "video": str(video_path)},
{"type": "text", "text": "Analyze this video."}
]
}
]
with torch.no_grad():
inputs = self.processor.apply_chat_template(
messages,
fps=self.fps,
add_generation_prompt=True,
tokenize=True,
return_dict=True,
return_tensors="pt",
)
inputs = {k: v.to(self.vision_model.device) for k, v in inputs.items()}
outputs = self.vision_model(**inputs, output_hidden_states=True)
hidden_states = outputs.hidden_states[-1]
# Average pooling over sequence dimension
embeddings = hidden_states.mean(dim=1).squeeze(0)
embeddings = embeddings.cpu().numpy()
return embeddings
def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
"""
Handle inference request
Args:
data: Input data containing either:
- "inputs": base64 encoded image or video file
- "video_url": URL to video file
Returns:
Dictionary with predictions for each emotion category
"""
try:
import base64
from PIL import Image
import io
file_path = None
is_image = False
# Handle different input formats
if "inputs" in data:
# Base64 encoded data
input_data = data["inputs"]
# Remove data URL prefix if present (e.g., "data:image/png;base64,")
if ',' in input_data and input_data.startswith('data:'):
input_data = input_data.split(',', 1)[1]
file_bytes = base64.b64decode(input_data)
# Try to detect if it's an image or video
try:
# Try to open as image
Image.open(io.BytesIO(file_bytes))
is_image = True
suffix = '.png'
except:
# Assume it's a video
is_image = False
suffix = '.avi'
# Save to temporary file
with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
tmp.write(file_bytes)
file_path = tmp.name
elif "video_url" in data:
# Download from URL
import requests
response = requests.get(data["video_url"])
with tempfile.NamedTemporaryFile(suffix='.avi', delete=False) as tmp:
tmp.write(response.content)
file_path = tmp.name
is_image = False
else:
return {"error": "No input provided. Use 'inputs' (base64) or 'video_url'"}
# Extract embeddings based on input type
if is_image:
embeddings = self.extract_image_embeddings(file_path)
else:
embeddings = self.extract_video_embeddings(file_path)
embeddings_tensor = torch.FloatTensor(embeddings).unsqueeze(0).to(self.device)
# Run classifiers
predictions = {}
with torch.no_grad():
for category, model in self.classifiers.items():
outputs = model(embeddings_tensor)
probabilities = torch.softmax(outputs, dim=1)
predicted_level = outputs.argmax(dim=1).item()
confidence = probabilities[0][predicted_level].item()
predictions[category] = {
"level": int(predicted_level),
"confidence": float(confidence),
"probabilities": probabilities[0].cpu().numpy().tolist()
}
# Clean up temporary file
if file_path:
Path(file_path).unlink(missing_ok=True)
return {
"success": True,
"predictions": predictions
}
except Exception as e:
return {
"success": False,
"error": str(e)
}
|