Engineering a Clinical-Grade AI Pipeline: From PyTorch to Production with an MLOps Flywheel

The lifecycle of a Machine Learning model shouldn’t end in a Jupyter Notebook. While achieving high validation accuracy on a static dataset is a great milestone, real-world value is only created when an AI system is accessible, explainable, and capable of continuous learning.

🚨 HIRING: Tech Talent 💰 $50–$120/hr | 🔥 Multiple Roles

Frontend • Backend • Full Stack • Mobile • AI/ML • DevOps 👉 **Apply Here**

Recently, I set out to bridge the gap between deep learning research and software engineering by building Pneumonia Vision AI — an automated, end-to-end screening tool for pediatric chest X-Rays.

My goal was not just to train a classifier, but to engineer a complete MLOps ecosystem. Here is a deep dive into the architecture, the code, and how I built a “Data Flywheel” to ensure the model gets smarter with every clinical interaction.

1. The Brain: Vision Pretraining with DenseNet121

Training deep convolutional networks from scratch on medical imagery is notoriously difficult due to limited dataset sizes. To combat this, I leveraged transfer learning using a pretrained PyTorch DenseNet121 architecture.

DenseNets connect each layer to every other layer in a feed-forward fashion. This drastically reduces the vanishing-gradient problem, strengthens feature propagation, and is highly parameter-efficient — making it exceptionally good at picking up the subtle, cloudy opacities associated with pneumonia.

I loaded the pretrained weights, froze the core feature-extraction layers, and replaced the final classification head for our binary task (Normal vs. Pneumonia):

import torch
import torch.nn as nn
from torchvision import models

# Load pretrained DenseNet121
model = models.densenet121(pretrained=True)

# Freeze core layers to prevent catastrophic forgetting
for param in model.parameters():
    param.requires_grad = False

# Replace the classifier for our 2 classes (Normal / Pneumonia)
num_ftrs = model.classifier.in_features
model.classifier = nn.Sequential(
    nn.Linear(num_ftrs, 256),
    nn.ReLU(),
    nn.Dropout(0.3),
    nn.Linear(256, 2) 
    # Notice: No Sigmoid or Softmax here!
)

An Architectural Gotcha: The Missing Softmax You might notice that the final layer stops at nn.Linear(256, 2) without an activation function like nn.Sigmoid() or nn.Softmax(). This is an intentional and highly critical PyTorch design pattern.

When using nn.CrossEntropyLoss to train the model, PyTorch automatically applies LogSoftmax and calculates the Negative Log Likelihood Loss under the hood. If I had manually added a Softmax layer in the model definition, CrossEntropyLoss would have applied it a second time, flattening the gradients and completely stalling the training process. The model must output raw, unnormalized numbers (logits) during training.

2. Algorithmic Optimization with Optuna

Guessing the perfect learning rate, weight decay, or dropout percentage is inefficient. To squeeze the maximum diagnostic accuracy out of the DenseNet architecture, I integrated Optuna, a hyperparameter optimization framework.

Instead of manual grid searches, I set up an Optuna study to algorithmically hunt for the optimal configuration by evaluating the validation loss across multiple trials:

import optuna
import torch.optim as optim

def objective(trial):
    # Suggest hyperparameters for the trial
    lr = trial.suggest_float("lr", 1e-5, 1e-2, log=True)
    optimizer_name = trial.suggest_categorical("optimizer", ["Adam", "RMSprop"])

    # Dynamically assign the optimizer
    optimizer = getattr(optim, optimizer_name)(model.parameters(), lr=lr)

    # Train for a few epochs and return validation accuracy...
    accuracy = train_and_evaluate(model, optimizer)
    return accuracy

study = optuna.create_study(direction="maximize")
study.optimize(objective, n_trials=20)

print("Best hyperparameters found: ", study.best_params)

By letting Optuna navigate the search space, the model converged on a hyperparameter combination that significantly boosted its sensitivity to pediatric pneumonia opacities.

3. Building Trust: Grad-CAM Explainability (XAI)

A “black box” algorithm is unacceptable in healthcare. If an AI predicts an 84% chance of pneumonia, the physician needs to know exactly what pixels influenced that decision.

To solve this, I integrated Grad-CAM (Gradient-weighted Class Activation Mapping). Grad-CAM uses the gradients of the target concept flowing into the final convolutional layer to produce a coarse localization map.

Instead of just returning a number, the model generates a visual colormap (red/yellow hotspots) overlaying the original X-Ray. This transforms the AI from a silent decision-maker into a collaborative diagnostic assistant.

4. The Engine: Serving the Model via FastAPI

