main.py

import cv2
from fer.fer import FER
import os
import pandas as pd
import yaml
import argparse
import sys

def load_config(config_path="config.yaml"):
    if os.path.exists(config_path):
        try:
            with open(config_path, 'r') as f:
                return yaml.safe_load(f)
        except Exception as e:
            print(f"Warning: Could not read config file '{config_path}': {e}")
    return {}

def main():
    parser = argparse.ArgumentParser(description="Visual Analysis Tool: Sentiment analysis from facial images.")
    parser.add_argument('--image_folder', type=str, default=None, help='Path to folder containing images')
    parser.add_argument('--output_file', type=str, default=None, help='Path to output Excel file')
    parser.add_argument('--config', type=str, default='config.yaml', help='Path to config file (default: config.yaml)')
    args = parser.parse_args()

    # Load config file
    config = load_config(args.config)

    # Use command-line args if provided, else config, else defaults
    image_folder = args.image_folder or config.get('image_folder', './example_images')
    output_file = args.output_file or config.get('output_file', './example_results.xlsx')

    # Generate a unique output filename if the file already exists
    output_file = args.output_file
    base, ext = os.path.splitext(output_file)
    counter = 1
    while os.path.exists(output_file):
        output_file = f"{base}_{counter}{ext}"
        counter += 1

    if not os.path.isdir(image_folder):
        print(f"Error: Image folder '{image_folder}' does not exist.")
        sys.exit(1)

    emotion_detector = FER(mtcnn=False)
    results = []

    for filename in os.listdir(image_folder):
        if filename.lower().endswith((".jpg", ".png", ".jpeg")):
            img_path = os.path.join(image_folder, filename)
            input_image = cv2.imread(img_path)
            if input_image is None:
                print(f"Warning: Could not read image '{filename}'. Skipping.")
                continue
            try:
                result = emotion_detector.detect_emotions(input_image)
                print(f"Processing {filename}: {len(result) if result else 0} faces detected")
            except Exception as e:
                print(f"Error processing '{filename}': {e}")
                results.append({
                    'Image': filename,
                    'Angry': 0,
                    'Disgust': 0,
                    'Fear': 0,
                    'Happy': 0,
                    'Sad': 0,
                    'Surprise': 0,
                    'Neutral': 0,
                    'Highest Emotion': f'ERROR: {str(e)}'
                })
                continue
            if result:
                emotions = result[0]['emotions']
                highest_emotion = max(emotions, key=emotions.get)
                results.append({
                    'Image': filename,
                    'Angry': emotions.get('angry', 0),
                    'Disgust': emotions.get('disgust', 0),
                    'Fear': emotions.get('fear', 0),
                    'Happy': emotions.get('happy', 0),
                    'Sad': emotions.get('sad', 0),
                    'Surprise': emotions.get('surprise', 0),
                    'Neutral': emotions.get('neutral', 0),
                    'Highest Emotion': highest_emotion
                })
                print(f"Emotion detected: {highest_emotion}")
            else:
                results.append({
                    'Image': filename,
                    'Angry': 0,
                    'Disgust': 0,
                    'Fear': 0,
                    'Happy': 0,
                    'Sad': 0,
                    'Surprise': 0,
                    'Neutral': 0,
                    'Highest Emotion': 'No face detected'
                })
                print(f"No face detected in {filename}")

    if not results:
        print("No images processed. No results to save.")
        sys.exit(0)

    df = pd.DataFrame(results)
    try:
        df.to_excel(output_file, index=False)
        print(f"Results saved to {output_file}")
        print(f"Total images: {len(results)}")
    except Exception as e:
        print(f"Error saving results to Excel: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()