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,15 +169,16 @@ def plot_to_result(self,
         plt.draw()
 
         try:
-            buf = fig.canvas.tostring_rgb()
+            buf = fig.canvas.buffer_rgba()
+            cols, rows = fig.canvas.get_width_height()
         except AttributeError:
             fig.canvas.draw()
-            buf = fig.canvas.tostring_rgb()
-        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)
+            buf = fig.canvas.buffer_rgba()
+            cols, rows = fig.canvas.get_width_height()
+        img_array = np.frombuffer(buf, dtype=np.uint8).reshape(rows, cols, 4)
+        img_array = cv2.cvtColor(img_array, cv2.COLOR_RGBA2BGR)
         plt.close()
-        return result
+        return img_array
 
     # Remark for refactoring: IMO a function should do one thing only, storing the image and plotting should be seperated and do not necessarily need to be class functions but standalone utility functions that the user can chain in his scripts to have more fine-grained control. 
     def plot(self,

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