Serving heavy PyTorch models requires a robust backend. I engineered a custom asynchronous FastAPI application, containerized via Docker, and hosted on Hugging Face Spaces.

This setup allows the frontend to securely send image data using multipart/form-data. Remember how we left out the Softmax function during training? During inference, we must apply it to convert the raw logits into readable confidence percentages for the UI:

from fastapi import FastAPI, UploadFile, File
import torch
import torch.nn.functional as F
import shutil

app = FastAPI()

@app.post("/api/v1/predict")
async def predict_xray(file: UploadFile = File(...)):
    if not file.content_type.startswith('image/'):
        raise HTTPException(status_code=400, detail='Invalid file format. Please upload an image.')

    try:
        image_bytes = await file.read()

        # Select a model with given weights
        assigned_model = random.choices(['DenseNet121_v1', 'ResNet50_Challenger'],
                        weights=[1, 0],  # To train another model in the future
                        k=1)[0]           # For now, we only have DenseNet121_v1, so it gets 100% of the traffic. This is where A/B testing logic would go in the future.


        ai_result = run_diagnostic_inference(image_bytes)

        original_url = upload_image_to_cloud(image_bytes, folder_name='raw_xrays')
        heatmap_url = upload_image_to_cloud(ai_result['heatmap_bytes'], folder_name='heatmaps')

        # Use Pydantic to validate the data before saving to MongoDB
        record = DiagnosticRecordSchema(
            original_image_url=original_url,
            heatmap_image_url=heatmap_url,
            prediction=ai_result['prediction'],
            confidence_score=ai_result['confidence'],
            probabilities=ai_result['probabilities'],
            model_variant=assigned_model
        )

        # model_dump() converts the Pydantic model safely to a dictionary for MongoDB
        result = await diagnostic_records.insert_one(record.model_dump())

        # Return a dictionary that strictly matches PredictionResponse
        return {
            'status': 'success',
            'record_id': str(result.inserted_id),
            'prediction': ai_result['prediction'],
            'confidence': ai_result['confidence'],
            'heatmap_url': heatmap_url,
            'original_url': original_url
        }
  except Exception as e:
        raise HTTPException(status_code=500, detail=f'Inference pipeline failed: {str(e)}')

5. The Interface: A Modern React & Tailwind Dashboard

The user experience is just as critical as the neural network. I built a lightning-fast React frontend, scaffolded with Vite and styled using the newly released Tailwind CSS v4.

Deployed globally on Vercel, the dashboard features a medical-grade drag-and-drop interface. When a scan is uploaded, Axios communicates with the Hugging Face backend, dynamically updating the UI to display the prediction and the Grad-CAM heatmap in real-time.

6. The Secret Sauce: The MLOps Data Flywheel

The most critical architectural feature of Pneumonia Vision AI is what happens after the prediction is made.

Directly below the Grad-CAM heatmap, the physician is presented with a simple A/B testing prompt: “Do you agree with this AI diagnosis?”

If the model makes a mistake, the physician clicks “Disagree.” The React frontend immediately fires a PUT request back to the FastAPI server, logging the interaction directly into a MongoDB database using the unique record_id.

// React Frontend: Capturing Clinical Edge Cases
const submitFeedback = async (isAgree) => {
    try {
        const payload = {
            physician_override: !isAgree, // True if the AI was wrong
            notes: isAgree ? "Agreed with AI" : "Physician disagreed with AI diagnosis"
        };

        // Update the MongoDB record via FastAPI
        await axios.put(`${API_URL}/api/v1/override/${result.record_id}`, payload);
        setFeedbackStatus('success');

    } catch (err) {
        console.error("Feedback Loop Error:", err);
    }
};

This creates a true Data Flywheel. Over time, I can query the database for every single scan where physician_override == true. That highly specific subset of edge-cases, false positives, and anomalies becomes the exact dataset I will use to train the next, more resilient iteration of the DenseNet model.

The Road Ahead

Building Pneumonia Vision AI was an incredible challenge in connecting the dots between deep learning layers, hyperparameter tuning with Optuna, API routing, web interfaces, and database management. By combining pretrained architectures with modern full-stack engineering, we can create tools that are not only highly accurate but are built to continuously learn in production.

• Live Demo: https://precision-diagnostics-xai.vercel.app/
• GitHub Repository: https://github.com/mobadara/portfolio-frontend

Thank you for being a part of the community

Before you go:

👉 Be sure to clap and follow the writer ️👏️️

👉 Follow us: **Linkedin| [Medium](https://medium.com/codetodeploy)**

👉 CodeToDeploy Tech Community is live on Discord — **Join now!**

Disclosure: This post includes affiliate and partnership links.