【LogoDetection 数据集处理】（4）提取每张图片的logo区域

参考：【LogoDetection 数据集处理】（3）将训练集按照类别划分为多个文件夹

之前我们已经将训练集按照类别划分为50个文件夹，每个文件夹存在一个类别的图片：

在很多时候，我们需要提取出图片中的logo区域。

原图：

提取的logo区域：

从划分好的50和文件夹中的图片提取logo区域，得到50个类别的logo分别放在50个文件夹中。

代码如下：

""" 提取出每张图片的logo区域。 """""

import os
import json
import shutil
import cv2
from tqdm import tqdm


categories_list = [
    {
    
        "id": 1,
        "name": "bingdundun"
    },
    {
    
        "id": 2,
        "name": "sanyo"
    },
    {
    
        "id": 3,
        "name": "Eifini"
    },
    {
    
        "id": 4,
        "name": "PSALTER"
    },
    {
    
        "id": 5,
        "name": "Beaster"
    },
    {
    
        "id": 6,
        "name": "ON"
    },
    {
    
        "id": 7,
        "name": "BYREDO"
    },
    {
    
        "id": 8,
        "name": "Ubras"
    },
    {
    
        "id": 9,
        "name": "Eternelle"
    },
    {
    
        "id": 10,
        "name": "PerfectDiary"
    },
    {
    
        "id": 11,
        "name": "huaxizi"
    },
    {
    
        "id": 12,
        "name": "Clarins"
    },
    {
    
        "id": 13,
        "name": "Loccitane"
    },
    {
    
        "id": 14,
        "name": "Versace"
    },
    {
    
        "id": 15,
        "name": "Mizuno"
    },
    {
    
        "id": 16,
        "name": "Lining"
    },
    {
    
        "id": 17,
        "name": "DoubleStar"
    },
    {
    
        "id": 18,
        "name": "YONEX"
    },
    {
    
        "id": 19,
        "name": "ToryBurch"
    },
    {
    
        "id": 20,
        "name": "Gucci"
    },
    {
    
        "id": 21,
        "name": "LouisVuitton"
    },
    {
    
        "id": 22,
        "name": "CARTELO"
    },
    {
    
        "id": 23,
        "name": "JORDAN"
    },
    {
    
        "id": 24,
        "name": "KENZO"
    },
    {
    
        "id": 25,
        "name": "UNDEFEATED"
    },
    {
    
        "id": 26,
        "name": "BoyLondon"
    },
    {
    
        "id": 27,
        "name": "TREYO"
    },
    {
    
        "id": 28,
        "name": "carhartt"
    },
    {
    
        "id": 29,
        "name": "jierou"
    },
    {
    
        "id": 30,
        "name": "Blancpain"
    },
    {
    
        "id": 31,
        "name": "GXG"
    },
    {
    
        "id": 32,
        "name": "ledin"
    },
    {
    
        "id": 33,
        "name": "Diadora"
    },
    {
    
        "id": 34,
        "name": "TUCANO"
    },
    {
    
        "id": 35,
        "name": "Loewe"
    },
    {
    
        "id": 36,
        "name": "GraniteGear"
    },
    {
    
        "id": 37,
        "name": "DESCENTE"
    },
    {
    
        "id": 38,
        "name": "OSPREY"
    },
    {
    
        "id": 39,
        "name": "Swatch"
    },
    {
    
        "id": 40,
        "name": "erke"
    },
    {
    
        "id": 41,
        "name": "MassimoDutti"
    },
    {
    
        "id": 42,
        "name": "PINKO"
    },
    {
    
        "id": 43,
        "name": "PALLADIUM"
    },
    {
    
        "id": 44,
        "name": "origins"
    },
    {
    
        "id": 45,
        "name": "Trendiano"
    },
    {
    
        "id": 46,
        "name": "yiner"
    },
    {
    
        "id": 47,
        "name": "MonsterGuardians"
    },
    {
    
        "id": 48,
        "name": "fuerjia"
    },
    {
    
        "id": 49,
        "name": "IPSA"
    },
    {
    
        "id": 50,
        "name": "Schwarzkopf"
    }
]


# 划分好类别的训练集路径
data_path = "TrainSet_50Classes"
annoations_path = "dataset/fewshotlogodetection_round1_train_202204/train/annotations/instances_train2017.json"
classdir_list = os.listdir(data_path)


# 读取annoations的json文件
with open(annoations_path, 'r', encoding='utf-8') as f:
    annoations_dict = json.load(f)
images_list = annoations_dict["images"]
annotations_list = annoations_dict["annotations"]


# 创建最外层文件夹
logoarea_dir = "LogoArea_50classes"
isExists_trainset_dir = os.path.exists(logoarea_dir)
if not isExists_trainset_dir:
    os.mkdir(logoarea_dir)


# 创建50个子文件夹,以及class_id与文件夹名称的字典。
class_dict = {
    }
for i in range(len(categories_list)):
    dir_name = categories_list[i]["name"]
    dir_name_path = os.path.join(logoarea_dir, dir_name)
    class_dict[i + 1] = dir_name_path
    isExists = os.path.exists(dir_name_path)
    if isExists:
        shutil.rmtree(dir_name_path)
    os.makedirs(dir_name_path, exist_ok=True)


# 提取logo区域
for classdir_name in tqdm(classdir_list,desc="process"):
    classdir_path = os.path.join(data_path, classdir_name)
    pic_list = os.listdir(classdir_path)
    logo_num = 0    # 用于图片命名的序号
    for pic_name in pic_list:
        pic_path = os.path.join(classdir_path, pic_name)
        image = cv2.imread(pic_path)
        for i in range(len(images_list)):
            if pic_name == images_list[i]["file_name"]:
                image_id = images_list[i]["id"]
                for j in range(len(annotations_list)):
                    if image_id == annotations_list[j]["image_id"]:
                        bbox = annotations_list[j]["bbox"]
                        x0, y0, w, h = bbox
                        logo_area = image[y0:y0 + h, x0:x0 + w]   # 提取logo区域

                        # 图片命名从 000 开始
                        if logo_num < 10:
                            save_pic_name = classdir_name + "_00" + str(logo_num) + ".jpg"
                        elif logo_num < 100:
                            save_pic_name = classdir_name + "_0" + str(logo_num) + ".jpg"
                        else:
                            save_pic_name = classdir_name + "_" + str(logo_num) + ".jpg"
                        save_pic_path = os.path.join(logoarea_dir, classdir_name, save_pic_name)
                        cv2.imwrite(save_pic_path, logo_area)
                        logo_num += 1

得到一个 LogoArea_50classes 文件夹如下：