Upload 18 files

app.py

@@ -0,0 +1,61 @@
+import gradio as gr
+from PIL import Image
+from io import BytesIO
+from src.pipeline import InferencePipeline
+from src.app.config import load_config
+
+# Load configuration and initialize the inference pipeline
+config = load_config()
+inference_pipeline = InferencePipeline(config)
+
+def process_image_from_bytes(file, apply_clahe_postprocess,apply_pre_contrast_adjustment,return_original_size):
+    """
+    Process the image bytes using the inference pipeline.
+
+    Args:
+        file_bytes: The image file in bytes.
+        apply_clahe_postprocess: Boolean indicating if CLAHE postprocessing should be applied.
+
+    Returns:
+        The processed image.
+    """
+    try:
+        # Perform super-resolution
+        sr_image = inference_pipeline.run(file, apply_pre_contrast_adjustment=apply_pre_contrast_adjustment, apply_clahe_postprocess=apply_clahe_postprocess,return_original_size=return_original_size)
+        return sr_image
+    except Exception as e:
+        return f"An exception occurred: {str(e)}"
+
+# Define the Gradio interface
+def gradio_interface():
+    with gr.Blocks() as demo:
+        gr.Markdown("""
+        # X-Ray Image Super-Resolution-Denoiser Demo
+        Provide image bytes to process and optionally apply CLAHE postprocessing.
+        """)
+
+        with gr.Row():
+            file_input = gr.File(label="Upload Image (PNG, JPEG, or DICOM)")
+            apply_clahe_checkbox = gr.Checkbox(label="Apply CLAHE Postprocessing", value=False)
+            apply_pre_contrast_adjustment_checkbox = gr.Checkbox(label="Apply PreContrast Adjustment", value=False)
+            return_original_size_checkbox = gr.Checkbox(label="Return Original Size", value=True)
+
+        process_button = gr.Button("Process Image")
+        output_image = gr.Image(label="Processed Image")
+
+
+
+        process_button.click(
+            process_image_from_bytes,
+            inputs=[file_input, apply_clahe_checkbox,apply_pre_contrast_adjustment_checkbox,return_original_size_checkbox],
+            outputs=output_image
+        )
+
+    return demo
+
+# Launch the Gradio interface
+demo = gradio_interface()
+
+demo.launch(
+    debug=True,
+)

configs/config.yaml

@@ -0,0 +1,21 @@
+model:
+  source: "huggingface"        # Options: "huggingface" or "local"
+  repo_id: "SerdarHelli/super_res_xray"   # Required if source is "huggingface"
+  filename: "net_g.pth" # Model weights filename in HF repo
+  weights: "/path/to/weights.pth"  # Optional if using local weights
+  scale: 4
+  device: "cuda"  # Options: "cuda", "cpu"
+
+preprocessing:
+  unsharping_mask:
+      kernel_size : 7
+      strength: 2
+  brightness:
+      factor : 1.2
+      
+postprocessing:
+  clahe:
+    clipLimit: 2
+    tileGridSize: 
+      - 16
+      - 16

requirements.txt

@@ -0,0 +1,12 @@
+numpy == 1.26.4
+opencv-python == 4.10.0.84
+Pillow == 11.0.0
+torch == 2.5.1
+torchvision == 0.20.1
+tqdm == 4.67.1
+pydicom == 3.0.1
+fastapi == 0.115.6
+uvicorn == 0.34.0
+scikit-image == 0.25.0
+python-multipart == 0.0.20
+huggingface-hub == 0.25.2

src/__init__.py

@@ -0,0 +1,3 @@
+from .network.model import RealESRGAN
+from .pipeline import InferencePipeline
+from .preprocess import *

src/app/config.py

@@ -0,0 +1,20 @@
+import yaml
+
+def load_config(config_path="configs/config.yaml"):
+    """
+    Load the configuration from a YAML file.
+
+    Args:
+        config_path: Path to the configuration file.
+
+    Returns:
+        dict: Configuration dictionary.
+    """
+    try:
+        with open(config_path, "r") as file:
+            config = yaml.safe_load(file)
+        return config
+    except FileNotFoundError:
+        raise Exception(f"Configuration file '{config_path}' not found.")
+    except yaml.YAMLError as e:
+        raise Exception(f"Error parsing configuration file: {e}")

src/app/exceptions.py

@@ -0,0 +1,31 @@
+class ModelLoadError(Exception):
+    """Raised when the model fails to load."""
+    def __init__(self, message="Failed to load the model."):
+        self.message = message
+        super().__init__(self.message)
+
+
+class PreprocessingError(Exception):
+    """Raised when an error occurs during preprocessing."""
+    def __init__(self, message="Error during image preprocessing."):
+        self.message = message
+        super().__init__(self.message)
+
+class PostprocessingError(Exception):
+    """Raised when an error occurs during postprocessing."""
+    def __init__(self, message="Error during image postprocessing."):
+        self.message = message
+        super().__init__(self.message)
+
+class InferenceError(Exception):
+    """Raised when an error occurs during inference."""
+    def __init__(self, message="Error during inference."):
+        self.message = message
+        super().__init__(self.message)
+
+
+class InputError(Exception):
+    """Raised when an error occurs during loading input."""
+    def __init__(self, message="Error loading input."):
+        self.message = message
+        super().__init__(self.message)

src/app/main.py

