Upload 26 files

.env

@@ -0,0 +1,5 @@
+SECRET_KEY=your-secret-key
+DATABASE_URL=postgresql://postgres:mysecretpassword@localhost:5432/piecefinder
+REDIS_URL=redis://localhost:6379/0
+YOLO_MODEL_PATH=/app/models/yolov8n.pt
+SAM_MODEL_PATH=/app/models/mobile_sam.pth

app.py

@@ -2,6 +2,5 @@ from app import create_app
 
 app = create_app()
 
-
 if __name__ == '__main__':
-    app.run()
+    app.run()

app/__init__.py

@@ -0,0 +1,23 @@
+from flask import Flask
+from flask_sqlalchemy import SQLAlchemy
+from flask_cors import CORS
+
+db = SQLAlchemy()
+
+def create_app():
+    app = Flask(__name__)
+    app.config['SQLALCHEMY_DATABASE_URI'] = 'postgresql://postgres:mysecretpassword@localhost:5432/piecefinder'
+    app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False
+    app.config['SECRET_KEY'] = 'your-secret-key'
+    app.config['REDIS_URL'] = 'redis://localhost:6379/0'
+
+    db.init_app(app)
+    CORS(app)  # Enable CORS
+
+    with app.app_context():
+        db.create_all()
+
+    from app.routes import bp
+    app.register_blueprint(bp)
+
+    return app

app/config.py

@@ -0,0 +1,8 @@
+import os
+
+class Config:
+    SECRET_KEY = os.environ.get('SECRET_KEY', 'default-secret-key')
+    DATABASE_URL = os.environ.get('DATABASE_URL', 'postgresql://user:pass@localhost/piecefinder')
+    REDIS_URL = os.environ.get('REDIS_URL', 'redis://localhost:6379/0')
+    YOLO_MODEL_PATH = os.environ.get('YOLO_MODEL_PATH', 'yolov8n.pt')
+    SAM_MODEL_PATH = os.environ.get('SAM_MODEL_PATH', 'mobile_sam.pth')

app/models/feedback_manager.py

@@ -0,0 +1,50 @@
+from sqlalchemy import create_engine
+import json
+from app import create_app
+
+class FeedbackManager:
+    def __init__(self):
+        app = create_app()
+        self.db = app.db
+
+        # Initialize feedback table if not exists
+        with self.db.connect() as conn:
+            conn.execute("""
+                CREATE TABLE IF NOT EXISTS feedback (
+                    id SERIAL PRIMARY KEY,
+                    data JSONB,
+                    timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+                )
+            """)
+
+    def record_feedback(self, piece_id: int, slot_id: int, suggestions: dict, user_choice: int, is_correct: bool):
+        feedback_entry = {
+            'piece_id': piece_id,
+            'slot_id': slot_id,
+            'suggestions': suggestions,
+            'user_choice': user_choice,
+            'is_correct': is_correct
+        }
+        with self.db.connect() as conn:
+            conn.execute(
+                "INSERT INTO feedback (data) VALUES (%s)",
+                (json.dumps(feedback_entry),)
+            )
+
+    def get_training_data(self) -> dict:
+        positive, negative = [], []
+        with self.db.connect() as conn:
+            results = conn.execute("SELECT data FROM feedback").fetchall()
+            for row in results:
+                entry = json.loads(row[0])
+                if entry['is_correct']:
+                    positive.append({
+                        'piece_id': entry['piece_id'],
+                        'slot_id': entry['slot_id']
+                    })
+                else:
+                    negative.append({
+                        'piece_id': entry['piece_id'],
+                        'slot_id': entry['slot_id']
+                    })
+        return {'positive': positive, 'negative': negative}

app/models/matcher.py

