add np

helpers.py

@@ -2,6 +2,7 @@ import gc
 import torch
 from scipy.signal import fftconvolve
 from PIL import Image
+import numpy as np
 
 def flush():
     gc.collect()

test.py

@@ -0,0 +1,50 @@
+
+
+class FondantInferenceModel:
+    """FondantInferenceModel class that abstracts the model loading and inference.
+    User needs to implement an inference, pre/postprocess step and pass the class to the FondantInferenceComponent.
+    The FondantInferenceComponent will then load the model and prepare it for inference.
+    The examples folder can then show examples for a pytorch / huggingface / tensorflow / ... model.
+    """
+    def __init__(self, device: str = "cpu"):
+        self.device = device
+        # load model
+        self.model = self.load_model()
+        # set model to eval mode
+        self.eval()
+        
+    def load_model(self):
+        # load model
+        ...
+    
+    def eval(self):
+        # prepare for inference
+        self.model = self.model.eval()
+        self.model = self.model.to(self.device)
+    
+    def preprocess(self, input):
+        # preprocess input
+        ...
+    
+    def postprocess(self, output):
+        # postprocess output
+        ...
+    
+    def __call__(self, *args, **kwargs):
+        processed_inputs = self.preprocess(*args, **kwargs)
+        outputs = self.model(*processed_inputs)
+        processed_outputs = self.postprocess(outputs)
+        return processed_outputs
+    
+
+class FondantInferenceComponent(FondantTransformComponent, FondantInferenceModel):
+    # loads the model and prepares it for inference
+    
+    def transform(
+        self, args: argparse.Namespace, dataframe: dd.DataFrame
+    ) -> dd.DataFrame:
+        # by using the InferenceComponent, the model is automatically loaded and prepared for inference
+        # you just need to call the infer method
+        # the self.infer method calls the model.__call__ method of the FondantInferenceModel
+        output = self.infer(args.image)
+