error resolved

app.py

@@ -4,28 +4,36 @@ import os
 from model import create_effnet_b2
 from timeit import default_timer as timer
 from typing import Tuple, Dict
+from PIL import Image
+import numpy as np
+
 
 class_names = ["pizza", "steak", "sushi"]
 
-effbet_b2_model , efftnet_b2_transform = create_effnet_b2()
+effnet_b2_model , effnet_b2_transform = create_effnet_b2()
 
-effbet_b2_model.load_state_dict(torch.load(f = "./effnet_b2.pt", map_location = torch.device("cpu")))
+effnet_b2_model.load_state_dict(torch.load(f = "./effnet_b2.pt", map_location = torch.device("cpu")))
 
-def predict(img)-> Tuple[Dict,float]:
-    
+def predict(img) -> Tuple[Dict, float]:
     start_time = timer()
-    img = efftnet_b2_transform(img).unsqueeze(0)
-    
-    effbet_b2_model.eval()
     
+    # Convert from NumPy array to PIL image
+    if isinstance(img, np.ndarray):
+        img = Image.fromarray(img.astype("uint8"), "RGB")
+
+    img = effnet_b2_transform(img).unsqueeze(0)
+
+    effnet_b2_model.eval()
     with torch.inference_mode():
-        pred_prob = torch.softmax(effbet_b2_model(img), 1)
-        
-    pred_label_probs = {class_names[i] : float(pred_prob[0][i]) for i in range(len(class_names))}
-    
+        pred_prob = torch.softmax(effnet_b2_model(img), dim=1)
+
+    pred_label_probs = {
+        class_names[i]: float(pred_prob[0][i]) for i in range(len(class_names))
+    }
+
     end_time = timer()
-    pred_time = round(end_time - start_time , 4)
-    
+    pred_time = round(end_time - start_time, 4)
+
     return pred_label_probs, pred_time
 
 
@@ -46,7 +54,7 @@ demo = gr.Interface(
     allow_flagging="never"
 )
 
-demo.launch(debug=False)
+demo.launch(debug=False, share=True)
     
 
     

model.py

@@ -1,4 +1,5 @@
 
+import torch
 from torchvision.models import EfficientNet_B2_Weights, efficientnet_b2
 from torch import nn