Severstal Dataset to COCO

简介

Kaggle上的钢板表面缺陷检测数据集Severstal给的是一个csv格式的标注文件，不是一般的coco或者voc格式。这里给出将其转为coco格式的代码。对于数据集的分析可以参考这篇文章：https://www.qixinbo.info/2020/02/15/kaggle-steel/

代码

# %%
import numpy as np
import pandas as pd
import os, cv2, json, random, shutil
from collections import defaultdict
from pycocotools.coco import COCO
from PIL import Image
# %% [markdown]

# # 读取数据
# %%

train_df = pd.read_csv("train.csv")

train_df.head()
# %% [markdown]

# # 有无缺陷划分
# %%

idx_no_defect = []

idx_defect = []
for col in range(0, len(train_df), 4):

img_names = [str(i).split("_")[0] for i in train_df.iloc[col:col+4, 0].values]

if not (img_names[0]  img_names[1]  img_names[2] == img_names[3]):

raise ValueError
labels <span class="token operator">=</span> train_df<span class="token punctuation">.</span>iloc<span class="token punctuation">[</span>col<span class="token punctuation">:</span>col<span class="token operator">+</span><span class="token number">4</span><span class="token punctuation">,</span> <span class="token number">1</span><span class="token punctuation">]</span>
<span class="token keyword">if</span> labels<span class="token punctuation">.</span>isna<span class="token punctuation">(</span><span class="token punctuation">)</span><span class="token punctuation">.</span><span class="token builtin">all</span><span class="token punctuation">(</span><span class="token punctuation">)</span><span class="token punctuation">:</span>
    idx_no_defect<span class="token punctuation">.</span>append<span class="token punctuation">(</span>col<span class="token punctuation">)</span>
<span class="token keyword">else</span><span class="token punctuation">:</span>
    idx_defect<span class="token punctuation">.</span>append<span class="token punctuation">(</span>col<span class="token punctuation">)</span>

# %% [markdown]

# # 标注掩码处理函数
# %%

def name_and_mask(start_idx):

col = start_idx

img_names = [str(i).split("_")[0] for i in train_df.iloc[col:col+4, 0].values]

if not (img_names[0]  img_names[1]  img_names[2] == img_names[3]):

raise ValueError
labels <span class="token operator">=</span> train_df<span class="token punctuation">.</span>iloc<span class="token punctuation">[</span>col<span class="token punctuation">:</span>col<span class="token operator">+</span><span class="token number">4</span><span class="token punctuation">,</span> <span class="token number">1</span><span class="token punctuation">]</span>
mask <span class="token operator">=</span> np<span class="token punctuation">.</span>zeros<span class="token punctuation">(</span><span class="token punctuation">(</span><span class="token number">256</span><span class="token punctuation">,</span> <span class="token number">1600</span><span class="token punctuation">,</span> <span class="token number">4</span><span class="token punctuation">)</span><span class="token punctuation">,</span> dtype<span class="token operator">=</span>np<span class="token punctuation">.</span>uint8<span class="token punctuation">)</span>

<span class="token keyword">for</span> idx<span class="token punctuation">,</span> label <span class="token keyword">in</span> <span class="token builtin">enumerate</span><span class="token punctuation">(</span>labels<span class="token punctuation">.</span>values<span class="token punctuation">)</span><span class="token punctuation">:</span>
    <span class="token keyword">if</span> label <span class="token keyword">is</span> <span class="token keyword">not</span> np<span class="token punctuation">.</span>nan<span class="token punctuation">:</span>
        mask_label <span class="token operator">=</span> np<span class="token punctuation">.</span>zeros<span class="token punctuation">(</span><span class="token number">1600</span><span class="token operator">*</span><span class="token number">256</span><span class="token punctuation">,</span> dtype<span class="token operator">=</span>np<span class="token punctuation">.</span>uint8<span class="token punctuation">)</span>
        label <span class="token operator">=</span> label<span class="token punctuation">.</span>split<span class="token punctuation">(</span><span class="token string">" "</span><span class="token punctuation">)</span>
        positions <span class="token operator">=</span> <span class="token builtin">map</span><span class="token punctuation">(</span><span class="token builtin">int</span><span class="token punctuation">,</span> label<span class="token punctuation">[</span><span class="token punctuation">:</span><span class="token punctuation">:</span><span class="token number">2</span><span class="token punctuation">]</span><span class="token punctuation">)</span>
        length <span class="token operator">=</span> <span class="token builtin">map</span><span class="token punctuation">(</span><span class="token builtin">int</span><span class="token punctuation">,</span> label<span class="token punctuation">[</span><span class="token number">1</span><span class="token punctuation">:</span><span class="token punctuation">:</span><span class="token number">2</span><span class="token punctuation">]</span><span class="token punctuation">)</span>
        <span class="token keyword">for</span> pos<span class="token punctuation">,</span> le <span class="token keyword">in</span> <span class="token builtin">zip</span><span class="token punctuation">(</span>positions<span class="token punctuation">,</span> length<span class="token punctuation">)</span><span class="token punctuation">:</span>
            mask_label<span class="token punctuation">[</span>pos<span class="token operator">-</span><span class="token number">1</span><span class="token punctuation">:</span>pos<span class="token operator">+</span>le<span class="token operator">-</span><span class="token number">1</span><span class="token punctuation">]</span> <span class="token operator">=</span> <span class="token number">1</span>
        mask<span class="token punctuation">[</span><span class="token punctuation">:</span><span class="token punctuation">,</span> <span class="token punctuation">:</span><span class="token punctuation">,</span> idx<span class="token punctuation">]</span> <span class="token operator">=</span> mask_label<span class="token punctuation">.</span>reshape<span class="token punctuation">(</span><span class="token number">256</span><span class="token punctuation">,</span> <span class="token number">1600</span><span class="token punctuation">,</span> order<span class="token operator">=</span><span class="token string">'F'</span><span class="token punctuation">)</span>

