app.py

# -*- coding: utf-8 -*-

import sys
import os
import uuid
from datetime import datetime
import httpx
from fastapi import FastAPI, File, UploadFile, Form, HTTPException, Request, status
import cv2
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
from modelscope.outputs import OutputKeys
import numpy as np
from starlette.staticfiles import StaticFiles
from starlette.templating import Jinja2Templates

app = FastAPI()

model_paths = {
    "universal": {'path': 'damo/cv_unet_universal-matting', 'task': Tasks.universal_matting},
    "people": {'path': 'damo/cv_unet_image-matting', 'task': Tasks.portrait_matting},
}

default_model = list(model_paths.keys())[0]
default_model_info = model_paths[default_model]
loaded_models = {default_model: pipeline(default_model_info['task'], model=default_model_info['path'])}

UPLOAD_FOLDER = "./upload"
OUTPUT_FOLDER = "./output"

os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(OUTPUT_FOLDER, exist_ok=True)


class ModelLoader:
    def __init__(self):
        self.loaded_models = {default_model: loaded_models[default_model]}

    def load_model(self, model_name):
        if model_name not in self.loaded_models:
            model_info = model_paths[model_name]
            if not model_info:
                raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid model selection")
            model_path = model_info['path']
            task_group = model_info['task']

            self.loaded_models[model_name] = pipeline(task_group, model=model_path)
        return self.loaded_models[model_name]


model_loader = ModelLoader()


def get_filename():
    filename = uuid.uuid4()
    original_image_filename = f"original_{filename}.png"
    image_filename = f"image_{filename}.png"
    mask_filename = f"mask_{filename}.png"
    return original_image_filename, image_filename, mask_filename


# remove excess transparent background and crop the image
def crop_image_by_alpha_channel(input_image: np.ndarray | str, output_path: str):
    img_array = cv2.imread(input_image, cv2.IMREAD_UNCHANGED) if isinstance(input_image, str) else input_image
    if img_array.shape[2] != 4:
        raise ValueError("Input image must have an alpha channel")

    alpha_channel = img_array[:, :, 3]
    bbox = cv2.boundingRect(alpha_channel)
    x, y, w, h = bbox
    cropped_img_array = img_array[y:y + h, x:x + w]
    cv2.imwrite(output_path, cropped_img_array)
    return output_path


def process_image(image_bytes: bytes):
    img = cv2.imdecode(np.frombuffer(image_bytes, np.uint8), cv2.IMREAD_UNCHANGED)
    final_img = convert_image_to_white_background(image=img)
    if final_img is None:
        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid image")
    return final_img


def convert_image_to_white_background(image_path: str = None, image: np.ndarray | None = None):
    try:
        if image_path is not None:
            img = cv2.imread(image_path, cv2.IMREAD_UNCHANGED)
        elif image is not None:
            img = image
        else:
            raise ValueError("Either image_path or image must be provided.")

        if img.shape[2] == 4:
            alpha_channel = img[:, :, 3]
            rgb_channels = img[:, :, :3]

            alpha_channel_3d = alpha_channel[:, :, np.newaxis] / 255.0
            alpha_channel_3d = np.repeat(alpha_channel_3d, 3, axis=2)

            white_background_image = np.ones_like(rgb_channels, dtype=np.uint8) * 255

            foreground = cv2.multiply(rgb_channels, alpha_channel_3d, dtype=cv2.CV_8UC3)
            background = cv2.multiply(white_background_image, 1 - alpha_channel_3d, dtype=cv2.CV_8UC3)

            final_img = cv2.add(foreground, background)
        else:
            final_img = img
        return final_img
    except Exception as e:
        print(f'Error: {e}')
        return None


@app.post("/switch_model/{new_model}")
async def switch_model(new_model: str):
    if new_model not in model_paths:
        return {"content": "Invalid model selection"}, status.HTTP_400_BAD_REQUEST
    model_info = model_paths[new_model]

    loaded_models[new_model] = pipeline(model_info['task'], model=model_info['path'])
    model_loader.loaded_models = loaded_models
    return {"content": f"Switched to model: {new_model}"}, status.HTTP_200_OK


