File size: 1,372 Bytes
c4e13a2 f95494c c4e13a2 |
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 |
import cv2
import numpy as np
from rembg import remove
from ultralytics import YOLO
class ImageProcessor:
def __init__(self, model_path):
self.model = YOLO(model_path)
self.class_names = {0: "upper_clothes", 1: "lower_clothes"}
def remove_background(self, image_bytes):
return remove(image_bytes)
def process_image(self, image_bytes):
# Background removal
bg_removed = self.remove_background(image_bytes)
# Convert to OpenCV format
nparr = np.frombuffer(bg_removed, np.uint8)
img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
# Segmentation
results = self.model.predict(img)
return self._process_masks(results, img)
def _process_masks(self, results, img):
segmented = {}
if results[0].masks is not None:
for mask, class_id in zip(results[0].masks.data, results[0].boxes.cls):
class_id = int(class_id.item())
mask_np = mask.cpu().numpy()
mask_resized = cv2.resize(mask_np, (img.shape[1], img.shape[0]), interpolation=cv2.INTER_NEAREST)
_, binary_mask = cv2.threshold(mask_resized, 0.5, 255, cv2.THRESH_BINARY)
binary_mask = binary_mask.astype(np.uint8)
segmented[self.class_names[class_id]] = binary_mask
return segmented |