|
from typing import Tuple |
|
from io import BytesIO |
|
|
|
from fastapi import FastAPI, UploadFile, File |
|
from fastapi.middleware.cors import CORSMiddleware |
|
from starlette.responses import StreamingResponse |
|
|
|
import numpy as np |
|
|
|
import cv2 |
|
|
|
import torch |
|
from torchvision import transforms |
|
|
|
from PIL import Image, ImageChops |
|
|
|
import segmentation_models_pytorch as smp |
|
|
|
|
|
constants: dict = { |
|
'encoder_name': 'resnet34', |
|
'encoder_weights': 'imagenet', |
|
|
|
'sigmoid_threshold': 0.55, |
|
|
|
'model_path': 'models/production/unetplusplus_resnet34.pth' |
|
} |
|
|
|
def load_model() -> smp.UnetPlusPlus: |
|
''' |
|
Returns: |
|
model: smp.UnetPlusPlus |
|
''' |
|
global model |
|
|
|
if model is None: |
|
model = smp.UnetPlusPlus(encoder_name=constants['encoder_name'], |
|
encoder_weights=constants['encoder_weights'], |
|
in_channels=3, |
|
classes=1).to(device) |
|
|
|
model.load_state_dict(torch.load(constants['model_path'], map_location=device)) |
|
|
|
return model |
|
|
|
def draw_bounding_boxes(mask: np.array) -> Tuple[np.array, float]: |
|
''' |
|
Arguments: |
|
mask: np.array (numpy) |
|
|
|
Returns: |
|
Tuple[np.array, float] |
|
''' |
|
mask = mask.astype(np.uint8) |
|
|
|
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) |
|
|
|
if len(contours) == 0: |
|
return mask, 0.0 |
|
|
|
mask_bgr = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR) |
|
mean_area = [] |
|
|
|
for contour in contours: |
|
x, y, w, h = cv2.boundingRect(contour) |
|
|
|
max_x = x + w |
|
max_y = y + h |
|
|
|
cv2.rectangle(mask_bgr, (x, y), (max_x, max_y), (255, 0, 0), 1) |
|
mean_area.append(abs(max_x - x) * abs(max_y - y)) |
|
|
|
return mask_bgr, sum(mean_area) / len(mean_area) |
|
|
|
transform = transforms.Compose([transforms.ToTensor(), |
|
transforms.Resize((256, 256), antialias=True)]) |
|
|
|
device = 'cuda' if torch.cuda.is_available() else 'cpu' |
|
|
|
model = None |
|
|
|
app = FastAPI(title='Brain MRI Medical Images Segmentation Hugging Face') |
|
|
|
app.add_middleware(CORSMiddleware, |
|
allow_headers=['*'], |
|
expose_headers=['has_tumor', 'average_pixel_area']) |
|
|
|
@app.post('/predict') |
|
async def predict(file: UploadFile = File(...)): |
|
input_image = Image.open(BytesIO(await file.read())).convert('RGB') |
|
array_image = np.array(input_image) |
|
transformed_image = transform(array_image) |
|
transformed_image = transformed_image.unsqueeze(0) |
|
transformed_image = transformed_image.to(device) |
|
|
|
model = load_model() |
|
model.eval() |
|
|
|
with torch.no_grad(): |
|
prediction = model(transformed_image) |
|
|
|
prediction = torch.sigmoid(prediction) |
|
prediction = (prediction > constants['sigmoid_threshold']).float() |
|
prediction = prediction.squeeze() |
|
prediction = prediction.cpu() |
|
prediction = prediction.numpy() |
|
prediction = prediction * 255 |
|
prediction, mean_area = draw_bounding_boxes(prediction) |
|
|
|
transformed_image = transformed_image.cpu() |
|
transformed_image = transformed_image[0].permute(1, 2, 0) |
|
transformed_image = transformed_image.numpy() * 255 |
|
transformed_image = transformed_image.astype(np.uint8) |
|
|
|
input_image = Image.fromarray(transformed_image).convert('RGBA') |
|
predicted_mask = Image.fromarray(prediction).convert('RGBA') |
|
result = ImageChops.screen(input_image, predicted_mask) |
|
|
|
bytes_io = BytesIO() |
|
result.save(bytes_io, format='PNG') |
|
bytes_io.seek(0) |
|
|
|
return StreamingResponse(bytes_io, media_type='image/png', headers={'has_tumor': f'{prediction.max() > 0}', 'average_pixel_area': f'{round(mean_area, 2)}'}) |
|
|
