app.py

import onnxruntime
import numpy as np
import cv2
from PIL import Image
import gradio as gr 
import tempfile
import os

from typing import Any

def preprocess_image(image: Image.Image, input_size=(640,640)):
    img = np.array(image)
    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    
    h, w, _ = img.shape
    scale = min(input_size[0] / h, input_size[1] / w)
    new_h, new_w = int(h * scale), int(w * scale)
    img_resized = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_LINEAR)

    #padding 
    img_padded = np.zeros((input_size[0], input_size[1], 3), dtype=np.uint8)
    img_padded[:new_h, :new_w] = img_resized

    #Normalize 
    img_padded = img_padded.astype(np.float32) / 255.0
    img_padded = np.transpose(img_padded, (2, 0, 1))  # HWC -> CHW format
    img_padded = np.expand_dims(img_padded, axis=0)

    return img_padded, scale, (new_h, new_w)

def nms(boxes, scores, iou_threshold=0.5):
    # boxes (num boxes, 4), "4" contain (x, y, w, h) in center format (cx, cy, w, h)
    x1 = boxes[:, 0] - boxes[:, 2] / 2
    y1 = boxes[:, 1] - boxes[:, 3] / 2
    x2 = boxes[:, 0] + boxes[:, 2] / 2
    y2 = boxes[:, 1] + boxes[:, 3] / 2

    areas = np.maximum(0.0, x2 - x1) * np.maximum(0.0, y2 - y1)
    order = scores.argsort()[::-1]

    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])

        w = np.maximum(0.0, xx2 - xx1)
        h = np.maximum(0.0, yy2 - yy1)
        inter = w * h
        iou = inter / (areas[i] + areas[order[1:]] - inter + 1e-10)

        mask = iou < iou_threshold
        order = order[1:][mask]

    return np.array(keep, dtype=np.int64)

def run_inference(model_path, input_tensor: np.ndarray) -> Any:
    available_providers = onnxruntime.get_available_providers()
    print(f"Available providers: {available_providers}")
    #
    providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
    try:
        session = onnxruntime.InferenceSession(model_path, providers=providers, sess_options=None)
        print(f"Using provider: {session.get_providers()[0]}")
    except Exception as e:
        print(f"Error initializing session: {e}")
        raise
    input_name = session.get_inputs()[0].name
    output_names = [output.name for output in session.get_outputs()]
    outputs = session.run(output_names, {input_name: input_tensor})
    return outputs

def _ensure_pred_layout(predictions: np.ndarray, num_classes: int, num_mask_coefficients: int) -> np.ndarray:
    """
    Ensure predictions have shape (1, N, C) where C = 4 + num_classes + num_mask_coefficients.
    ONNX exports use (1, C, N).
    """
    if predictions.ndim != 3:
        raise ValueError(f"Unexpected predictions ndim: {predictions.ndim}, expected 3")
    # Determine C either in axis=1 or axis=2
    C_expected = 4 + num_classes + num_mask_coefficients
    b, a1, a2 = predictions.shape
    if a1 == C_expected and a2 != C_expected:
        # (1, C, N) -> (1, N, C)
        predictions = np.transpose(predictions, (0, 2, 1))
    elif a2 == C_expected and a1 != C_expected:
        # already (1, N, C)
        pass
    else:
        if a1 < a2:
            predictions = np.transpose(predictions, (0, 2, 1))
    return predictions

def _ensure_proto_layout(proto: np.ndarray, num_mask_coefficients: int) -> np.ndarray:
    """
    Ensure proto has shape (1, num_mask_coefficients, H, W).
    Some exports use (1, H, W, num_mask_coefficients).
    """
    if proto.ndim != 4:
        raise ValueError(f"Unexpected proto ndim: {proto.ndim}, expected 4")
    b, d1, d2, d3 = proto.shape
    # If channels-last
    if d3 == num_mask_coefficients and d1 != num_mask_coefficients:
        proto = np.transpose(proto, (0, 3, 1, 2))
    return proto