@app.post("/matting")
async def matting(image: UploadFile = File(...), model: str = Form(default=default_model, alias="model")):
    try:
        image_bytes = await image.read()
        img = cv2.imdecode(np.frombuffer(image_bytes, np.uint8), cv2.IMREAD_UNCHANGED)

        if model not in model_paths:
            return {"content": "Invalid model selection"}, status.HTTP_400_BAD_REQUEST

        selected_model = model_loader.load_model(model)
        original_image_filename, image_filename, mask_filename = get_filename()
        cv2.imwrite(os.path.join(UPLOAD_FOLDER, original_image_filename), img)

        final_img = convert_image_to_white_background(image=img)
        if final_img is None:
            raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid image")
        result = selected_model(final_img)

        cv2.imwrite(os.path.join(OUTPUT_FOLDER, image_filename), result[OutputKeys.OUTPUT_IMG])
        cv2.imwrite(os.path.join(OUTPUT_FOLDER, mask_filename), result[OutputKeys.OUTPUT_IMG][:, :, 3])

        response_data = {
            "code": 0,
            "result_image_url": f"/output/{image_filename}",
            "mask_image_url": f"/output/{mask_filename}",
            "original_image_size": {"width": img.shape[1], "height": img.shape[0]},
            "generation_time": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        }
        return response_data
    except HTTPException as e:
        return {"error": str(e)}, e.status_code
    except Exception as e:
        return {"error": str(e)}, status.HTTP_500_INTERNAL_SERVER_ERROR


@app.post("/matting/url")
async def matting_url(request: Request, model: str = Form(default=default_model, alias="model")):
    try:
        json_data = await request.json()
        image_url = json_data.get("image_url")
    except Exception as e:
        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=f"Error parsing JSON data: {str(e)}")

    if not image_url:
        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Image URL is required")

    try:
        async with httpx.AsyncClient() as client:
            response = await client.get(image_url)
            response.raise_for_status()
            img_array = np.frombuffer(response.content, dtype=np.uint8)
            img = cv2.imdecode(img_array, cv2.IMREAD_UNCHANGED)
    except httpx.RequestError as e:
        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=f"Failed to fetch image from URL: {str(e)}")

    if model not in model_paths:
        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid model selection")

    selected_model = model_loader.load_model(model)
    original_image_filename, image_filename, mask_filename = get_filename()
    cv2.imwrite(os.path.join(UPLOAD_FOLDER, original_image_filename), img)

    final_img = convert_image_to_white_background(image=img)
    if final_img is None:
        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid image")
    result = selected_model(final_img)

    cv2.imwrite(os.path.join(OUTPUT_FOLDER, image_filename), result[OutputKeys.OUTPUT_IMG])
    cv2.imwrite(os.path.join(OUTPUT_FOLDER, mask_filename), result[OutputKeys.OUTPUT_IMG][:, :, 3])

    response_data = {
        "code": 0,
        "result_image_url": f"/output/{image_filename}",
        "mask_image_url": f"/output/{mask_filename}",
        "original_image_size": {"width": img.shape[1], "height": img.shape[0]},
        "generation_time": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
    }

    return response_data


templates = Jinja2Templates(directory="web")
app.mount("/static", StaticFiles(directory="./web/static"), name="static")
app.mount("/output", StaticFiles(directory="./output"), name="output")
app.mount("/upload", StaticFiles(directory="./upload"), name="upload")


@app.get("/")
async def read_index(request: Request):
    return templates.TemplateResponse(
        "index.html", {
            "request": request,
            "default_model": default_model,
            "available_models": list(model_paths.keys())
        })


if __name__ == "__main__":
    import uvicorn

    default_bind_host = "0.0.0.0" if sys.platform != "win32" else "127.0.0.1"
    uvicorn.run(app, host=default_bind_host, port=8000)