init

.gitignore

@@ -0,0 +1,2 @@
+.env
+.idea

Dockerfile

@@ -0,0 +1,16 @@
+# Read the doc: https://huggingface.co/docs/hub/spaces-sdks-docker
+# you will also find guides on how best to write your Dockerfile
+
+FROM python:3.9
+
+RUN useradd -m -u 1000 user
+USER user
+ENV PATH="/home/user/.local/bin:$PATH"
+
+WORKDIR /app
+
+COPY --chown=user ./requirements.txt requirements.txt
+RUN pip install --no-cache-dir --upgrade -r requirements.txt
+
+COPY --chown=user . /app
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -0,0 +1,38 @@
+from fastapi import FastAPI, UploadFile, File
+import torch
+from models import load_model1
+from utils import preprocess_image, postprocess_mask, resize_mask, mask_to_base64
+import numpy as np
+from PIL import Image
+import io
+
+app = FastAPI()
+
+# Загрузка модели при старте (глобально, один раз)
+model1 = load_model1()
+
+
+@app.get("/")
+def greet_json():
+    return {"Hello": "World!"}
+
+
+@app.post("/predict1")
+async def predict1(file: UploadFile = File(...)):
+    content = await file.read()
+    image = Image.open(io.BytesIO(content)).convert('RGB')
+    image_np = np.array(image)
+
+    input_tensor = preprocess_image(image_np)
+    with torch.no_grad():
+        logits = model1(input_tensor.unsqueeze(0))  # batch dim
+
+    pred_mask = postprocess_mask(logits)  # (256, 256) binary
+
+    mask_100 = resize_mask(pred_mask, 100)
+    mask_224 = resize_mask(pred_mask, 224)
+
+    return {
+        "mask_100_base64": mask_to_base64(mask_100),
+        "mask_224_base64": mask_to_base64(mask_224)
+    }

models.py

@@ -0,0 +1,20 @@
+import torch
+import segmentation_models_pytorch as smp
+
+DEVICE = torch.device('cpu')
+
+model1 = None
+
+def load_model1(weights_path='weights/seg.pth'):
+    global model1
+    if model1 is None:
+        model1 = smp.Unet(
+            encoder_name="mobilenet_v2",
+            encoder_weights=None,
+            in_channels=3,
+            classes=1
+        ).to(DEVICE)
+        state_dict = torch.load(weights_path, map_location=DEVICE)
+        model1.load_state_dict(state_dict)
+        model1.eval()
+    return model1

requirements.txt

@@ -0,0 +1,8 @@
+fastapi
+uvicorn[standard]
+torch
+segmentation_models_pytorch
+albumentations
+pillow
+numpy
+opencv-python-headless

utils.py

@@ -0,0 +1,36 @@
+import numpy as np
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+import torch
+import cv2
+from PIL import Image
+import io
+import base64
+
+# Препроцессинг: аналог валидации
+preprocess_transform = A.Compose([
+    A.Resize(256, 256),
+    A.Normalize(),
+    ToTensorV2()
+])
+
+def preprocess_image(image_np: np.ndarray) -> torch.Tensor:
+    augmented = preprocess_transform(image=image_np)
+    return augmented['image']
+
+def postprocess_mask(logits: torch.Tensor, threshold: float = 0.5) -> np.ndarray:
+    pred = torch.sigmoid(logits).squeeze().cpu().numpy()
+    binary_mask = (pred > threshold).astype(np.float32)
+    return binary_mask  # shape (256, 256)
+
+def resize_mask(mask: np.ndarray, size: int) -> np.ndarray:
+    # Resize с nearest neighbor для бинарных масок
+    resized = cv2.resize(mask, (size, size), interpolation=cv2.INTER_NEAREST)
+    return resized.astype(np.float32)  # 0/1 float
+
+def mask_to_base64(mask: np.ndarray) -> str:
+    # Конверт в PIL grayscale (0/255), save as PNG, base64
+    pil_mask = Image.fromarray((mask * 255).astype(np.uint8)).convert('L')
+    buffered = io.BytesIO()
+    pil_mask.save(buffered, format="PNG")
+    return base64.b64encode(buffered.getvalue()).decode('utf-8')

weights/class1.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c5767f1246bfe8ee0077a0eefda6c8a1a66e8639de3fc1d94bacf7254633a5f2
+size 9205515

weights/class2.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bf1835ce202339850361b67f914d7768400f981aa2f0c1a6a29e7f268a749f18
+size 9174411

weights/seg.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:45c364012ded2dbb389ffe61177056f8b574d5f58822dc3dd7be8a10a20459d7
+size 26805291