@@ -0,0 +1,61 @@
+import numpy as np
+import faiss
+from sklearn.metrics.pairwise import cosine_similarity
+import cv2
+
+class PuzzleMatchingEngine:
+    def __init__(self):
+        self.weight_geometric = 0.35
+        self.weight_color = 0.30
+        self.weight_edge = 0.35
+        self.faiss_index = faiss.IndexFlatL2(10)  # Assume 10D edge features
+
+    def find_matches(self, pieces: list, slots: list, top_k: int = 3) -> dict:
+        matches = {}
+        if not pieces or not slots:
+            return matches
+
+        # Prepare FAISS index with piece edge features
+        piece_edge_features = np.array([p['features']['edge'] for p in pieces], dtype=np.float32)
+        self.faiss_index.add(piece_edge_features)
+
+        for slot in slots:
+            slot_matches = []
+            slot_edge_features = np.array([slot['features']['edge']], dtype=np.float32)
+            distances, indices = self.faiss_index.search(slot_edge_features, top_k * 2)
+
+            for idx in indices[0]:
+                piece = pieces[idx]
+                geometric_score = self._calculate_geometric_similarity(piece['features']['geometric'], slot['features']['geometric'])
+                color_score = self._calculate_color_similarity(piece['features']['color'], slot['features']['color'])
+                edge_score = self._calculate_edge_similarity(piece['features']['edge'], slot['features']['edge'])
+                confidence = (
+                    self.weight_geometric * geometric_score +
+                    self.weight_color * color_score +
+                    self.weight_edge * edge_score
+                )
+                slot_matches.append({
+                    'piece_id': piece['id'],
+                    'slot_id': slot['id'],
+                    'confidence': float(confidence),
+                    'geometric_score': float(geometric_score),
+                    'color_score': float(color_score),
+                    'edge_score': float(edge_score)
+                })
+
+            slot_matches.sort(key=lambda x: x['confidence'], reverse=True)
+            matches[slot['id']] = slot_matches[:top_k]
+
+        return matches
+
+    def _calculate_geometric_similarity(self, piece_features: dict, slot_features: dict) -> float:
+        # Compare area, perimeter, aspect ratio, etc.
+        area_diff = abs(piece_features['area'] - slot_features['area']) / max(piece_features['area'], slot_features['area'])
+        perimeter_diff = abs(piece_features['perimeter'] - slot_features['perimeter']) / max(piece_features['perimeter'], slot_features['perimeter'])
+        return 1.0 - (0.5 * area_diff + 0.5 * perimeter_diff)
+
+    def _calculate_color_similarity(self, piece_histogram: np.ndarray, slot_histogram: np.ndarray) -> float:
+        return cv2.compareHist(piece_histogram, slot_histogram, cv2.HISTCMP_CORREL)
+
+    def _calculate_edge_similarity(self, piece_features: np.ndarray, slot_features: np.ndarray) -> float:
+        return cosine_similarity(piece_features.reshape(1, -1), slot_features.reshape(1, -1))[0, 0]

app/models/segmenter.py

@@ -0,0 +1,47 @@
+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

app/routes/count.py

@@ -0,0 +1,31 @@
+from flask import Blueprint, request, jsonify
+from app.models.segmenter import PuzzlePieceSegmenter
+import base64
+import numpy as np
+import cv2
+
+count_bp = Blueprint('count', __name__)
+
+@count_bp.route('/count', methods=['POST'])
+def count_pieces():
+    try:
+        data = request.get_json()
+        image_base64 = data.get('image_base64')
+        if not image_base64:
+            return jsonify({'error': 'No image provided'}), 400
+
+        # Decode base64 image
+        image_data = base64.b64decode(image_base64)
+        nparr = np.frombuffer(image_data, np.uint8)
+        image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+
+        # Initialize segmenter
+        segmenter = PuzzlePieceSegmenter()
+
+        # Perform segmentation and count
+        pieces = segmenter.segment_pieces(image)
+        count = len(pieces)
+
+        return jsonify({'count': count}), 200
+    except Exception as e:
+        return jsonify({'error': str(e)}), 500

app/routes/feedback.py

@@ -0,0 +1,27 @@
+from flask import Blueprint, request, jsonify
+from app.models.feedback_manager import FeedbackManager
+
+feedback_bp = Blueprint('feedback', __name__)
+
+@feedback_bp.route('/feedback', methods=['POST'])
+def submit_feedback():
+    try:
+        data = request.get_json()
+        piece_id = data.get('pieceId')
+        slot_id = data.get('slotId')
+        suggestions = data.get('suggestions')
+        user_choice = data.get('userChoice')
+        is_correct = data.get('isCorrect')
+
+        if not all([piece_id, slot_id, suggestions, user_choice is not None, is_correct is not None]):
+            return jsonify({'error': 'Missing required fields'}), 400
+
+        # Initialize feedback manager
+        feedback_manager = FeedbackManager()
+
+        # Record feedback
+        feedback_manager.record_feedback(piece_id, slot_id, suggestions, user_choice, is_correct)
+
+        return jsonify({'message': 'Feedback recorded successfully'}), 200
+    except Exception as e:
+        return jsonify({'error': str(e)}), 500

