HYTOS

"""

Training For Small Object

"""

import os

import argparse

import torch.nn as nn

from torch.utils.data import DataLoader

from src.doftech_dataset import DoftechDataset

from src.utils import *

from src.loss import YoloLoss

from src.yolo_net import Yolo

from src.yolo_doftech import YoloD

import shutil

import visdom

import cv2

import pickle

import numpy as np

from src.vis_utils import array_tool as at

from src.vis_utils.vis_tool import visdom_bbox

loss_data = {'X': [], 'Y': [], 'legend_U':['total', 'coord', 'conf', 'cls']}

# 형상 CLASS

DOFTECH_CLASSES= ['gate', 'globe', 'butterfly', 'check', 'ball', 'relief',

                  '3way_solenoid', 'gate_pressure', 'globe_pressure', 'butterfly_pressure', 'ball_shutoff', 'ball_pressure','ball_motor', 'plug_pressure',

                  'circle', 'inst_console', 'inst_console_dcs', 'inst_console_sih', 'logic_dcs', 'utility', 'specialty_items', 'logic', 'logic_local_console_dcs',

                  'reducer', 'blind_spectacle_open', 'blind_insertion_open', 'blind_spectacle_close', 'blind_insertion_close',

                  'strainer_basket', 'strainer_conical', 'fitting_capillary_tubing', 'meter_ultrasonic', 'strainer_y', 'tube_pitot'

                  ,'opc']

use_voc_model = True

use_visdom = False

if use_visdom:

    visdom = visdom.Visdom(port='8088')

    #visdom = visdom.Visdom()

