initial commit

main.py

@@ -0,0 +1,127 @@
+import io
+from fastapi import FastAPI, File, UploadFile
+from fastapi.responses import JSONResponse, StreamingResponse
+from PIL import Image, UnidentifiedImageError
+from routes.segmentation import segment_chess_board
+from routes.detection import detect_pieces
+from routes.fen_generator import gen_fen
+from typing import List, Dict, Any
+from pydantic import BaseModel
+
+app = FastAPI()
+
+class DetectionResults(BaseModel):
+    boxes: list
+    confidences: list
+    classes: list
+
+class FenRequest(BaseModel):
+    detections: DetectionResults
+    perspective: str
+
+@app.get("/")
+async def read_root():
+    return {
+        "name": "Narendra",
+        "age": 20,
+        "Gender": "Male"
+    }
+
+@app.post("/getSeg")
+async def get_seg(file: UploadFile = File(...)):
+    print(f'Image received: {file.filename}')
+
+    try:
+        image_content = await file.read()
+        if not image_content:
+            return JSONResponse(content={"error": "Empty file uploaded"}, status_code=400)
+
+        try:
+            image = Image.open(io.BytesIO(image_content))
+        except UnidentifiedImageError:
+            return JSONResponse(content={"error": "Invalid image format"}, status_code=400)
+
+        # If segment_chess_board is async, use `await`, otherwise remove `await`
+        segmented_image = await segment_chess_board(image)
+
+        if isinstance(segmented_image, dict):
+            return JSONResponse(content=segmented_image, status_code=400)
+
+        # Save to in-memory bytes
+        img_bytes = io.BytesIO()
+        segmented_image.save(img_bytes, format="PNG")
+        img_bytes.seek(0)
+
+        print("Image successfully processed and returned")
+        return StreamingResponse(
+            img_bytes,
+            media_type="image/png",
+            headers={"Content-Disposition": "inline; filename=output.png"}
+        )
+
+
+    except Exception as e:
+        return JSONResponse(content={"error": str(e)}, status_code=500)
+    
+
+@app.post("/getCoords")
+async def get_coords(file: UploadFile = File(...)):
+    try:
+        image_content = await file.read()
+
+        if not image_content:
+            print("No image found")
+            return JSONResponse(content={"error": "Empty file uploaded"}, status_code=400)
+
+        try:
+            image = Image.open(io.BytesIO(image_content))
+        except UnidentifiedImageError:
+            return JSONResponse(content={"error": "Invalid image format"}, status_code=400)
+
+        detection_results = await detect_pieces(image)
+
+        if "error" in detection_results:
+            return JSONResponse(content=detection_results, status_code=400)
+
+        print("Image successfully processed and returned")
+        return JSONResponse(content={"detections": detection_results}, status_code=200)
+
+    except Exception as e:
+        print(f"Unexpected error: {str(e)}")
+        return JSONResponse(content={"error": "Unexpected error occurred", "details": str(e)}, status_code=500)
+    
+
+@app.post("/getFen")
+async def get_fen(request: FenRequest):
+    results = request.detections
+    perspective = request.perspective
+    try:
+        if perspective not in ["w", "b"]:
+            return JSONResponse(
+                content={"error": "Perspective must be 'w' (white) or 'b' (black)"},
+                status_code=400
+            )
+
+        if not results.boxes or not results.confidences or not results.classes:
+            return JSONResponse(
+                content={"error": "Invalid input", "details": "Missing required fields"},
+                status_code=400
+            )
+        
+        print(results.model_dump())
+
+        fen = gen_fen(results.model_dump(), perspective)
+
+        if not fen:
+            return JSONResponse(
+                content={"error": "FEN generation failed", "details": "Invalid input data"},
+                status_code=500
+            )
+
+        return JSONResponse(content={"FEN": fen}, status_code=200)
+
+    except Exception as e:
+        return JSONResponse(
+            content={"error": "Unexpected error occurred", "details": str(e)},
+            status_code=500
+        )

routes/detection.py

@@ -0,0 +1,29 @@
+from ultralytics import YOLO
+from PIL import Image
+detect_model = YOLO(r'D:\venv\chess-vision\models\chessDetection3d.pt')
+
+async def detect_pieces(image : Image):
+    if image is None:
+        print("No image is there")
+        return {"error" : "No image detected"}
+    
+    results = detect_model.predict(image)
+
+    if not results or len(results) == 0:
+        print("No results are there")
+        return {"error" : "No results found"}
+    
+    boxes = results[0].boxes.xyxy.tolist()  
+    confidences = results[0].boxes.conf.tolist()  
+    classes = results[0].boxes.cls.tolist()  
+
+    class_names = []
+
+    for idx in classes:
+        class_names.append(detect_model.names[idx])
+
+    return {
+        "boxes": boxes,
+        "confidences": confidences,
+        "classes": class_names
+    }

