Spaces:
Sleeping
Sleeping
File size: 3,500 Bytes
25429d7 6231b6a 25429d7 6231b6a 25429d7 6231b6a 25429d7 6231b6a 25429d7 6231b6a 25429d7 6231b6a 25429d7 6231b6a 25429d7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | from fastapi import APIRouter, File, UploadFile, HTTPException
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
# Create the router instance
router = APIRouter(
prefix='/api/v1',
tags=['Diagnostics']
)
# 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')
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)}')
@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(
{'_id': ObjectId(record_id)},
{'$set': {
'physician_override': override_data.physician_override,
'override_notes': override_data.notes
}}
)
if update_result.modified_count == 0:
raise HTTPException(status_code=404, detail='Diagnostic record not found.')
return {'status': 'success', 'message': 'Feedback recorded. A/B metrics updated.'}
except Exception as e:
raise HTTPException(status_code=500, detail=f'Database update failed: {str(e)}') |