<span class="token keyword">return</span> img_names<span class="token punctuation">[</span><span class="token number">0</span><span class="token punctuation">]</span><span class="token punctuation">,</span> mask

def mask_to_rle(mask):

'''

Convert a mask into RLE
Parameters: 
mask (numpy.array): binary mask of numpy array where 1 - mask, 0 - background

Returns: 
sring: run length encoding 
'''</span>
pixels<span class="token operator">=</span> mask<span class="token punctuation">.</span>T<span class="token punctuation">.</span>flatten<span class="token punctuation">(</span><span class="token punctuation">)</span><span class="token comment">#转置之后再flatten同上</span>
pixels <span class="token operator">=</span> np<span class="token punctuation">.</span>concatenate<span class="token punctuation">(</span><span class="token punctuation">[</span><span class="token punctuation">[</span><span class="token number">0</span><span class="token punctuation">]</span><span class="token punctuation">,</span> pixels<span class="token punctuation">,</span> <span class="token punctuation">[</span><span class="token number">0</span><span class="token punctuation">]</span><span class="token punctuation">]</span><span class="token punctuation">)</span><span class="token comment">#加0是为了后面错位-</span>
runs <span class="token operator">=</span> np<span class="token punctuation">.</span>where<span class="token punctuation">(</span>pixels<span class="token punctuation">[</span><span class="token number">1</span><span class="token punctuation">:</span><span class="token punctuation">]</span> <span class="token operator">!=</span> pixels<span class="token punctuation">[</span><span class="token punctuation">:</span><span class="token operator">-</span><span class="token number">1</span><span class="token punctuation">]</span><span class="token punctuation">)</span><span class="token punctuation">[</span><span class="token number">0</span><span class="token punctuation">]</span> <span class="token operator">+</span> <span class="token number">1</span><span class="token comment">#这里要+1</span>
<span class="token comment">#其实就是再计算后面位置和前面位置的值相等；不等就输出索引号+1</span>
runs<span class="token punctuation">[</span><span class="token number">1</span><span class="token punctuation">:</span><span class="token punctuation">:</span><span class="token number">2</span><span class="token punctuation">]</span> <span class="token operator">-=</span> runs<span class="token punctuation">[</span><span class="token punctuation">:</span><span class="token punctuation">:</span><span class="token number">2</span><span class="token punctuation">]</span>
<span class="token comment">#每一个不相等索引号相减，后一个-前一个就是长度</span>
<span class="token keyword">return</span> <span class="token string">' '</span><span class="token punctuation">.</span>join<span class="token punctuation">(</span><span class="token builtin">str</span><span class="token punctuation">(</span>x<span class="token punctuation">)</span> <span class="token keyword">for</span> x <span class="token keyword">in</span> runs<span class="token punctuation">)</span>

#主要生成mask，根据rle_string(这玩意就是上面的EncodedPixels，图像标签)

# from https://www.kaggle.com/robertkag/rle-to-mask-converter

def rle_to_mask(rle_string,height,width):