app/routes/match.py

@@ -0,0 +1,53 @@
+from flask import Blueprint, request, jsonify
+from app.models.matcher import PuzzleMatchingEngine
+from app.models.segmenter import PuzzlePieceSegmenter
+import base64
+import numpy as np
+import cv2
+
+match_bp = Blueprint('match', __name__)
+
+@match_bp.route('/match', methods=['POST'])
+def match_pieces():
+    try:
+        data = request.get_json()
+        image_base64 = data.get('image_base64')
+        if not image_base64:
+            return jsonify({'error': 'No image provided'}), 400
+
+        # Decode base64 image
+        image_data = base64.b64decode(image_base64)
+        nparr = np.frombuffer(image_data, np.uint8)
+        image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+
+        # Initialize segmenter and matcher
+        segmenter = PuzzlePieceSegmenter()
+        matcher = PuzzleMatchingEngine()
+
+        # Segment pieces and slots (assuming slots are from a reference image)
+        pieces = segmenter.segment_pieces(image)
+        # For demo, assume slots are pre-segmented or same as pieces
+        slots = pieces  # Replace with actual slot segmentation logic
+
+        # Perform matching
+        matches = matcher.find_matches(pieces, slots, top_k=3)
+
+        # Serialize matches
+        serialized_matches = []
+        for slot_id, slot_matches in matches.items():
+            for match in slot_matches:
+                serialized_matches.append({
+                    'piece_id': match['piece_id'],
+                    'slot_id': match['slot_id'],
+                    'confidence': match['confidence'],
+                    'geometric_score': match['geometric_score'],
+                    'color_score': match['color_score'],
+                    'edge_score': match['edge_score']
+                })
+
+        return jsonify({
+            'matches': serialized_matches,
+            'slots': [{'id': slot['id']} for slot in slots]
+        }), 200
+    except Exception as e:
+        return jsonify({'error': str(e)}), 500

app/routes/segment.py

@@ -0,0 +1,44 @@
+from flask import Blueprint, request, jsonify
+from app.models.segmenter import PuzzlePieceSegmenter
+import base64
+import numpy as np
+import cv2
+
+segment_bp = Blueprint('segment', __name__)
+
+@segment_bp.route('/segment', methods=['POST'])
+def segment_image():
+    try:
+        data = request.get_json()
+        image_base64 = data.get('image_base64')
+        if not image_base64:
+            return jsonify({'error': 'No image provided'}), 400
+
+        # Decode base64 image
+        image_data = base64.b64decode(image_base64)
+        nparr = np.frombuffer(image_data, np.uint8)
+        image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+
+        # Initialize segmenter
+        segmenter = PuzzlePieceSegmenter()
+
+        # Perform segmentation
+        pieces = segmenter.segment_pieces(image)
+
+        # Serialize pieces
+        serialized_pieces = [
+            {
+                'id': piece['id'],
+                'image': base64.b64encode(piece['image']).decode('utf-8'),
+                'features': {
+                    'geometric': piece['features']['geometric'],
+                    'color': piece['features']['color'],
+                    'edge': piece['features']['edge']
+                }
+            }
+            for piece in pieces
+        ]
+
+        return jsonify({'pieces': serialized_pieces}), 200
+    except Exception as e:
+        return jsonify({'error': str(e)}), 500

app/utils/feature_extraction.py

@@ -0,0 +1,34 @@
+import cv2
+import numpy as np
+
+def extract_piece_features(image: np.ndarray, contour: np.ndarray) -> tuple:
+    # Create mask for the piece
+    mask = np.zeros(image.shape[:2], dtype=np.uint8)
+    cv2.drawContours(mask, [contour], -1, 255, -1)
+    
+    # Extract piece image
+    piece_image = cv2.bitwise_and(image, image, mask=mask)
+    
+    # Geometric features
+    area = cv2.contourArea(contour)
+    perimeter = cv2.arcLength(contour, True)
+    rect = cv2.minAreaRect(contour)
+    aspect_ratio = rect[1][0] / rect[1][1] if rect[1][1] != 0 else 1.0
+    
+    # Color histogram
+    color_histogram = cv2.calcHist([piece_image], [0, 1, 2], mask, [8, 8, 8], [0, 256, 0, 256, 0, 256])
+    color_histogram = cv2.normalize(color_histogram, color_histogram).flatten()
+    
+    # Edge features (simplified)
+    gray = cv2.cvtColor(piece_image, cv2.COLOR_BGR2GRAY)
+    edges = cv2.Canny(gray, 100, 200)
+    edge_features = np.histogram(edges[mask == 255], bins=10, range=(0, 255))[0].astype(np.float32)
+    edge_features /= edge_features.sum() + 1e-10
+    
+    features = {
+        'geometric': {'area': float(area), 'perimeter': float(perimeter), 'aspect_ratio': float(aspect_ratio)},
+        'color': color_histogram,
+        'edge': edge_features
+    }
+    
+    return piece_image, features

