Spaces:
Build error
Build error
| 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 |