Fix compatibility with PyTorch 2.6+ and matplotlib >=3.6

- Fix torch.load() weights_only default change (PyTorch 2.6+)
  Add weights_only=False to load official model checkpoints
- Fix FigureCanvas.tostring_rgb() removal (matplotlib >=3.6)
  Replace with buffer_rgba() and update color channel handling
- Add missing psutil dependency to requirements.txt

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: Xyc2016

fastsam/prompt.py

@@ -169,13 +169,13 @@ def plot_to_result(self,
         plt.draw()
 
         try:
-            buf = fig.canvas.tostring_rgb()
+            buf = fig.canvas.buffer_rgba()
         except AttributeError:
             fig.canvas.draw()
-            buf = fig.canvas.tostring_rgb()
+            buf = fig.canvas.buffer_rgba()
         cols, rows = fig.canvas.get_width_height()
-        img_array = np.frombuffer(buf, dtype=np.uint8).reshape(rows, cols, 3)
-        result = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)
+        img_array = np.frombuffer(buf, dtype=np.uint8).reshape(rows, cols, 4)
+        result = cv2.cvtColor(img_array, cv2.COLOR_RGBA2BGR)
         plt.close()
         return result
 

requirements.txt

@@ -8,6 +8,7 @@ scipy>=1.4.1
 torch>=1.7.0
 torchvision>=0.8.1
 tqdm>=4.64.0
+psutil
 
 pandas>=1.1.4
 seaborn>=0.11.0

ultralytics/nn/tasks.py

@@ -515,7 +515,7 @@ def torch_safe_load(weight):
     check_suffix(file=weight, suffix='.pt')
     file = attempt_download_asset(weight)  # search online if missing locally
     try:
-        return torch.load(file, map_location='cpu'), file  # load
+        return torch.load(file, map_location='cpu', weights_only=False), file  # load
     except ModuleNotFoundError as e:  # e.name is missing module name
         if e.name == 'models':
             raise TypeError(
@@ -530,7 +530,7 @@ def torch_safe_load(weight):
                        f"run a command with an official YOLOv8 model, i.e. 'yolo predict model=yolov8n.pt'")
         check_requirements(e.name)  # install missing module
 
-        return torch.load(file, map_location='cpu'), file  # load
+        return torch.load(file, map_location='cpu', weights_only=False), file  # load
 
 
 def attempt_load_weights(weights, device=None, inplace=True, fuse=False):

utils/tools.py

@@ -180,14 +180,14 @@ def fast_process(
     plt.draw()
 
     try:
-        buf = fig.canvas.tostring_rgb()
+        buf = fig.canvas.buffer_rgba()
     except AttributeError:
         fig.canvas.draw()
-        buf = fig.canvas.tostring_rgb()
-    
+        buf = fig.canvas.buffer_rgba()
+
     cols, rows = fig.canvas.get_width_height()
-    img_array = np.fromstring(buf, dtype=np.uint8).reshape(rows, cols, 3)
-    cv2.imwrite(os.path.join(save_path, result_name), cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR))
+    img_array = np.frombuffer(buf, dtype=np.uint8).reshape(rows, cols, 4)
+    cv2.imwrite(os.path.join(save_path, result_name), cv2.cvtColor(img_array, cv2.COLOR_RGBA2BGR))
 
 
 # CPU post process