@@ -0,0 +1,83 @@
+from fastapi import FastAPI,Request,status
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse
+from src.app.routes import inference
+from src.app.exceptions import ModelLoadError, PreprocessingError, InferenceError,InputError, PostprocessingError
+import torch
+import os
+import sys
+
+app = FastAPI(title="Super Resolution Dental X-ray API", version="1.0.0")
+app.add_middleware(CORSMiddleware, allow_origins=["*"])
+
+# Include API routes
+app.include_router(inference.router, prefix="/inference", tags=["Inference"])
+
+@app.get("/")
+def read_root():
+    return {"message": "Welcome to the Super Resolution Dental X-ray API"}
+
+@app.get("/health", tags=["Health"])
+async def health_check():
+    """
+    Health check endpoint to ensure the API and CUDA are running.
+
+    Returns:
+        dict: Status message indicating the API and CUDA availability.
+    """
+    def bash(command):
+        return os.popen(command).read()
+
+    # Check CUDA status
+
+    # Construct response
+    return {
+        "status": "Healthy",
+        "message": "API is running successfully.",
+        "cuda": {
+            "sys.version": sys.version,
+            "torch.__version__": torch.__version__,
+            "torch.cuda.is_available()": torch.cuda.is_available(),
+            "torch.version.cuda": torch.version.cuda,
+            "torch.backends.cudnn.version()": torch.backends.cudnn.version(),
+            "torch.backends.cudnn.enabled": torch.backends.cudnn.enabled,
+            "nvidia-smi": bash('nvidia-smi')
+        }
+    }
+
+# Custom exception handlers
+@app.exception_handler(ModelLoadError)
+async def model_load_error_handler(request: Request, exc: ModelLoadError):
+    return JSONResponse(
+        status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+        content={"error": "ModelLoadError", "message": exc.message},
+    )
+
+@app.exception_handler(PreprocessingError)
+async def preprocessing_error_handler(request: Request, exc: PreprocessingError):
+    return JSONResponse(
+        status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+        content={"error": "PreprocessingError", "message": exc.message},
+    )
+
+@app.exception_handler(PostprocessingError)
+async def postprocessing_error_handler(request: Request, exc: PostprocessingError):
+    return JSONResponse(
+        status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+        content={"error": "PostprocessingError", "message": exc.message},
+    )
+
+@app.exception_handler(InferenceError)
+async def inference_error_handler(request: Request, exc: InferenceError):
+    return JSONResponse(
+        status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+        content={"error": "InferenceError", "message": exc.message},
+    )
+
+
+@app.exception_handler(InputError)
+async def input_load_error_handler(request: Request, exc: InputError):
+    return JSONResponse(
+        status_code=status.HTTP_400_BAD_REQUEST,
+        content={"error": "InputError", "message": exc.message},
+    )

src/app/routes/inference.py

@@ -0,0 +1,87 @@
+from fastapi import APIRouter, UploadFile, File
+from io import BytesIO
+from src.app.config import load_config
+from src.pipeline import InferencePipeline
+# Define the router
+router = APIRouter()
+
+# Load configuration
+config = load_config()
+inference_pipeline = InferencePipeline(config)
+
+from fastapi import APIRouter, UploadFile, File, HTTPException
+from io import BytesIO
+import os
+from fastapi import HTTPException
+from fastapi.responses import FileResponse
+
+@router.post("/predict")
+async def process_image(
+    file: UploadFile = File(...),
+    apply_clahe_postprocess: bool = False,
+    apply_pre_contrast_adjustment: bool = True,
+    return_original_size: bool = True
+
+):
+    """
+    API endpoint to process and super-resolve an image.
+
+    Args:
+        file: Image file to process (PNG, JPEG, or DICOM).
+        apply_clahe_postprocess: Boolean indicating if CLAHE should be applied post-processing.
+
+    Returns:
+        FileResponse: Processed image file or error message.
+    """
+    try:
+
+        # Validate apply_clahe_postprocess parameter
+        if not isinstance(apply_clahe_postprocess, bool):
+            raise HTTPException(
+                status_code=400,
+                detail="The 'apply_clahe_postprocess' parameter must be a boolean."
+            )
+        if not isinstance(apply_pre_contrast_adjustment, bool):
+            raise HTTPException(
+                status_code=400,
+                detail="The 'apply_pre_contrast_adjustment' parameter must be a boolean."
+            )
+        if not isinstance(return_original_size, bool):
+                raise HTTPException(
+                    status_code=400,
+                    detail="The 'return_original_size' parameter must be a boolean."
+                )
+
+        # Read the uploaded file into memory
+        file_bytes = await file.read()
+
+        # Perform inference with the pipeline
+        sr_image = inference_pipeline.run(BytesIO(file_bytes), apply_clahe_postprocess=apply_clahe_postprocess, 
+                                          apply_pre_contrast_adjustment = apply_pre_contrast_adjustment,
+                                          return_original_size = return_original_size
+                                          )
+
+        # Save the processed image to a temporary file
+        output_file_path = "output_highres.png"
+        sr_image.save(output_file_path, format="PNG")
+
+        # Return the file as a response
+        return FileResponse(
+            path=output_file_path,
+            media_type="image/png",
+            filename="processed_image.png"
+        )
+
+    except HTTPException as e:
+        raise e
+
+    except Exception as e:
+        raise HTTPException(
+            status_code=500,
+            detail=f"An error occurred during processing: {str(e)}"
+        )
+    finally:
+        # Cleanup temporary file if it exists
+        if os.path.exists("output_highres.png"):
+            os.remove("output_highres.png")
+

src/network/arch_utils.py

