ok

main.py

@@ -0,0 +1,84 @@
+from fastapi import FastAPI, Form, Depends, Request, File, UploadFile
+from fastapi.encoders import jsonable_encoder
+from fastapi.responses import JSONResponse
+from fastapi.middleware.cors import CORSMiddleware
+
+import segmentation_models_pytorch as smp
+import torch
+import numpy as np
+import cv2
+import os
+from torch.utils.data import Dataset, DataLoader
+from PIL import Image
+from io import BytesIO
+
+DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+model = smp.PAN(encoder_name="resnext50_32x4d", in_channels=3, classes=1)
+model.to(DEVICE).load_state_dict(torch.load("/model/pan_resnext50_32x4d_adam_lr001_batch16_epoch_50.ckpt", map_location=DEVICE))
+
+app = FastAPI()
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # Replace with the list of allowed origins for production
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+image_dataset = []
+for file in os.listdir("/images"):
+    image_dataset.append(cv2.resize(cv2.imread('images/' + file), (160, 544)))
+
+@app.get("/")
+async def root():
+    return {"message": "Hello World"}
+
+class CustomDataset(Dataset):
+    def __init__(self, data, transform=None):
+        self.data = data
+        self.transform = transform
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        sample = {
+            'image': self.data[idx],
+        }
+
+        if self.transform:
+            sample = self.transform(sample)
+
+        return sample
+
+@app.post("/segmentation")
+async def segmentation(file: UploadFile = File(...)):
+    contents = await file.read()
+    image = Image.open(BytesIO(contents))
+    open_cv_image = np.array(image)
+    open_cv_image = cv2.resize(open_cv_image, (160, 544))
+
+    image_dataset.insert(0, open_cv_image)
+    dataset = CustomDataset(image_dataset)
+    dataloader = DataLoader(dataset, batch_size=1, shusffle=False, num_workers=0)
+
+    try:
+        for batch in dataloader:
+            temp_image = batch['image'].to(DEVICE)
+            output = model(temp_image)
+        output[0] = (output[0] > 0.5)
+        output = output[0].squeeze().cpu().numpy()
+
+    except Exception as e:
+        print(e)
+        return JSONResponse(status_code=500, content={"error": str(e)})
+    else:
+        return JSONResponse(status_code=200, content={"result": 'good'})
+    
+@app.post("/predict")
+async def predict():
+    return None
+
+        

model/pan_resnext50_32x4d_adam_lr001_batch16_epoch_50.ckpt

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

requirements.txt

@@ -0,0 +1,6 @@
+segmentation-models-pytorch
+# pandas
+numpy
+torch
+opencv-python
+fastapi==0.103.2