def postprocess_output(output, conf_threshold=0.25, iou_threshold=0.5,
                        input_size=(640, 640), scale=1.0, padded_size=None, img_shape=None):
    predictions = output[0]
    proto = output[1]

    num_classes = 21
    num_mask_coefficients = 32

    predictions = _ensure_pred_layout(predictions, num_classes, num_mask_coefficients)
    proto = _ensure_proto_layout(proto, num_mask_coefficients)

    boxes = predictions[:, :, :4]  # (1, N, 4) in (cx, cy, w, h)
    scores = predictions[:, :, 4:4 + num_classes]
    mask_coefficients = predictions[:, :, 4 + num_classes:]

    class_ids = np.argmax(scores, axis=-1).squeeze(0)
    confidences = np.max(scores, axis=-1).squeeze(0)

    keep_mask = confidences > conf_threshold
    boxes = boxes.squeeze(0)[keep_mask]  # (K, 4)
    confidences = confidences[keep_mask]
    class_ids = class_ids[keep_mask]
    mask_coefficients = mask_coefficients.squeeze(0)[keep_mask]

    boxes = boxes / max(scale, 1e-12)

    if padded_size is not None:
        ph, pw = padded_size
        boxes[:, 0] = np.clip(boxes[:, 0], 0, pw)  # cx
        boxes[:, 1] = np.clip(boxes[:, 1], 0, ph)  # cy
        boxes[:, 2] = np.clip(boxes[:, 2], 0, pw)  # w
        boxes[:, 3] = np.clip(boxes[:, 3], 0, ph)  # h

    if boxes.size == 0:
        return boxes, confidences, class_ids, np.array([])

    indices = nms(boxes, confidences, iou_threshold)
    boxes = boxes[indices]
    confidences = confidences[indices]
    class_ids = class_ids[indices]
    mask_coefficients = mask_coefficients[indices]

    # Generate masks with boxes and img_shape
    masks = generate_segmentation_mask(mask_coefficients, proto, input_size, scale, padded_size, boxes, img_shape)

    return boxes, confidences, class_ids, masks

def generate_segmentation_mask(mask_coefficients, proto, input_size, scale, padded_size, boxes, img_shape):
    proto = proto[0]  # (32, H, W)
    mask_h, mask_w = proto.shape[1:]
    masks = []
    orig_h, orig_w = img_shape

    for i, coeff in enumerate(mask_coefficients):
        mask = np.dot(coeff, proto.reshape(proto.shape[0], -1)).reshape(mask_h, mask_w)
        mask = 1.0 / (1.0 + np.exp(-mask))  # Sigmoid

        mask_up = cv2.resize(mask, input_size, interpolation=cv2.INTER_LINEAR)
        mask_valid = mask_up[:padded_size[0], :padded_size[1]]
        mask_final = cv2.resize(mask_valid, (orig_w, orig_h), interpolation=cv2.INTER_LINEAR)

        box = boxes[i]  # (cx, cy, w, h)
        cx, cy, bw, bh = box
        x1 = max(0, int(cx - bw / 2))
        y1 = max(0, int(cy - bh / 2))
        x2 = min(orig_w, int(cx + bw / 2))
        y2 = min(orig_h, int(cy + bh / 2))

        clipped_mask = np.zeros(img_shape, dtype=np.float32)
        clipped_mask[y1:y2, x1:x2] = mask_final[y1:y2, x1:x2]
        clipped_mask = (clipped_mask > 0.5).astype(np.uint8)
        masks.append(clipped_mask.astype(bool))

    return np.array(masks, dtype=bool)

def visualize_results(img: np.ndarray, boxes, confidences, class_ids, masks, class_names):
    print(f"Image shape: {img.shape}, Masks shape: {masks.shape if len(masks) > 0 else 'No masks'}")  # Debug
    for mask in masks:
        mask_color = np.random.randint(0, 255, (3,), dtype=np.uint8)
        img[mask] = img[mask] * 0.5 + mask_color * 0.5
    return Image.fromarray(img)

def read_class_names(file_path):
    with open(file_path, "r") as file:
        class_names = [ line.strip() for line in file if line.strip()]
    return class_names

def predict(image):
    class_names = read_class_names("car_part_class_names.txt")
    onnx_model_path = "car_part_segment.onnx"

    input_size=(640, 640)
    #Preprocess
    input_data, scale, padded_size = preprocess_image(image, input_size=input_size)
    img_shape = (image.height, image.width)

    #Inference
    outputs = run_inference(onnx_model_path, input_data)

    boxes, confidences, class_ids, masks = postprocess_output(
        outputs,
        conf_threshold=0.25,
        iou_threshold=0.45,
        input_size=input_size,
        scale=scale,
        padded_size=padded_size,
        img_shape=img_shape
    )

    img_np = np.array(image) # PIL to numpy RGB
    output_img = visualize_results(img_np, boxes, confidences, class_ids, masks, class_names)

    return output_img

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Input Image"),
    outputs=gr.Image(type="pil", label="Segmented Output"),
    title="Car Part Segmentation",
    description="Upload an image of a car part to segment it using a ONNX model.",
    examples=[
        "huyndai.jpg"
    ]
)

demo.launch()