File size: 1,287 Bytes
203a0a4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
import cv2
import numpy as np
from app.utils import preprocess_image, extract_features

class SimpleSegmenter:
    def __init__(self):
        self.min_piece_area = 500
        self.max_piece_area = 50000

    def segment_pieces(self, image: np.ndarray) -> list:
        # Preprocess image
        processed = preprocess_image(image)
        gray = cv2.cvtColor(processed, cv2.COLOR_BGR2GRAY)
        _, thresh = cv2.threshold(gray, 240, 255, cv2.THRESH_BINARY_INV)

        # Find contours
        contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        pieces = []
        piece_id = 0

        for contour in contours:
            area = cv2.contourArea(contour)
            if self.min_piece_area <= area <= self.max_piece_area:
                mask = np.zeros(image.shape[:2], dtype=np.uint8)
                cv2.drawContours(mask, [contour], -1, 255, -1)
                piece_image = cv2.bitwise_and(image, image, mask=mask)
                features = extract_features(piece_image, mask)
                pieces.append({
                    'id': piece_id,
                    'image': cv2.imencode('.jpg', piece_image)[1].tobytes(),
                    'features': features
                })
                piece_id += 1

        return pieces