如何在Python中按指定规则顺序保存检测对象的图片文件名？

我需要将检测到的对象按规则生成图片文件名保存：同类对象序号连续递增，不受检测顺序影响。例如检测到对象列表[chair, tv, bed, chair, bed, chair, chair]，期望生成文件名：
[chair.png, chair1.png, chair2.png, chair3.png, bed.png, bed1.png, tv.png]

但当前得到的结果是：
[bed.png, bed2.png, chair.png, chair1.png, chair3.png, chair4.png, tv.png]

推测是count变量设置错误，但无法定位问题。

我的代码：

%cd /content/drive/MyDrive/Now_fine/yolov7-mask/binary_masks_images
folder_name = 'my_folder_' + time.strftime("%Y_%m_%d_%H_%M_%S")
os.mkdir(folder_name)
count = 0
for one_mask, bbox, cls, conf in zip(pred_masks_np, nbboxes, pred_cls, pred_conf):
    if conf < 0.25:
        continue
    else:
        label = names[int(cls)]
        print(label) 
        if os.path.exists('/content/drive/MyDrive/Now_fine/yolov7-mask/binary_masks_images/'+folder_name+'/'+label+'.png'):
            count+=1
            plt.imsave('/content/drive/MyDrive/Now_fine/yolov7-mask/binary_masks_images/'+folder_name+'/'+label+str(count)+'.png', one_mask, cmap='gray')
        else:
            plt.imsave('/content/drive/MyDrive/Now_fine/yolov7-mask/binary_masks_images/'+folder_name+'/'+label+'.png', one_mask, cmap='gray')

问题原因 #

你用了全局的count变量，所有类别共享同一个计数器，导致不同类别的编号互相干扰。比如先处理chair时count累加，后续处理bed时，count已经不是0，直接生成带大序号的文件名，出现跳号。

解决方案 #

给每个类别单独维护计数器，用字典存储每个类别的已保存数量，确保每个类别自己的编号从0开始连续递增。

修改后的代码：

%cd /content/drive/MyDrive/Now_fine/yolov7-mask/binary_masks_images
import os
import time
import matplotlib.pyplot as plt

folder_name = 'my_folder_' + time.strftime("%Y_%m_%d_%H_%M_%S")
os.mkdir(folder_name)
# 用字典存储每个类别的计数，每个类别独立统计
class_counter = {}

for one_mask, bbox, cls, conf in zip(pred_masks_np, nbboxes, pred_cls, pred_conf):
    if conf < 0.25:
        continue
    label = names[int(cls)]
    print(label)
    
    # 初始化该类别的计数器，不存在则设为0
    if label not in class_counter:
        class_counter[label] = 0
    
    current_count = class_counter[label]
    # 生成文件名：计数为0时用label.png，否则用label+序号.png
    filename = f"{label}.png" if current_count == 0 else f"{label}{current_count}.png"
    
    save_path = os.path.join(folder_name, filename)
    plt.imsave(save_path, one_mask, cmap='gray')
    
    # 对应类别计数加1
    class_counter[label] += 1

代码说明 #

1. class_counter字典：每个键对应一个类别，值是该类别已保存的图片数量
2. 处理每个对象时，先检查类别是否在字典中，不存在则初始化为0
3. 根据当前计数生成文件名，计数为0时不加序号，否则添加序号
4. 保存完成后，对应类别的计数加1，保证下一个同类别对象的序号连续

内容的提问来源于stack exchange，提问作者Sharif Bhatti