Update app.py

app.py

@@ -67,112 +67,21 @@ def predict_next_move(fen, stockfish):
         return "Invalid FEN notation!"
 
     best_move = stockfish.get_best_move()
-    ans = transform_string(best_move)
-    return f"The predicted next move is: {ans}" if best_move else "No valid move found (checkmate/stalemate)."
+    return f"The predicted next move is: {best_move}" if best_move else "No valid move found (checkmate/stalemate)."
+
         
-def process_image(image_path):
-    # Ensure output directory exists
-    if not os.path.exists('output'):
-        os.makedirs('output')
-
-    # Load the segmentation model
-    segmentation_model = YOLO("segmentation.pt")
-
-    # Run inference to get segmentation results
-    results = segmentation_model.predict(
-        source=image_path,
-        conf=0.8  # Confidence threshold
-    )
-
-    # Initialize variables for the segmented mask and bounding box
-    segmentation_mask = None
-    bbox = None
-
-    for result in results:
-        if result.boxes.conf[0] >= 0.8:  # Filter results by confidence
-            segmentation_mask = result.masks.data.cpu().numpy().astype(np.uint8)[0]
-            bbox = result.boxes.xyxy[0].cpu().numpy()  # Get the bounding box coordinates
-            break
-
-    if segmentation_mask is None:
-        print("No segmentation mask with confidence above 0.8 found.")
-        return None
-
-    # Load the image
-    image = cv2.imread(image_path)
-
-    # Convert the image to RGB format
-    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-
-    # Resize segmentation mask to match the input image dimensions
-    segmentation_mask_resized = cv2.resize(segmentation_mask, (image.shape[1], image.shape[0]))
-
-    # Extract bounding box coordinates
-    if bbox is not None:
-        x1, y1, x2, y2 = bbox
-        # Crop the segmented region based on the bounding box
-        cropped_segment = image[int(y1):int(y2), int(x1):int(x2)]
-
-        # Convert the cropped segment to RGB
-        cropped_segment_rgb = cv2.cvtColor(cropped_segment, cv2.COLOR_BGR2RGB)
-
-        # Save the cropped segmented image
-        cropped_image_path = 'output/cropped_segment.jpg'
-        cv2.imwrite(cropped_image_path, cropped_segment)
-        print(f"Cropped segmented image saved to {cropped_image_path}")
-
-        # Display the image in Streamlit
-        st.image(cropped_segment_rgb, caption="Uploaded Image (Cropped)", use_column_width=True)
-
-        # Return the cropped RGB image
-        return cropped_segment_rgb
-
-
-def transform_string(input_str):
-    # Remove extra spaces and convert to lowercase
-    input_str = input_str.strip().lower()
-
-    # Check if input is valid
-    if len(input_str) != 4 or not input_str[0].isalpha() or not input_str[1].isdigit() or \
-       not input_str[2].isalpha() or not input_str[3].isdigit():
-        return "Invalid input"
-
-    # Define mappings
-    letter_mapping = {
-        'a': 'h', 'b': 'g', 'c': 'f', 'd': 'e',
-        'e': 'd', 'f': 'c', 'g': 'b', 'h': 'a'
-    }
-    number_mapping = {
-        '1': '8', '2': '7', '3': '6', '4': '5',
-        '5': '4', '6': '3', '7': '2', '8': '1'
-    }
-
-    # Transform string
-    result = ""
-    for i, char in enumerate(input_str):
-        if i % 2 == 0:  # Letters
-            result += letter_mapping.get(char, "Invalid")
-        else:  # Numbers
-            result += number_mapping.get(char, "Invalid")
-    
-    # Check for invalid transformations
-    if "Invalid" in result:
-        return "Invalid input"
-
-    return result
-
-
-
-# Streamlit app
-def main():
-    st.title("Chessboard Position Detection and Move Prediction")
-    
-    os.chmod("/home/user/app/stockfish-ubuntu-x86-64-sse41-popcnt", 0o755)
 
-   
 
 
 
+
+
+def main():
+    st.title("Chessboard Position Detection and Move Prediction")
+
+    # Set permissions for the Stockfish engine binary
+    os.chmod("/home/user/app/stockfish-ubuntu-x86-64-sse41-popcnt", 0o755)
+
     # User uploads an image or captures it from their camera
     image_file = st.camera_input("Capture a chessboard image") or st.file_uploader("Upload a chessboard image", type=["jpg", "jpeg", "png"])
 
@@ -184,96 +93,84 @@ def main():
         with open(temp_file_path, "wb") as f:
             f.write(image_file.getbuffer())
 