def train(name=None, classes=None, bigs=None, root_path=None, pre_trained_model_path=None):

    global DOFTECH_CLASSES

    DOFTECH_CLASSES = classes

    parser = argparse.ArgumentParser("You Only Look Once: Unified, Real-Time Object Detection")

    parser.add_argument("--image_size", type=int, default=512, help="The common width and height for all images")

    parser.add_argument("--batch_size", type=int, default=10, help="The number of images per batch")

    # Training 기본 Setting

    parser.add_argument("--momentum", type=float, default=0.9)

    parser.add_argument("--decay", type=float, default=0.0005)

    parser.add_argument("--dropout", type=float, default=0.5)

    parser.add_argument("--num_epoches", type=int, default=205)

    parser.add_argument("--test_interval", type=int, default=20, help="Number of epoches between testing phases")

    parser.add_argument("--object_scale", type=float, default=1.0)

    parser.add_argument("--noobject_scale", type=float, default=0.5)

    parser.add_argument("--class_scale", type=float, default=1.0)

    parser.add_argument("--coord_scale", type=float, default=5.0)

    parser.add_argument("--reduction", type=int, default=32)

    parser.add_argument("--es_min_delta", type=float, default=0.0,

                        help="Early stopping's parameter: minimum change loss to qualify as an improvement")

    parser.add_argument("--es_patience", type=int, default=0,

                        help="Early stopping's parameter: number of epochs with no improvement after which training will be stopped. Set to 0 to disable this technique.")

    # 확인해야 하는 PATH

    parser.add_argument("--data_path", type=str, default=os.path.join(root_path, 'training'), help="the root folder of dataset") # 학습 데이터 경로 -> image와 xml의 상위 경로 입력

    parser.add_argument("--data_path_test", type=str, default=os.path.join(root_path, 'test'), help="the root folder of dataset") # 테스트 데이터 경로 -> test할 이미지만 넣으면 됨

    #parser.add_argument("--pre_trained_model_type", type=str, choices=["model", "params"], default="model")

    parser.add_argument("--pre_trained_model_path", type=str, default=pre_trained_model_path) # Pre-training 된 모델 경로

    parser.add_argument("--saved_path", type=str, default=os.path.join(root_path, 'checkpoint')) # training 된 모델 저장 경로

    parser.add_argument("--conf_threshold", type=float, default=0.35)

    parser.add_argument("--nms_threshold", type=float, default=0.5)

    opt = parser.parse_args()

    if not os.path.isdir(opt.saved_path):

        os.mkdir(opt.saved_path)

    # 학습할 클래스들을 저장하고 인식 시 불러와 사용합니다.

    with open(os.path.join(opt.saved_path, name + "_info.info"), 'w') as stream:

        con = str(len(DOFTECH_CLASSES))

        names = '\n'.join(DOFTECH_CLASSES)

        bigs = '\n'.join(bigs)

        con = con + '\n' + names + '\n' + '***bigs***' + '\n' + bigs

        stream.write(con)

    if torch.cuda.is_available():

        torch.cuda.manual_seed(123)

    else:

        torch.manual_seed(123)

    learning_rate_schedule = {"0": 1e-5, "5": 1e-4,

                              "80": 1e-5, "110": 1e-6}

    training_params = {"batch_size": 1,#opt.batch_size,

                       "shuffle": True,

                       "drop_last": True,

                       "collate_fn": custom_collate_fn}

    test_params = {"batch_size": opt.batch_size,

                   "shuffle": False,

                   "drop_last": False,

                   "collate_fn": custom_collate_fn}

    training_set = DoftechDataset(opt.data_path, opt.image_size, is_training=True, classes=DOFTECH_CLASSES)

    training_generator = DataLoader(training_set, **training_params)

    test_set = DoftechDataset(opt.data_path_test, opt.image_size, is_training=False, classes=DOFTECH_CLASSES)

    test_generator = DataLoader(test_set, **test_params)

    # BUILDING MODEL =======================================================================

    if use_voc_model :

        pre_model = Yolo(20).cuda()

        pre_model.load_state_dict(torch.load(opt.pre_trained_model_path), strict=False)

        model = YoloD(pre_model, training_set.num_classes).cuda()

    else :

        model = Yolo(training_set.num_classes).cuda()

    nn.init.normal_(list(model.modules())[-1].weight, 0, 0.01)

    criterion = YoloLoss(training_set.num_classes, model.anchors, opt.reduction)

    optimizer = torch.optim.SGD(model.parameters(), lr=1e-5, momentum=opt.momentum, weight_decay=opt.decay)

    best_loss = 1e10

    best_epoch = 0

    num_iter_per_epoch = len(training_generator)

    loss_step = 0

    # ======================================================================================

    # TRAINING =============================================================================

    save_count = 0

    model.train()

    for epoch in range(opt.num_epoches):

        if str(epoch) in learning_rate_schedule.keys():

            for param_group in optimizer.param_groups:

                param_group['lr'] = learning_rate_schedule[str(epoch)]

        for iter, batch in enumerate(training_generator):

            image, label, image2 = batch

            image = Variable(image.cuda(), requires_grad=False)

            if torch.cuda.is_available():

                image = Variable(image.cuda(), requires_grad=True)

                origin = Variable(image2.cuda(), requires_grad=False)

            else:

                image = Variable(image, requires_grad=True)

            optimizer.zero_grad()

            logits = model(image)

            loss, loss_coord, loss_conf, loss_cls = criterion(logits, label)

            loss.backward()

            optimizer.step()

            if iter % (opt.test_interval * 5) == 0:

                print("Epoch: {}/{}, Iteration: {}/{}, Lr: {}, Loss:{:.5f} (Coord:{:.5f} Conf:{:.5f} Cls:{:.5f})".format

                    (epoch + 1, opt.num_epoches, iter + 1, num_iter_per_epoch, optimizer.param_groups[0]['lr'], loss,

                    loss_coord,loss_conf,loss_cls))

                if use_visdom:

                    predictions = post_processing(logits, opt.image_size, DOFTECH_CLASSES, model.anchors, opt.conf_threshold,

                                                  opt.nms_threshold)

                    gt_image = at.tonumpy(image[0])

                    gt_image = visdom_bbox(gt_image, label[0])

                    visdom.image(gt_image, opts=dict(title='gt_box_image'), win=3)

                    #

                    origin_image = at.tonumpy(origin[0])

                    origin_image = visdom_bbox(origin_image, [])

                    visdom.image(origin_image, opts=dict(title='origin_box_image'), win=4)

                    image = at.tonumpy(image[0])

                    if len(predictions) != 0:

                        box_image = visdom_bbox(image, predictions[0])

                        visdom.image(box_image, opts=dict(title='box_image'), win=2)

                    elif len(predictions) == 0:

                        box_image = visdom_bbox(image, [])

                        visdom.image(box_image, opts=dict(title='box_image'), win=2)

                    loss_dict = {

                        'total' : loss.item(),

                        'coord' : loss_coord.item(),

                        'conf' : loss_conf.item(),

                        'cls' : loss_cls.item()

                    }

                    visdom_loss(visdom, loss_step, loss_dict)

                    loss_step = loss_step + 1

        if epoch % opt.test_interval == 0:

            model.eval()

            loss_ls = []

            loss_coord_ls = []

            loss_conf_ls = []

            loss_cls_ls = []

            for te_iter, te_batch in enumerate(test_generator):

                te_image, te_label, _ = te_batch

                num_sample = len(te_label)

                if torch.cuda.is_available():

                    te_image = te_image.cuda()

                with torch.no_grad():

                    te_logits = model(te_image)

                    batch_loss, batch_loss_coord, batch_loss_conf, batch_loss_cls = criterion(te_logits, te_label)

                loss_ls.append(batch_loss * num_sample)

                loss_coord_ls.append(batch_loss_coord * num_sample)

                loss_conf_ls.append(batch_loss_conf * num_sample)

                loss_cls_ls.append(batch_loss_cls * num_sample)

            te_loss = sum(loss_ls) / test_set.__len__()

            te_coord_loss = sum(loss_coord_ls) / test_set.__len__()

            te_conf_loss = sum(loss_conf_ls) / test_set.__len__()

            te_cls_loss = sum(loss_cls_ls) / test_set.__len__()

            print("Test>> Epoch: {}/{}, Lr: {}, Loss:{:.5f} (Coord:{:.5f} Conf:{:.5f} Cls:{:.5f})".format(

                epoch + 1, opt.num_epoches, optimizer.param_groups[0]['lr'], te_loss, te_coord_loss, te_conf_loss, te_cls_loss))

            model.train()

            if te_loss + opt.es_min_delta < best_loss:

                save_count += 1

                best_loss = te_loss

                best_epoch = epoch

                print("SAVE MODEL")

                # for debug for each loss

                #torch.save(model.state_dict(), os.path.join(opt.saved_path, name + "_only_params_" + str(save_count) + "_" + "{:.5f}".format(best_loss) + ".pth"))

                #torch.save(model, os.path.join(opt.saved_path, name + "_whole_model_" + str(save_count) + "_" + "{:.5f}".format(best_loss) + ".pth"))

                # save

                torch.save(model.state_dict(), os.path.join(opt.saved_path, name + "_only_params.pth"))

                torch.save(model, os.path.join(opt.saved_path, name + "_whole_model.pth"))

            else:

                save_count += 1

                # for debug for each loss

                #torch.save(model.state_dict(), os.path.join(opt.saved_path, name + "_only_params_" + str(save_count) + "_" + "{:.5f}".format(te_loss) + ".pth"))

                #torch.save(model, os.path.join(opt.saved_path, name + "_whole_model_" + str(save_count) + "_" + "{:.5f}".format(te_loss) + ".pth"))

            # Early stopping

            if epoch - best_epoch > opt.es_patience > 0:

                print("Stop training at epoch {}. The lowest loss achieved is {}".format(epoch, te_loss))

                break

