masktococo

最新推荐文章于 2025-08-25 22:33:19 发布
原创最新推荐文章于 2025-08-25 22:33:19 发布 · 260 阅读
1 ·
CC 4.0 BY-SA版权
文章标签：
#python
python 专栏收录该内容
112 篇文章
订阅专栏
from convertrle import torle
import json
import os
import cv2
import numpy as np

def xyxy2xywh(x):
    # Convert nx4 boxes from [x1, y1, x2, y2] to [x, y, w, h] where xy1=top-left, xy2=bottom-right
    y = np.zeros_like(x)
    y[:, 0] = (x[:, 0] + x[:, 2]) / 2  # x center
    y[:, 1] = (x[:, 1] + x[:, 3]) / 2  # y center
    y[:, 2] = x[:, 2] - x[:, 0]  # width
    y[:, 3] = x[:, 3] - x[:, 1]  # height
    return y

def getarea(points):
    return cv2.contourArea(np.array(points))

CocoDIct = {}
CocoDIct['info'] = {'description': 'COCO 2017 Dataset', 'url': 'http://cocodataset.org', 'version': '1.0', 'year': 2017, 'contributor': 'COCO Consortium', 'date_created': '2017/09/01'}
CocoDIct['licenses'] = [{'url': 'http://creativecommons.org/licenses/by-nc-sa/2.0/', 'id': 1, 'name': 'Attribution-NonCommercial-ShareAlike License'}, {'url': 'http://creativecommons.org/licenses/by-nc/2.0/', 'id': 2, 'name': 'Attribution-NonCommercial License'}, {'url': 'http://creativecommons.org/licenses/by-nc-nd/2.0/', 'id': 3, 'name': 'Attribution-NonCommercial-NoDerivs License'}, {'url': 'http://creativecommons.org/licenses/by/2.0/', 'id': 4, 'name': 'Attribution License'}, {'url': 'http://creativecommons.org/licenses/by-sa/2.0/', 'id': 5, 'name': 'Attribution-ShareAlike License'}, {'url': 'http://creativecommons.org/licenses/by-nd/2.0/', 'id': 6, 'name': 'Attribution-NoDerivs License'}, {'url': 'http://flickr.com/commons/usage/', 'id': 7, 'name': 'No known copyright restrictions'}, {'url': 'http://www.usa.gov/copyright.shtml', 'id': 8, 'name': 'United States Government Work'}]
CocoDIct['images'] = []
CocoDIct['annotations'] = []
CocoDIct['categories'] = [{'supercategory': 'icecream', 'id': 1, 'name': 'clear'},{'supercategory': 'icecream', 'id': 2, 'name': 'not_so_clear'},{'supercategory': 'icecream', 'id': 3, 'name': 'interference'}]
categorydict = {'clear':1,'not_so_clear':2,'interference':3}
image_id = 1
id = 1
for dir in os.listdir('/home/lixuan/workspace/dataset/icecreamtwo/labelme_train/masks'):
    image = {}
    image['license'] = 1
    image['file_name'] = dir + '.jpg'
    image['coco_url'] = 'http://creativecommons.org/{}'.format(dir+'.jpg')
    img = cv2.imread('/home/lixuan/workspace/dataset/icecreamtwo/labelme_train/jpg/{}'.format(dir + '.jpg'))
    height,width,_ = img.shape
    image['height'] = height
    image['width'] = width
    image['date_captured'] = '2013-11-25 14:00:10'
    image['flickr_url'] = 'http://farm6.staticflickr.com/5533/10257288534_c916fafd78_z.jpg'
    image['id'] = image_id
    CocoDIct['images'].append(image)
    for file in os.listdir('/home/lixuan/workspace/dataset/icecreamtwo/labelme_train/masks/{}'.format(dir)):
        annotation = {}
        if not 'None' in file:
            segmentation = {}
            try:
                counts = torle('/home/lixuan/workspace/dataset/icecreamtwo/labelme_train/masks/{}/{}'.format(dir,file))
                segmentation['counts'] = counts
                segmentation['size'] = [height,width]
                annotation['segmentation'] = segmentation
                Points = np.load('/home/lixuan/workspace/dataset/icecreamtwo/labelme_train/txts/{}/{}'.format(dir,file.replace('jpg','npy')), allow_pickle=True)
                area = 0
                for points in Points:
                    area += getarea(points)
                annotation['area'] = area
                annotation['iscrowd'] = 1
                annotation['image_id'] = image_id
                mask = cv2.imread('/home/lixuan/workspace/dataset/icecreamtwo/labelme_train/masks/{}/{}'.format(dir,file),0)
                ys,xs = np.where(mask > 0)
                box = [min(xs),min(ys),max(xs),max(ys)]
                bbox = xyxy2xywh(np.array([box]))[0]
                annotation['bbox'] = bbox.tolist()
                try:
                    annotation['category_id'] = categorydict[file.split('icecream_')[1].split('-')[0].split('.')[0]]
                except:
                    annotation['category_id'] = 3
                annotation['id'] = id
                CocoDIct['annotations'].append(annotation)
                id += 1
            except Exception as e:
                pass
        else:
            annotation['iscrowd'] = 0
            annotation['image_id'] = image_id
            mask = cv2.imread('/home/lixuan/workspace/dataset/icecreamtwo/labelme_train/masks/{}/{}'.format(dir,file),0)
            ys, xs = np.where(mask > 0)
            box = [min(xs), min(ys), max(xs), max(ys)]
            bbox = xyxy2xywh(np.array([box]))[0]
            annotation['bbox'] = bbox.tolist()
            points = np.load('/home/lixuan/workspace/dataset/icecreamtwo/labelme_train/txts/{}/{}'.format(dir,file.replace('jpg','npy')))
            area = getarea(points)
            segmentation = [points.flatten().tolist()]
            annotation['segmentation'] = segmentation
            try:
                annotation['category_id'] = categorydict[file.split('icecream_')[1].split('-')[0].split('.')[0]]
            except:
                annotation['category_id'] = 3
            annotation['id'] = id
            annotation['area'] = area
            CocoDIct['annotations'].append(annotation)
            id += 1
    image_id += 1
with open('icecream_coco.json','w') as f:
    json.dump(CocoDIct,f)