Spaces:
Sleeping
Sleeping
fixed path to data
Browse files
agentcourse_unit4/tools/chess_board_recognizer.py
CHANGED
|
@@ -1,3 +1,4 @@
|
|
|
|
|
| 1 |
import tempfile
|
| 2 |
from typing import Optional
|
| 3 |
|
|
@@ -38,8 +39,11 @@ class ChessBoardRecognizerTool(Tool):
|
|
| 38 |
|
| 39 |
def __init__(self):
|
| 40 |
super().__init__()
|
| 41 |
-
|
| 42 |
-
|
|
|
|
|
|
|
|
|
|
| 43 |
|
| 44 |
def forward(self, image_path: str, is_white_turn: Optional[bool] = None) -> str:
|
| 45 |
processed_image = self._process_image(image_path)
|
|
|
|
| 1 |
+
import os
|
| 2 |
import tempfile
|
| 3 |
from typing import Optional
|
| 4 |
|
|
|
|
| 39 |
|
| 40 |
def __init__(self):
|
| 41 |
super().__init__()
|
| 42 |
+
|
| 43 |
+
parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
|
| 44 |
+
|
| 45 |
+
self.model_std = YOLO(f"{parent_dir}/data/standard.pt")
|
| 46 |
+
self.model_seg = YOLO(f"{parent_dir}/data/segmentation.pt")
|
| 47 |
|
| 48 |
def forward(self, image_path: str, is_white_turn: Optional[bool] = None) -> str:
|
| 49 |
processed_image = self._process_image(image_path)
|