def main():
yolo = Create_Yolo(input_size=416, CLASSES="tools/labels.txt")
checkpoints_path = "checkpoints/yolov3_custom_Tiny"
yolo.load_weights(checkpoints_path)
# yolo.save('save/yolov3')
yolo.save('save/yolov3.h5')
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
import sys
foldername = os.path.basename(os.getcwd())
if foldername == "tools":
os.chdir("..")
sys.path.insert(1, os.getcwd())
import tensorflow as tf
from yolov3.yolov4 import Create_Yolo
from yolov3.utils import load_yolo_weights, Load_Yolo_model
from yolov3.configs import *
if YOLO_TYPE == "yolov4":
Darknet_weights = YOLO_V4_TINY_WEIGHTS if TRAIN_YOLO_TINY else YOLO_V4_WEIGHTS
if YOLO_TYPE == "yolov3":
Darknet_weights = YOLO_V3_TINY_WEIGHTS if TRAIN_YOLO_TINY else YOLO_V3_WEIGHTS
if YOLO_CUSTOM_WEIGHTS == False:
yolo = Create_Yolo(input_size=YOLO_INPUT_SIZE, CLASSES=YOLO_COCO_CLASSES)
load_yolo_weights(yolo, Darknet_weights) # use Darknet weights
else:
#yolo = Create_Yolo(input_size=YOLO_INPUT_SIZE, CLASSES=TRAIN_CLASSES)
#yolo.load_weights(YOLO_CUSTOM_WEIGHTS) # use custom weights
yolo = Load_Yolo_model()
yolo.summary()
yolo.save(f'./checkpoints/{YOLO_TYPE}-{YOLO_INPUT_SIZE}')
print(f"model saves to /checkpoints/{YOLO_TYPE}-{YOLO_INPUT_SIZE}")