app/utils/preprocess.py

@@ -0,0 +1,14 @@
+import cv2
+import numpy as np
+
+def preprocess_image(image: np.ndarray) -> np.ndarray:
+    # Convert to LAB color space and apply CLAHE
+    lab = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)
+    clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
+    lab[:, :, 0] = clahe.apply(lab[:, :, 0])
+    enhanced = cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)
+
+    # Apply bilateral filtering
+    denoised = cv2.bilateralFilter(enhanced, 9, 75, 75)
+
+    return denoised

app/utils/visualize.py

@@ -0,0 +1,19 @@
+import cv2
+import numpy as np
+
+def visualize_pieces(image: np.ndarray, pieces: list) -> np.ndarray:
+    vis_image = image.copy()
+    for piece in pieces:
+        contour = piece.get('contour')  # Assume contour is stored
+        if contour is not None:
+            cv2.drawContours(vis_image, [contour], -1, (0, 255, 0), 2)
+            cv2.putText(
+                vis_image,
+                f"ID: {piece['id']}",
+                (int(contour[:, :, 0].min()), int(contour[:, :, 1].min()) - 10),
+                cv2.FONT_HERSHEY_SIMPLEX,
+                0.5,
+                (0, 0, 255),
+                1
+            )
+    return vis_image

deployments/Dockerfile

@@ -0,0 +1,15 @@
+FROM python:3.10-slim
+
+WORKDIR /app
+
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY . .
+
+ENV FLASK_APP=main.py
+ENV FLASK_RUN_HOST=0.0.0.0
+
+EXPOSE 5000
+
+CMD ["flask", "run"]

deployments/kubernetes.yaml

@@ -0,0 +1,43 @@
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: piecefinder-backend
+spec:
+  replicas: 2
+  selector:
+    matchLabels:
+      app: piecefinder-backend
+  template:
+    metadata:
+      labels:
+        app: piecefinder-backend
+    spec:
+      containers:
+      - name: piecefinder-backend
+        image: piecefinder-backend:latest
+        ports:
+        - containerPort: 5000
+        env:
+        - name: DATABASE_URL
+          valueFrom:
+            secretKeyRef:
+              name: db-credentials
+              key: database-url
+        - name: REDIS_URL
+          valueFrom:
+            secretKeyRef:
+              name: redis-credentials
+              key: redis-url
+---
+apiVersion: v1
+kind: Service
+metadata:
+  name: piecefinder-backend
+spec:
+  selector:
+    app: piecefinder-backend
+  ports:
+  - protocol: TCP
+    port: 80
+    targetPort: 5000
+  type: ClusterIP

deployments/requirements.txt

@@ -0,0 +1,12 @@
+flask
+opencv-python
+numpy
+ultralytics
+segment-anything
+faiss-cpu
+tensorflow
+scikit-image
+sqlalchemy
+psycopg2-binary
+redis
+flask-cors

main.py

@@ -0,0 +1,6 @@
+from app import create_app
+
+app = create_app()
+
+if __name__ == '__main__':
+    app.run(debug=True, host='0.0.0.0', port=5000)

models/mobile_sam.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6dbb90523a35330fedd7f1d3dfc66f995213d81b29a5ca8108dbcdd4e37d6c2f
+size 40728226

models/yolov8n.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f59b3d833e2ff32e194b5bb8e08d211dc7c5bdf144b90d2c8412c47ccfc83b36
+size 6549796

scripts/data_augmentation.py

