File size: 5,284 Bytes
08a38c0 9b4fb85 08a38c0 | 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 | """
Simple emotion classification with a Deit-Base model hosted on the Hugging Face Hub.
Install:
pip install torch timm safetensors huggingface_hub pillow
Quick use:
from emotion_classifier import EmotionClassifier
clf = EmotionClassifier()
result = clf.predict("path/to/image.jpg")
print(result)
# {'prediction': 'Neutral', 'confidence': 0.95, 'Optimistic': 0.01,
# 'Pessimistic': 0.01, 'Hostile': 0.02, 'Neutral': 0.95}
Command line use:
python emotion_classifier.py path/to/image_or_folder
python emotion_classifier.py path/to/folder --output results.csv
"""
import argparse
from pathlib import Path
import torch
import torch.nn.functional as F
import timm
from PIL import Image
from torchvision import transforms
from safetensors.torch import load_file
from huggingface_hub import hf_hub_download
# ---------------------------------------------------------------------------
# Fixed model settings (must match how the checkpoint was trained)
# ---------------------------------------------------------------------------
REPO_ID = "Simon14142/emotion-recognition-4class-deit"
FILENAME = "DeiT-Base_checkpoint_final.safetensors"
MODEL_NAME = "deit_base_patch16_224"
IMG_SIZE = 224
EMOTION_NAMES = ["Optimistic", "Pessimistic", "Hostile", "Neutral"]
IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD = [0.229, 0.224, 0.225]
IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp", ".bmp", ".tiff"}
class EmotionClassifier:
"""Loads the Deit-Base emotion model straight from the Hugging Face Hub."""
def __init__(self, repo_id: str = REPO_ID, filename: str = FILENAME,
device: str | None = None, precision: str = "fp32"):
self.device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
checkpoint_path = hf_hub_download(repo_id=repo_id, filename=filename)
self.model = timm.create_model(MODEL_NAME, pretrained=False, num_classes=len(EMOTION_NAMES))
state_dict = load_file(checkpoint_path, device="cpu")
missing, unexpected = self.model.load_state_dict(state_dict, strict=True)
if missing:
print(f"Warning - missing keys: {missing}")
if unexpected:
print(f"Warning - unexpected keys: {unexpected}")
self.precision = precision
if precision == "bf16":
if self.device.type == "cuda" and torch.cuda.is_bf16_supported():
self.model = self.model.to(torch.bfloat16)
else:
print("BF16 not supported here, falling back to FP32.")
self.precision = "fp32"
self.model = self.model.to(self.device).eval()
self.transform = transforms.Compose([
transforms.Resize((IMG_SIZE, IMG_SIZE)),
transforms.ToTensor(),
transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
])
@torch.no_grad()
def predict(self, image, threshold: float = 0.4) -> dict:
"""image: a file path (str/Path) or a PIL.Image. Returns a result dict."""
if not isinstance(image, Image.Image):
image = Image.open(image).convert("RGB")
tensor = self.transform(image).unsqueeze(0).to(self.device)
if self.precision == "bf16":
tensor = tensor.to(torch.bfloat16)
logits = self.model(tensor)
probs = F.softmax(logits.float(), dim=1).squeeze(0).cpu()
pred_idx = int(probs.argmax())
conf = float(probs[pred_idx])
label = EMOTION_NAMES[pred_idx] if conf >= threshold else "Unknown"
return {
"prediction": label,
"confidence": round(conf, 4),
**{name: round(float(probs[i]), 4) for i, name in enumerate(EMOTION_NAMES)},
}
def _collect_image_paths(input_path: str) -> list[str]:
p = Path(input_path)
if p.is_file():
return [str(p)]
if p.is_dir():
paths = sorted(str(f) for f in p.iterdir() if f.suffix.lower() in IMAGE_EXTENSIONS)
if not paths:
raise FileNotFoundError(f"No images found in: {input_path}")
return paths
raise FileNotFoundError(f"Path does not exist: {input_path}")
def main():
parser = argparse.ArgumentParser(description="Emotion classification with Deit-Base")
parser.add_argument("input", help="Path to an image file or a folder of images")
parser.add_argument("--threshold", type=float, default=0.4)
parser.add_argument("--output", help="Optional CSV path to save results")
args = parser.parse_args()
clf = EmotionClassifier()
paths = _collect_image_paths(args.input)
rows = []
for path in paths:
result = clf.predict(path, threshold=args.threshold)
result["image_path"] = path
rows.append(result)
print(f"{Path(path).name:<40} -> {result['prediction']:<12} ({result['confidence']:.3f})")
if args.output:
import csv
with open(args.output, "w", newline="", encoding="utf-8") as f:
writer = csv.DictWriter(f, fieldnames=["image_path", "prediction", "confidence", *EMOTION_NAMES])
writer.writeheader()
writer.writerows(rows)
print(f"\nSaved {len(rows)} results to {args.output}")
if __name__ == "__main__":
main()
|