Update app.py

made changes to user our segmentation model instead of roboflow one

app.py

@@ -95,63 +95,29 @@ def main():
             f.write(image_file.getbuffer())
 
         # Load the YOLO models
-        model = YOLO("fine_tuned_on_all_data.pt")  # Replace with your trained model weights file
+        model = YOLO("ChessDetection3d")  # Replace with your trained model weights file
 
         
-        # seg_model = YOLO("segmentation.pt")
+        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]
-
-
-
-        
-        image = cv2.imread(temp_file_path)
-        
-        result = CLIENT.infer(image, model_id="chessboard-segmentation/1")
-        
-        if image is None:
-            st.write("Error: Image not loaded.")
-        
-        
-        prediction_data = result
+        # Load and process the image
+        img = cv2.imread(temp_file_path)
+        r = seg_model.predict(source=temp_file_path)
         
-        
-        if not prediction_data.get('predictions'):
-            st.write("No board found.")
-            return 
-        else:
-            for prediction in prediction_data.get('predictions', []):
-                if 'x' in prediction and 'y' in prediction and 'width' in prediction and 'height' in prediction:
-                    x, y, w, h = prediction['x'], prediction['y'], prediction['width'], prediction['height']
-                    # print(f"Bounding box coordinates: ({x}, {y}), width={w}, height={h}")
-            
-                    x1, y1 = int(x - w / 2), int(y - h / 2)
-                    x2, y2 = int(x + w / 2), int(y + h / 2)
-            
-                    src_pts = np.array([[x1, y1], [x2, y1], [x2, y2], [x1, y2]], dtype="float32")
-                    # print(f"Source Points: {src_pts}")
-            
-                    chessboard_size = 600
-                    dst_pts = np.array([[0, 0], [chessboard_size - 1, 0], [chessboard_size - 1, chessboard_size - 1], [0, chessboard_size - 1]], dtype="float32")
-                    # print(f"Destination Points: {dst_pts}")
+        if len(r) == 0 || len(r) > 1:
+            if len(r) == 0:
+                st.write("NO BOARD IN THE IMAGE")
+            elif len(r) > 1:
+                st.write("Multiple boards are there in the image, please take only at a time")
+            return
             
-                    matrix = cv2.getPerspectiveTransform(src_pts, dst_pts)
-            
-                    # Apply the perspective warp
-                    transformed_chessboard = cv2.warpPerspective(image, matrix, (chessboard_size, chessboard_size))
-                    
-            
-                    # Convert images to RGB for display
-                    new_img = cv2.cvtColor(transformed_chessboard, cv2.COLOR_BGR2RGB)
-                    
+        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))
         st.image(image, caption="Segmented Chessboard", use_container_width=True)