Add config.json with model_type

config.json

@@ -1,3 +1,10 @@
 {
-    "model_type": "yolos"
+    "model_type": "u2net",
+    "architectures": [
+        "U2NET"
+    ],
+    "task": "image-segmentation",
+    "pipeline_class": "U2NetPipeline",
+    "model_file": "u2net.py",
+    "weights_file": "u2net.pth"
 }

inference.py

@@ -1,34 +0,0 @@
-import torch
-from u2net import U2NET
-from torchvision import transforms
-import numpy as np
-from PIL import Image
-import torch.nn.functional as F
-import data_transforms
-
-# Load the model
-def load_model():
-    model = U2NET(3, 1)
-    model.load_state_dict(torch.load("u2net.pth", map_location="cpu"))
-    model.eval()
-    return model
-
-# Preprocessing function (same as you defined locally)
-def preprocess(image):
-    transform = transforms.Compose([data_transforms.RescaleT(320), data_transforms.ToTensorLab(flag=0)])
-    label_3 = np.zeros(image.shape)
-    label = np.zeros(label_3.shape[0:2])
-    sample = transform({"imidx": np.array([0]), "image": image, "label": label})
-    return sample
-
-# Inference function
-def infer(model, image):
-    input_size = [1024, 1024]
-    im_shp = image.shape[0:2]
-    im_tensor = torch.tensor(image, dtype=torch.float32).permute(2, 0, 1)
-    im_tensor = F.upsample(torch.unsqueeze(im_tensor, 0), input_size, mode="bilinear").type(torch.uint8)
-    image = torch.divide(im_tensor, 255.0)
-    result = model(image)
-    result = torch.squeeze(F.upsample(result[0][0], im_shp, mode='bilinear'), 0)
-    result = (result - result.min()) / (result.max() - result.min())
-    return result.numpy()

model_info.json

@@ -1,4 +0,0 @@
-{
-    "model_type": "image-segmentation",
-    "task": "image-segmentation"
-}

requirements.txt

@@ -1,7 +1,8 @@
 torch
 torchvision
 numpy
-pillow
+Pillow
 matplotlib
 scikit-image
 huggingface-hub
+transformers

u2net_pipeline.py

@@ -0,0 +1,54 @@
+import torch
+from torchvision import transforms
+from PIL import Image
+import numpy as np
+import torch.nn.functional as F
+from u2net import U2NET
+import data_transforms
+from transformers import Pipeline
+
+class U2NetPipeline(Pipeline):
+    def __init__(self, model, **kwargs):
+        super().__init__(model=model, **kwargs)
+        self.model = U2NET(3, 1)
+        self.model.load_state_dict(torch.load(model, map_location="cpu"))
+        self.model.eval()
+
+    def _sanitize_parameters(self, **kwargs):
+        return {}, {}, {}
+
+    def preprocess(self, image):
+        if isinstance(image, str):
+            image = Image.open(image).convert("RGB")
+        elif isinstance(image, Image.Image):
+            image = image.convert("RGB")
+        else:
+            raise ValueError("Input must be a PIL Image or a path to an image file")
+
+        image = np.array(image)
+        transform = transforms.Compose([data_transforms.RescaleT(320), data_transforms.ToTensorLab(flag=0)])
+        sample = transform({"imidx": np.array([0]), "image": image, "label": np.zeros(image.shape[:2])})
+
+        input_size = [1024, 1024]
+        im_tensor = sample['image'].unsqueeze(0)
+        im_tensor = F.interpolate(im_tensor, input_size, mode="bilinear")
+        image = torch.divide(im_tensor, 255.0)
+        image = transforms.Normalize([0.5, 0.5, 0.5], [1.0, 1.0, 1.0])(image)
+
+        return {"image": image, "original_size": image.shape[2:]}
+
+    def _forward(self, model_inputs):
+        with torch.no_grad():
+            outputs = self.model(model_inputs["image"])
+        return {"outputs": outputs, "original_size": model_inputs["original_size"]}
+
+    def postprocess(self, model_outputs):
+        result = model_outputs["outputs"][0][0]
+        result = F.interpolate(result, size=model_outputs["original_size"], mode='bilinear', align_corners=False)
+        result = result.squeeze().cpu().numpy()
+        ma, mi = result.max(), result.min()
+        result = (result - mi) / (ma - mi)
+        return (result * 255).astype(np.uint8)
+
+def load_model():
+    return U2NetPipeline("u2net.pth")