Fix bug in issue . ()

* Fix bug in .
When run 'python.exe segment/predict.py --visualize' will throw AttributeError: 'tuple' object has no attribute 'shape'

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
pull/10752/head^2
tiancailin 2024-01-03 15:07:19 +08:00 committed by GitHub
parent 52db52b927
commit bd1a8295fb
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 3 additions and 1 deletions

View File

@ -59,7 +59,9 @@ def feature_visualization(x, module_type, stage, n=32, save_dir=Path('runs/detec
n: Maximum number of feature maps to plot
save_dir: Directory to save results
"""
if 'Detect' not in module_type:
if ('Detect'
not in module_type) and ('Segment'
not in module_type): # 'Detect' for Object Detect task,'Segment' for Segment task
batch, channels, height, width = x.shape # batch, channels, height, width
if height > 1 and width > 1:
f = save_dir / f"stage{stage}_{module_type.split('.')[-1]}_features.png" # filename