Narendra9009 commited on
Commit
8d083d0
·
1 Parent(s): 7a1529d

changed get fen function for receiving image as input

Browse files
Files changed (1) hide show
  1. main.py +27 -31
main.py CHANGED
@@ -5,7 +5,7 @@ from PIL import Image, UnidentifiedImageError
5
  from routes.segmentation import segment_chess_board
6
  from routes.detection import detect_pieces
7
  from routes.fen_generator import gen_fen
8
- from typing import List, Dict, Any
9
  from pydantic import BaseModel
10
 
11
  app = FastAPI()
@@ -15,10 +15,6 @@ class DetectionResults(BaseModel):
15
  confidences: list
16
  classes: list
17
 
18
- class FenRequest(BaseModel):
19
- detections: DetectionResults
20
- perspective: str
21
-
22
  @app.get("/")
23
  async def read_root():
24
  return {
@@ -92,36 +88,36 @@ async def get_coords(file: UploadFile = File(...)):
92
 
93
 
94
  @app.post("/getFen")
95
- async def get_fen(request: FenRequest):
96
- results = request.detections
97
- perspective = request.perspective
 
 
98
  try:
99
- if perspective not in ["w", "b"]:
100
- return JSONResponse(
101
- content={"error": "Perspective must be 'w' (white) or 'b' (black)"},
102
- status_code=400
103
- )
104
-
105
- if not results.boxes or not results.confidences or not results.classes:
106
- return JSONResponse(
107
- content={"error": "Invalid input", "details": "Missing required fields"},
108
- status_code=400
109
- )
110
-
111
- print(results.model_dump())
112
 
113
- fen = gen_fen(results.model_dump(), perspective)
114
 
 
 
 
 
 
115
  if not fen:
116
- return JSONResponse(
117
- content={"error": "FEN generation failed", "details": "Invalid input data"},
118
- status_code=500
119
- )
120
 
121
  return JSONResponse(content={"FEN": fen}, status_code=200)
122
-
123
  except Exception as e:
124
- return JSONResponse(
125
- content={"error": "Unexpected error occurred", "details": str(e)},
126
- status_code=500
127
- )
 
5
  from routes.segmentation import segment_chess_board
6
  from routes.detection import detect_pieces
7
  from routes.fen_generator import gen_fen
8
+ from typing import List, Dict, Any, Union
9
  from pydantic import BaseModel
10
 
11
  app = FastAPI()
 
15
  confidences: list
16
  classes: list
17
 
 
 
 
 
18
  @app.get("/")
19
  async def read_root():
20
  return {
 
88
 
89
 
90
  @app.post("/getFen")
91
+ async def get_fen(file : UploadFile = File(), perspective : str = "w" | "b"):
92
+
93
+ if perspective not in ["w" , "b"]:
94
+ return JSONResponse(content={"error" : "Perspective should be w (white) or b (black)"}, status_code=500)
95
+
96
  try:
97
+ image_content = await file.read()
98
+ if not image_content:
99
+ return JSONResponse(content={"error": "Empty file uploaded"}, status_code=400)
100
+
101
+ try:
102
+ image = Image.open(io.BytesIO(image_content))
103
+ except UnidentifiedImageError:
104
+ return JSONResponse(content={"error": "Invalid image format"}, status_code=400)
105
+
106
+ segmented_image = await segment_chess_board(image)
107
+ if isinstance(segmented_image, dict):
108
+ return JSONResponse(content=segmented_image, status_code=400)
 
109
 
110
+ segmented_image = segmented_image.resize((224, 224))
111
 
112
+ detection_results = await detect_pieces(segmented_image)
113
+ if "error" in detection_results:
114
+ return JSONResponse(content=detection_results, status_code=400)
115
+
116
+ fen = gen_fen(detection_results, perspective)
117
  if not fen:
118
+ return JSONResponse(content={"error": "FEN generation failed", "details": "Invalid input data"}, status_code=500)
 
 
 
119
 
120
  return JSONResponse(content={"FEN": fen}, status_code=200)
121
+
122
  except Exception as e:
123
+ return JSONResponse(content={"error": "Unexpected error occurred", "details": str(e)}, status_code=500)