utils.py

from PIL import Image
import io
import numpy as np
import pydicom
import torch
from fastapi import HTTPException, UploadFile
from skimage.transform import resize


def load_image(image):
    image = image.convert("RGB")
    image_array = np.asarray(image) / 255.0  # Normalize to [0,1]
    image_array = resize(image_array, (224, 224))
    image_tensor = torch.tensor(image_array, dtype=torch.float32).permute(2, 0, 1)  # CxHxW
    
    mean = torch.tensor([0.485, 0.456, 0.406])
    std = torch.tensor([0.229, 0.224, 0.225])
    image_tensor = (image_tensor - mean[:, None, None]) / std[:, None, None]
    
    return image_tensor.unsqueeze(0)  # Add batch dimension

async def convert_to_png(file: UploadFile) -> Image.Image:
    """Converts JPG, PNG, or DICOM to a PNG format"""
    image_data = await file.read()
    if file.content_type in ["image/jpeg", "image/png", "image/jpg"]:
        image = Image.open(io.BytesIO(image_data))
        return image
    
    if file.content_type == "application/dicom" or file.filename.endswith(".dcm") or file.filename.endswith(".dicom"):
        dicom_data = pydicom.dcmread(io.BytesIO(image_data))
        pixel_array = dicom_data.pixel_array

        if pixel_array.dtype != np.uint8:
            pixel_array = (pixel_array / pixel_array.max() * 255).astype(np.uint8)
        
        image = Image.fromarray(pixel_array).convert("RGB")
        return image
    
    raise HTTPException(status_code=400, detail="Unsupported media type")