coco/tools.py
2020-11-23 15:22:58 +01:00

317 lines
No EOL
11 KiB
Python

import pycocotools.coco
import argparse
import logging
import os
import pprint
import sqlite3
from coco.storage import COCOStorage, Annotation, Segment
import cv2
import mahotas
import subprocess
import tqdm
import numpy as np
import ast
import svgwrite
from svgwrite.extensions import Inkscape
from xml.etree import ElementTree
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("coco")
argParser = argparse.ArgumentParser(description='Create shape SVG\'s')
argParser.add_argument(
'--annotations',
type=str,
default='../../datasets/COCO/annotations/instances_val2017.json'
)
argParser.add_argument(
'--categories',
action='store_true',
help='Show categories'
)
argParser.add_argument(
'--db',
type=str,
metavar='DATABASE',
help='SQLite db filename, will be created if not existing'
)
argParser.add_argument(
'--propagate',
action='store_true',
help='Store annotation data in sqlite db'
)
argParser.add_argument(
'--zerkine',
action='store_true',
help='Find and store annotation Zerkine moments for those that do not have it yet'
)
argParser.add_argument(
'--similar',
type=int,
metavar="ANNOTATION_ID",
help='Find similar shapes for annotation'
)
argParser.add_argument(
'--stickers',
type=str,
metavar="SVG_FILENAME",
help="""
Create an SVG with sticker pages (afterwards convert to EPS: \"for f in *; do echo $f; inkscape -f $f --export-eps $f.eps; done\")
To PDF:
for f in stickers-v1/*.svg; do echo $f; inkscape $f -o pdf/$f.pdf; done
pdfunite stickers-v1/*.pdf stickers-v1.pdf
"""
)
argParser.add_argument(
'--no-texture',
action='store_true',
help="""
Don't fill with texture, but only print cut lines
"""
)
args = argParser.parse_args()
logger.info(f"Load {args.annotations}")
coco = pycocotools.coco.COCO(args.annotations)
if args.categories:
cats = {}
for id, cat in coco.cats.items():
if cat['supercategory'] not in cats:
cats[cat['supercategory']] = []
cats[cat['supercategory']].append(cat)
# pp = pprint.PrettyPrinter(indent=4)
pprint.pprint(cats, sort_dicts=False)
storage = None
if args.db:
storage = COCOStorage(args.db)
con = storage.con
if args.propagate:
logger.info("Create categories")
cur = con.cursor()
cur.executemany('INSERT OR IGNORE INTO categories(id, supercategory, name) VALUES (:id, :supercategory, :name)', coco.cats.values())
con.commit()
logger.info("Images...")
cur.executemany('''
INSERT OR IGNORE INTO images(id, flickr_url, coco_url, width, height, date_captured)
VALUES (:id, :flickr_url, :coco_url, :width, :height, :date_captured)
''', coco.imgs.values())
con.commit()
logger.info("Annotations...")
def annotation_generator():
for c in coco.anns.values():
ann = c.copy()
ann['bbox_top'] = ann['bbox'][1]
ann['bbox_left'] = ann['bbox'][0]
ann['bbox_width'] = ann['bbox'][2]
ann['bbox_height'] = ann['bbox'][3]
yield ann
cur.executemany('''
INSERT OR IGNORE INTO annotations(id, image_id, category_id, iscrowd, area, bbox_top, bbox_left, bbox_width, bbox_height)
VALUES (:id, :image_id, :category_id, :iscrowd, :area, :bbox_top, :bbox_left, :bbox_width, :bbox_height)
''', annotation_generator())
con.commit()
logger.info("Segments...")
def segment_generator():
for ann in coco.anns.values():
for i, seg in enumerate(ann['segmentation']):
yield {
'id': ann['id']*10 + i, # create a uniqe segment id, supports max 10 segments per annotation
'annotation_id': ann['id'],
'points': str(seg)[1:-1],
}
cur.executemany('''
INSERT OR IGNORE INTO segments(id, annotation_id, points)
VALUES (:id, :annotation_id, :points)
''', segment_generator())
con.commit()
logger.info("Done...")
if args.zerkine:
nr = storage.countAnnotationsWithoutZerkine()
for i in tqdm.tqdm(range(nr)):
annotation = storage.getAnnotationWithoutZerkine()
normAnn = annotation.getNormalised(100, 100)
filenameRoot = '/tmp/tmp_ann_to_convert'
dwg = normAnn.asSvg(filenameRoot + '.svg', square=True, bg='black')
dwg.save()
# convert to rasterised
subprocess.call([
'inkscape',
'-f', filenameRoot + '.svg',
'-e', filenameRoot + '.png',
],
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
)
# read with opencv
image = cv2.imread(filenameRoot + '.png', cv2.IMREAD_GRAYSCALE)
moments = mahotas.features.zernike_moments(image, 21)
storage.storeZerkineForAnnotation(annotation, moments, delayCommit = True)
if not i % 100:
storage.con.commit()
storage.con.commit()
if args.similar:
# todo find similar
annotation = storage.getAnnotationById(args.similar, withZerkine=True)
dwg = annotation.asSvg(f'tmp/source.svg', square=True)
dwg.save()
shapeA = np.array(annotation.segments[0].points)
annMoments = np.array(annotation.zerkine_moment)
distances = []
print(annotation)
# Stupid, this seems to have been superfluous
# zerkines = storage.getZerkines()
# for zerkine in tqdm.tqdm(zerkines):
# if annotation.id == zerkine['id']:
# continue
#
# diff = annMoments - np.array(Annotation.parseZerkineFromDB(zerkine['zerkine_moment']))
# distance = np.linalg.norm(diff)
# distances.append((zerkine['id'], distance))
anns = storage.getAllAnnotationPoints()
for ann in tqdm.tqdm(anns):
try:
shapeB = np.array(Segment.asCoordinates(ast.literal_eval('['+ann['points']+']')))
# fourth param is require, but according to docs does nothing
distance= cv2.matchShapes(shapeA, shapeB, cv2.CONTOURS_MATCH_I2, 1)
distances.append((ann['id'], distance))
except Exception as e:
logger.critical(f"Exception comparing {annotation.id} to {ann['id']}, points: {ann['points']}")
logger.exception(e)
distances = sorted(distances, key=lambda d: d[1])
for i in range(10):
similarAnnotation = storage.getAnnotationById(distances[i][0])
print(coco.cats[similarAnnotation.category_id])
dwg = similarAnnotation.asSvg(f'tmp/result_{i}.svg', square=True)
dwg.save()
if args.stickers:
grid = (3, 4) # items in the grid x,y
size = (105, 148) # in mm
sizeFactor = 5 # influences the size of the patterns
viewBoxSize = (size[0] * sizeFactor, size[1] * sizeFactor)
margin = 10
textMargin = 5
gridSize = (
int((viewBoxSize[0]-((grid[0]+1)*margin))/grid[0]),
int((viewBoxSize[1]-((grid[1]+1)*margin))/grid[1])
)
# see also textures.xml
textureIds = ["#siqwx","#wjnbs","#pnfez","#ejtxy","#obabs","#hehoj","#mrwjs","#ryjbw","#rkkau","#vbjcl","#zzehx","#mumke","#brhhk","#gujvh","#hfgqa","#lrbsh","#bndby","#bfnxk","#ydler","#pnxdr","#htqlj","#nunnt","#tidaw","#tcdum","#kwwja","#hgdkl","#nvkwz","#uzdqb","#fgshk","#vknil","#yeenr","#mslkw","#eibaw","#meama","#akuvz","#khkpp","#ibnow","#wivvx","#svksy","#xhmew","#jmiqu","#gfcer","#iueil","#iufvt","#ugkud","#dchzd","#nejks","#dqseb","#yhrwm","#bmiet","#qovkk","#hxoiq","#jfguh","#kbpkl","#ikarj","#nucap","#qfsqn","#bboqt","#pxkjn","#lbnx","#nxkmp","#snojb","#oioil","#hvldz","#qpscp","#oborh","#crobu","#ydhwn","#geanf","#sdfeo","#cgtma","#rjfrc","#uhcys","#lrgem","#osiho","#etssd","#esxcs","#hczhr","#nnhxw","#wrlbu"]
nr = 0
total_nr = len(coco.cats)
max_cat_id = max([cid for cid, _ in coco.cats.items()])
for category_id, cat in coco.cats.items():
nr+=1
filename = os.path.join(
args.stickers,
f"{category_id}_{cat['supercategory']}_{cat['name']}.svg")
dwg = svgwrite.Drawing(
filename,
size=(f'{size[0]}mm', f'{size[1]}mm'),
viewBox=f"0 0 {viewBoxSize[0]} {viewBoxSize[1]}"
)
annotations = storage.getRandomAnnotations(
limit = grid[0]*grid[1],
category_id = category_id
)
inkscape = Inkscape(dwg)
# Layers/groups to cut stickers at 123stickers.nl
contourG = inkscape.layer(label='Snijlijnen')
drawingG = inkscape.layer(label='Shapes')
# dwg.add(svgwrite.container.Defs())
dwg.add(drawingG)
dwg.add(contourG)
font_size = 10
# COCO cat id's go until 90, whereas there's only 80 categories.
# for consistency, we stick with id instead of the nr.
text = dwg.text(
f"{category_id:02d}/{max_cat_id}",
insert=(textMargin, textMargin+font_size), font_size=font_size, fill='black', font_family='Arial',
)
drawingG.add(text)
text = dwg.text(
f"{cat['supercategory']} - {cat['name']}",
insert=(viewBoxSize[0]-textMargin, textMargin+font_size), font_size=font_size, fill='black', font_family='Arial',
style='text-anchor:end;')
drawingG.add(text)
text = dwg.text(
f"Common Objects In Context",
insert=(textMargin, viewBoxSize[1]-textMargin), font_size=font_size, fill='black', font_family='Arial',
)
drawingG.add(text)
text = dwg.text(
f"http://plottingd.at/a",
insert=(viewBoxSize[0]-textMargin, viewBoxSize[1]-textMargin), font_size=font_size, fill='black', font_family='Arial',
style='text-anchor:end;')
drawingG.add(text)
for i, annotation in enumerate(annotations):
normAnn = annotation.getNormalised(gridSize[0], gridSize[1])
translation = normAnn.getTranslationToCenter()
# print(translation)
pX = i%grid[0]
pY = int(i/grid[0])
posX = pX*gridSize[0] + (pX+1)*margin - translation[0]
posY = pY*gridSize[1] + (pY+1)*margin - translation[1]
# print(i, posX, posY, gridSize)
positionG = svgwrite.container.Group(transform=f'translate({posX}, {posY})')
normAnn.writeToDrawing(positionG, stroke='#2FEE2F', stroke_width='1pt', fill_opacity="0")
contourG.add(positionG)
position2G = svgwrite.container.Group(transform=f'translate({posX}, {posY})')
pattern_id = textureIds[category_id % len(textureIds)]
normAnn.writeToDrawing(position2G, fill=f'url({pattern_id})', stroke='blue', stroke_width='0')
drawingG.add(position2G)
xml = dwg.get_xml()
if not args.no_texture:
with open('textures.xml', 'r') as fp:
textureTree = ElementTree.fromstring(fp.read())
defsTree = xml.find('defs')
for pattern in textureTree:
defsTree.append(pattern)
xmlString = ElementTree.tostring(xml)
with open(filename, 'wb') as fp:
# print(xmlString)
fp.write(xmlString)
logger.info(f"Wrote to {filename}")