routes/fen_generator.py

@@ -0,0 +1,114 @@
+import json
+
+FEN_MAPPING = {
+    "black-pawn": "p", "black-rook": "r", "black-knight": "n", "black-bishop": "b", "black-queen": "q", "black-king": "k",
+    "white-pawn": "P", "white-rook": "R", "white-knight": "N", "white-bishop": "B", "white-queen": "Q", "white-king": "K"
+}
+
+# there are some issues in the code i,e in the line 87 , possible modification for that is   rank = int(grid_position[1]) - 1
+
+# Grid settings
+border = 0  
+grid_size = 224  
+block_size = grid_size // 8  
+
+x_labels = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h']  
+y_labels = [8, 7, 6, 5, 4, 3, 2, 1]  
+
+def get_grid_coordinate(pixel_x, pixel_y, perspective):
+    try:
+        adjusted_x = pixel_x - border
+        adjusted_y = pixel_y - border
+
+        if adjusted_x < 0 or adjusted_y < 0 or adjusted_x >= grid_size or adjusted_y >= grid_size:
+            return None  
+
+        x_index = adjusted_x // block_size
+        y_index = adjusted_y // block_size
+
+        if x_index < 0 or x_index >= len(x_labels) or y_index < 0 or y_index >= len(y_labels):
+            return None  
+
+        if perspective == "b":
+            x_index = 7 - x_index  
+            y_index = 7 - y_index  
+
+        file = x_labels[x_index]
+        rank = y_labels[y_index]
+
+        return f"{file}{rank}"
+    except Exception as e:
+        print(f"Error in get_grid_coordinate: {e}")
+        return None  
+
+def gen_fen(result: dict, p: str):
+    try:
+        if not isinstance(result, dict):
+            print("Error: Expected a dictionary for result")
+            return None
+
+        boxes = result.get("boxes", [])
+        classes = result.get("classes", [])
+
+        if not boxes or not classes:
+            print("Error: Missing 'boxes' or 'classes' in input")
+            return None
+
+        if len(boxes) != len(classes):
+            print("Error: Mismatch between bounding boxes and class labels")
+            return None
+
+        height, width = 224, 224  
+        board = [["8"] * 8 for _ in range(8)]  
+
+        for box, class_name in zip(boxes, classes):
+            if not isinstance(box, (list, tuple)) or len(box) != 4:
+                print(f"Skipping invalid box: {box}")
+                continue  
+
+            fen_piece = FEN_MAPPING.get(class_name, None)
+            if not fen_piece:
+                print(f"Skipping unrecognized piece: {class_name}")
+                continue  
+
+            try:
+                x_min, y_min, x_max, y_max = map(int, box)
+            except ValueError:
+                print(f"Skipping box with invalid values: {box}")
+                continue  
+
+            center_x, center_y = (x_min + x_max) / 2, (y_min + y_max) / 2
+            pixel_x = int(center_x)
+            pixel_y = int(height - center_y)  
+
+            grid_position = get_grid_coordinate(pixel_x, pixel_y, p)
+            if grid_position:
+                file = ord(grid_position[0]) - ord('a')
+                rank = int(grid_position[1]) - 1
+
+                if 0 <= rank < 8 and 0 <= file < 8:
+                    board[rank][file] = fen_piece
+                else:
+                    print(f"Skipping out-of-bounds grid position: {grid_position}")
+
+        fen_rows = []
+        for row in board:
+            fen_row = ""
+            empty_count = 0
+            for cell in row:
+                if cell == "8":
+                    empty_count += 1
+                else:
+                    if empty_count > 0:
+                        fen_row += str(empty_count)
+                        empty_count = 0
+                    fen_row += cell
+            if empty_count > 0:
+                fen_row += str(empty_count)
+            fen_rows.append(fen_row)  # FIXED: Ensured last row is added
+
+        return "/".join(fen_rows)
+
+    except Exception as e:
+        print(f"Error in gen_fen: {e}")
+        return None  

routes/segmentation.py

@@ -0,0 +1,23 @@
+from ultralytics import YOLO
+from PIL import Image
+
+seg_model = YOLO(r'D:\venv\chess-vision\models\SegModel (1).pt')
+
+async def segment_chess_board(image : Image):
+    if image is None:
+        return {"error" : "No image found" }    
+    
+    results = seg_model.predict(image)
+
+    if not results or len(results) == 0:
+        return {"error" : "No chessboard detected"}
+
+    if len(results) > 1:
+        return {"error" : "Multiple  chess boards found in the image"}
+    
+    xywh = results[0].boxes.xyxy[0].tolist()
+    x_min, y_min, x_max, y_max = map(int, xywh)
+    
+    segmented_image = image.crop((x_min, y_min, x_max, y_max)) 
+
+    return segmented_image