# emotion detection with artemis # given a CSV with image filepaths in a column 'image_file', # append a column 'grounding_emotion' and a column 'emotions' with the detected emotion. # 'emotions' contains detected the emotions ordered by descending probability. # 'grounding_emotion' contains the main emotion (the most probable one) import pandas as pd import numpy as np from PIL import Image import torch from tqdm import tqdm #### checkpoints def save_checkpoint(checkpoint_path,df, grounding_emo_list, emotions_list): output_df = df.copy() output_df['grounding_emotion'] = grounding_emo_list output_df['emotions'] = emotions_list output_df.to_csv(checkpoint_path, index= False, # don't write a new 'Index' column ) print("Saved checkpoint!") def load_checkpoint_grounding_emotion(checkpoint_path): try: print("reading checkpoint at ", checkpoint_path) df = pd.read_csv(checkpoint_path) cached_grounding_emotion= { row['image_file']: row['grounding_emotion'] for _, row in df.iterrows() } print(f"Checkpoint loaded succesfully to cache: {len(cached_grounding_emotion)} processed files") return cached_grounding_emotion except: print("Checkpoint was not loaded") return cached_grounding_emotion_dict def load_checkpoint_emotions(checkpoint_path): try: print("reading checkpoint at ", checkpoint_path) df = pd.read_csv(checkpoint_path) cached_emotions= { row['image_file']: row['emotions'] for _, row in df.iterrows() } print(f"Checkpoint loaded succesfully to cache: {len(cached_emotions)} processed files") return cached_emotions except: print("Checkpoint was not loaded") return cached_emotions_dict def get_checkpoint_path(output_path): #checkpoint_path = "checkpoint" + os.path.basename(output_path) #checkpoint_path = os.path.join( os.path.dirname(output_path), checkpoint_path) #return checkpoint_path return output_path cached_grounding_emotion_dict = {} # to avoid recomputing cached_emotions_dict = {} # to avoid recomputing def get_all_emotions(filepath, model, cached_grounding_emotion_dict, cached_emotions_dict): emotions = cached_emotions_dict.get(filepath) grounding_emotion = cached_grounding_emotion_dict.get(filepath) if emotions is None: grounding_emotion, emotions = get_all_emotions_in_image(filepath, model) cached_grounding_emotion_dict[filepath] = grounding_emotion cached_emotions_dict[filepath] = emotions return grounding_emotion, emotions def get_all_emotions_in_image(filepath, model): with Image.open(filepath).convert('RGB') as img: img = transformation(img).unsqueeze(0)# unsqueeze to add artificial first dimension # emotion detection emotion_vector = model(img) # apply the model emotion_vector = np.exp(emotion_vector.detach().numpy()) # calculate probabilities sorted_indices = (-emotion_vector).argsort()[0] # sort from most to least likely emotions = [(IDX_TO_EMOTION[ind], emotion_vector[0][ind]) for ind in sorted_indices] #construct the csv line #emotions_df['image_file'].append(filename) grounding_emotion = emotions[0][0] return grounding_emotion, emotions ################### utilities #### Artemis emotions: ARTEMIS_EMOTIONS = ['amusement', 'awe', 'contentment', 'excitement', 'anger', 'disgust', 'fear', 'sadness', 'something else'] EMOTION_TO_IDX = {e: i for i, e in enumerate(ARTEMIS_EMOTIONS)} IDX_TO_EMOTION = {EMOTION_TO_IDX[e]: e for e in EMOTION_TO_IDX} ### Artemis image preprocessing #from artemis.in_out.neural_net_oriented import image_transformation import torchvision.transforms as transforms image_net_mean = [0.485, 0.456, 0.406] image_net_std = [0.229, 0.224, 0.225] def image_transformation(img_dim, lanczos=True): """simple transformation/pre-processing of image data.""" if lanczos: resample_method = Image.LANCZOS else: resample_method = Image.BILINEAR normalize = transforms.Normalize(mean=image_net_mean, std=image_net_std) img_transforms = dict() img_transforms['train'] = transforms.Compose([transforms.Resize((img_dim, img_dim), resample_method), transforms.ToTensor(), normalize]) # Use same transformations as in train (since no data-augmentation is applied in train) img_transforms['test'] = img_transforms['train'] img_transforms['val'] = img_transforms['train'] img_transforms['rest'] = img_transforms['train'] return img_transforms transformation = image_transformation(255)['train'] ### Artemis load model #from artemis.in_out.neural_net_oriented import torch_load_model import warnings def torch_load_model(checkpoint_file, map_location=None): """ Wrap torch.load to catch standard warning of not finding the nested implementations. :param checkpoint_file: :param map_location: :return: """ with warnings.catch_warnings(): warnings.simplefilter("ignore") model = torch.load(checkpoint_file, map_location=map_location) return model ######## ######### performing emotion detection def artemis_emotions_detection(input_file, output_file): checkpoint_path = get_checkpoint_path(output_file) cached_grounding_emotion_dict = load_checkpoint_grounding_emotion(checkpoint_path) cached_emotions_dict = load_checkpoint_emotions(checkpoint_path) #load the model # device = torch.device("cuda:" + str(0) if torch.cuda.is_available() else "cpu") # it is not working with gpus, only cpu :/ device =torch.device("cpu") model = torch_load_model(img2emo_model_path, map_location=device) recognized_grounding_emotion_per_image = [] recognized_emotions_per_image = [] processed_files = set(cached_emotions_dict.keys()) df = pd.read_csv(input_file) iterable_list = list(enumerate( df['image_file'])) for elem in tqdm(iterable_list): idx = elem[0] filepath = elem[1] #save checkpoint every 50 files if (not (len(processed_files) % 49) ): print(f"Images processed: {len(processed_files)}") save_checkpoint( checkpoint_path, df.iloc[:idx], recognized_grounding_emotion_per_image, recognized_emotions_per_image ) grounding_emotion, emotions = get_all_emotions( filepath, model, cached_grounding_emotion_dict, cached_emotions_dict ) recognized_grounding_emotion_per_image.append(grounding_emotion) recognized_emotions_per_image.append(emotions) processed_files.add(filepath) recognized_grounding_emotion_per_image = pd.Series(recognized_grounding_emotion_per_image) recognized_emotions_per_image = pd.Series(recognized_emotions_per_image) return recognized_grounding_emotion_per_image, recognized_emotions_per_image import argparse if __name__ == "__main__": parser = argparse.ArgumentParser(prog="Artemis emotion detection", description='Recognizes the emotions per image in an image list') parser.add_argument("--input_file", "-in", metavar='in', type=str, nargs=1, help='input file containing images-paths for emotion detection.', #default=[default_painting_folder] ) parser.add_argument("--output_file", "-out", metavar='out', type=str, nargs=1, help='output file containing images-paths + grounding (main) emotion + ranked emotions' #default=[default_interpretation_folder] ) parser.add_argument("--model_path", "-mp", metavar='mp', type=str, nargs=1, help='artemis img2emo model path' #default=[default_interpretation_folder] ) args = parser.parse_args() input_csv_file = args.input_file[0] output_csv_file = args.output_file[0] #where is the image to emotion classifier? img2emo_model_path = args.model_path[0] print(">>> input file: " , input_csv_file) print(">>> output file: ", output_csv_file) print(">>> img to emotion model path: ", img2emo_model_path) # perform object recognition recognized_grounding_emotion_per_image, recognized_emotions_per_image = artemis_emotions_detection(input_csv_file, output_csv_file) # add a column with the recognized objects output_df = pd.read_csv(input_csv_file) output_df['grounding_emotion'] = recognized_grounding_emotion_per_image output_df['emotions'] = recognized_emotions_per_image output_df.to_csv(output_csv_file, index= False, # don't write a new 'Index' column )