removed osail_utils

app.py

@@ -3,7 +3,6 @@ import os
 import gradio as gr
 import matplotlib.pyplot as plt
 import numpy as np
-import osail_utils
 import pandas as pd 
 import skimage
 from mediffusion import DiffusionModule
@@ -29,6 +28,22 @@ BASELINE_NOISE = torch.randn(1, 1, 256, 256).cuda().half()
 
 # Model helper functions
 
+class LoadImageD(mn.transforms.Transform):
+    def __init__(self, keys, transpose=False, normalize=False):
+        self.keys = keys
+        self.transpose = transpose
+        self.normalize = normalize
+    def __call__(self, data):
+        for key in self.keys:
+            img = skimage.io.imread(data[key])
+            if self.transpose:
+                img = img.transpose(0, 1)
+            if self.normalize:
+                img -= img.min()
+                img /= (img.max()+1e-6)
+            data[key] = img
+        return data
+
 def create_ds(img_paths):
     if type(img_paths) == str:
         img_paths = [img_paths]
@@ -36,7 +51,7 @@ def create_ds(img_paths):
         
     # Get the transforms            
     Ts_list = [
-            osail_utils.io.LoadImageD(keys=["img"], transpose=True, normalize=True),
+            LoadImageD(keys=["img"], transpose=True, normalize=True),
             mn.transforms.EnsureChannelFirstD(
                 keys=["img"], channel_dim="no_channel"
             ),