'''

convert RLE(run length encoding) string to numpy array
Parameters: 
rleString (str): Description of arg1 
height (int): height of the mask
width (int): width of the mask 

Returns: 
numpy.array: numpy array of the mask
'''</span>
rows<span class="token punctuation">,</span> cols <span class="token operator">=</span> height<span class="token punctuation">,</span> width
<span class="token keyword">if</span> rle_string <span class="token operator">==</span> <span class="token operator">-</span><span class="token number">1</span><span class="token punctuation">:</span>
    <span class="token keyword">return</span> np<span class="token punctuation">.</span>zeros<span class="token punctuation">(</span><span class="token punctuation">(</span>height<span class="token punctuation">,</span> width<span class="token punctuation">)</span><span class="token punctuation">)</span>
rleNumbers <span class="token operator">=</span> <span class="token punctuation">[</span><span class="token builtin">int</span><span class="token punctuation">(</span>numstring<span class="token punctuation">)</span> <span class="token keyword">for</span> numstring <span class="token keyword">in</span> rle_string<span class="token punctuation">.</span>split<span class="token punctuation">(</span><span class="token string">' '</span><span class="token punctuation">)</span><span class="token punctuation">]</span>
<span class="token comment">#变换为多列格式(像素位置，长度)</span>
rlePairs <span class="token operator">=</span> np<span class="token punctuation">.</span>array<span class="token punctuation">(</span>rleNumbers<span class="token punctuation">)</span><span class="token punctuation">.</span>reshape<span class="token punctuation">(</span><span class="token operator">-</span><span class="token number">1</span><span class="token punctuation">,</span><span class="token number">2</span><span class="token punctuation">)</span>
img <span class="token operator">=</span> np<span class="token punctuation">.</span>zeros<span class="token punctuation">(</span>rows<span class="token operator">*</span>cols<span class="token punctuation">,</span>dtype<span class="token operator">=</span>np<span class="token punctuation">.</span>uint8<span class="token punctuation">)</span>
<span class="token comment">#创建mask，此时还是一行</span>
<span class="token keyword">for</span> index<span class="token punctuation">,</span>length <span class="token keyword">in</span> rlePairs<span class="token punctuation">:</span>
    index <span class="token operator">-=</span> <span class="token number">1</span>
    img<span class="token punctuation">[</span>index<span class="token punctuation">:</span>index<span class="token operator">+</span>length<span class="token punctuation">]</span> <span class="token operator">=</span> <span class="token number">255</span>
<span class="token comment">#以上画竖线，给出的+1了</span>
img <span class="token operator">=</span> img<span class="token punctuation">.</span>reshape<span class="token punctuation">(</span>cols<span class="token punctuation">,</span>rows<span class="token punctuation">)</span><span class="token comment">#因为画的是竖线，所以cols为行</span>
<span class="token keyword">return</span> img<span class="token punctuation">.</span>T

# %% [markdown]

# # 生成数据集
# %%

def get_dataset(idxs, src_img_path, dst_img_path, json_file):

# 添加类别定义

categories = [{'id': 1, 'name': 'class1', 'supercategory': 'object'},

{'id': 2, 'name': 'class2', 'supercategory': 'object'},

{'id': 3, 'name': 'class3', 'supercategory': 'object'},

{'id': 4, 'name': 'class4', 'supercategory': 'object'}]