@@ -0,0 +1,197 @@
+import math
+import torch
+from torch import nn as nn
+from torch.nn import functional as F
+from torch.nn import init as init
+from torch.nn.modules.batchnorm import _BatchNorm
+
+@torch.no_grad()
+def default_init_weights(module_list, scale=1, bias_fill=0, **kwargs):
+    """Initialize network weights.
+
+    Args:
+        module_list (list[nn.Module] | nn.Module): Modules to be initialized.
+        scale (float): Scale initialized weights, especially for residual
+            blocks. Default: 1.
+        bias_fill (float): The value to fill bias. Default: 0
+        kwargs (dict): Other arguments for initialization function.
+    """
+    if not isinstance(module_list, list):
+        module_list = [module_list]
+    for module in module_list:
+        for m in module.modules():
+            if isinstance(m, nn.Conv2d):
+                init.kaiming_normal_(m.weight, **kwargs)
+                m.weight.data *= scale
+                if m.bias is not None:
+                    m.bias.data.fill_(bias_fill)
+            elif isinstance(m, nn.Linear):
+                init.kaiming_normal_(m.weight, **kwargs)
+                m.weight.data *= scale
+                if m.bias is not None:
+                    m.bias.data.fill_(bias_fill)
+            elif isinstance(m, _BatchNorm):
+                init.constant_(m.weight, 1)
+                if m.bias is not None:
+                    m.bias.data.fill_(bias_fill)
+
+
+def make_layer(basic_block, num_basic_block, **kwarg):
+    """Make layers by stacking the same blocks.
+
+    Args:
+        basic_block (nn.module): nn.module class for basic block.
+        num_basic_block (int): number of blocks.
+
+    Returns:
+        nn.Sequential: Stacked blocks in nn.Sequential.
+    """
+    layers = []
+    for _ in range(num_basic_block):
+        layers.append(basic_block(**kwarg))
+    return nn.Sequential(*layers)
+
+
+class ResidualBlockNoBN(nn.Module):
+    """Residual block without BN.
+
+    It has a style of:
+        ---Conv-ReLU-Conv-+-
+         |________________|
+
+    Args:
+        num_feat (int): Channel number of intermediate features.
+            Default: 64.
+        res_scale (float): Residual scale. Default: 1.
+        pytorch_init (bool): If set to True, use pytorch default init,
+            otherwise, use default_init_weights. Default: False.
+    """
+
+    def __init__(self, num_feat=64, res_scale=1, pytorch_init=False):
+        super(ResidualBlockNoBN, self).__init__()
+        self.res_scale = res_scale
+        self.conv1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1, bias=True)
+        self.conv2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1, bias=True)
+        self.relu = nn.ReLU(inplace=True)
+
+        if not pytorch_init:
+            default_init_weights([self.conv1, self.conv2], 0.1)
+
+    def forward(self, x):
+        identity = x
+        out = self.conv2(self.relu(self.conv1(x)))
+        return identity + out * self.res_scale
+
+
+class Upsample(nn.Sequential):
+    """Upsample module.
+
+    Args:
+        scale (int): Scale factor. Supported scales: 2^n and 3.
+        num_feat (int): Channel number of intermediate features.
+    """
+
+    def __init__(self, scale, num_feat):
+        m = []
+        if (scale & (scale - 1)) == 0:  # scale = 2^n
+            for _ in range(int(math.log(scale, 2))):
+                m.append(nn.Conv2d(num_feat, 4 * num_feat, 3, 1, 1))
+                m.append(nn.PixelShuffle(2))
+        elif scale == 3:
+            m.append(nn.Conv2d(num_feat, 9 * num_feat, 3, 1, 1))
+            m.append(nn.PixelShuffle(3))
+        else:
+            raise ValueError(f'scale {scale} is not supported. ' 'Supported scales: 2^n and 3.')
+        super(Upsample, self).__init__(*m)
+
+
+def flow_warp(x, flow, interp_mode='bilinear', padding_mode='zeros', align_corners=True):
+    """Warp an image or feature map with optical flow.
+
+    Args:
+        x (Tensor): Tensor with size (n, c, h, w).
+        flow (Tensor): Tensor with size (n, h, w, 2), normal value.
+        interp_mode (str): 'nearest' or 'bilinear'. Default: 'bilinear'.
+        padding_mode (str): 'zeros' or 'border' or 'reflection'.
+            Default: 'zeros'.
+        align_corners (bool): Before pytorch 1.3, the default value is
+            align_corners=True. After pytorch 1.3, the default value is
+            align_corners=False. Here, we use the True as default.
+
+    Returns:
+        Tensor: Warped image or feature map.
+    """
+    assert x.size()[-2:] == flow.size()[1:3]
+    _, _, h, w = x.size()
+    # create mesh grid
+    grid_y, grid_x = torch.meshgrid(torch.arange(0, h).type_as(x), torch.arange(0, w).type_as(x))
+    grid = torch.stack((grid_x, grid_y), 2).float()  # W(x), H(y), 2
+    grid.requires_grad = False
+
+    vgrid = grid + flow
+    # scale grid to [-1,1]
+    vgrid_x = 2.0 * vgrid[:, :, :, 0] / max(w - 1, 1) - 1.0
+    vgrid_y = 2.0 * vgrid[:, :, :, 1] / max(h - 1, 1) - 1.0
+    vgrid_scaled = torch.stack((vgrid_x, vgrid_y), dim=3)
+    output = F.grid_sample(x, vgrid_scaled, mode=interp_mode, padding_mode=padding_mode, align_corners=align_corners)
+
+    # TODO, what if align_corners=False
+    return output
+
+
+def resize_flow(flow, size_type, sizes, interp_mode='bilinear', align_corners=False):
+    """Resize a flow according to ratio or shape.
+
+    Args:
+        flow (Tensor): Precomputed flow. shape [N, 2, H, W].
+        size_type (str): 'ratio' or 'shape'.
+        sizes (list[int | float]): the ratio for resizing or the final output
+            shape.
+            1) The order of ratio should be [ratio_h, ratio_w]. For
+            downsampling, the ratio should be smaller than 1.0 (i.e., ratio
+            < 1.0). For upsampling, the ratio should be larger than 1.0 (i.e.,
+            ratio > 1.0).
+            2) The order of output_size should be [out_h, out_w].
+        interp_mode (str): The mode of interpolation for resizing.
+            Default: 'bilinear'.
+        align_corners (bool): Whether align corners. Default: False.
+
+    Returns:
+        Tensor: Resized flow.
+    """
+    _, _, flow_h, flow_w = flow.size()
+    if size_type == 'ratio':
+        output_h, output_w = int(flow_h * sizes[0]), int(flow_w * sizes[1])
+    elif size_type == 'shape':
+        output_h, output_w = sizes[0], sizes[1]
+    else:
+        raise ValueError(f'Size type should be ratio or shape, but got type {size_type}.')
+
+    input_flow = flow.clone()
+    ratio_h = output_h / flow_h
+    ratio_w = output_w / flow_w
+    input_flow[:, 0, :, :] *= ratio_w
+    input_flow[:, 1, :, :] *= ratio_h
+    resized_flow = F.interpolate(
+        input=input_flow, size=(output_h, output_w), mode=interp_mode, align_corners=align_corners)
+    return resized_flow
+
+
+# TODO: may write a cpp file
+def pixel_unshuffle(x, scale):
+    """ Pixel unshuffle.
+
+    Args:
+        x (Tensor): Input feature with shape (b, c, hh, hw).
+        scale (int): Downsample ratio.
+
+    Returns:
+        Tensor: the pixel unshuffled feature.
+    """
+    b, c, hh, hw = x.size()
+    out_channel = c * (scale**2)
+    assert hh % scale == 0 and hw % scale == 0
+    h = hh // scale
+    w = hw // scale
+    x_view = x.view(b, c, h, scale, w, scale)
+    return x_view.permute(0, 1, 3, 5, 2, 4).reshape(b, out_channel, h, w)

src/network/model.py

@@ -0,0 +1,65 @@
+import os
+import torch
+from PIL import Image
+import numpy as np
+
+from .rrdbnet_arch import RRDBNet
+from .utils import pad_reflect, split_image_into_overlapping_patches, stich_together, \
+                   unpad_image
+
+
+
+
+class RealESRGAN:
+    def __init__(self, device, scale=4):
+        self.device = device
+        self.scale = scale
+        self.model = RRDBNet(
+            num_in_ch=3, num_out_ch=3, num_feat=64, 
+            num_block=23, num_grow_ch=32, scale=scale
+        )
+        
+    def load_weights(self, model_path):
+        
+        loadnet = torch.load(model_path)
+        if 'params' in loadnet:
+            self.model.load_state_dict(loadnet['params'], strict=True)
+        elif 'params_ema' in loadnet:
+            self.model.load_state_dict(loadnet['params_ema'], strict=True)
+        else:
+            self.model.load_state_dict(loadnet, strict=True)
+        self.model.eval()
+        self.model.to(self.device)
+        
+    @torch.cuda.amp.autocast()
+    def predict(self, lr_image, batch_size=4, patches_size=192,
+                padding=24, pad_size=15):
+        scale = self.scale
+        device = self.device
+        lr_image = np.array(lr_image)
+        lr_image = pad_reflect(lr_image, pad_size)
+
+        patches, p_shape = split_image_into_overlapping_patches(
+            lr_image, patch_size=patches_size, padding_size=padding
+        )
+        img = torch.FloatTensor(patches/255).permute((0,3,1,2)).to(device).detach()
+
+        with torch.no_grad():
+            res = self.model(img[0:batch_size])
+            for i in range(batch_size, img.shape[0], batch_size):
+                res = torch.cat((res, self.model(img[i:i+batch_size])), 0)
+
+        sr_image = res.permute((0,2,3,1)).clamp_(0, 1).cpu()
+        np_sr_image = sr_image.numpy()
+
+        padded_size_scaled = tuple(np.multiply(p_shape[0:2], scale)) + (3,)
+        scaled_image_shape = tuple(np.multiply(lr_image.shape[0:2], scale)) + (3,)
+        np_sr_image = stich_together(
+            np_sr_image, padded_image_shape=padded_size_scaled, 
+            target_shape=scaled_image_shape, padding_size=padding * scale
+        )
+        sr_img = (np_sr_image*255).astype(np.uint8)
+        sr_img = unpad_image(sr_img, pad_size*scale)
+        sr_img = Image.fromarray(sr_img)
+
+        return sr_img

src/network/rrdbnet_arch.py

@@ -0,0 +1,121 @@
+import torch
+from torch import nn as nn
+from torch.nn import functional as F
+
+from .arch_utils import default_init_weights, make_layer, pixel_unshuffle
+
+
+class ResidualDenseBlock(nn.Module):
+    """Residual Dense Block.
+
+    Used in RRDB block in ESRGAN.
+
+    Args:
+        num_feat (int): Channel number of intermediate features.
+        num_grow_ch (int): Channels for each growth.
+    """
+
+    def __init__(self, num_feat=64, num_grow_ch=32):
+        super(ResidualDenseBlock, self).__init__()
+        self.conv1 = nn.Conv2d(num_feat, num_grow_ch, 3, 1, 1)
+        self.conv2 = nn.Conv2d(num_feat + num_grow_ch, num_grow_ch, 3, 1, 1)
+        self.conv3 = nn.Conv2d(num_feat + 2 * num_grow_ch, num_grow_ch, 3, 1, 1)
+        self.conv4 = nn.Conv2d(num_feat + 3 * num_grow_ch, num_grow_ch, 3, 1, 1)
+        self.conv5 = nn.Conv2d(num_feat + 4 * num_grow_ch, num_feat, 3, 1, 1)
+
+        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+
+        # initialization
+        default_init_weights([self.conv1, self.conv2, self.conv3, self.conv4, self.conv5], 0.1)
+
+    def forward(self, x):
+        x1 = self.lrelu(self.conv1(x))
+        x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1)))
+        x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1)))
+        x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1)))
+        x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))
+        # Emperically, we use 0.2 to scale the residual for better performance
+        return x5 * 0.2 + x
+
+
+class RRDB(nn.Module):
+    """Residual in Residual Dense Block.
+
+    Used in RRDB-Net in ESRGAN.
+
+    Args:
+        num_feat (int): Channel number of intermediate features.
+        num_grow_ch (int): Channels for each growth.
+    """
+
+    def __init__(self, num_feat, num_grow_ch=32):
+        super(RRDB, self).__init__()
+        self.rdb1 = ResidualDenseBlock(num_feat, num_grow_ch)
+        self.rdb2 = ResidualDenseBlock(num_feat, num_grow_ch)
+        self.rdb3 = ResidualDenseBlock(num_feat, num_grow_ch)
+
+    def forward(self, x):
+        out = self.rdb1(x)
+        out = self.rdb2(out)
+        out = self.rdb3(out)
+        # Emperically, we use 0.2 to scale the residual for better performance
+        return out * 0.2 + x
+
+
+class RRDBNet(nn.Module):
+    """Networks consisting of Residual in Residual Dense Block, which is used
+    in ESRGAN.
+
+    ESRGAN: Enhanced Super-Resolution Generative Adversarial Networks.
+
+    We extend ESRGAN for scale x2 and scale x1.
+    Note: This is one option for scale 1, scale 2 in RRDBNet.
+    We first employ the pixel-unshuffle (an inverse operation of pixelshuffle to reduce the spatial size
+    and enlarge the channel size before feeding inputs into the main ESRGAN architecture.
+
+    Args:
+        num_in_ch (int): Channel number of inputs.
+        num_out_ch (int): Channel number of outputs.
+        num_feat (int): Channel number of intermediate features.
+            Default: 64
+        num_block (int): Block number in the trunk network. Defaults: 23
+        num_grow_ch (int): Channels for each growth. Default: 32.
+    """
+
+    def __init__(self, num_in_ch, num_out_ch, scale=4, num_feat=64, num_block=23, num_grow_ch=32):
+        super(RRDBNet, self).__init__()
+        self.scale = scale
+        if scale == 2:
+            num_in_ch = num_in_ch * 4
+        elif scale == 1:
+            num_in_ch = num_in_ch * 16
+        self.conv_first = nn.Conv2d(num_in_ch, num_feat, 3, 1, 1)
+        self.body = make_layer(RRDB, num_block, num_feat=num_feat, num_grow_ch=num_grow_ch)
+        self.conv_body = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
+        # upsample
+        self.conv_up1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
+        self.conv_up2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
+        if scale == 8:
+            self.conv_up3 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
+        self.conv_hr = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
+        self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
+
+        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+
+    def forward(self, x):
+        if self.scale == 2:
+            feat = pixel_unshuffle(x, scale=2)
+        elif self.scale == 1:
+            feat = pixel_unshuffle(x, scale=4)
+        else:
+            feat = x
+        feat = self.conv_first(feat)
+        body_feat = self.conv_body(self.body(feat))
+        feat = feat + body_feat
+        # upsample
+        feat = self.lrelu(self.conv_up1(F.interpolate(feat, scale_factor=2, mode='nearest')))
+        feat = self.lrelu(self.conv_up2(F.interpolate(feat, scale_factor=2, mode='nearest')))
+        if self.scale == 8:
+            feat = self.lrelu(self.conv_up3(F.interpolate(feat, scale_factor=2, mode='nearest')))
+        out = self.conv_last(self.lrelu(self.conv_hr(feat)))
+        return out

src/network/utils.py

