Spaces:
Sleeping
Sleeping
File size: 1,633 Bytes
8fd2d1a 7cc2c5a 8fd2d1a 7cc2c5a 8fd2d1a 7cc2c5a 8fd2d1a 7cc2c5a 8fd2d1a 7cc2c5a 8fd2d1a 7cc2c5a 8fd2d1a |
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 44 45 46 47 48 49 50 51 52 53 |
import gradio as gr
from PIL import Image
from pathlib import Path
import numpy as np
from ultralytics import YOLO
MODEL_WEIGHTS_PATH = Path("weights/best.pt")
VERSION_PATH = Path("VERSION")
# Read version string from VERSION file
try:
VERSION = VERSION_PATH.read_text().strip()
except Exception:
VERSION = "unknown"
model = None
def get_model() -> YOLO:
"""
Returns the YOLO model instance.
"""
global model
if model is None:
if not MODEL_WEIGHTS_PATH.exists():
raise FileNotFoundError(f"Model weights not found at {MODEL_WEIGHTS_PATH}. Please deploy weights before running.")
model = YOLO(str(MODEL_WEIGHTS_PATH))
return model
def segment(image: Image.Image) -> tuple[Image.Image, str]:
"""
Returns a tuple: (segmentation mask PIL.Image, model version string)
"""
model = get_model()
img_np = np.array(image)
results = model(img_np)
if not results or not hasattr(results[0], "masks") or results[0].masks is None:
mask_img = Image.new("L", image.size, 0)
else:
mask = results[0].masks.data[0].cpu().numpy()
mask_img = Image.fromarray((mask * 255).astype(np.uint8))
mask_img = mask_img.resize(image.size)
return mask_img, str(VERSION)
iface = gr.Interface(
fn=segment,
inputs=gr.Image(type="pil"),
outputs=[gr.Image(type="pil", label="Segmentation Mask"), gr.Textbox(label="Model Version")],
title=f"YOLO Segmentation Model (version: {VERSION})",
description=f"Upload an image to get a segmentation mask. Model version: {VERSION}"
)
if __name__ == "__main__":
iface.launch()
|