feat: implement universal image loader to support various input types and update preprocessing functions accordingly

app.py

@@ -5,6 +5,9 @@ import numpy as np
 import os
 import time
 import logging
+import requests
+import io
+import tempfile
 
 # Assuming these are available from your utils and agents directories
 # You might need to adjust paths or copy these functions/classes if they are not directly importable.
@@ -66,16 +69,57 @@ CLASS_NAMES = {
     "model_7": ['Fake', 'Real'],
 }
 
+# Universal image loader
+def universal_image_loader(img_input):
+    """
+    Accepts a PIL Image, NumPy array, file path, or URL.
+    Returns a PIL Image.
+    """
+    if isinstance(img_input, Image.Image):
+        return img_input
+    if isinstance(img_input, np.ndarray):
+        return Image.fromarray(img_input)
+    if isinstance(img_input, str):
+        if img_input.startswith('http://') or img_input.startswith('https://'):
+            try:
+                response = requests.get(img_input)
+                response.raise_for_status()
+                return Image.open(io.BytesIO(response.content)).convert('RGB')
+            except Exception as e:
+                logger.error(f"Failed to load image from URL: {img_input} | Error: {e}")
+                raise
+        elif os.path.exists(img_input):
+            try:
+                return Image.open(img_input).convert('RGB')
+            except Exception as e:
+                logger.error(f"Failed to load image from file: {img_input} | Error: {e}")
+                raise
+        else:
+            logger.error(f"String input is not a valid file path or URL: {img_input}")
+            raise ValueError(f"Invalid image input: {img_input}")
+    logger.error(f"Unsupported image input type: {type(img_input)}")
+    raise TypeError(f"Unsupported image input type: {type(img_input)}")
+
+# Update all preprocessors to use universal_image_loader
+
 def preprocess_resize_256(image):
+    image = universal_image_loader(image)
     if image.mode != 'RGB':
         image = image.convert('RGB')
     return transforms.Resize((256, 256))(image)
 
 def preprocess_resize_224(image):
+    image = universal_image_loader(image)
     if image.mode != 'RGB':
         image = image.convert('RGB')
     return transforms.Resize((224, 224))(image)
 
+def preprocess_256(image):
+    image = universal_image_loader(image)
+    if image.mode != 'RGB':
+        image = image.convert('RGB')
+    return transforms.Resize((256, 256))(image)
+
 def postprocess_pipeline(prediction, class_names):
     # Assumes HuggingFace pipeline output
     return {pred['label']: pred['score'] for pred in prediction}
@@ -109,10 +153,6 @@ register_model_with_metadata(
 
 feature_extractor_3 = AutoFeatureExtractor.from_pretrained(MODEL_PATHS["model_3"], device=device)
 model_3 = AutoModelForImageClassification.from_pretrained(MODEL_PATHS["model_3"]).to(device)
-def preprocess_256(image):
-    if image.mode != 'RGB':
-        image = image.convert('RGB')
-    return transforms.Resize((256, 256))(image)
 def postprocess_logits_model3(outputs, class_names):
     logits = outputs.logits.cpu().numpy()[0]
     probabilities = softmax(logits)
@@ -171,8 +211,7 @@ register_model_with_metadata(
 )
 
 def preprocess_simple_prediction(image):
-    # The simple_prediction function expects a PIL image (filepath is handled internally)
-    return image
+    return universal_image_loader(image)
 
 def postprocess_simple_prediction(result, class_names):
     scores = {name: 0.0 for name in class_names}
@@ -184,10 +223,15 @@ def postprocess_simple_prediction(result, class_names):
     return scores
 
 def simple_prediction(img):
+    img = universal_image_loader(img)
+    # Save PIL image to a temporary file
+    with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp:
+        img.save(tmp, format="JPEG")
+        tmp_path = tmp.name
     client = Client("aiwithoutborders-xyz/OpenSight-Community-Forensics-Preview")
     result = client.predict(
-            input_image=handle_file(img),
-            api_name="/simple_predict"
+        input_image=tmp_path,
+        api_name="/simple_predict"
     )
     return result