Update predict_utils.py

predict_utils.py

@@ -3,8 +3,7 @@ import random
 import torch
 from PIL import Image
 
-from config import IMAGE_SIZE
-from data_utils import get_transform, load_charcoal_dataset
+from data_utils import get_eval_transform, prepare_splits, get_class_names
 from train_utils import load_model, get_runtime_device
 
 
@@ -19,7 +18,7 @@ def predict_uploaded_image(model_name: str, image: Image.Image):
     model, meta = load_model(model_name, device)
 
     class_names = meta["config"]["class_names"]
-    transform = get_transform()
+    transform = get_eval_transform()
 
     image = image.convert("RGB")
     tensor = transform(image).unsqueeze(0).to(device)
@@ -48,17 +47,23 @@ def test_random_sample(model_name: str):
     device = get_runtime_device()
     model, meta = load_model(model_name, device)
 
-    raw, class_names = load_charcoal_dataset()
-    test_dataset = raw["test"]
+    splits = prepare_splits()
+    class_names = get_class_names()
+    test_dataset = splits["test"]
 
     idx = random.randint(0, len(test_dataset) - 1)
     item = test_dataset[idx]
 
-    image = item["image"].convert("RGB")
+    image = item["image"]
+    if not isinstance(image, Image.Image):
+        image = Image.open(image)
+
+    image = image.convert("RGB")
+
     label = int(item["label"])
     label_name = class_names[label]
 
-    transform = get_transform()
+    transform = get_eval_transform()
     tensor = transform(image).unsqueeze(0).to(device)
 
     with torch.no_grad():