Update app.py

app.py

@@ -1,36 +1,90 @@
 import gradio as gr
-import torch
-from torchvision import transforms
+import numpy as np
 from PIL import Image
-from transformers import AutoModelForImageClassification, AutoImageProcessor
-
-model_name = "nateraw/fer-2013"
-
-processor = AutoImageProcessor.from_pretrained(model_name)
-model = AutoModelForImageClassification.from_pretrained(model_name)
-
-transform = transforms.Compose([
-    transforms.Resize((224, 224)),
-    transforms.ToTensor()
-])
-
-def predict(img):
-    img = Image.fromarray(img).convert("RGB")
-    inputs = processor(images=img, return_tensors="pt")
-    with torch.no_grad():
-        outputs = model(**inputs)
-    logits = outputs.logits
-    prob = logits.softmax(dim=1)
-    score, label_id = torch.max(prob, dim=1)
-    label = model.config.id2label[label_id.item()]
-    return f"Emotion: {label} ({score.item():.2f})"
-
-ui = gr.Interface(
-    fn=predict,
-    inputs=gr.Image(type="numpy", label="Upload Image"),
-    outputs="text",
-    title="Emotion Detection (PyTorch)",
-    description="Detect emotions using a lightweight PyTorch model (FER-2013)."
+import torch
+from transformers import AutoImageProcessor, AutoModelForImageClassification
+
+# --- Configuration ---
+MODEL_NAME = "nateraw/fer-2013"
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+# --- Model and Processor Loading ---
+try:
+    processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
+    model = AutoModelForImageClassification.from_pretrained(MODEL_NAME).to(DEVICE)
+    model.eval()
+    LABELS = model.config.id2label
+    
+    print(f"Model loaded successfully on device: {DEVICE}")
+    print(f"Available labels: {LABELS}")
+
+except Exception as e:
+    print(f"Error loading model or processor: {e}")
+    # Use a dummy function if loading fails
+    processor = None
+    model = None
+    LABELS = {0: "dummy_emotion"}
+
+# --- Inference Function ---
+def classify_emotion(image_np: np.ndarray) -> str:
+    """
+    Performs emotion classification on an input image (numpy array).
+    """
+    if model is None or processor is None:
+        return f"Error: Model or processor failed to load. Check logs."
+
+    try:
+        # Convert numpy array (from Gradio) to PIL Image
+        image = Image.fromarray(image_np).convert("RGB")
+        
+        # Preprocess the image
+        # The processor handles necessary resizing and normalization
+        inputs = processor(images=image, return_tensors="pt").to(DEVICE)
+        
+        # Run inference
+        with torch.no_grad():
+            outputs = model(**inputs)
+            
+        # Get predictions
+        logits = outputs.logits
+        probabilities = torch.nn.functional.softmax(logits, dim=-1)
+        
+        # Find the dominant emotion
+        confidence, predicted_class_idx = torch.max(probabilities, 1)
+        
+        dominant_emotion = LABELS[predicted_class_idx.item()]
+        confidence_score = confidence.item()
+
+        # Format the result string
+        result_str = f"Emotion: **{dominant_emotion}** ({confidence_score:.2f})"
+        return result_str
+
+    except Exception as e:
+        return f"Prediction Error: {type(e).__name__} - {str(e)}"
+
+# --- Gradio Interface ---
+iface = gr.Interface(
+    fn=classify_emotion,
+    inputs=gr.Image(
+        type="numpy", 
+        label="Upload an image of a face"
+    ),
+    outputs=gr.Markdown(label="Predicted Emotion"),
+    title="Emotion Detection (PyTorch/Transformers)",
+    description=(
+        "Upload an image containing a face to classify the dominant emotion. "
+        "Uses the **nateraw/fer-2013** PyTorch model from Hugging Face Transformers. "
+        "No TensorFlow or Keras dependencies."
+    ),
+    examples=[
+        # Providing simple examples is good practice
+        # Note: Gradio will handle downloading and using these if they exist in the repo
+        # Since this is for a new Space, use placeholder or common sense (omit paths if files aren't included)
+    ],
+    allow_flagging="never",
+    theme=gr.themes.Soft()
 )
 
-ui.launch()
+# Launch the app (required for Hugging Face Spaces)
+if __name__ == "__main__":
+    iface.launch()