simulated A/B testing

backend/ml_model.py

@@ -58,6 +58,72 @@ except Exception as e:
 
 
 # Image processing transformation pipeline
+transform_pipeline = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+])
 
+def run_diagnostic_inference(image_bytes: bytes) -> dict:
+    """
+    uns an image buffer through the DenseNet pipeline, generates diagnostic probabilities,
+    and constructs a visual Grad-CAM interpretability overlay.
+    
+    Args:
+        image_bytes (bytes): The input image in bytes format.
+        
+    Returns:
+        dict: A dictionary containing the predicted class, confidence score, and the Grad-CAM overlay image in bytes format.
+    """
+    if model is None:
+        raise RuntimeError('Model is not loaded. Cannot perform inference.') 
+    
+    # Open image from raw stream byte and convert to standard RGB format
+    image = Image.open(io.BytesIO(image_bytes)).convert('RGB')
+    # Generate tensor batch for PyTorch model execution
+    tensor_img = transform_pipeline(image).unsqueeze(0).to(DEVICE)
+    # Perform class inference
+    with torch.no_grad():
+        outputs = model(tensor_img)
+        probabilities = torch.softmax(outputs, dim=1).cpu().numpy()[0]
+        
+        # Class 0: Normal  |  Class 1: Pneumonia
+        normal_prob = float(probabilities[0])
+        pneumonia_prob = float(probabilities[1])
+        
+    prediction_class = 'Pneumonia' if pneumonia_prob > normal_prob else 'Normal'
+    confidence_score = max(normal_prob, pneumonia_prob)
+    
+    logging.info(f'Inference completed: Predicted class - {prediction_class}, Confidence score - {confidence_score:.4f}')
+    
+    # Execute Grad-CAM calculation
+    # Target the last dense block layer of DenseNet-121 for tracking applications
+    
+    target_layers = [model.features.denseblock4.denselayer16]
+    # Prepare standard float array of the image for overlay mapping
+    rgb_img_np = np.array(image.resize((224, 224)), dtype=np.float32) / 255.0
+    
+    cam = GradCAM(model=model, target_layers=target_layers)
+    
+    # Target the predicted class index specifically to see its activation layout
+    target_category = 1 if prediction_class == "Pneumonia" else 0
+    targets = [ClassifierOutputTarget(target_category)]
+    
+    # Compute visual grayscale activation map
+    grayscale_cam = cam(input_tensor=tensor_img, targets=targets)[0, :]
+    
+    # Blend the grayscale map into a colored heatmap directly onto our structural input image
+    cam_image_overlay = show_cam_on_image(rgb_img_np, grayscale_cam, use_rgb=True)
+    
+    # Step C: Encode Heatmap into binary buffer for Cloudinary output conversion
+    # Convert RGB array back to standard BGR image array layout for OpenCV conversion
+    cam_image_bgr = cv2.cvtColor((cam_image_overlay * 255).astype(np.uint8), cv2.COLOR_RGB2BGR)
+    _, buffer = cv2.imencode('.jpg', cam_image_bgr)
+    heatmap_bytes = buffer.tobytes()
     
-    
+    return {
+        'prediction': prediction_class,
+        'confidence': confidence_score,
+        'probabilities': {'Normal': normal_prob, 'Pneumonia': pneumonia_prob},
+        'heatmap_bytes': heatmap_bytes
+    }

backend/routers/predict.py

@@ -1,5 +1,7 @@
 from fastapi import APIRouter, File, UploadFile, HTTPException
-from ..schema import PredictionResponse, DiagnosticRecordSchema
+from bson import ObjectId
+import random
+from ..schema import PredictionResponse, DiagnosticRecordSchema, OverrideRequest
 from ..database import upload_image_to_cloud, diagnostic_records
 from ..ml_model import run_diagnostic_inference
 
@@ -12,15 +14,22 @@ router = APIRouter(
 # Enforce strict typing on the response using response_model
 @router.post('/predict', response_model=PredictionResponse)
 async def predict_xray(file: UploadFile = File(...)):
-    '''
+    """
     Analyzes an uploaded X-ray, uploads to Cloudinary, saves to MongoDB, 
     and returns strictly typed results.
-    '''
+    """
     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')
@@ -32,7 +41,8 @@ async def predict_xray(file: UploadFile = File(...)):
             heatmap_image_url=heatmap_url,
             prediction=ai_result['prediction'],
             confidence_score=ai_result['confidence'],
-            probabilities=ai_result['probabilities']
+            probabilities=ai_result['probabilities'],
+            model_variant=assigned_model
         )
 
         # model_dump() converts the Pydantic model safely to a dictionary for MongoDB
@@ -54,9 +64,9 @@ async def predict_xray(file: UploadFile = File(...)):
     
 @router.put('/override/{record_id}')
 async def submit_physician_override(record_id: str, override_data: OverrideRequest):
-    '''
+    """
     Updates a specific diagnostic record with the physician's feedback.
-    '''
+    """
     try:
         # Update the MongoDB document matching the ID
         update_result = await diagnostic_records.update_one(