Fix compatibility with PyTorch 2.6+ and matplotlib >=3.6#283
Fix compatibility with PyTorch 2.6+ and matplotlib >=3.6#283Xyc2016 wants to merge 1 commit intoCASIA-LMC-Lab:mainfrom
Conversation
There was a problem hiding this comment.
Pull request overview
Note
Copilot was unable to run its full agentic suite in this review.
This PR updates FastSAM to run on newer PyTorch (2.6+) and matplotlib (3.6+) by addressing breaking API changes and a missing runtime dependency.
Changes:
- Update
torch.load()calls to explicitly setweights_only=Falsefor PyTorch 2.6+ compatibility. - Replace deprecated/removed matplotlib canvas RGB extraction with RGBA buffer extraction and update numpy buffer conversion.
- Add missing
psutildependency torequirements.txt.
Reviewed changes
Copilot reviewed 3 out of 4 changed files in this pull request and generated 5 comments.
| File | Description |
|---|---|
| utils/tools.py | Switch figure-to-image extraction to buffer_rgba() + RGBA→BGR conversion and use np.frombuffer. |
| fastsam/prompt.py | Update plot rendering pipeline to use buffer_rgba() and RGBA→BGR conversion. |
| ultralytics/nn/tasks.py | Force torch.load(..., weights_only=False) to load checkpoints that include custom classes. |
| requirements.txt | Add psutil to avoid ModuleNotFoundError on fresh installs. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| 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 |
There was a problem hiding this comment.
torch_safe_load is intended to be a safer loader, but setting weights_only=False re-enables full pickle deserialization, which can lead to arbitrary code execution when loading untrusted .pt files. Recommendation: attempt torch.load(..., weights_only=True) first and only fall back to weights_only=False for a narrowly-scoped set of trusted/official weights (or gate it behind an explicit user opt-in / clear warning), or use PyTorch safe loading mechanisms (e.g., allowlisting needed globals) so third-party weights remain safe by default.
| 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 |
There was a problem hiding this comment.
torch_safe_load is intended to be a safer loader, but setting weights_only=False re-enables full pickle deserialization, which can lead to arbitrary code execution when loading untrusted .pt files. Recommendation: attempt torch.load(..., weights_only=True) first and only fall back to weights_only=False for a narrowly-scoped set of trusted/official weights (or gate it behind an explicit user opt-in / clear warning), or use PyTorch safe loading mechanisms (e.g., allowlisting needed globals) so third-party weights remain safe by default.
| 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)) |
There was a problem hiding this comment.
Catching AttributeError around buffer_rgba() is misleading here: if buffer_rgba truly doesn’t exist, calling it again in the except will still fail. If the intent is to ensure the canvas is rendered before reading the buffer, it’s clearer and more deterministic to always call fig.canvas.draw() (or fig.canvas.draw_idle() + flush) before buffer_rgba() and remove the try/except.
| 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 |
There was a problem hiding this comment.
Same issue as utils/tools.py: the AttributeError fallback won’t help if buffer_rgba is absent, and the duplicated get_width_height() is unnecessary. Calling fig.canvas.draw() unconditionally before reading the buffer allows you to remove the try/except and compute cols, rows once, simplifying the control flow.
Xyc2016
force-pushed
the
fix/compat-pytorch-matplotlib
branch
from
9389011 to
cc2cae1
Compare
- 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>
Xyc2016
force-pushed
the
fix/compat-pytorch-matplotlib
branch
from
cc2cae1 to
fb4c852
Compare
Summary
Fix three compatibility issues that prevent FastSAM from running with current PyTorch (2.6+) and matplotlib (3.6+).
Changes
1. PyTorch 2.6+ model loading
PyTorch 2.6+ changed
torch.load()defaultweights_onlyfromFalsetoTrue, breaking loading of official checkpoints that contain custom ultralytics classes.Approach: Try
weights_only=Truefirst, catchpickle.UnpicklingError/RuntimeError, then fall back toweights_only=Falsewith a warning.2. matplotlib >=3.6 rendering
matplotlib 3.6+ removed
FigureCanvas.tostring_rgb(). Replaced withbuffer_rgba()(available since matplotlib 3.1) in bothfastsam/prompt.pyandutils/tools.py. Also removed the broken try/except pattern that retried the same failing call.3. Missing psutil dependency
ultralytics/yolo/utils/checks.pyimportspsutilbut it was missing fromrequirements.txt.Files changed
fastsam/prompt.pytostring_rgb()→buffer_rgba(), RGB→RGBAutils/tools.pytostring_rgb()→buffer_rgba(),np.fromstring→np.frombuffer, RGB→RGBAultralytics/nn/tasks.pyweights_only=Truewith safe fallback toFalserequirements.txtpsutil