coco_dataset = {

'licenses': [],

'images': [],

'annotations': [],

'info': {

'year': 2023,

'version': '1.0',

'description': 'COCO dataset for instance segmentation',

'contributor': 'Your Name',

'url': 'https://yourwebsite.com',

'date_created': '2023-04-06',

},

'categories': categories,

}
<span class="token comment"># 遍历图像文件，并添加图像信息和标注信息</span>
<span class="token keyword">for</span> col <span class="token keyword">in</span> idxs<span class="token punctuation">:</span>
    img_names <span class="token operator">=</span> <span class="token punctuation">[</span><span class="token builtin">str</span><span class="token punctuation">(</span>i<span class="token punctuation">)</span><span class="token punctuation">.</span>split<span class="token punctuation">(</span><span class="token string">"_"</span><span class="token punctuation">)</span><span class="token punctuation">[</span><span class="token number">0</span><span class="token punctuation">]</span> <span class="token keyword">for</span> i <span class="token keyword">in</span> train_df<span class="token punctuation">.</span>iloc<span class="token punctuation">[</span>col<span class="token punctuation">:</span>col<span class="token operator">+</span><span class="token number">4</span><span class="token punctuation">,</span> <span class="token number">0</span><span class="token punctuation">]</span><span class="token punctuation">.</span>values<span class="token punctuation">]</span>
    <span class="token comment"># 添加图像信息</span>
    img_id <span class="token operator">=</span> <span class="token builtin">len</span><span class="token punctuation">(</span>coco_dataset<span class="token punctuation">[</span><span class="token string">'images'</span><span class="token punctuation">]</span><span class="token punctuation">)</span> <span class="token operator">+</span> <span class="token number">1</span>
    img_path <span class="token operator">=</span> os<span class="token punctuation">.</span>path<span class="token punctuation">.</span>join<span class="token punctuation">(</span><span class="token string">'train_images'</span><span class="token punctuation">,</span> img_names<span class="token punctuation">[</span><span class="token number">0</span><span class="token punctuation">]</span><span class="token punctuation">)</span>
    img <span class="token operator">=</span> Image<span class="token punctuation">.</span><span class="token builtin">open</span><span class="token punctuation">(</span>img_path<span class="token punctuation">)</span>
    img_width<span class="token punctuation">,</span> img_height <span class="token operator">=</span> img<span class="token punctuation">.</span>size
    img_info <span class="token operator">=</span> <span class="token punctuation">{<!-- --></span>
        <span class="token string">'id'</span><span class="token punctuation">:</span> img_id<span class="token punctuation">,</span>
        <span class="token string">'file_name'</span><span class="token punctuation">:</span> img_names<span class="token punctuation">[</span><span class="token number">0</span><span class="token punctuation">]</span><span class="token punctuation">,</span>
        <span class="token string">'width'</span><span class="token punctuation">:</span> img_width<span class="token punctuation">,</span>
        <span class="token string">'height'</span><span class="token punctuation">:</span> img_height<span class="token punctuation">,</span>
        <span class="token string">'license'</span><span class="token punctuation">:</span> <span class="token boolean">None</span><span class="token punctuation">,</span>
        <span class="token string">'flickr_url'</span><span class="token punctuation">:</span> <span class="token boolean">None</span><span class="token punctuation">,</span>
        <span class="token string">'coco_url'</span><span class="token punctuation">:</span> <span class="token boolean">None</span><span class="token punctuation">,</span>
        <span class="token string">'date_captured'</span><span class="token punctuation">:</span> <span class="token boolean">None</span>
    <span class="token punctuation">}</span>
    coco_dataset<span class="token punctuation">[</span><span class="token string">'images'</span><span class="token punctuation">]</span><span class="token punctuation">.</span>append<span class="token punctuation">(</span>img_info<span class="token punctuation">)</span>
    name<span class="token punctuation">,</span> mask <span class="token operator">=</span> name_and_mask<span class="token punctuation">(</span>col<span class="token punctuation">)</span>
    
    shutil<span class="token punctuation">.</span>copy<span class="token punctuation">(</span>os<span class="token punctuation">.</span>path<span class="token punctuation">.</span>join<span class="token punctuation">(</span>src_img_path<span class="token punctuation">,</span> name<span class="token punctuation">)</span><span class="token punctuation">,</span> os<span class="token punctuation">.</span>path<span class="token punctuation">.</span>join<span class="token punctuation">(</span>dst_img_path<span class="token punctuation">,</span> name<span class="token punctuation">)</span><span class="token punctuation">)</span>
    
    <span class="token keyword">for</span> ch <span class="token keyword">in</span> <span class="token builtin">range</span><span class="token punctuation">(</span><span class="token number">4</span><span class="token punctuation">)</span><span class="token punctuation">:</span>
        contours<span class="token punctuation">,</span> _ <span class="token operator">=</span> cv2<span class="token punctuation">.</span>findContours<span class="token punctuation">(</span>mask<span class="token punctuation">[</span><span class="token punctuation">:</span><span class="token punctuation">,</span> <span class="token punctuation">:</span><span class="token punctuation">,</span> ch<span class="token punctuation">]</span><span class="token punctuation">,</span> cv2<span class="token punctuation">.</span>RETR_EXTERNAL<span class="token punctuation">,</span> cv2<span class="token punctuation">.</span>CHAIN_APPROX_NONE<span class="token punctuation">)</span>
        <span class="token keyword">for</span> contour <span class="token keyword">in</span> contours<span class="token punctuation">:</span>
            x<span class="token punctuation">,</span> y<span class="token punctuation">,</span> w<span class="token punctuation">,</span> h <span class="token operator">=</span> cv2<span class="token punctuation">.</span>boundingRect<span class="token punctuation">(</span>contour<span class="token punctuation">)</span>
            polygon <span class="token operator">=</span> contour<span class="token punctuation">.</span>squeeze<span class="token punctuation">(</span><span class="token punctuation">)</span><span class="token punctuation">.</span>tolist<span class="token punctuation">(</span><span class="token punctuation">)</span>
            polygon<span class="token punctuation">.</span>append<span class="token punctuation">(</span>polygon<span class="token punctuation">[</span><span class="token number">0</span><span class="token punctuation">]</span><span class="token punctuation">)</span>  <span class="token comment"># 添加首尾闭合</span>
            polygon <span class="token operator">=</span> np<span class="token punctuation">.</span>array<span class="token punctuation">(</span>polygon<span class="token punctuation">)</span>
            polygon <span class="token operator">=</span> polygon<span class="token punctuation">.</span>flatten<span class="token punctuation">(</span><span class="token punctuation">)</span><span class="token punctuation">.</span>tolist<span class="token punctuation">(</span><span class="token punctuation">)</span>
            ann_id <span class="token operator">=</span> <span class="token builtin">len</span><span class="token punctuation">(</span>coco_dataset<span class="token punctuation">[</span><span class="token string">'annotations'</span><span class="token punctuation">]</span><span class="token punctuation">)</span> <span class="token operator">+</span> <span class="token number">1</span>
            ann_info <span class="token operator">=</span> <span class="token punctuation">{<!-- --></span>
                <span class="token string">'id'</span><span class="token punctuation">:</span> ann_id<span class="token punctuation">,</span>
                <span class="token string">'image_id'</span><span class="token punctuation">:</span> img_id<span class="token punctuation">,</span>
                <span class="token string">'category_id'</span><span class="token punctuation">:</span> ch<span class="token operator">+</span><span class="token number">1</span><span class="token punctuation">,</span>
                <span class="token string">'segmentation'</span><span class="token punctuation">:</span> <span class="token punctuation">[</span>polygon<span class="token punctuation">]</span><span class="token punctuation">,</span>
                <span class="token string">'bbox'</span><span class="token punctuation">:</span> <span class="token punctuation">[</span>x<span class="token punctuation">,</span> y<span class="token punctuation">,</span> w<span class="token punctuation">,</span> h<span class="token punctuation">]</span><span class="token punctuation">,</span>
                <span class="token string">'area'</span><span class="token punctuation">:</span> w <span class="token operator">*</span> h<span class="token punctuation">,</span>
                <span class="token string">'iscrowd'</span><span class="token punctuation">:</span> <span class="token number">0</span>
            <span class="token punctuation">}</span>
            coco_dataset<span class="token punctuation">[</span><span class="token string">'annotations'</span><span class="token punctuation">]</span><span class="token punctuation">.</span>append<span class="token punctuation">(</span>ann_info<span class="token punctuation">)</span>
