File size: 1,419 Bytes
2a2e2dd 159ad53 2a2e2dd 159ad53 1d57c26 159ad53 2a2e2dd 159ad53 2a2e2dd 159ad53 2a2e2dd 159ad53 1d57c26 159ad53 1d57c26 159ad53 1d57c26 159ad53 1d57c26 159ad53 2a2e2dd 1d57c26 b0a26d5 2a2e2dd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 |
import gradio as gr
import onnxruntime as ort
import numpy as np
from PIL import Image
import requests
from torchvision import transforms
# Load ONNX model
ort_session = ort.InferenceSession("model.onnx") # Ensure model.onnx is in your app folder
# Preprocessing function
def preprocess(image):
image = image.resize((320, 320)).convert("RGB")
image_np = np.array(image).astype(np.float32) / 255.0
image_np = image_np.transpose(2, 0, 1) # HWC -> CHW
image_np = np.expand_dims(image_np, axis=0) # Add batch dimension
return image_np
# Inference + Postprocessing
def segment_dress(image):
input_tensor = preprocess(image)
inputs = {ort_session.get_inputs()[0].name: input_tensor}
outputs = ort_session.run(None, inputs)
pred = outputs[0][0][0]
pred = (pred - pred.min()) / (pred.max() - pred.min())
pred_img = Image.fromarray((pred * 255).astype(np.uint8)).resize(image.size)
# Apply mask to image
image_np = np.array(image.convert("RGB"))
mask = np.array(pred_img).astype(np.float32) / 255.0
masked = (image_np * mask[..., None]).astype(np.uint8)
return Image.fromarray(masked)
# Gradio app
gr.Interface(
fn=segment_dress,
inputs=gr.Image(type="pil", label="Upload Image"),
outputs=gr.Image(type="pil", label="Segmented Dress"),
title="Background Removal",
description="Upload an image and Remove the Background"
).launch() |