def create_model(self, skip_detect_layer=True):
model = Yolonet(n_classes=len(self._model_config['labels']),
arch=self._arch)
# tf_format = self._pretrained_config['tf_format']
# if os.path.exists(tf_format):
# model.load_weights(tf_format)
# print('TF pretrained weights loaded from {}'.format(tf_format))
keras_weights = self._pretrained_config['keras_format']
if os.path.exists(keras_weights):
self._load_weights(model, keras_weights)
# model.load_weights(keras_weights, by_name=True)
print('Keras pretrained weights loaded from {}'.format(
keras_weights))
elif self._arch == 'darknet' and not os.path.exists(keras_weights):
download_if_not_exists(
self._pretrained_config['darknet_format'],
'https://pjreddie.com/media/files/yolov3.weights')
model.load_darknet_params(
self._pretrained_config['darknet_format'], skip_detect_layer)
print('Original yolov3 weights loaded')
return model
def _initYolonet(self):
if self._yolonet:
return self._yolonet
nClasses = len(self.labelNames)
self._yolonet = Yolonet(n_classes=nClasses)
if self._weights:
if self._weights.endswith('.h5'): # keras format
self._yolonet.load_weights(self._weights)
else: # darknet format. skip_detect_layer=True because where is no detection with initial weights
self._yolonet.load_darknet_params(self._weights,
skip_detect_layer=True)
return self._yolonet
class CounterScreenModel:
labelNames = ["counter", "counter_screen"]
anchors = [
10, 13, 16, 30, 33, 23, 30, 61, 62, 45, 59, 119, 116, 90, 156, 198,
373, 326
]
net_size = 416
def __init__(self, weights):
self._yolonet = None
self._weights = weights
def yolonet(self):
return self._initYolonet()
def _initYolonet(self):
if self._yolonet:
return self._yolonet
nClasses = len(self.labelNames)
self._yolonet = Yolonet(n_classes=nClasses)
if self._weights:
if self._weights.endswith('.h5'): # keras format
self._yolonet.load_weights(self._weights)
else: # darknet format. skip_detect_layer=True because where is no detection with initial weights
self._yolonet.load_darknet_params(self._weights,
skip_detect_layer=True)
return self._yolonet
def detect(self, image, threshold=.5):
detector = YoloDetector(self.yolonet(),
self.anchors,
net_size=self.net_size)
return detector.detect(image, threshold) # boxes, labels, probs
def detectScreen(self, image):
counterScreenLabel = 1
boxes, labels, probs = self.detect(image, 0.8)
# filter only screens
onlyScreens = ((b, l, p) for b, l, p in zip(boxes, labels, probs)
if l == counterScreenLabel)
boxes, labels, probs = unzip(onlyScreens, [], [], [])
if len(boxes) == 0:
return None
return boxes[0]
def create_model(self, skip_detect_layer=True):
model = Yolonet(n_classes=len(self._model_config["labels"]))
keras_weights = self._pretrained_config["keras_format"]
if os.path.exists(keras_weights):
model.load_weights(keras_weights)
print("Keras pretrained weights loaded from {}!!".format(keras_weights))
else:
download_if_not_exists(self._pretrained_config["darknet_format"],
"https://pjreddie.com/media/files/yolov3.weights")
model.load_darknet_params(self._pretrained_config["darknet_format"], skip_detect_layer)
print("Original yolov3 weights loaded!!")
return model
def create_model(self, skip_detect_layer=True):
model = Yolonet(n_classes=len(self._model_config["labels"]))
keras_weights = self._pretrained_config["keras_format"]
if os.path.exists(keras_weights):
model.load_weights(keras_weights)
# *disable print statement*
# print("Keras pretrained weights loaded from {}!!".format(keras_weights))
# *disable print statement*
else:
# *inserted to prevent download*
raise Exception(
'Weights file not found, check \"keras_format\"" in \"kitti.json\" file'
)
# *inserted to prevent download*
download_if_not_exists(
self._pretrained_config["darknet_format"],
"https://pjreddie.com/media/files/yolov3.weights")
model.load_darknet_params(
self._pretrained_config["darknet_format"], skip_detect_layer)
print("Original yolov3 weights loaded!!")
return model
labels_naming=config["model"]["labels"],
anchors=config["model"]["anchors"],
jitter=config["train"]["jitter"])
valid_generator = BatchGenerator(ann_fnames,
config["train"]["train_image_folder"],
batch_size=config["train"]["batch_size"],
labels_naming=config["model"]["labels"],
anchors=config["model"]["anchors"],
jitter=False,
shuffle=False)
print(train_generator.steps_per_epoch)
# 2. create model
model = Yolonet(n_classes=len(config["model"]["labels"]))
model.load_darknet_params(config["pretrained"]["darknet_format"], skip_detect_layer=True)
# 4. traini
train_fn(model,
train_generator,
valid_generator,
learning_rate=config["train"]["learning_rate"],
save_dname=config["train"]["save_folder"],
num_epoches=config["train"]["num_epoch"])
# 5. prepare sample images
img_fnames = glob.glob(os.path.join(config["train"]["train_image_folder"], "*.*"))
imgs = [cv2.imread(fname)[:,:,::-1] for fname in img_fnames]
# 6. create new model & load trained weights
optimizer.apply_gradients(zip(grads, model.variables))
if i == 0 or (i + 1) % verbose == 0:
logits = model(images_tensor)
loss_value = loss_fn(list_y_trues, logits)
print("{}-th loss = {}".format(i, loss_value))
if min_loss_value > loss_value:
print(" update weight {}".format(loss_value))
min_loss_value = loss_value
model.save_weights("{}.h5".format(fname))
if __name__ == '__main__':
tf.enable_eager_execution()
from yolo.samples import sample_images, sample_list_y_trues, sample_true_boxes
# 1. setup dataset
images_tensor = tf.constant(sample_images())
true_boxes = tf.constant(sample_true_boxes())
list_y_trues = [tf.constant(arr) for arr in sample_list_y_trues()]
# 2. create model
model = Yolonet(18)
# 3. define optimizer
optimizer = tf.train.AdamOptimizer(learning_rate=1e-4)
# 4. training
# train(images_tensor, list_y_trues, true_boxes, optimizer, model, 2, 1)
# -*- coding: utf-8 -*-
import tensorflow as tf
tf.enable_eager_execution()
from yolo.net import Yolonet
from yolo.train import train
YOLOV3_WEIGHTS = "yolov3.weights"
if __name__ == '__main__':
import os
from yolo.dataset.generator import create_generator
from yolo import PROJECT_ROOT
# 1. create generator
ann_dir = os.path.join(PROJECT_ROOT, "samples", "anns")
img_dir = os.path.join(PROJECT_ROOT, "samples", "imgs")
generator = create_generator(img_dir, ann_dir)
# 2. create model
model = Yolonet(n_classes=1)
model.load_darknet_params(YOLOV3_WEIGHTS, True)
# 3. define optimizer
optimizer = tf.train.AdamOptimizer(learning_rate=1e-4)
# 4. training
train(generator, optimizer, model, 100, 1)
class DigitsOnScreenModel:
labelNames = ["0", "1", "2", "3", "4", "5", "6", "7", "8", "9"]
anchors = [
10, 13, 16, 30, 33, 23, 30, 61, 62, 45, 59, 119, 116, 90, 156, 198,
373, 326
]
net_size = 320
def __init__(self, weights):
self._yolonet = None
self._weights = weights
def yolonet(self):
return self._initYolonet()
def _initYolonet(self):
if self._yolonet:
return self._yolonet
nClasses = len(self.labelNames)
self._yolonet = Yolonet(n_classes=nClasses)
if self._weights:
if self._weights.endswith('.h5'): # keras format
self._yolonet.load_weights(self._weights)
else: # darknet format. skip_detect_layer=True because where is no detection with initial weights
self._yolonet.load_darknet_params(self._weights,
skip_detect_layer=True)
return self._yolonet
def detect(self, image, threshold=.5):
detector = YoloDetector(self.yolonet(),
self.anchors,
net_size=self.net_size)
return detector.detect(image, threshold) # boxes, labels, probs
class _io:
pad = 40
def preprocess(self, image):
gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
gray = cv2.medianBlur(gray, 3, gray)
inverted = imInvert(gray, out=gray)
binarized = binarizeSauvola(inverted, windowSize=41, k=.1)
binarized = imInvert(binarized, out=binarized)
binarized = cv2.cvtColor(binarized, cv2.COLOR_GRAY2RGB)
# imshowWait(DEBUG=binarized)
return binarized
def postprocess(self, boxes):
pad = self.pad
boxes = [(x1 - pad, y1 - pad, x2 - pad, y2 - pad)
for x1, y1, x2, y2 in boxes]
return boxes
def extendDetectionBox(self, proposedBox):
x1, y1, x2, y2 = proposedBox
pad = self.pad
return x1 - pad, y1 - pad, x2 + pad, y2 + pad
def detectDigits(self, image, proposedBox):
io = self._io()
viewImage = imageByBox(image, proposedBox)
preprocessed = io.preprocess(viewImage)
boxes, labels, probs = self.detect(preprocessed, .9)
# TODO: sort result by box.x1 (from left to right)
# boxes = io.postprocess(boxes)
return viewImage, boxes, labels, probs
@classmethod
def createWithLastWeights(cls):
weights = "./weights/weights_7_3.263.h5"
return cls(weights)