In [1]:
import os
import cv2
from shutil import copyfile, rmtree
from sklearn.model_selection import train_test_split
from django.db.models import Q
from django.contrib.auth.models import User
from app.models import Label,Image,Batch, Comment, STATUS_CHOICES
from app.models import TODO, TAGGING, REVIEWING, DONE
In [2]:
dataset_path = os.path.join('.', 'static', 'dataset')
yolo_path = os.path.join('.', 'yolo')
yolo_anno_dir = os.path.join(yolo_path, 'annotate')
obj_path = os.path.join(yolo_anno_dir, 'obj')
if os.path.exists(obj_path):
rmtree(obj_path)
os.makedirs(obj_path)
In [40]:
classcnt = 0
classidx = dict()
stop = ['/.', '.', './', '\.', '\.']
def get_class(label):
global classidx, classcnt
info = [label.brand, label.model]
cls = "{}.{}".format(*info).lower().strip()
if not any(info) or cls in stop:
cls = "unknown"
if not cls in classidx:
classidx[cls] = classcnt
classcnt += 1
return cls
def get_coord(label):
return label.x, label.y, label.width, label.height
def get_yolo_anno(mat, label):
imh, imw = mat.shape[:2]
x, y, w, h = get_coord(label)
print(x, y, w, h)
cls = get_class(label)
dw = 1.0/imw
dh = 1.0/imh
x = (x+(x+w))/2.0
y = (y+(y+h))/2.0
x = x*dw
w = w*dw
y = y*dh
h = h*dh
# x, y, w, h = get_coord(label)
# rx = (x+(w/2.0))/imw
# ry = (y+(h/2.0))/imh
# rw = w/2.0/imw
# rh = h/2.0/imh
# cls = get_class(label)
return classidx[cls], x, y, w, h
def recheck_ratio(mat, labels, filename='test.png'):
imh, imw = mat.shape[:2]
print(imh)
tmp = mat.copy()
for label in labels:
cls, rx, ry, rw, rh = label
w = int(rw*imw)
h = int(rh*imh)
x1 = int(rx*imw-w/2.0)
y1 = int(ry*imh-h/2.0)
x2 = int(x1+w)
y2 = int(y1+h)
tmp = cv2.rectangle(tmp, (x1, y1), (x2, y2), (0, 255, 0))
w = int(rw*imw)
h = int(rh*imh)
x = int(rx*imw)
y = int(ry*imh)
tmp = cv2.circle(tmp, (x, y), h, (0,0,255))
cv2.imwrite(filename, tmp)
In [41]:
images = Image.objects.filter(batch__status=REVIEWING)\
.exclude(label__isnull=True)
for i, image in enumerate(images):
labels = list()
labelled = Label.objects.filter(image=image)
image_path = os.path.join(dataset_path, image.src_path)
fname, ext = os.path.splitext(image.src_path)
fname = image.id
obj_img_path = os.path.join(obj_path, "{}.png".format(fname))
obj_txt_path = os.path.join(obj_path, "{}.txt".format(fname))
mat = cv2.imread(image_path)
for label in labelled:
ret = get_yolo_anno(mat, label)
labels.append(ret)
# !!!DEBUG
# recheck_ratio(mat, labels, "./yolo/{}.png".format(i))
# copyfile(image_path, obj_img_path)
# with open(obj_txt_path, 'w+') as fp:
# for label in labels:
# fp.write(" ".join([str(x) for x in label]))
# fp.write(os.linesep)
In [ ]:
namespath = os.path.join(yolo_anno_dir, 'obj.names')
with open(namespath, 'w+') as fp:
for k, v in sorted( classidx.items(), key=lambda x:x[1] ):
fp.write(k)
fp.write(os.linesep)
datapath = os.path.join(yolo_anno_dir, 'obj.data')
context = """classes= {}
train = yolo/annotate/train.txt
valid = yolo/annotate/test.txt
names = yolo/annotate/obj.names
backup = backup/""".format(len(classidx))
with open(datapath, 'w+') as fp:
fp.write(context)
all_images = images.all().values_list('id')
train, test = train_test_split(all_images, test_size=0.1)
print len(train), len(test)
with open(os.path.join(yolo_anno_dir, 'train.txt'), 'w+') as fp:
for x in train:
fname = "{}.jpg".format(x[0])
fp.write(os.path.join(obj_path, fname)+os.linesep)
with open(os.path.join(yolo_anno_dir, 'test.txt'), 'w+') as fp:
for x in test:
fname = "{}.jpg".format(x[0])
fp.write(os.path.join(obj_path, fname)+os.linesep)
In [ ]:
def create_cfg(args) :
return """[net]
batch={}
subdivisions={}
width=608
height=608
channels=3
momentum=0.9
decay=0.0005
angle=0
saturation = 1.5
exposure = 1.5
hue=.1
learning_rate=0.001
burn_in=1000
max_batches = 500200
policy=steps
steps=400000,450000
scales=.1,.1
[convolutional]
batch_normalize=1
filters=32
size=3
stride=1
pad=1
activation=leaky
[maxpool]
size=2
stride=2
[convolutional]
batch_normalize=1
filters=64
size=3
stride=1
pad=1
activation=leaky
[maxpool]
size=2
stride=2
[convolutional]
batch_normalize=1
filters=128
size=3
stride=1
pad=1
activation=leaky
[convolutional]
batch_normalize=1
filters=64
size=1
stride=1
pad=1
activation=leaky
[convolutional]
batch_normalize=1
filters=128
size=3
stride=1
pad=1
activation=leaky
[maxpool]
size=2
stride=2
[convolutional]
batch_normalize=1
filters=256
size=3
stride=1
pad=1
activation=leaky
[convolutional]
batch_normalize=1
filters=128
size=1
stride=1
pad=1
activation=leaky
[convolutional]
batch_normalize=1
filters=256
size=3
stride=1
pad=1
activation=leaky
[maxpool]
size=2
stride=2
[convolutional]
batch_normalize=1
filters=512
size=3
stride=1
pad=1
activation=leaky
[convolutional]
batch_normalize=1
filters=256
size=1
stride=1
pad=1
activation=leaky
[convolutional]
batch_normalize=1
filters=512
size=3
stride=1
pad=1
activation=leaky
[convolutional]
batch_normalize=1
filters=256
size=1
stride=1
pad=1
activation=leaky
[convolutional]
batch_normalize=1
filters=512
size=3
stride=1
pad=1
activation=leaky
[maxpool]
size=2
stride=2
[convolutional]
batch_normalize=1
filters=1024
size=3
stride=1
pad=1
activation=leaky
[convolutional]
batch_normalize=1
filters=512
size=1
stride=1
pad=1
activation=leaky
[convolutional]
batch_normalize=1
filters=1024
size=3
stride=1
pad=1
activation=leaky
[convolutional]
batch_normalize=1
filters=512
size=1
stride=1
pad=1
activation=leaky
[convolutional]
batch_normalize=1
filters=1024
size=3
stride=1
pad=1
activation=leaky
#######
[convolutional]
batch_normalize=1
size=3
stride=1
pad=1
filters=1024
activation=leaky
[convolutional]
batch_normalize=1
size=3
stride=1
pad=1
filters=1024
activation=leaky
[route]
layers=-9
[convolutional]
batch_normalize=1
size=1
stride=1
pad=1
filters=64
activation=leaky
[reorg]
stride=2
[route]
layers=-1,-4
[convolutional]
batch_normalize=1
size=3
stride=1
pad=1
filters=1024
activation=leaky
[convolutional]
size=1
stride=1
pad=1
filters={}
activation=linear
[region]
anchors = 0.57273, 0.677385, 1.87446, 2.06253, 3.33843, 5.47434, 7.88282, 3.52778, 9.77052, 9.16828
bias_match=1
classes={}
coords=4
num=5
softmax=1
jitter=.3
rescore=1
object_scale=5
noobject_scale=1
class_scale=1
coord_scale=1
absolute=1
thresh = .6
random=1""".format(*args)
In [ ]:
nclass = len(classidx)
train_cfg = (64, 8, (nclass+5)*5, nclass)
test_cfg = (1, 1, (nclass+5)*5, nclass)
path = os.path.join(yolo_path, 'yolo-obj-train.cfg')
with open(path, 'w+') as fp:
fp.write(create_cfg(train_cfg))
path = os.path.join(yolo_path, 'yolo-obj-test.cfg')
with open(path, 'w+') as fp:
fp.write(create_cfg(test_cfg))
In [ ]:
# ./darknet detector train yolo/annotate/obj.data cfg/yolo-obj.cfg darknet19_448.conv.23
len(images)
In [ ]: