def main(param_dir):
cfg_path = "../4Others/yolo.cfg"
test_root_dir = "../2CvTrain"
test_label_csv_mame = '../2CvTrain/label.csv'
blocks = parse_cfg(cfg_path)
label_csv_mame = '../2CvTrain/label.csv'
params = prep_params(param_dir, label_csv_mame)
from yolo_v3 import yolo_v3
model = yolo_v3(params, blocks)
test_transform = transforms.Compose([
transforms.Resize(model.params["height"]),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
test_data = CustData(test_label_csv_mame, transform=test_transform)
test_loader = DataLoader(test_data,
shuffle=False,
batch_size=model.params["batch_size"],
collate_fn=my_collate,
num_workers=0)
start = time.time()
best_map, best_ap, best_conf, specific_conf_map, specific_conf_ap,\
map_frame = get_map(model, test_loader, train=False, loop_conf=False)
print(time.time() - start)
return best_map, best_ap, best_conf, specific_conf_map, specific_conf_ap, \
map_frame
def conf_map_frame(iou_conf, conf_list):
cuda = True
specific_conf = 0.9
cfg_path = "../4Others/color_ball.cfg"
test_root_dir = "../1TestData"
test_label_csv_mame = '../1TestData/label.csv'
classes = load_classes('../4Others/color_ball.names')
blocks = parse_cfg(cfg_path)
model = yolo_v3(blocks)
checkpoint_path = "../4TrainingWeights/experiment/2018-11-30_05_19_48.404150/2018-11-30_05_59_33.719706_model.pth"
checkpoint = torch.load(checkpoint_path)
model.load_state_dict(checkpoint)
model = model.cuda()
test_transform = transforms.Compose([
transforms.Resize(model.net["height"]),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
test_data = CustData(test_label_csv_mame,
test_root_dir,
transform=test_transform)
test_loader = DataLoader(test_data,
shuffle=False,
batch_size=model.net["batch"],
collate_fn=my_collate,
num_workers=6)
best_map, best_ap, best_conf, specific_conf_map, specific_conf_ap,\
map_frame = get_map(model, test_loader, cuda, conf_list, iou_conf,
classes, False, specific_conf, True)
return best_map, best_ap, best_conf, specific_conf_map, specific_conf_ap, \
map_frame
def compare():
date_time_now = str(datetime.datetime.now()).replace(" ", "_").replace(
":", "_")
compare_path = f"../5Compare/individual_train/{date_time_now}/"
if not os.path.exists(compare_path):
os.makedirs(compare_path)
config_name = "exp_config.p"
conf_list = np.arange(start=0.1, stop=0.95, step=0.025)
seed_range = range(420, 425)
classes = load_classes('../4Others/color_ball.names')
cfg_path = "../4Others/color_ball.cfg"
blocks = parse_cfg(cfg_path)
model = yolo_v3(blocks)
model.load_weights("../4Weights/yolov3.weights", cust_train_zero=True)
with open(compare_path + config_name, "wb") as fp:
pickle.dump(model.net, fp, protocol=pickle.HIGHEST_PROTOCOL)
time_taken_df = pd.DataFrame(columns=list(seed_range))
for cls_index, cls in enumerate(classes):
to_path = f"../{cls}/"
for index, seed in enumerate(seed_range):
model.load_weights("../4Weights/yolov3.weights",
cust_train_zero=True)
x_start = time.time()
random.seed(seed)
if not os.path.exists(to_path):
os.makedirs(to_path)
sub_name = f"{cls}_seed_{seed}_"
name_list = ["img_name", "c", "gx", "gy", "gw", "gh"]
# Original label names
label_csv_mame = '../color_balls/label.csv'
img_txt_path = "../color_balls/*.txt"
prep_labels(img_txt_path, name_list, label_csv_mame)
# sub sampled label names
sub_sample_csv_name = to_path + "label.csv"
sub_sample_txt_path = to_path + "*.txt"
prep_labels(sub_sample_txt_path, name_list, sub_sample_csv_name)
# label_csv_mame = '../1TestData/label.csv'
# img_txt_path = "../1TestData/*.txt"
move_images(label_name=label_csv_mame,
to_path=to_path,
action_fn=action_fn,
cat_nums=[cls_index])
best_map, best_ap, best_conf, specific_conf_map, specific_conf_ap,\
map_frame = main(model,
classes, conf_list,
sub_sample_csv_name,
sub_sample_txt_path,
to_path,
cuda=True,
specific_conf=0.5,
sub_name=sub_name,
selected_cls=[str(x) for x in [cls_index]])
map_frame.to_csv(f"{compare_path+sub_name}.csv", index=True)
shutil.rmtree(to_path)
time_taken_df.loc[cls, seed] = time.time() - x_start
# if you change the csv_name, pls change accordingly in the
# visualization part
time_taken_df.to_csv(compare_path + 'time_taken.csv', index=True)
def compare():
date_time_now = str(datetime.datetime.now()).replace(" ", "_").replace(
":", "_")
compare_path = f"../5Compare/input_size/{date_time_now}/"
if not os.path.exists(compare_path):
os.makedirs(compare_path)
config_name = "exp_config.p"
num_anchors = 3
conf_list = np.arange(start=0.1, stop=0.95, step=0.025)
seed_range = range(424, 428)
input_sizes = [448, 480]
classes = load_classes('../4Others/color_ball.names')
cfg_path = "../4Others/color_ball.cfg"
blocks = parse_cfg(cfg_path)
model = yolo_v3(blocks)
model.load_weights("../4Weights/yolov3.weights", cust_train_zero=True)
model.net['seed_range'] = seed_range
label_csv_mame = '../color_balls/label.csv'
img_txt_path = "../color_balls/*.txt"
root_dir = "../color_balls"
# save the list of input sizes into the conf dictionary
model.net['width'] = input_sizes
model.net['height'] = input_sizes
# save the dictionary into local pickle file
with open(compare_path + config_name, "wb") as fp:
pickle.dump(model.net, fp, protocol=pickle.HIGHEST_PROTOCOL)
time_taken_df = pd.DataFrame(columns=list(seed_range))
for input_size in input_sizes:
# now parse the size into the model
model.net['width'] = input_size
model.net['height'] = input_size
set_anchors_to_model(model, num_anchors, label_csv_mame, input_size,
input_size)
for index, seed in enumerate(seed_range):
model.load_weights("../4Weights/yolov3.weights",
cust_train_zero=True)
x_start = time.time()
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
sub_name = f"{input_size}_seed_{seed}_"
# label_csv_mame = '../1TestData/label.csv'
# img_txt_path = "../1TestData/*.txt"
best_map, best_ap, best_conf, specific_conf_map, specific_conf_ap,\
map_frame = main(model,
classes, conf_list,
label_csv_mame,
img_txt_path,
root_dir,
cuda=True,
specific_conf=0.5,
sub_name=sub_name)
map_frame.to_csv(f"{compare_path+sub_name}.csv", index=True)
time_taken_df.loc[input_size, seed] = time.time() - x_start
time_taken_df.to_csv(compare_path + 'time_taken.csv', index=True)
def compare():
date_time_now = str(datetime.datetime.now()).replace(" ", "_").replace(
":", "_")
compare_path = f"../5Compare/one_anchor/448_480/{date_time_now}/"
if not os.path.exists(compare_path):
os.makedirs(compare_path)
num_anchors = 3
conf_list = np.arange(start=0.1, stop=0.95, step=0.025)
seed_range = list(range(1227, 1242))
input_sizes = [448, 480]
classes = load_classes('../4Others/color_ball.names')
cfg_path = "../4Others/color_ball_one_anchor.cfg"
blocks = parse_cfg(cfg_path)
model = yolo_v3(blocks)
model.load_weights("../4Weights/yolov3.weights", cust_train_zero=True)
model.net['seed_range'] = seed_range
label_csv_mame = '../color_balls/label.csv'
img_txt_path = "../color_balls/*.txt"
root_dir = "../color_balls"
time_taken_df = pd.DataFrame(columns=list(seed_range))
for input_size in input_sizes:
# now parse the size into the model
model.net['width'] = input_size
model.net['height'] = input_size
set_anchors_to_model(model, num_anchors, label_csv_mame, input_size,
input_size)
yolo_layer = [(model.layer_type_dic['yolo'][i]) for i in range(3)]
model.net['anchors'] = [
model.module_list[yolo_layer[i]][0].anchors.tolist()
for i in range(3)
]
for index, seed in enumerate(seed_range):
model.load_weights("../4Weights/yolov3.weights",
cust_train_zero=True)
x_start = time.time()
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
sub_name = f"{input_size}_seed_{seed}_"
# label_csv_mame = '../1TestData/label.csv'
# img_txt_path = "../1TestData/*.txt"
best_map, best_ap, best_conf, specific_conf_map, specific_conf_ap,\
map_frame = main(model,
classes, conf_list,
label_csv_mame,
img_txt_path,
root_dir,
cuda=True,
specific_conf=0.5,
sub_name=sub_name)
map_frame.to_csv(f"{compare_path+sub_name}.csv", index=True)
time_taken_df.loc[input_size, seed] = time.time() - x_start
time_taken_df.to_csv(compare_path + 'time_taken.csv', index=True)
def compare():
date_time_now = str(datetime.datetime.now()).replace(" ", "_").replace(
":", "_")
compare_path = f"../5Compare/multiple_train/{date_time_now}/"
if not os.path.exists(compare_path):
os.makedirs(compare_path)
config_name = "exp_config.p"
conf_list = np.arange(start=0.1, stop=0.95, step=0.025)
seed_range = range(1218, 1221)
classes = load_classes('../4Others/color_ball.names')
cfg_path = "../4Others/color_ball_one_anchor.cfg"
blocks = parse_cfg(cfg_path)
model = yolo_v3(blocks)
model.load_weights("../4Weights/yolov3.weights", cust_train_zero=True)
model.net['seed_used'] = list(seed_range)
with open(compare_path + config_name, "wb") as fp:
pickle.dump(model.net, fp, protocol=pickle.HIGHEST_PROTOCOL)
time_taken_df = pd.DataFrame(columns=list(seed_range))
for index, seed in enumerate(seed_range):
to_path = f"../{seed}/"
model.load_weights("../4Weights/yolov3.weights", cust_train_zero=True)
x_start = time.time()
random.seed(seed)
sub_name = f"_seed_{seed}_"
# Original label names
label_csv_mame = '../color_balls/label.csv'
img_txt_path = "../color_balls/*.txt"
# sub sampled label names
# label_csv_mame = '../1TestData/label.csv'
# img_txt_path = "../1TestData/*.txt"
best_map, best_ap, best_conf, specific_conf_map, specific_conf_ap,\
map_frame = main(model,
classes, conf_list,
label_csv_mame,
img_txt_path,
to_path,
cuda=True,
specific_conf=0.5,
sub_name=sub_name)
map_frame.to_csv(f"{compare_path+sub_name}.csv", index=True)
time_taken_df.loc[0, seed] = time.time() - x_start
# if you change the csv_name, pls change accordingly in the
# visualization part
time_taken_df.to_csv(compare_path + 'time_taken.csv', index=True)
def get_single_img_map(cfg_path, det, label_csv_mame, params):
map_frame = pd.DataFrame(columns=['mAP', 'ap'])
assert params["height"] % 32 == 0
assert params["height"] > 32
blocks = parse_cfg(cfg_path)
model = yolo_v3(params, blocks)
# model.load_weights("../4Weights/yolov3.weights")
if params['cuda']:
model = model.cuda()
for parameter in model.parameters():
parameter.requires_grad = False
# yolo v3 down size the imput images 32 strides, therefore the input needs
# to be a multiplier of 32 and > 32
# put to evaluation mode so that gradient wont be calculated
model.eval()
transform = transforms.Compose([
LetterBoxImage_cv([params['height'], params['height']]),
ImgToTensorCv(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
data = CustDataCV(label_csv_mame,
transform=transform,
detection_phase=True)
data_loader = DataLoader(data,
shuffle=False,
batch_size=1,
collate_fn=my_collate_detection,
num_workers=0,
worker_init_fn=worker_init_fn)
for original_imgs, images, imlist, im_dim_list, labels in data_loader:
if model.params['cuda']:
images = images.cuda()
prediction = model(images)
output = filter_results(prediction, params)
mAP, ap = simgle_img_map(model, output, labels)
map_frame.loc[imlist[-1], 'mAP'] = mAP
map_frame.loc[imlist[-1], 'ap'] = ap
return map_frame
def main():
cuda = True
cfg_path = "../4Others/color_ball.cfg"
test_root_dir = "../1TestData"
test_label_csv_mame = '../1TestData/label.csv'
classes = load_classes('../4Others/color_ball.names')
blocks = parse_cfg(cfg_path)
model = yolo_v3(blocks)
conf_list = np.arange(start=0.2, stop=0.95, step=0.025)
checkpoint_path = "../4TrainingWeights/2018-11-07_23_13_38.391465/2018-11-08_02_45_20.195250_model.pth"
checkpoint = torch.load(checkpoint_path)
model.load_state_dict(checkpoint)
model = model.cuda()
test_transform = transforms.Compose([
transforms.Resize(model.net["height"]),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
test_data = CustData(test_label_csv_mame,
test_root_dir,
transform=test_transform)
test_loader = DataLoader(test_data,
shuffle=True,
batch_size=1,
collate_fn=my_collate,
num_workers=4)
start = time.time()
sum_ap, sum_map, ap, mAP, map_frame = get_map(model,
test_loader,
cuda,
conf_list,
0.5,
classes,
train=False)
print(time.time() - start)
return sum_ap, sum_map, ap, mAP, map_frame
# change cuda to True if you wish to use gpu
cuda = False
colors = pkl.load(open("../4Others/pallete", "rb"))
# images = "../1RawData/"
# det = "../2ProcessedData/"
classes = load_classes('../4Others/coco.names')
batch_size = 12
# Used in filter_results to apply a confidence mask
confidence = float(0.5)
# Used in filter_results to remove points with IoU < nms_thesh
nms_thesh = float(0.4)
start = 0
num_classes = 80
cfg_path = "../4Others/yolov3.cfg"
blocks = parse_cfg(cfg_path)
# setup our model object from the yolo_v3.py file
model = yolo_v3(blocks)
model.load_weights("../4Weights/yolov3.weights")
if cuda:
model = model.cuda()
for params in model.parameters():
params.requires_grad = False
if version == 'V1':
model.layer_type_dic['net_info']["height"] = 416
model.layer_type_dic['net_info']["width"] = 416
elif version == 'V2':
model.layer_type_dic['net_info']["height"] = 32 * 7
model.layer_type_dic['net_info']["width"] = 32 * 7
inp_dim = int(model.layer_type_dic['net_info']["height"])
def run_webcam(cfg_path,
param_path,
colors,
count_delay,
nms_thesh=False,
confidence=False,
count_time=10):
with open(param_path) as fp:
params = json.load(fp)
if type(params['anchors']) == list:
params['anchors'] = np.array(params['anchors'])
blocks = parse_cfg(cfg_path)
model = yolo_v3(params, blocks)
for parameter in model.parameters():
parameter.requires_grad = False
if params['cuda']:
model = model.cuda()
transform = transforms.Compose([
LetterboxImage_cv([params['height'], params['height']]),
ImgToTensorCv(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
assert params['height'] % 32 == 0
assert params['height'] > 32
cap = cv2.VideoCapture(0)
fps_list = []
add = False
print(f"this is height {params['height']}")
if params['height'] > 480:
# 720 p
cap.set(3, 1280)
else:
# 480 p
cap.set(3, 800)
print(f"width :{cap.get(3)}")
print(f"height :{cap.get(4)}")
assert cap.isOpened(), 'Cannot capture source'
start = time.time()
count_start_time = start + count_delay
if not confidence:
confidence = params['confidence']
if not nms_thesh:
nms_thesh = params['nms_thesh']
while cap.isOpened():
ret, frame = cap.read()
frame = cv2.flip(frame, 1)
if ret:
# we need to unsqueeze to add batch size dimension to tensor in
# order to fit into pytorch data format
img = transform(frame).unsqueeze(0)
im_dim = frame.shape[1], frame.shape[0]
im_dim = torch.FloatTensor(im_dim).repeat(1, 2)
if params['cuda']:
im_dim = im_dim.cuda()
img = img.cuda()
output = model(img)
output = filter_results(output, params)
# when there is no dection
if type(output) == int:
cv2.imshow("frame", frame)
fps = (1 / (time.time() - start))
# print(f"FPS of the video is {fps:5.4f}")
if add:
fps_list.append(fps)
if (time.time() - count_start_time) > count_time:
print(f"avg_fps: {np.mean(fps_list):5.4f}")
cv2.destroyAllWindows()
cap.release()
return np.mean(fps_list)
break
elif time.time() > count_start_time:
count_start_time = time.time()
fps_list.append(fps)
add = True
start = time.time()
key = cv2.waitKey(1) & 0xFF
if key == 27:
break
continue
im_dim = im_dim.repeat(output.size(0), 1)
scaling_factor = torch.min(params['height'] / im_dim,
1)[0].view(-1, 1)
output[:,
[1, 3]] -= (params['height'] -
scaling_factor * im_dim[:, 0].view(-1, 1)) / 2
output[:,
[2, 4]] -= (params['height'] -
scaling_factor * im_dim[:, 1].view(-1, 1)) / 2
output[:, 1:5] /= scaling_factor
for i in range(output.shape[0]):
output[i, [1, 3]] = torch.clamp(output[i, [1, 3]], 0.0,
im_dim[i, 0])
output[i, [2, 4]] = torch.clamp(output[i, [2, 4]], 0.0,
im_dim[i, 1])
list(map(lambda x: write(x, frame, params['classes']), output))
cv2.imshow("frame", frame)
key = cv2.waitKey(1) & 0xff
fps = 1 / (time.time() - start)
# print(f"FPS of the video is {fps:5.4f}")
if add:
fps_list.append(fps)
if (time.time() - count_start_time) > count_time:
print(f"avg_fps: {np.mean(fps_list):5.4f}")
cv2.destroyAllWindows()
cap.release()
return np.mean(fps_list)
break
elif time.time() > count_start_time:
count_start_time = time.time()
fps_list.append(fps)
add = True
start = time.time()
if key == 27:
break
else:
break
def compare():
date_time_now = str(datetime.datetime.now()).replace(" ", "_").replace(
":", "_")
compare_path = f"../5Compare/epoch_experiment/{date_time_now}/"
if not os.path.exists(compare_path):
os.makedirs(compare_path)
config_name = "exp_config.p"
conf_list = np.arange(start=0.2, stop=0.95, step=0.025)
seed_range = range(420, 425)
to_path_list = {
"../250_to_300_imgs/": 1.05 / 7,
"../400_to_450_imgs/": 1 / 4,
"../550_to_600_imgs/": 3 / 7,
"../700_to_750_imgs/": 4.7 / 7
}
epochs_list = [25, 30, 35, 40]
classes = load_classes('../4Others/color_ball.names')
cfg_path = "../4Others/color_ball.cfg"
blocks = parse_cfg(cfg_path)
model = yolo_v3(blocks)
model.load_weights("../4Weights/yolov3.weights", cust_train_zero=True)
model.net['epochs_list'] = epochs_list
model.net['img_sampling_info'] = to_path_list
with open(compare_path + config_name, "wb") as fp:
pickle.dump(model.net, fp, protocol=pickle.HIGHEST_PROTOCOL)
for to_path, train_size in to_path_list.items():
time_taken_df = pd.DataFrame(columns=list(epochs_list))
file_name = to_path.strip(".").strip("/")
for index, seed in enumerate(seed_range):
x_start = time.time()
random.seed(seed)
if not os.path.exists(to_path):
os.makedirs(to_path)
name_list = ["img_name", "c", "gx", "gy", "gw", "gh"]
# Original label names
label_csv_mame = '../color_balls/label.csv'
img_txt_path = "../color_balls/*.txt"
prep_labels(img_txt_path, name_list, label_csv_mame)
# sub sampled label names
sub_sample_csv_name = to_path + "label.csv"
sub_sample_txt_path = to_path + "*.txt"
# label_csv_mame = '../1TestData/label.csv'
# img_txt_path = "../1TestData/*.txt"
move_images(label_name=label_csv_mame,
to_path=to_path,
action_fn=split_train,
train_size=train_size,
random_state=seed)
for epoch in epochs_list:
model.load_weights("../4Weights/yolov3.weights",
cust_train_zero=True)
model.net['epochs'] = epoch
sub_name = f"{file_name}_seed_{seed}_epoch_{epoch}_"
best_map, best_ap, best_conf, specific_conf_map,\
specific_conf_ap, map_frame = main(model,
classes, conf_list,
sub_sample_csv_name,
sub_sample_txt_path,
to_path,
cuda=True,
specific_conf=0.5,
sub_name=sub_name)
map_frame.to_csv(f"{compare_path+sub_name}.csv", index=True)
time_taken_df.loc[file_name, epoch] = time.time() - x_start
# if you change the csv_name, pls change accordingly in the
# visualization part
time_taken_df.to_csv(compare_path + f"time_taken_{seed}.csv",
index=True)
shutil.rmtree(to_path)
def detection(cfg_path, det, label_csv_mame, params):
assert params["height"] % 32 == 0
assert params["height"] > 32
blocks = parse_cfg(cfg_path)
model = yolo_v3(params, blocks)
# model.load_weights("../4Weights/yolov3.weights")
if params['cuda']:
model = model.cuda()
for parameter in model.parameters():
parameter.requires_grad = False
# yolo v3 down size the imput images 32 strides, therefore the input needs
# to be a multiplier of 32 and > 32
# put to evaluation mode so that gradient wont be calculated
model.eval()
transform = transforms.Compose([
LetterBoxImage_cv([params['height'], params['height']]),
ImgToTensorCv(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
data = CustDataCV(label_csv_mame,
transform=transform,
detection_phase=True)
data_loader = DataLoader(data,
shuffle=False,
batch_size=model.params['batch_size'],
collate_fn=my_collate_detection,
num_workers=0,
worker_init_fn=worker_init_fn)
for original_imgs, images, imlist, im_dim_list, labels in data_loader:
im_dim_list = torch.FloatTensor(im_dim_list).repeat(1, 2)
if model.params['cuda']:
images = images.cuda()
im_dim_list = im_dim_list.cuda()
prediction = model(images)
model.params['num_classes']
output = filter_results(prediction, params)
if type(output) != int:
im_dim_list = torch.index_select(im_dim_list, 0, output[:,
0].long())
scaling_factor = torch.min(params["height"] / im_dim_list,
1)[0].view(-1, 1)
# Clamp all elements in input into the range [ min, max ] and
# return a resulting tensor:
# it is to make sure the min max of width and height are not larger
# or smaller than the image boundary
output[:, [1, 3]] -= (params["height"] - scaling_factor *
im_dim_list[:, 0].view(-1, 1)) / 2
output[:, [2, 4]] -= (params["height"] - scaling_factor *
im_dim_list[:, 1].view(-1, 1)) / 2
output[:, 1:5] /= scaling_factor
for i in range(output.shape[0]):
output[i, [1, 3]] = torch.clamp(output[i, [1, 3]], 0.0,
im_dim_list[i, 0])
output[i, [2, 4]] = torch.clamp(output[i, [2, 4]], 0.0,
im_dim_list[i, 1])
list(
map(
lambda x: detection_write(x, original_imgs, params[
"classes"]), output))
det_names = pd.Series(imlist).apply(
lambda x: "{}/{}".format(det,
x.split("/")[-1]))
list(map(cv2.imwrite, det_names, original_imgs))
if model.params['cuda']:
torch.cuda.empty_cache()
def run_training(params,
label_csv_mame,
name_list,
test_label_csv_mame,
test_img_txt_path,
cfg_path,
params_dir,
valid_label_csv_mame=False,
valid_img_txt_path=False):
random.seed(params['seed'])
np.random.seed(params['seed'])
torch.manual_seed(params['seed'])
torch.cuda.manual_seed(params['seed'])
logging.basicConfig(level=logging.DEBUG,
format="[%(asctime)s %(filename)s] %(message)s")
# prepare test lables
prep_labels(test_img_txt_path, name_list, test_label_csv_mame)
# parse model architect from architect config file
blocks = parse_cfg(cfg_path)
# training transformation
pre_trans = RandomCrop(jitter=params['rand_crop'],
inp_dim=params["height"])
train_transform = transforms.Compose([
transforms.ColorJitter(brightness=params["exposure"],
saturation=params["saturation"]),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
test_pre_trans = OfssetLabels(resize=True, input_dim=params['height'])
test_transform = transforms.Compose([
ImgToTensorCv(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
train_data = CustData(label_csv_mame,
pre_trans=pre_trans,
transform=train_transform)
test_data = CustDataCV(test_label_csv_mame,
pre_trans=test_pre_trans,
transform=test_transform)
train_loader = DataLoader(train_data,
shuffle=True,
batch_size=params["batch_size"],
collate_fn=my_collate,
num_workers=params['num_workers'],
worker_init_fn=worker_init_fn)
test_loader = DataLoader(test_data,
shuffle=False,
batch_size=params["batch_size"],
collate_fn=my_collate,
num_workers=params['num_workers'],
worker_init_fn=worker_init_fn)
# initiate model
model = yolo_v3(params, blocks)
# Initiate validation data loader if there is any
if valid_label_csv_mame:
valid_data = CustDataCV(valid_label_csv_mame, transform=test_transform)
valid_loader = DataLoader(valid_data,
shuffle=False,
batch_size=params["batch_size"],
collate_fn=my_collate,
num_workers=params['num_workers'],
worker_init_fn=worker_init_fn)
# Start training
best_map, best_ap, best_conf, specific_conf_map, specific_conf_ap,\
map_frame = model.fit(train_loader, valid_loader, test_loader)
else:
# Start training
best_map, best_ap, best_conf, specific_conf_map, specific_conf_ap, \
map_frame = model.fit(train_loader, test_loader)
return best_map, best_ap, best_conf, specific_conf_map, specific_conf_ap, \
map_frame
