Update app.py

app.py

@@ -4,6 +4,7 @@ from torchvision import models, transforms
 from PIL import Image
 import torch.nn.functional as F
 import torch.nn as nn
+import numpy as np
 
 # Class labels
 class_names = ['fake', 'real']
@@ -35,19 +36,26 @@ model, device = load_model()
 
 # Inference function
 def predict(image):
+    # Convert numpy array to PIL Image if needed
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
+    
+    # Convert RGBA to RGB if needed
     if image.mode == "RGBA":
         image = image.convert("RGB")
-
+    
+    # Apply transforms
     image = data_transforms(image).unsqueeze(0).to(device)
+    
     with torch.no_grad():
         outputs = model(image)
         probs = F.softmax(outputs, dim=1)
         conf, pred = torch.max(probs, 1)
-
-    label = class_names[pred.item()]
-    confidence = f"{conf.item() * 100:.2f}%"
+        label = class_names[pred.item()]
+        confidence = f"{conf.item() * 100:.2f}%"
+    
     return f"{label} ({confidence})"
 
 # Gradio interface
 demo = gr.Interface(fn=predict, inputs="image", outputs="text", api_name="predict")
-demo.launch()
+demo.launch()