MbModel / app.py
MobvxtR's picture
Update app.py
4fcc79c verified
import gradio as gr
import cv2
import numpy as np
from utils.preprocessing import ImageProcessor
# Initialize processor
processor = ImageProcessor("models/best.pt")
def process_image(input_image):
if input_image is None:
raise gr.Error("Please upload an image first!")
# Convert Gradio Image to bytes
_, img_bytes = cv2.imencode(".png", input_image)
# Process image
results = processor.process_image(img_bytes.tobytes())
# Format outputs
return {
class_name: (mask * 255).astype(np.uint8)
for class_name, mask in results.items()
}
# Gradio interface
with gr.Blocks(title="Fashion Segmenter") as demo:
gr.Markdown("# 🧥 Fashion Item Segmenter")
with gr.Row():
input_image = gr.Image(label="Upload Clothing Image", type="numpy")
output_gallery = gr.Gallery(label="Segmented Items", columns=2)
with gr.Row():
run_btn = gr.Button("Process Image", variant="primary")
examples = gr.Examples(
examples=["sample1.jpg", "sample2.jpg"],
inputs=[input_image],
label="Example Images"
)
run_btn.click(
fn=process_image,
inputs=[input_image],
outputs=[output_gallery],
show_progress=True
)
if __name__ == "__main__":
demo.launch()