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