@@ -0,0 +1,133 @@
+import numpy as np
+import torch
+from PIL import Image
+import os
+import io
+
+def pad_reflect(image, pad_size):
+    imsize = image.shape
+    height, width = imsize[:2]
+    new_img = np.zeros([height+pad_size*2, width+pad_size*2, imsize[2]]).astype(np.uint8)
+    new_img[pad_size:-pad_size, pad_size:-pad_size, :] = image
+    
+    new_img[0:pad_size, pad_size:-pad_size, :] = np.flip(image[0:pad_size, :, :], axis=0) #top
+    new_img[-pad_size:, pad_size:-pad_size, :] = np.flip(image[-pad_size:, :, :], axis=0) #bottom
+    new_img[:, 0:pad_size, :] = np.flip(new_img[:, pad_size:pad_size*2, :], axis=1) #left
+    new_img[:, -pad_size:, :] = np.flip(new_img[:, -pad_size*2:-pad_size, :], axis=1) #right
+    
+    return new_img
+
+def unpad_image(image, pad_size):
+    return image[pad_size:-pad_size, pad_size:-pad_size, :]
+
+
+def process_array(image_array, expand=True):
+    """ Process a 3-dimensional array into a scaled, 4 dimensional batch of size 1. """
+    
+    image_batch = image_array / 255.0
+    if expand:
+        image_batch = np.expand_dims(image_batch, axis=0)
+    return image_batch
+
+
+def process_output(output_tensor):
+    """ Transforms the 4-dimensional output tensor into a suitable image format. """
+    
+    sr_img = output_tensor.clip(0, 1) * 255
+    sr_img = np.uint8(sr_img)
+    return sr_img
+
+
+def pad_patch(image_patch, padding_size, channel_last=True):
+    """ Pads image_patch with with padding_size edge values. """
+    
+    if channel_last:
+        return np.pad(
+            image_patch,
+            ((padding_size, padding_size), (padding_size, padding_size), (0, 0)),
+            'edge',
+        )
+    else:
+        return np.pad(
+            image_patch,
+            ((0, 0), (padding_size, padding_size), (padding_size, padding_size)),
+            'edge',
+        )
+
+
+def unpad_patches(image_patches, padding_size):
+    return image_patches[:, padding_size:-padding_size, padding_size:-padding_size, :]
+
+
+def split_image_into_overlapping_patches(image_array, patch_size, padding_size=2):
+    """ Splits the image into partially overlapping patches.
+    The patches overlap by padding_size pixels.
+    Pads the image twice:
+        - first to have a size multiple of the patch size,
+        - then to have equal padding at the borders.
+    Args:
+        image_array: numpy array of the input image.
+        patch_size: size of the patches from the original image (without padding).
+        padding_size: size of the overlapping area.
+    """
+    
+    xmax, ymax, _ = image_array.shape
+    x_remainder = xmax % patch_size
+    y_remainder = ymax % patch_size
+    
+    # modulo here is to avoid extending of patch_size instead of 0
+    x_extend = (patch_size - x_remainder) % patch_size
+    y_extend = (patch_size - y_remainder) % patch_size
+    
+    # make sure the image is divisible into regular patches
+    extended_image = np.pad(image_array, ((0, x_extend), (0, y_extend), (0, 0)), 'edge')
+    
+    # add padding around the image to simplify computations
+    padded_image = pad_patch(extended_image, padding_size, channel_last=True)
+    
+    xmax, ymax, _ = padded_image.shape
+    patches = []
+    
+    x_lefts = range(padding_size, xmax - padding_size, patch_size)
+    y_tops = range(padding_size, ymax - padding_size, patch_size)
+    
+    for x in x_lefts:
+        for y in y_tops:
+            x_left = x - padding_size
+            y_top = y - padding_size
+            x_right = x + patch_size + padding_size
+            y_bottom = y + patch_size + padding_size
+            patch = padded_image[x_left:x_right, y_top:y_bottom, :]
+            patches.append(patch)
+    
+    return np.array(patches), padded_image.shape
+
+
+def stich_together(patches, padded_image_shape, target_shape, padding_size=4):
+    """ Reconstruct the image from overlapping patches.
+    After scaling, shapes and padding should be scaled too.
+    Args:
+        patches: patches obtained with split_image_into_overlapping_patches
+        padded_image_shape: shape of the padded image contructed in split_image_into_overlapping_patches
+        target_shape: shape of the final image
+        padding_size: size of the overlapping area.
+    """
+    
+    xmax, ymax, _ = padded_image_shape
+    patches = unpad_patches(patches, padding_size)
+    patch_size = patches.shape[1]
+    n_patches_per_row = ymax // patch_size
+    
+    complete_image = np.zeros((xmax, ymax, 3))
+    
+    row = -1
+    col = 0
+    for i in range(len(patches)):
+        if i % n_patches_per_row == 0:
+            row += 1
+            col = 0
+        complete_image[
+        row * patch_size: (row + 1) * patch_size, col * patch_size: (col + 1) * patch_size,:
+        ] = patches[i]
+        col += 1
+    return complete_image[0: target_shape[0], 0: target_shape[1], :]

src/pipeline.py