def visdom_loss(visdom, loss_step, loss_dict):

    loss_data['X'].append(loss_step)

    loss_data['Y'].append([loss_dict[k] for k in loss_data['legend_U']])

    visdom.line(

        X=np.stack([np.array(loss_data['X'])] * len(loss_data['legend_U']), 1),

        Y=np.array(loss_data['Y']),

        win=30,

        opts=dict(xlabel='Step',

                  ylabel='Loss',

                  title='YOLO_V2',

                  legend=loss_data['legend_U']),

        update='append'

    )

if __name__ == "__main__":

    datas = ['gate', 'globe', 'butterfly', 'check', 'ball', 'relief',

                  '3way_solenoid', 'gate_pressure', 'globe_pressure', 'butterfly_pressure', 'ball_shutoff', 'ball_pressure','ball_motor', 'plug_pressure',

                  'inst', 'func_valve', 'inst_console', 'inst_console_dcs', 'inst_console_sih', 'logic_dcs', 'utility', 'specialty_items', 'logic', 'logic_local_console_dcs',

                  'reducer', 'blind_spectacle_open', 'blind_insertion_open', 'blind_spectacle_close', 'blind_insertion_close',

                  'strainer_basket', 'strainer_conical', 'fitting_capillary_tubing', 'meter_ultrasonic', 'strainer_y', 'tube_pitot',

                  'opc']

    data_path = os.path.join(os.path.dirname(os.path.realpath(__file__)) + '\\Data\\', 'VnV')

    train(name='VnV', classes=datas, root_path=data_path, pre_trained_model_path=os.path.dirname(os.path.realpath(

                                                       __file__)) + '\\pre_trained_model\\only_params_trained_yolo_voc')