StyleSavvy-ai / utils /detector.py
hank1229's picture
Update utils/detector.py
ca23413 verified
raw
history blame contribute delete
809 Bytes
from PIL import Image
from io import BytesIO
from models.vision import VisionModel
from utils.bg_removal import remove_background
vision = VisionModel()
def detect_clothing(image_input, do_bg_remove: bool = False):
# 1) If filepath, load; else assume PIL.Image
if isinstance(image_input, str):
img = Image.open(image_input).convert("RGB")
else:
img = image_input
# 2) Optional background removal
if do_bg_remove:
buf = BytesIO()
img.save(buf, format="JPEG")
img_bytes = buf.getvalue()
img = remove_background(img_bytes)
else:
img = img.convert("RGB")
# 3) Run detection
detections = vision.detect(img)
if not detections:
detections = [{"label": "outfit", "score": 1.0, "box": []}]
return detections