from PIL import Image
import numpy as np
from matplotlib import pyplot as plt
import re
import math
import random
from rknn.api import RKNN

modelConfiguration = "Yolo-Fastest-voc/yolov3-tiny.cfg"   
modelWeights = "Yolo-Fastest-voc/yolov3-tiny.weights"  


# modelConfiguration = "Yolo-Fastest-voc/yolo-fastest-xl.cfg"   
# modelWeights = "Yolo-Fastest-voc/yolo-fastest-xl.weights"  

# modelConfiguration = "Yolo-Fastest-coco/yolo-fastest-xl.cfg"   
# modelWeights = "Yolo-Fastest-coco/yolo-fastest-xl.weights" 

# modelConfiguration = "backup/yolo-fastest.cfg"   
# modelWeights = "backup/yolo-fastest_best.weights"  


if __name__ == '__main__':
	rknn=RKNN()
	
	rknn.config(mean_values=[0,0,0],std_values=[255,255,255],target_platform='rk3566')
	

	print('load model...')
	ret = rknn.load_darknet(model=modelConfiguration, weight=modelWeights)
	print("ret")
	print(ret)
	if ret != 0:
		print('load err...')
		exit(ret)
	print('done')
	

	print('building...')
	ret = rknn.build(do_quantization=True, dataset='./dataset.txt')
	if ret != 0:
		print('build fail!')
		exit(ret)
	print('done')
	
	ret = rknn.export_rknn('./yolov3-tiny_3566.rknn')
	if ret != 0:
		print('export fail!')
		exit(ret)
	exit(0)