@@ -0,0 +1,24 @@
+import cv2
+import numpy as np
+import os
+
+def generate_synthetic_puzzle(num_images: int, output_dir: str):
+    os.makedirs(output_dir, exist_ok=True)
+    
+    for i in range(num_images):
+        # Create blank image
+        image = np.ones((1000, 1000, 3), dtype=np.uint8) * 255
+        
+        # Add random puzzle pieces (simplified)
+        for _ in range(np.random.randint(10, 50)):
+            x = np.random.randint(100, 900)
+            y = np.random.randint(100, 900)
+            size = np.random.randint(50, 200)
+            color = (np.random.randint(0, 256), np.random.randint(0, 256), np.random.randint(0, 256))
+            cv2.rectangle(image, (x, y), (x + size, y + size), color, -1)
+        
+        # Save image
+        cv2.imwrite(os.path.join(output_dir, f'puzzle_{i}.jpg'), image)
+
+if __name__ == '__main__':
+    generate_synthetic_puzzle(10, 'synthetic_puzzles')

scripts/train.py

@@ -0,0 +1,30 @@
+from app.models.feedback_manager import FeedbackManager
+from app.models.matcher import PuzzleMatchingEngine
+
+def retrain_model():
+    feedback_manager = FeedbackManager()
+    training_data = feedback_manager.get_training_data()
+    
+    # Dummy retraining logic (adjust weights based on feedback)
+    matcher = PuzzleMatchingEngine()
+    positive_samples = training_data['positive']
+    negative_samples = training_data['negative']
+    
+    # Example: Increase weights for correct matches
+    if positive_samples:
+        matcher.weight_geometric += 0.01
+        matcher.weight_color += 0.01
+        matcher.weight_edge += 0.01
+        total = matcher.weight_geometric + matcher.weight_color + matcher.weight_edge
+        matcher.weight_geometric /= total
+        matcher.weight_color /= total
+        matcher.weight_edge /= total
+    
+    print("Model weights updated:", {
+        'geometric': matcher.weight_geometric,
+        'color': matcher.weight_color,
+        'edge': matcher.weight_edge
+    })
+
+if __name__ == '__main__':
+    retrain_model()

tests/test_feedback.py

@@ -0,0 +1,21 @@
+import unittest
+from app.models.feedback_manager import FeedbackManager
+
+class TestFeedback(unittest.TestCase):
+    def setUp(self):
+        self.feedback_manager = FeedbackManager()
+
+    def test_record_feedback(self):
+        self.feedback_manager.record_feedback(
+            piece_id=1,
+            slot_id=1,
+            suggestions={'piece_id': 1, 'confidence': 0.9},
+            user_choice=1,
+            is_correct=True
+        )
+        training_data = self.feedback_manager.get_training_data()
+        self.assertIn('positive', training_data)
+        self.assertGreaterEqual(len(training_data['positive']), 1)
+
+if __name__ == '__main__':
+    unittest.main()

tests/test_matching.py

@@ -0,0 +1,26 @@
+import unittest
+import numpy as np
+from app.models.matcher import PuzzleMatchingEngine
+
+class TestMatching(unittest.TestCase):
+    def setUp(self):
+        self.matcher = PuzzleMatchingEngine()
+        self.dummy_piece = {
+            'id': 0,
+            'features': {
+                'geometric': {'area': 1000, 'perimeter': 400, 'aspect_ratio': 1.0},
+                'color': np.zeros(512, dtype=np.float32),
+                'edge': np.zeros(10, dtype=np.float32)
+            }
+        }
+
+    def test_find_matches(self):
+        pieces = [self.dummy_piece]
+        slots = [self.dummy_piece]
+        matches = self.matcher.find_matches(pieces, slots)
+        self.assertIsInstance(matches, dict)
+        self.assertIn(0, matches)
+        self.assertGreaterEqual(len(matches[0]), 1)
+
+if __name__ == '__main__':
+    unittest.main()

tests/test_segmentation.py

@@ -0,0 +1,20 @@
+import unittest
+import cv2
+import numpy as np
+from app.models.segmenter import PuzzlePieceSegmenter
+
+class TestSegmentation(unittest.TestCase):
+    def setUp(self):
+        self.segmenter = PuzzlePieceSegmenter()
+        self.test_image = np.zeros((1000, 1000, 3), dtype=np.uint8)  # Dummy image
+
+    def test_segment_pieces(self):
+        pieces = self.segmenter.segment_pieces(self.test_image)
+        self.assertIsInstance(pieces, list)
+        for piece in pieces:
+            self.assertIn('id', piece)
+            self.assertIn('image', piece)
+            self.assertIn('features', piece)
+
+if __name__ == '__main__':
+    unittest.main()