In this repository I'm trying to do image segmentations on KITTI Road dataset using UNET built from scratch with PyTorch and serve the trained model over HTTP RESTful API for hosting using Docker containers.

Architecture U-NET:

from : U-Net: Convolutional Networks for Biomedical Image Segmentation https://arxiv.org/abs/1505.04597

LOSS function IOU loss :

IOU (Intersection over Union) loss is a valuable tool for optimizing segmentation models. It measures the overlap between predicted and ground truth segmentations, encouraging accurate boundary capture. By minimizing the IOU loss during training, models generate more precise segmentations, improving performance in tasks like medical image analysis and object detection.

Note : binary cross-entropy loss is suitable for pixel-wise classification tasks, while IOU loss is beneficial for evaluating and optimizing segmentation models that require accurate boundary capture. and it made more sense to choose a IOU loss over BCE loss

class IoULoss(nn.Module):
    def __init__(self):
        super(IoULoss, self).__init__()
        self.eps = 1e-6

    def forward(self, y_pred, y_true):
        # Flatten the input tensors
        y_pred = y_pred.view(-1)        
        y_true = y_true.view(-1)
        # Calculate the confusion matrix
        intersection = (y_pred * y_true).sum()
        union = y_pred.sum() + y_true.sum() - intersection

        # Calculate the IoU and return the complement as the loss
        iou = intersection / (union + self.eps)
        return 1 - iou

Results :

Support on Kaggle : https://www.kaggle.com/code/hossamemamo/kitti-road-segmentation-pytorch-unet-from-scratch

Disclaimer: This project is still work-in-progress

Running Server Locally

You will need Docker installed for this to work.

• Clone this repo and unzip the best_model_state.zip file (this contains the trained parameters).

• Take the output file best_model_state.bin and place it inside the www/ directory.

• cd into the www/ directory and run docker build -t torch-server .

• After it finishes, run docker run -d -p 5000:5000 torch-server

• Done! you have the model exposed over localhost:5000.

• Send your requests containing the image in the request form data with the name file using (postman/insomnia/curl).

Sending Requests

using curl:

curl -X POST -H "Content-Type: multipart/form-data" -F file="@my_img_file.png" "localhost:5000/upload" -o prediction_image.jpg

Cautions

follow this for consistent behavior as this software is still in early development.

• your input image to the server should be a png.
• the output server prediction image will be jpg.

Development Plan

• Download dataset and preprocess it.
• Built UNET architecture and train using PyTorch on GPU.
• Expose the model over an HTTP API.
• Build the server docker image.
• Deploy the docker image to docker hub.
• Build web front-end interface.
• Host containers on a cloud provider acting as backend for the front-end user interface.
• Add support for Real-Time.

Documentation Todos

• Add documentation for running the server locally without using Docker.

License

MIT