Spaces:
Sleeping
Sleeping
| """ | |
| CLI entrypoint for FER inference. | |
| Usage: | |
| python predict.py --image photo.jpg | |
| python predict.py --image photo.jpg --detect-face | |
| python predict.py --image img1.jpg img2.jpg img3.jpg | |
| python predict.py --folder ./test_images/ | |
| python predict.py --image photo.jpg --detect-face --save-output result.jpg | |
| python predict.py --image photo.jpg --weights ../models/model_weights.pth | |
| """ | |
| import argparse | |
| import os | |
| import sys | |
| from pathlib import Path | |
| import cv2 | |
| import numpy as np | |
| from PIL import Image | |
| # Allow running from any directory | |
| sys.path.insert(0, str(Path(__file__).parent)) | |
| from inference import FERPredictor | |
| from utils import visualize_prediction, draw_face_predictions | |
| SUPPORTED_EXTS = {'.jpg', '.jpeg', '.png', '.bmp', '.webp', '.tiff', '.tif'} | |
| EMOTION_ORDER = ['happy', 'neutral', 'surprise', 'sad', 'angry', 'fear', 'disgust'] | |
| BAR_MAX_WIDTH = 20 | |
| def _bar(prob: float) -> str: | |
| filled = round(prob * BAR_MAX_WIDTH) | |
| return '█' * filled + ' ' * (BAR_MAX_WIDTH - filled) | |
| def _print_result(image_path: str, result: dict): | |
| print(f"\nImage: {image_path}") | |
| print('─' * 45) | |
| print(f"Emotion : {result['emotion']}") | |
| print(f"Confidence : {result['confidence']*100:.1f}%") | |
| print('─' * 45) | |
| print("All emotions:") | |
| sorted_probs = sorted(result['probabilities'].items(), key=lambda x: x[1], reverse=True) | |
| for emotion, prob in sorted_probs: | |
| bar = _bar(prob) | |
| print(f" {emotion:<10} {bar} {prob*100:.1f}%") | |
| print() | |
| def _print_face_result(image_path: str, face_results: list[dict]): | |
| print(f"\nImage: {image_path} [{len(face_results)} face(s) detected]") | |
| for res in face_results: | |
| idx = res.get('face_index', 0) | |
| bbox = res.get('bbox') | |
| bbox_str = f" bbox: {bbox}" if bbox else " (no face detected, ran on full image)" | |
| print(f"\n Face #{idx + 1}{bbox_str}") | |
| print(f" Emotion : {res['emotion']}") | |
| print(f" Confidence : {res['confidence']*100:.1f}%") | |
| sorted_probs = sorted(res['probabilities'].items(), key=lambda x: x[1], reverse=True) | |
| print(" Probabilities:") | |
| for emotion, prob in sorted_probs: | |
| bar = _bar(prob) | |
| print(f" {emotion:<10} {bar} {prob*100:.1f}%") | |
| print() | |
| def _collect_images(args) -> list[str]: | |
| paths = [] | |
| if args.image: | |
| paths.extend(args.image) | |
| if args.folder: | |
| folder = Path(args.folder) | |
| if not folder.is_dir(): | |
| print(f"[ERROR] Folder not found: {args.folder}") | |
| sys.exit(1) | |
| for p in sorted(folder.iterdir()): | |
| if p.suffix.lower() in SUPPORTED_EXTS: | |
| paths.append(str(p)) | |
| if not paths: | |
| print(f"[WARNING] No supported image files found in: {args.folder}") | |
| return paths | |
| def _save_annotated(image_path: str, face_results: list[dict], save_path: str): | |
| img_bgr = cv2.imread(image_path) | |
| if img_bgr is None: | |
| print(f"[WARNING] Could not reload image for annotation: {image_path}") | |
| return | |
| annotated = draw_face_predictions(img_bgr, face_results) | |
| cv2.imwrite(save_path, annotated) | |
| print(f"[INFO] Annotated image saved to: {save_path}") | |
| def main(): | |
| parser = argparse.ArgumentParser( | |
| description='FER Inference — predict facial expressions from images.' | |
| ) | |
| parser.add_argument('--image', nargs='+', help='Path(s) to input image(s).') | |
| parser.add_argument('--folder', help='Folder of images to process.') | |
| parser.add_argument('--detect-face', action='store_true', | |
| help='Run face detection before prediction.') | |
| parser.add_argument('--face-method', default='mtcnn', choices=['mtcnn', 'haar'], | |
| help='Face detection method (default: mtcnn).') | |
| parser.add_argument('--weights', default='../models/model_weights.pth', | |
| help='Path to model_weights.pth (default: ../models/model_weights.pth).') | |
| parser.add_argument('--device', default='auto', | |
| help='Device: auto | cpu | cuda | cuda:0 (default: auto).') | |
| parser.add_argument('--save-output', help='Save annotated image to this path.') | |
| parser.add_argument('--save-plot', help='Save probability bar chart to this path.') | |
| args = parser.parse_args() | |
| if not args.image and not args.folder: | |
| parser.print_help() | |
| sys.exit(1) | |
| image_paths = _collect_images(args) | |
| if not image_paths: | |
| print("[ERROR] No images to process.") | |
| sys.exit(1) | |
| predictor = FERPredictor(weights_path=args.weights, device=args.device) | |
| for img_path in image_paths: | |
| if not os.path.exists(img_path): | |
| print(f"[WARNING] File not found, skipping: {img_path}") | |
| continue | |
| if args.detect_face: | |
| face_results = predictor.predict_with_face_detection( | |
| img_path, method=args.face_method | |
| ) | |
| _print_face_result(img_path, face_results) | |
| if args.save_output: | |
| save_path = args.save_output if len(image_paths) == 1 else \ | |
| f"{Path(args.save_output).stem}_{Path(img_path).stem}{Path(args.save_output).suffix}" | |
| _save_annotated(img_path, face_results, save_path) | |
| if args.save_plot and face_results: | |
| # Plot probabilities of first face | |
| from utils import plot_emotion_bars | |
| plot_emotion_bars(face_results[0]['probabilities'], | |
| title=f"{face_results[0]['emotion']} ({face_results[0]['confidence']*100:.1f}%)", | |
| save_path=args.save_plot) | |
| else: | |
| result = predictor.predict_image(img_path) | |
| _print_result(img_path, result) | |
| if args.save_output: | |
| img = Image.open(img_path) | |
| save_path = args.save_output if len(image_paths) == 1 else \ | |
| f"{Path(args.save_output).stem}_{Path(img_path).stem}{Path(args.save_output).suffix}" | |
| visualize_prediction(img, result, save_path=save_path) | |
| if args.save_plot: | |
| from utils import plot_emotion_bars | |
| plot_emotion_bars(result['probabilities'], | |
| title=f"{result['emotion']} ({result['confidence']*100:.1f}%)", | |
| save_path=args.save_plot) | |
| if __name__ == '__main__': | |
| main() | |