<span class="token keyword">with</span> <span class="token builtin">open</span><span class="token punctuation">(</span>json_file<span class="token punctuation">,</span> <span class="token string">'w'</span><span class="token punctuation">)</span> <span class="token keyword">as</span> f<span class="token punctuation">:</span>
    json<span class="token punctuation">.</span>dump<span class="token punctuation">(</span>coco_dataset<span class="token punctuation">,</span> f<span class="token punctuation">)</span>

# %%

# 定义数据集的根目录

dataset_dir = 'Severstal'

# 创建保存训练集、验证集和测试集的目录

annotations_dir = os.path.join(dataset_dir, 'annotations')

train_dir_all = os.path.join(dataset_dir, 'train2017all')

train_dir = os.path.join(dataset_dir, 'train2017')

val_dir = os.path.join(dataset_dir, 'val2017')

test_dir = os.path.join(dataset_dir, 'test2017')

train_json_all = os.path.join(dataset_dir, 'annotations', 'instances_train2017all.json')

train_json = os.path.join(dataset_dir, 'annotations', 'instances_train2017.json')

test_json = os.path.join(dataset_dir, 'annotations', 'instances_test2017.json')

val_json = os.path.join(dataset_dir, 'annotations', 'instances_val2017.json')
folders = [annotations_dir, train_dir_all, train_dir, val_dir, test_dir]
for f in folders:

if os.path.exists(f):

shutil.rmtree(f)

os.mkdir(f)
# 划分训练测试

test_ratio = 0.1
random.shuffle(idx_defect)

num_test = int(len(idx_defect) * test_ratio)
test_images_idx = idx_defect[:num_test]

train_images_idx_all = idx_defect[num_test:]

val_images_idx = idx_defect[num_test:num_test3]

train_images_idx = idx_defect[num_test3:]
get_dataset(test_images_idx, 'train_images', test_dir, test_json)

get_dataset(train_images_idx_all, 'train_images', train_dir_all, train_json_all)

get_dataset(train_images_idx, 'train_images', train_dir, train_json)

get_dataset(val_images_idx, 'train_images', val_dir, val_json)