-        # Process the image using its file path
-        processed_image = process_image(temp_file_path)
-
-        if processed_image is not None:
-            # Resize the image to 224x224
-            processed_image = cv2.resize(processed_image, (224, 224))
-            height, width, _ = processed_image.shape
-
-            # Initialize the YOLO model
-            model = YOLO("fine_tuned_on_all_data.pt")  # Replace with your trained model weights file
-
-            # Run detection
-            results = model.predict(source=processed_image, save=False, save_txt=False, conf=0.5)
-
-            # Initialize the board for FEN (empty rows represented by "8")
-            board = [["8"] * 8 for _ in range(8)]
-
-            # Extract predictions and map to FEN board
-            for result in results[0].boxes:
-                x1, y1, x2, y2 = result.xyxy[0].tolist()
-                class_id = int(result.cls[0])
-                class_name = model.names[class_id]
-
-                # Convert class_name to FEN notation
-                fen_piece = FEN_MAPPING.get(class_name, None)
-                if not fen_piece:
-                    continue
-
-                # Calculate the center of the bounding box
-                center_x = (x1 + x2) / 2
-                center_y = (y1 + y2) / 2
-
-                # Convert to integer pixel coordinates
-                pixel_x = int(center_x)
-                pixel_y = int(height - center_y)  # Flip Y-axis for generic coordinate system
-
-                # Get grid coordinate
-                grid_position = get_grid_coordinate(pixel_x, pixel_y)
-
-                if grid_position != "Pixel outside grid bounds":
-                    file = ord(grid_position[0]) - ord('a')  # Column index (0-7)
-                    rank = int(grid_position[1]) - 1  # Row index (0-7)
-
-                    # Place the piece on the board
-                    board[7 - rank][file] = fen_piece  # Flip rank index for FEN
-
-            # Generate the FEN string
-            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)
-
-            position_fen = "/".join(fen_rows)
-
-            # Ask the user for the next move side
-            move_side = st.selectbox("Select the side to move:", ["w (White)", "b (Black)"])
-            move_side = "w" if move_side.startswith("w") else "b"
-
-            # Append the full FEN string continuation
-            fen_notation = f"{position_fen} {move_side} - - 0 0"
-
-            st.subheader("Generated FEN Notation:")
-            st.code(fen_notation)
-            
-            # Initialize the Stockfish engine
-            stockfish_path = os.path.join(os.getcwd(), "stockfish-ubuntu-x86-64-sse41-popcnt")
-            stockfish = Stockfish(
-            path=stockfish_path,
-            depth=15,
-            parameters={"Threads": 2, "Minimum Thinking Time": 30}
-            )
-
-            # Predict the next move
-            next_move = predict_next_move(fen_notation, stockfish)
-            st.subheader("Stockfish Recommended Move:")
-            st.write(next_move)
-
-        else:
-            st.error("Failed to process the image. Please try again.")
-
-if __name__ == "__main__":
-    main()
+        # Load the YOLO models
+        model = YOLO("fine_tuned_on_all_data.pt")  # Replace with your trained model weights file
+        seg_model = YOLO("segmentation.pt")
+
+        # Load and process the image
+        img = cv2.imread(temp_file_path)
+        r = seg_model.predict(source=temp_file_path)
+        xyxy = r[0].boxes.xyxy
+        x_min, y_min, x_max, y_max = map(int, xyxy[0])
+        new_img = img[y_min:y_max, x_min:x_max]
+
+        # Resize the image to 224x224
+        image = cv2.resize(new_img, (224, 224))
+        height, width, _ = image.shape
+
+        # Get user input for perspective
+        p = st.radio("Select perspective:", ["b (Black)", "w (White)"])
+        p = p[0].lower()
+
+        # Initialize the board for FEN (empty rows represented by "8")
+        board = [["8"] * 8 for _ in range(8)]
+
+        # Run detection
+        results = model.predict(source=image, save=False, save_txt=False, conf=0.6)
+
+        # Extract predictions and map to FEN board
+        for result in results[0].boxes:
+            x1, y1, x2, y2 = result.xyxy[0].tolist()
+            class_id = int(result.cls[0])
+            class_name = model.names[class_id]
+
+            fen_piece = FEN_MAPPING.get(class_name, None)
+            if not fen_piece:
+                continue
+
+            center_x = (x1 + x2) / 2
+            center_y = (y1 + y2) / 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 != "Pixel outside grid bounds":
+                file = ord(grid_position[0]) - ord('a')
+                rank = int(grid_position[1]) - 1
+                board[rank][file] = fen_piece
+
+        # Generate the FEN string
+        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)
+
+        position_fen = "/".join(fen_rows)
+        move_side = st.radio("Select the side to move:", ["w (White)", "b (Black)"])[0].lower()
+        fen_notation = f"{position_fen} {move_side} - - 0 0"
+        st.subheader("Generated FEN Notation:")
+        st.code(fen_notation)
+
+        # Initialize the Stockfish engine
+        stockfish_path = os.path.join(os.getcwd(), "stockfish-ubuntu-x86-64-sse41-popcnt")
+        stockfish = Stockfish(
+        path=stockfish_path,
+        depth=15,
+        parameters={"Threads": 2, "Minimum Thinking Time": 30}
+        )
+
+        # Predict the next move
+        next_move = predict_next_move(fen_notation, stockfish)
+        st.subheader("Stockfish Recommended Move:")
+        st.write(next_move)