@@ -0,0 +1,236 @@
+import torch
+from PIL import Image
+import numpy as np
+from io import BytesIO
+from huggingface_hub import hf_hub_download
+from pathlib import Path
+
+from src.preprocess import read_xray, enhance_exposure, unsharp_masking, apply_clahe, resize_pil_image, increase_brightness
+from src.network.model import RealESRGAN
+from src.app.exceptions import InputError, ModelLoadError, PreprocessingError, InferenceError,PostprocessingError
+
+class ModelLoadError(Exception):
+    pass
+
+class InferencePipeline:
+    def __init__(self, config):
+        """
+        Initialize the inference pipeline using configuration.
+
+        Args:
+            config: Configuration dictionary.
+        """
+        self.config = config
+        self.device = config["model"].get("device", "cuda" if torch.cuda.is_available() else "cpu")
+        self.scale = config["model"].get("scale", 4)
+
+        model_source = config["model"].get("source", "local")
+        self.model = RealESRGAN(self.device, scale=self.scale)
+
+        print(f"Using device: {self.device}")
+        
+        try:
+            if model_source == "huggingface":
+                repo_id = config["model"]["repo_id"]
+                filename = config["model"]["filename"]
+                local_path = hf_hub_download(repo_id=repo_id, filename=filename)
+                self.load_weights(local_path)
+            else:
+                local_path = config["model"]["weights"]
+                self.load_weights(local_path)
+        except Exception as e:
+            raise ModelLoadError(f"Failed to load the model: {str(e)}")
+
+    def load_weights(self, model_weights):
+        """
+        Load the model weights.
+
+        Args:
+            model_weights: Path to the model weights file.
+        """
+        try:
+            self.model.load_weights(model_weights)
+        except FileNotFoundError:
+            raise ModelLoadError(f"Model weights not found at '{model_weights}'.")
+        except Exception as e:
+            raise ModelLoadError(f"Error loading weights: {str(e)}")
+    def preprocess(self, image_path_or_bytes, apply_pre_contrast_adjustment=True, is_dicom=False):
+        """
+        Preprocess the input image.
+
+        Args:
+            image_path: Path to the input image file.
+            is_dicom: Boolean indicating if the input is a DICOM file.
+
+        Returns: 
+            PIL Image: Preprocessed image.
+        """
+        try:
+            if is_dicom:
+                img = read_xray(image_path_or_bytes)
+            else:
+                img = Image.open(image_path_or_bytes)
+            
+            if apply_pre_contrast_adjustment:
+                img = enhance_exposure(np.array(img))
+                
+            if isinstance(img,np.ndarray):
+                img = Image.fromarray(((img / np.max(img))*255).astype(np.uint8))   
+                
+            if img.mode not in ['RGB']:
+                img = img.convert('RGB')
+                
+            img = unsharp_masking(
+                img,
+                self.config["preprocessing"]["unsharping_mask"].get("kernel_size", 7),
+                self.config["preprocessing"]["unsharping_mask"].get("strength", 2)
+            )
+            img = increase_brightness(
+                        img,
+                        self.config["preprocessing"]["brightness"].get("factor", 1.2),
+                    )
+
+            
+            if img.mode not in ['RGB']:
+                img = img.convert('RGB')
+                
+        
+            return img, img.size
+        except Exception as e:
+            raise PreprocessingError(f"Error during preprocessing: {str(e)}")
+
+    def postprocess(self, image_array):
+        """
+        Postprocess the output from the model.
+
+        Args:
+            image_array: PIL.Image output from the model.
+
+        Returns:
+            PIL Image: Postprocessed image.
+        """
+        try:
+            return apply_clahe(
+                image_array,
+                self.config["postprocessing"]["clahe"].get("clipLimit", 2.0),
+                tuple(self.config["postprocessing"]["clahe"].get("tileGridSize", [16, 16]))
+            )
+        except Exception as e:
+            raise PostprocessingError(f"Error during postprocessing: {str(e)}")
+
+    def is_dicom(self, file_path_or_bytes):
+        """
+        Check if the input file is a DICOM file.
+
+        Args:
+            file_path_or_bytes (str or bytes or BytesIO): Path to the file, byte content, or BytesIO object.
+
+        Returns:
+            bool: True if the file is a DICOM file, False otherwise.
+        """
+        try:
+            if isinstance(file_path_or_bytes, str):
+                # Check the file extension
+                file_extension = Path(file_path_or_bytes).suffix.lower()
+                if file_extension in ['.dcm', '.dicom']:
+                    return True
+
+                # Open the file and check the header
+                with open(file_path_or_bytes, 'rb') as file:
+                    header = file.read(132)
+                    return header[-4:] == b'DICM'
+
+            elif isinstance(file_path_or_bytes, BytesIO):
+                file_path_or_bytes.seek(0)
+                header = file_path_or_bytes.read(132)
+                file_path_or_bytes.seek(0)  # Reset the stream position
+                return header[-4:] == b'DICM'
+
+            elif isinstance(file_path_or_bytes, bytes):
+                header = file_path_or_bytes[:132]
+                return header[-4:] == b'DICM'
+
+        except Exception as e:
+            print(f"Error during DICOM validation: {e}")
+            return False
+
+        return False
+        
+    def validate_input(self, input_data):
+        """
+        Validate the input data to ensure it is suitable for processing.
+
+        Args:
+            input_data: Path to the input file, bytes content, or BytesIO object.
+
+        Returns:
+            bool: True if the input is valid, raises InputError otherwise.
+        """
+        if isinstance(input_data, str):
+            # Check if the file exists
+            if not Path(input_data).exists():
+                raise InputError(f"Input file '{input_data}' does not exist.")
+
+            # Check if the file type is supported
+            file_extension = Path(input_data).suffix.lower()
+            if file_extension not in ['.png', '.jpeg', '.jpg', '.dcm', '.dicom']:
+                raise InputError(f"Unsupported file type '{file_extension}'. Supported types are PNG, JPEG, and DICOM.")
+        
+        elif isinstance(input_data, BytesIO):
+            # Check if BytesIO data is not empty
+            if input_data.getbuffer().nbytes == 0:
+                raise InputError("Input BytesIO data is empty.")
+
+        else:
+            raise InputError("Unsupported input type. Must be a file path, byte content, or BytesIO object.")
+
+        return True
+    
+    def infer(self, input_image):
+        """
+        Perform inference on a single image.
+
+        Args:
+            input_image: PIL Image to be processed.
+
+        Returns:
+            PIL Image: Super-resolved image.
+        """
+        try:
+            # Perform inference
+            input_array = np.array(input_image)
+            sr_array = self.model.predict(input_array)
+            return sr_array
+        
+        except Exception as e:
+            raise InferenceError(f"Error during inference: {str(e)}")
+        
+    def run(self, input_path,  apply_pre_contrast_adjustment = True, apply_clahe_postprocess=False, return_original_size = True):
+        """
+        Process a single image and save the output.
+
+        Args:
+            input_path: Path to the input image file.
+            is_dicom: Boolean indicating if the input is a DICOM file.
+            apply_clahe_postprocess: Boolean indicating if CLAHE should be applied post-processing.
+        """
+        # Validate the input
+        self.validate_input(input_path)
+
+        is_dicom =self.is_dicom(input_path)
+
+        img, original_size = self.preprocess(input_path, is_dicom=is_dicom, apply_pre_contrast_adjustment = apply_pre_contrast_adjustment)
+
+        if img is None:
+            raise InputError(f"Invalid Input")
+ 
+
+        sr_image = self.infer(img)
+
+        if apply_clahe_postprocess:
+            sr_image = self.postprocess(sr_image)
+
+        if return_original_size:
+            sr_image = resize_pil_image(sr_image, target_shape = original_size)
+            
+        return sr_image

src/preprocess.py

