Create utils/preprocessing.py

utils/preprocessing.py

@@ -0,0 +1,34 @@
+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()
+                # ... [your mask processing logic here] ...
+                segmented[self.class_names[class_id]] = processed_mask
+        return segmented