import os
import numpy as np
import keras.backend as K
from keras.layers import Input, Lambda
from keras.models import Model
from keras.optimizers import Adam
from keras.callbacks import TensorBoard, ModelCheckpoint, ReduceLROnPlateau, EarlyStopping
# 指定那幾塊GPU是對於本程序可見的 其實有沒有都沒有什麼區別 tensor默認佔用全部的
# 一般用在忽略哪塊GPU 多GPU只能用一塊那種任務
from yolo3.model import preprocess_true_boxes, yolo_body, tiny_yolo_body, yolo_loss
from yolo3.utils import get_random_data
# 多Gpu訓練
from keras.utils import multi_gpu_model
import keras.backend.tensorflow_backend as KTF
import tensorflow as tf
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
def _main():
# 訓練前記得修改yolov3.cfg和model_data/voc_class.txt
annotation_path = '1.14/train/label1.txt' # 這裏是要的由VOC生成的txt,裏面包括image地址和框和種類 需修改
log_dir = '1.14/weights/' # 是存放生成好的新的weights的
classes_path = '1.14/train/classes.txt' # 這個我只標記了4類,所以要修改,原來voc是有20類的
anchors_path = '1.14/train/anchors.txt' # 這裏的anchor暫時先不修改(最好從新kmean算anchors)
class_names = get_classes(classes_path)
num_classes = len(class_names)
anchors = get_anchors(anchors_path)
# 修改大小,按照圖片最小的邊來修改
input_shape = (544,960) # multiple of 32, hw 這裏是否需要修改
config = tf.ConfigProto()
config.gpu_options.allow_growth=True #不全部佔滿顯存, 按需分配
sess = tf.Session(config=config)
KTF.set_session(sess)
is_tiny_version = len(anchors)==6 # default setting
if is_tiny_version:
model = create_tiny_model(input_shape, anchors, num_classes,
freeze_body=2, weights_path='model_data/tiny_yolo_weights.h5')
else:
# 這裏要修改,這裏的權重是基礎權重,在這個基礎上繼續訓練
model = create_model(input_shape, anchors, num_classes,
freeze_body=2, weights_path='12.5/weigths/dog_detect_544x960.h5') # make sure you know what you freeze