@@ -0,0 +1,187 @@
+import pydicom
+import numpy as np
+from pydicom.pixels import apply_voi_lut
+from skimage import exposure
+from PIL import Image,ImageEnhance
+import cv2
+
+def read_xray(path, voi_lut=True, fix_monochrome=True):
+    """
+    Read and preprocess a DICOM X-ray image.
+
+    Parameters:
+    - path: Path to the DICOM file.
+    - voi_lut: Apply VOI LUT if available.
+    - fix_monochrome: Fix inverted monochrome images.
+
+    Returns:
+    - NumPy array: Preprocessed X-ray image.
+    """
+    dicom = pydicom.dcmread(path)
+
+    # Apply VOI LUT if available
+    if voi_lut:
+        data = apply_voi_lut(dicom.pixel_array, dicom)
+    else:
+        data = dicom.pixel_array
+
+    # Fix inverted monochrome images
+    if fix_monochrome and dicom.PhotometricInterpretation == "MONOCHROME1":
+        data = np.amax(data) - data
+
+    # Normalize data to start from 0
+    data = data - np.min(data)
+
+    return data
+    
+def resize_pil_image(image: Image.Image, target_shape: tuple) -> Image.Image:
+    """
+    Resizes a PIL image based on a target shape.
+
+    Args:
+        image: Input PIL image.
+        target_shape: Desired shape for resizing. It can be a 2D tuple (height, width)
+                       or a 3D tuple (height, width, channels), where channels will be ignored.
+                       
+    Returns:
+        Resized PIL image.
+    """
+    # Convert image to a numpy array
+    np_image = np.array(image)
+
+    # Extract the original height and width from the numpy array
+    height, width = np_image.shape[:2]
+
+    # If the target shape is 2D (height, width)
+    if len(target_shape) == 2:
+        new_width, new_height = target_shape
+    elif len(target_shape) == 3:
+        # If the target shape is 3D (height, width, channels), only change the first two dimensions
+        new_width, new_height = target_shape[:2]
+    else:
+        raise ValueError("Target shape must be either 2D or 3D.")
+
+    # Resize the image using cv2 or PIL's in-built resizing (no channels affected)
+    pil_resized_image = Image.fromarray(np_image).resize((new_width, new_height), Image.LANCZOS)
+
+    return pil_resized_image
+    
+def enhance_exposure(img):
+    """
+    Enhance image exposure using histogram equalization.
+
+    Parameters:
+    - img: Input image as a NumPy array.
+
+    Returns:
+    - PIL.Image: Exposure-enhanced image.
+    """
+    img = exposure.equalize_hist(img)
+    img = exposure.equalize_adapthist(img / np.max(img))
+    img = (img * 255).astype(np.uint8)
+    return Image.fromarray(img)
+
+def unsharp_masking(image, kernel_size=5, strength=0.25):
+    """
+    Apply unsharp masking to enhance image sharpness.
+
+    Parameters:
+    - image: Input image as a NumPy array or PIL.Image.
+    - kernel_size: Size of the Gaussian blur kernel.
+    - strength: Strength of the high-pass filter.
+
+    Returns:
+    - PIL.Image: Sharpened image.
+    """
+    image = np.array(image)
+
+    # Convert to grayscale if needed
+    if len(image.shape) == 3:
+        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+    else:
+        gray = image
+
+    # Apply Gaussian blur and calculate high-pass filter
+    blurred = cv2.GaussianBlur(gray, (kernel_size, kernel_size), 0)
+    high_pass = cv2.subtract(gray, blurred)
+
+    # Combine high-pass with original image
+    sharpened = cv2.addWeighted(gray, 1, high_pass, strength, 0)
+
+    return Image.fromarray(sharpened)
+
+def increase_contrast(image: Image.Image, factor: float) -> Image.Image:
+    """
+    Increases the contrast of the input PIL image by a given factor.
+    
+    Args:
+        image: Input PIL image.
+        factor: Factor by which to increase the contrast. 
+                A factor of 1.0 means no change, values greater than 1.0 increase contrast,
+                values between 0.0 and 1.0 decrease contrast.
+    
+    Returns:
+        Image with increased contrast.
+    """
+    if image.mode not in ['RGB', 'L']:
+        image = image.convert('RGB')
+        
+    enhancer = ImageEnhance.Contrast(image)
+    image_enhanced = enhancer.enhance(factor)
+    return image_enhanced
+
+def increase_brightness(image: Image.Image, factor: float) -> Image.Image:
+    """
+    Increases the brightness of the input PIL image by a given factor.
+    
+    Args:
+        image: Input PIL image.
+        factor: Factor by which to increase the brightness. 
+                A factor of 1.0 means no change, values greater than 1.0 increase brightness,
+                values between 0.0 and 1.0 decrease brightness.
+    
+    Returns:
+        Image with increased brightness.
+    """
+    if image.mode not in ['RGB', 'L']:
+        image = image.convert('RGB')
+    
+    enhancer = ImageEnhance.Brightness(image)
+    image_enhanced = enhancer.enhance(factor)
+    return image_enhanced
+
+def apply_clahe(image, clipLimit=2.0, tileGridSize=(8, 8)):
+    """
+    Apply CLAHE (Contrast Limited Adaptive Histogram Equalization) to an image.
+
+    Parameters:
+    - image: Input image as a PIL.Image.
+    - clipLimit: Threshold for contrast limiting.
+    - tileGridSize: Size of the grid for histogram equalization.
+
+    Returns:
+    - Processed image in the same format as the input (PIL.Image).
+    """
+
+    image_np = np.array(image)
+    image_np = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
+
+
+    # Apply CLAHE based on image type
+    if len(image_np.shape) == 2:
+        # Grayscale image
+        clahe = cv2.createCLAHE(clipLimit=clipLimit, tileGridSize=tileGridSize)
+        processed = clahe.apply(image_np)
+    else:
+        # Color image: Apply CLAHE on the L channel in LAB space
+        lab = cv2.cvtColor(image_np, cv2.COLOR_BGR2LAB)
+        L, A, B = cv2.split(lab)
+
+        clahe = cv2.createCLAHE(clipLimit=clipLimit, tileGridSize=tileGridSize)
+        L_clahe = clahe.apply(L)
+        lab_clahe = cv2.merge((L_clahe, A, B))
+        processed = cv2.cvtColor(lab_clahe, cv2.COLOR_LAB2BGR)
+
+    processed_rgb = cv2.cvtColor(processed, cv2.COLOR_BGR2RGB)
+    return Image.fromarray(processed_rgb)
+

tests/test_inference_pipeline.py

