added fastapi app and config scripts

app.py

@@ -0,0 +1,147 @@
+import cv2
+import json
+import torch
+import base64
+import logging
+import numpy as np
+from fastapi import FastAPI, File, UploadFile, Form
+
+from config import settings
+
+from image_colourization_cgan.image_utils import *
+from image_colourization_cgan.model import ImageToImageConditionalGAN
+
+
+def activate_dropout(m):
+    if type(m) == torch.nn.Dropout:
+        m.train()
+    return
+
+
+app = FastAPI()
+logging.basicConfig(level=logging.INFO)
+
+device = settings.device
+image_size = settings.image_size
+
+file_model_local = f"./artifacts/colorizer_cgan_90.pt"
+file_model_cont = f"/data/models/colorizer_cgan_90.pt"
+
+colour_model_cgan = ImageToImageConditionalGAN(device)
+colour_model_cgan.eval()
+
+try:
+    logging.info(f"loading model from {file_model_local}")
+    colour_model_cgan.load_state_dict(torch.load(file_model_local, map_location=device))
+except:
+    logging.info(f"loading model from {file_model_cont}")
+    colour_model_cgan.load_state_dict(torch.load(file_model_cont, map_location=device))
+colour_model_cgan.to(device)
+colour_model_cgan.net_gen.apply(activate_dropout)
+
+
+def get_prediction(img_arr: np.ndarray) -> np.ndarray:
+    """
+    ---------
+    Arguments
+    ---------
+    img_arr: ndarray
+        a numpy array of the image
+
+    -------
+    Returns
+    -------
+    img_gen_rgb : ndarray
+        a numpy representing the generated colourized image
+    """
+
+    img_gray_resized = resize_image(img_arr, (image_size, image_size))
+    # resized grayscale is in [0, 1]
+
+    img_l = rescale_grayscale_image_l_channel(img_gray_resized)
+    # L channel is in [0, 100]
+
+    # apply pre-processing on L channel image
+    img_l_preprocessed = apply_image_l_pre_processing(img_l)
+
+    # repeat L channel 3 times because ResNet needs a 3 channel input
+    img_l_preprocessed = np.repeat(
+        np.expand_dims(img_l_preprocessed, axis=-1), 3, axis=-1
+    )
+    img_l_preprocessed = np.expand_dims(img_l_preprocessed, axis=0)
+
+    # NCHW format
+    img_l_preprocessed = np.transpose(img_l_preprocessed, (0, 3, 1, 2))
+
+    img_l_tensor = torch.tensor(img_l_preprocessed).float()
+    img_l_tensor = img_l_tensor.to(device, dtype=torch.float)
+
+    gen_img_ab_tensor = colour_model_cgan.net_gen(img_l_tensor)
+    gen_img_ab = gen_img_ab_tensor.detach().cpu().numpy()
+    gen_img_ab = np.squeeze(gen_img_ab)
+    gen_img_ab = np.transpose(gen_img_ab, [1, 2, 0])
+    gen_img_ab_postprocessed = apply_image_ab_post_processing(gen_img_ab)
+
+    # concat L and Generator network generated ab channels
+    gen_img_lab = np.concatenate(
+        (np.expand_dims(img_l, axis=-1), gen_img_ab_postprocessed), axis=-1
+    )
+    # convert Lab to RGB
+    img_gen_rgb = convert_lab2rgb(gen_img_lab)
+    img_gen_rgb = img_gen_rgb * 255
+    img_gen_rgb = img_gen_rgb.astype(np.uint8)
+
+    return img_gen_rgb
+
+
+@app.get("/info")
+def get_app_info() -> dict:
+    """
+    -------
+    Returns
+    -------
+    dict_info : dict
+        a dictionary with info to be sent as a response to get request
+    """
+    dict_info = {"app_name": settings.app_name, "version": settings.version}
+    return dict_info
+
+
+@app.post("/predict")
+def _file_upload(image_file: UploadFile = File(...)) -> dict:
+    """
+    ---------
+    Arguments
+    ---------
+    image_file: object
+        an object of type UploadFile
+
+    -------
+    Returns
+    -------
+    response_json : dict
+        a dict as a response json for the post request
+    """
+    try:
+        # if the file is sent via post request with open()
+        img_str = image_file.file.read()
+        img_decoded = cv2.imdecode(np.frombuffer(img_str, np.uint8), 0)
+    except:
+        # if the file is sent via post request from streamlit
+        img_decoded = cv2.imdecode(np.frombuffer(image_file.getvalue(), np.uint8), 0)
+
+    logging.info(image_file)
+
+    img_gen_rgb = get_prediction(img_decoded)
+    image_sum = np.sum(img_gen_rgb)
+    logging.info(f"image_sum: {image_sum}")
+    _, img_encoded = cv2.imencode(".PNG", img_gen_rgb)
+    img_encoded = base64.b64encode(img_encoded)
+
+    response_json = {
+        "name": image_file.filename,
+        "image_sum": str(image_sum),
+        "encoded_image": img_encoded,
+    }
+    # logging.info(response_json)
+    return response_json

config.py

@@ -0,0 +1,11 @@
+from pydantic_settings import BaseSettings
+
+
+class Settings(BaseSettings):
+    app_name: str = "CGAN Image Colourization API"
+    version: str = "2024.04.15"
+    image_size: int = 320
+    device: str = "cpu"
+
+
+settings = Settings()