Update README.md

README.md

@@ -44,23 +44,45 @@ To use the model for inference:
 from torchvision import transforms
 import torch
 from PIL import Image
+
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
 # Load model
 model = torch.load("EmoSet_clip_Lora_16.0R_8.0alphaLora_32_batch_0.0001_headmlp.pth").to(device).eval()
-# Load an image
-image = Image.open("image_path.jpg").convert("RGB")
-# Preprocess and predict
-def Emo_preprocess():
+
+# Emotion label mapping
+idx2label = {
+    0: "amusement",
+    1: "awe",
+    2: "contentment",
+    3: "excitement",
+    4: "anger",
+    5: "disgust",
+    6: "fear",
+    7: "sadness"
+}
+
+# Preprocessing function
+def emo_preprocess():
     transform = transforms.Compose([
-    transforms.Resize(224),
-    transforms.CenterCrop(size=(224,224)),  
-    transforms.ToTensor(),
-    # Note: The model normalizes the image inside the forward pass
-    # using mean = (0.48145466, 0.4578275, 0.40821073) and 
-    # std = (0.26862954, 0.26130258, 0.27577711)
-])
+        transforms.Resize(224),
+        transforms.CenterCrop(size=(224, 224)),
+        transforms.ToTensor(),
+        # Note: The model normalizes the image inside the forward pass
+        # using mean = (0.48145466, 0.4578275, 0.40821073) and 
+        # std = (0.26862954, 0.26130258, 0.27577711)
+    ])
     return transform
-image = Emo_preprocess()(image).unsqueeze(0).to(device)
+
+# Load an image
+image = Image.open("image_path.jpg").convert("RGB")
+image = emo_preprocess()(image).unsqueeze(0).to(device)
+
+# Run inference
 with torch.no_grad():
-    emo_label = model(image).item()
-print(f"Predicted Emotion: {emo_label}")
+    outputs = model(image)
+    _, predicted = outputs.max(1)  # Get the class index
+
+# Get emotion label
+predicted_emotion = idx2label[predicted.item()]
+print(f"Predicted Emotion: {predicted_emotion}")