如何解决Python中调用next(os.walk())时出现的StopIteration错误？

Hey there, let's break down why you're hitting that StopIteration error and get your code working smoothly.

First, Why the StopIteration Error Happens #

The error comes from next(os.walk('FYP/Fig_Dataset'))[2] — os.walk() returns a generator that yields tuples of (root directory, subdirectories, files) as it traverses the path. When you call next() on it, you're asking for the first result. But if the path you specified doesn't exist, or there's nothing inside it (no files/subdirs), the generator runs out of items to yield, hence the StopIteration exception.

Looking at your error trace, you used FYP/Fig_Dataset in the failing line, but your code later uses MyDrive/FYP/Fig_Dataset — this mismatch is probably the main culprit. If you're running this in Google Colab, you also need to make sure you've mounted your Drive correctly, otherwise the path won't be accessible.

Step 1: Fix the Path Issues #

First, confirm your paths are correct:

1. If you're in Colab, mount your Drive first (you might have forgotten this step):

   from google.colab import drive
   drive.mount('/content/drive')
   

2. Use the full absolute path to your directories. In Colab, Drive files live under /content/drive/MyDrive/, so your paths should be:

   img_path = '/content/drive/MyDrive/FYP/Fig_Dataset'
   msk_path = '/content/drive/MyDrive/FYP/Ground_Truth'
   

3. Verify the paths exist before proceeding to avoid errors:

   import os
   if not os.path.exists(img_path):
       print(f"Oops! Can't find the image folder at {img_path}")
   if not os.path.exists(msk_path):
       print(f"Oops! Can't find the mask folder at {msk_path}")
   

Step 2: Safely Get Files with os.walk() #

Instead of using next() directly (which can fail), use a loop to get the files from the root directory (this does the same thing as next() but is safer):

img_files = []
for root, dirs, files in os.walk(img_path):
    img_files = files  # Grab all files in the root folder
    break  # Stop after the first iteration (we don't need subdirs)

msk_files = []
for root, dirs, files in os.walk(msk_path):
    msk_files = files
    break

Alternatively, if you still want to use next(), wrap it in a try-except block to handle empty paths:

try:
    img_files = next(os.walk(img_path))[2]
except StopIteration:
    print(f"No files found in {img_path}")
    img_files = []

Step 3: Fix Other Small Bugs in Your Code #

Your code has a few other issues that will cause problems once you fix the path:

1. Incorrect sort parameter: img_files.sort(2) is invalid — list.sort() doesn't take 2 as an argument. If you want to sort filenames alphabetically, just use img_files.sort(). If you need numeric sorting (e.g., 1.jpg, 2.jpg instead of 10.jpg, 1.jpg), use:

   img_files.sort(key=lambda x: int(x.split('.')[0]))
   

2. Bad initialization for X and Y: X = [2] and Y = [2] will add an integer 2 as the first element of your data lists, which will clash with the image arrays you append later. Initialize them as empty lists instead:

   X = []
   Y = []
   

3. Path concatenation issues: Using 'MyDrive//FYP/Fig_Dataset/{}' can lead to messy paths. Use os.path.join() to safely build paths across operating systems:

   img = cv2.imread(os.path.join(img_path, img_fl))
   msk_filename = f"{img_fl.split('.')[0]}_segmentation.png"
   msk = cv2.imread(os.path.join(msk_path, msk_filename), cv2.IMREAD_GRAYSCALE)
   

4. Check if masks exist: Add a check to make sure the corresponding mask file exists before trying to process it, so you don't get unexpected errors:

   if msk is not None:
       resized_msk = cv2.resize(msk, (256, 192), interpolation=cv2.INTER_CUBIC)
       Y.append(resized_msk)
   else:
       print(f"Warning: Mask not found for {img_fl}")
   

Full Fixed Code #

Here's the complete corrected code that addresses all these issues:

import os
import cv2
from tqdm import tqdm
from google.colab import drive

# Mount Google Drive (only needed for Colab)
drive.mount('/content/drive')

# Define correct paths
img_path = '/content/drive/MyDrive/FYP/Fig_Dataset'
msk_path = '/content/drive/MyDrive/FYP/Ground_Truth'

# Validate paths
if not os.path.exists(img_path):
    raise FileNotFoundError(f"Image directory not found: {img_path}")
if not os.path.exists(msk_path):
    raise FileNotFoundError(f"Mask directory not found: {msk_path}")

# Get files from directories
img_files = next(os.walk(img_path))[2]
msk_files = next(os.walk(msk_path))[2]

# Sort files (adjust key if you need numeric sorting)
img_files.sort()
msk_files.sort()

print(f"Found {len(img_files)} image files")
print(f"Found {len(msk_files)} mask files")

# Initialize data lists
X = []
Y = []

# Process each image and mask
for img_fl in tqdm(img_files):
    if img_fl.split('.')[-1] == 'jpg':
        # Load and process image
        img = cv2.imread(os.path.join(img_path, img_fl))
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        resized_img = cv2.resize(img, (256, 192), interpolation=cv2.INTER_CUBIC)
        X.append(resized_img)
        
        # Load and process corresponding mask
        msk_filename = f"{img_fl.split('.')[0]}_segmentation.png"
        msk = cv2.imread(os.path.join(msk_path, msk_filename), cv2.IMREAD_GRAYSCALE)
        if msk is not None:
            resized_msk = cv2.resize(msk, (256, 192), interpolation=cv2.INTER_CUBIC)
            Y.append(resized_msk)
        else:
            print(f"Skipping {img_fl} - no matching mask found")

内容的提问来源于stack exchange，提问作者NURNAJWA ERWINA MD ROSLI