@@ -0,0 +1,204 @@
+import pytest
+from fastapi.testclient import TestClient
+import numpy as np
+from PIL import Image
+from io import BytesIO
+import pydicom
+from pydicom.dataset import Dataset, FileDataset
+import tempfile
+import os
+
+import sys
+from pathlib import Path
+
+# Add the src directory to the Python path
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent ))
+
+from src.app.main import app
+from src.pipeline import InferencePipeline
+
+# Initialize test client
+client = TestClient(app)
+
+@pytest.fixture
+def pipeline_config():
+    return {
+        "model": {
+            "weights": "weights/model.pth",
+            "scale": 4,
+            "device": "cpu"
+        },
+        "preprocessing": {
+            "unsharping_mask": {
+                "kernel_size": 7,
+                "strength": 0.5
+            }
+        },
+        "postprocessing": {
+            "clahe": {
+                "clipLimit": 2,
+                "tileGridSize": [16, 16]
+            }
+        }
+    }
+
+@pytest.fixture
+def pipeline(pipeline_config):
+    return InferencePipeline(pipeline_config)
+
+def create_dummy_dicom():
+    """Create a dummy DICOM file for testing."""
+    meta = Dataset()
+    meta.MediaStorageSOPClassUID = "1.2.840.10008.5.1.4.1.1.2"
+    meta.MediaStorageSOPInstanceUID = "1.2.3"
+    meta.TransferSyntaxUID = pydicom.uid.ExplicitVRLittleEndian
+
+    ds = FileDataset("", {}, file_meta=meta, preamble=b"\x00" * 128)
+    
+    # Required Patient and Image Information
+    ds.PatientName = "Test"
+    ds.PatientID = "12345"
+    ds.Modality = "CT"
+    ds.StudyInstanceUID = "1.2.3.4.5.6.7.8.9.10"
+    ds.SeriesInstanceUID = "1.2.3.4.5.6.7.8.9.11"
+    ds.SOPInstanceUID = "1.2.3.4.5.6.7.8.9.12"
+    ds.StudyDate = "20240101"
+    ds.StudyTime = "120000"
+    ds.Manufacturer = "TestManufacturer"
+
+    # Required Image Data Information
+    ds.PhotometricInterpretation = "MONOCHROME2"
+    ds.Rows = 128
+    ds.Columns = 128
+    ds.BitsAllocated = 16
+    ds.BitsStored = 16  # Add missing Bits Stored
+    ds.HighBit = 15  # Highest bit set
+    ds.PixelRepresentation = 0  # Unsigned integer
+    ds.SamplesPerPixel = 1  # Single-channel (grayscale)
+    ds.PixelData = (np.random.rand(128, 128) * 65535).astype(np.uint16).tobytes()
+
+    # Save to a temporary file
+    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".dcm")
+    ds.save_as(temp_file.name)
+    return temp_file.name
+
+    
+def test_is_dicom(pipeline):
+    dicom_path = create_dummy_dicom()
+    
+    # Test with file path
+    assert pipeline.is_dicom(dicom_path) is True
+    
+    # Test with BytesIO
+    with open(dicom_path, "rb") as f:
+        dicom_bytes = BytesIO(f.read())
+    assert pipeline.is_dicom(dicom_bytes) is True
+
+    # Test with invalid BytesIO (non-DICOM content)
+    non_dicom_bytes = BytesIO()
+    non_dicom_bytes.write(b"\x89PNG\r\n\x1a\n" + b"\x00" * 128)  # Write invalid header
+    non_dicom_bytes.seek(0)
+    assert pipeline.is_dicom(non_dicom_bytes) is False
+
+    os.remove(dicom_path)
+
+
+def test_is_dicom(pipeline):
+    dicom_path = create_dummy_dicom()
+    
+    # Test with file path
+    assert pipeline.is_dicom(dicom_path) is True, "DICOM file path should be recognized as DICOM"
+    
+    # Test with BytesIO
+    with open(dicom_path, "rb") as f:
+        dicom_bytes = BytesIO(f.read())
+    assert pipeline.is_dicom(dicom_bytes) is True, "BytesIO DICOM content should be recognized as DICOM"
+
+
+
+    # Test with invalid BytesIO (non-DICOM content)
+    non_dicom_bytes = BytesIO()
+    non_dicom_bytes.write(b"\x89PNG\r\n\x1a\n" + b"\x00" * 128)  # Write invalid header
+    non_dicom_bytes.seek(0)
+    assert pipeline.is_dicom(non_dicom_bytes) is False, "Non-DICOM BytesIO should not be recognized as DICOM"
+
+    # Test with invalid raw bytes
+    invalid_raw_bytes = b"\x89PNG\r\n\x1a\n" + b"\x00" * 128
+    assert pipeline.is_dicom(invalid_raw_bytes) is False, "Invalid raw bytes should not be recognized as DICOM"
+
+    os.remove(dicom_path)
+
+
+
+def test_preprocess_normal_image(pipeline):
+    # Create a dummy image
+    image = Image.new("RGB", (128, 128), color="red")
+    
+    # Test with BytesIO
+    image_bytes = BytesIO()
+    image.save(image_bytes, format="JPEG")
+    image_bytes.seek(0)
+
+    processed_image_bytes = pipeline.preprocess(image_bytes, is_dicom=False)
+    assert isinstance(processed_image_bytes, Image.Image)
+    
+    # Test with file path
+    temp_image_path = tempfile.NamedTemporaryFile(delete=False, suffix=".jpg").name
+    image.save(temp_image_path)
+    
+    processed_image_path = pipeline.preprocess(temp_image_path, is_dicom=False)
+    assert isinstance(processed_image_path, Image.Image)
+
+    os.remove(temp_image_path)
+
+
+def test_infer(pipeline):
+    # Create a dummy image
+    image = Image.new("RGB", (128, 128), color="red")
+
+    # Perform inference
+    result = pipeline.infer(image)
+    assert isinstance(result, Image.Image)
+
+
+def test_postprocess(pipeline):
+
+    image = Image.new("RGB", (128, 128), color="red")
+    result = pipeline.postprocess(image)
+
+    assert isinstance(result, Image.Image)
+
+
+def test_api_predict_normal_image():
+    # Create a dummy image
+    image = Image.new("RGB", (128, 128), color="red")
+    image_bytes = BytesIO()
+    image.save(image_bytes, format="JPEG")
+    image_bytes.seek(0)
+
+    response = client.post(
+        "/inference/predict",  # Adjusted to include the prefix
+        files={"file": ("test.jpg", image_bytes, "image/jpeg")},
+        data={"apply_clahe_postprocess": "false"}  # Ensure proper boolean conversion
+    )
+    assert response.status_code == 200, response.text
+    assert response.headers["content-type"] == "image/png"
+
+
+def test_api_predict_dicom():
+    dicom_path = create_dummy_dicom()
+    
+    # Use BytesIO for testing
+    with open(dicom_path, "rb") as f:
+        dicom_bytes = BytesIO(f.read())
+
+    response = client.post(
+        "/inference/predict",  # Adjusted to include the prefix
+        files={"file": ("test.dcm", dicom_bytes, "application/dicom")},
+        data={"apply_clahe_postprocess": "false"}  # Ensure proper boolean conversion
+    )
+    assert response.status_code == 200, response.text
+    assert response.headers["content-type"] == "image/png"
+
+    os.remove(dicom_path)
+