piecefinder / app /segmenter.py
Steph254's picture
Upload 12 files
203a0a4 verified
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