Spaces:
Runtime error
Runtime error
| import cv2 | |
| import numpy as np | |
| from ultralytics import YOLO | |
| from segment_anything import SamPredictor, sam_model_registry | |
| from app.utils.preprocess import preprocess_image | |
| from app.utils.feature_extraction import extract_piece_features | |
| class PuzzlePieceSegmenter: | |
| def __init__(self): | |
| # Load models | |
| self.yolo = YOLO("yolov8n.pt") # Lightweight YOLOv8 model | |
| self.sam = sam_model_registry["vit_b"](checkpoint="mobile_sam.pth") | |
| self.predictor = SamPredictor(self.sam) | |
| self.min_piece_area = 500 | |
| self.max_piece_area = 50000 | |
| def segment_pieces(self, image: np.ndarray) -> list: | |
| # Preprocess image | |
| processed = preprocess_image(image) | |
| # YOLOv8 for coarse detection | |
| results = self.yolo(processed) | |
| boxes = results[0].boxes.xyxy.cpu().numpy() | |
| # SAM for fine segmentation | |
| self.predictor.set_image(processed) | |
| pieces = [] | |
| piece_id = 0 | |
| for box in boxes: | |
| masks, _, _ = self.predictor.predict(box_coordinates=box, multimask_output=False) | |
| for mask in masks: | |
| contours, _ = cv2.findContours(mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) | |
| for contour in contours: | |
| area = cv2.contourArea(contour) | |
| if self.min_piece_area <= area <= self.max_piece_area: | |
| # Extract features and image for this piece | |
| piece_image, features = extract_piece_features(processed, contour) | |
| if features: | |
| pieces.append({ | |
| 'id': piece_id, | |
| 'image': cv2.imencode('.jpg', piece_image)[1].tobytes(), | |
| 'features': features | |
| }) | |
| piece_id += 1 | |
| return pieces |