Esempio n. 1
0
def main():
    torch.set_grad_enabled(False)
    torch.backends.cudnn.benchmark = True

    test_dir = "./input/deepfake-detection-challenge/test_videos"
    csv_path = "./input/deepfake-detection-challenge/sample_submission.csv"

    face_detector = FaceDetector()
    face_detector.load_checkpoint(
        "./input/dfdc-pretrained-2/RetinaFace-Resnet50-fixed.pth")
    loader = DFDCLoader(test_dir, face_detector, T.ToTensor())

    model1 = xception(num_classes=2, pretrained=False)
    ckpt = torch.load("./input/dfdc-pretrained-2/xception-hg-2.pth",
                      map_location=torch.device('cpu'))
    model1.load_state_dict(ckpt["state_dict"])
    model1 = model1.cpu()
    model1.eval()

    model2 = WSDAN(num_classes=2, M=8, net="xception", pretrained=False).cpu()
    ckpt = torch.load("./input/dfdc-pretrained-2/ckpt_x.pth",
                      map_location=torch.device('cpu'))
    model2.load_state_dict(ckpt["state_dict"])
    model2.eval()

    model3 = WSDAN(num_classes=2, M=8, net="efficientnet",
                   pretrained=False).cpu()
    ckpt = torch.load("./input/dfdc-pretrained-2/ckpt_e.pth",
                      map_location=torch.device('cpu'))
    model3.load_state_dict(ckpt["state_dict"])
    model3.eval()

    zhq_nm_avg = torch.Tensor([.4479, .3744, .3473]).view(1, 3, 1, 1).cpu()
    zhq_nm_std = torch.Tensor([.2537, .2502, .2424]).view(1, 3, 1, 1).cpu()

    for batch in loader:
        batch = batch.cpu()

        i1 = F.interpolate(batch, size=299, mode="bilinear")
        i1.sub_(0.5).mul_(2.0)
        o1 = model1(i1).softmax(-1)[:, 1].cpu().numpy()

        i2 = (batch - zhq_nm_avg) / zhq_nm_std
        o2, _, _ = model2(i2)
        o2 = o2.softmax(-1)[:, 1].cpu().numpy()

        i3 = F.interpolate(i2, size=300, mode="bilinear")
        o3, _, _ = model3(i3)
        o3 = o3.softmax(-1)[:, 1].cpu().numpy()

        out = 0.2 * o1 + 0.7 * o2 + 0.1 * o3
        loader.feedback(out)

    with open(csv_path) as fin, open("submission.csv", "w") as fout:
        fout.write(next(fin))
        for line in fin:
            fname = line.split(",", 1)[0]
            pred = loader.score[fname]
            print("%s,%.6f" % (fname, pred), file=fout)
Esempio n. 2
0
def select_device(device='', apex=False, batch_size=None):
    # device = 'cpu' or '0' or '0,1,2,3'
    cpu_request = device.lower() == 'cpu'
    if device and not cpu_request:  # if device requested other than 'cpu'
        os.environ['CUDA_VISIBLE_DEVICES'] = device  # set environment variable
        assert torch.cuda.is_available(
        ), 'CUDA unavailable, invalid device %s requested' % device  # check availablity

    cuda = False if cpu_request else torch.cuda.is_available()
    if cuda:
        c = 1024**2  # bytes to MB
        ng = torch.cuda.device_count()
        if ng > 1 and batch_size:  # check that batch_size is compatible with device_count
            assert batch_size % ng == 0, 'batch-size %g not multiple of GPU count %g' % (
                batch_size, ng)
        x = [torch.cuda.get_device_properties(i) for i in range(ng)]
        s = 'Using CUDA ' + (
            'Apex ' if apex else ''
        )  # apex for mixed precision https://github.com/NVIDIA/apex
        for i in range(0, ng):
            if i == 1:
                s = ' ' * len(s)
            print(
                "%sdevice%g _CudaDeviceProperties(name='%s', total_memory=%dMB)"
                % (s, i, x[i].name, x[i].total_memory / c))
    else:
        print('Using CPU')

    print('')  # skip a line
    return torch.device('cuda:0' if cuda else 'cpu')
Esempio n. 3
0
def _loss_func(preds, targs):
    # Loss multipliers are printed at the end of train_run, to balance the number of annotated pixels
    preds_b = preds[:,1,:,:].unsqueeze(1)
    preds_c = preds[:,0,:,:].unsqueeze(1)
    mask_benign = targs == 0
    masked_preds_benign = preds_c[mask_benign].float() - preds_b[mask_benign].float()
    masked_targs_benign = targs[mask_benign].float()
    if len(masked_targs_benign) > 0:
        loss_benign = L1Loss()(masked_preds_benign, masked_targs_benign) * 1
    else:
        loss_benign = torch.full((), 0., device=torch.device("cuda"))

    mask_clinsig = targs > 0
    masked_preds_clinsig = preds_c[mask_clinsig].float() - preds_b[mask_clinsig].float()
    masked_targs_clinsig = targs[mask_clinsig].float()
    if len(masked_preds_clinsig) > 0:
        loss_clinsig = L1Loss()(masked_preds_clinsig, masked_targs_clinsig) * 5
    else:
        loss_clinsig = torch.full((), 0., device=torch.device("cuda"))
    return loss_benign + loss_clinsig
Esempio n. 4
0
    def gpu_status(self, gpu_arg):
        """
        Checks on the status of a GPU for model device assignment.
        INPUTS:
            1. Command line GPU switch argument:    <bool>
        RETURNS:
            1. PyTorch device type available:       <device object>
        """
        # check if the GPU is currently available and set device flag appropriately
        self.device_location = "cuda:0" if torch.cuda.is_available() else "cpu"

        # if GPU switch was selected during program initiation
        if gpu_arg == True and self.device_location == "cuda:0":
            return torch.device(self.device_location)
        
        # GPU switch not requested of device not available
        else:
            if gpu_arg:
                print("Sorry, GPU not available, using CPU instead...")
            self.device_location = "cpu"
            return torch.device(self.device_location)
Esempio n. 5
0
def _acc_dice(preds, targs):
    preds_b = preds[:,1,:,:].unsqueeze(1)
    preds_c = preds[:,0,:,:].unsqueeze(1)
    preds = preds_c - preds_b
    global acc_benign_vals
    mask = (targs == 1) | (preds > 0.5)
    masked_preds = preds[mask].float()
    masked_targs = targs[mask].float()
    if len(masked_targs) == 0:
        out = torch.full((), np.average(acc_benign_vals), device=torch.device("cuda"))
    else:
        iou = ((masked_preds - masked_targs).abs() <= 0.5).float().mean()
        out = 2 * iou / (1 + iou)
        acc_benign_vals.append(out)
        if len(acc_benign_vals) > 200:
            acc_benign_vals = acc_benign_vals[-100:]
    return out
 def init(state, im_x, total_num):
     for i in range(len(state)):
         target_pos.append(state[i][:2])
         target_sz.append(state[i][2:4])
         tracking_index.append(index * 100 + total_num + i)
         im_z_crop, _ = get_crop(im_x,
                                 target_pos[i],
                                 target_sz[i],
                                 z_size,
                                 avg_chans=avg_chans,
                                 context_amount=context_amount,
                                 func_get_subwindow=get_subwindow_tracking)
         im_z_crops.append(im_z_crop)
         array = torch.from_numpy(
             np.ascontiguousarray(
                 im_z_crops[i].transpose(2, 0, 1)[np.newaxis, ...],
                 np.float32)).to(torch.device("cuda"))
         lost.append(0)
         with torch.no_grad():
             features.append(Model(array, phase=phase))
Esempio n. 7
0
def setup(model: nn.Module) -> nn.Module:
    '''配置模型的GPU/CPU环境
    
    Args:
        model(nn.Module): 需要设置的模型

    Return:
        model(nn.Module): 配置好的模型
        device(torch.device): 模型可用的环境
    '''
    use_cuda = torch.cuda.is_available()
    device = torch.device('cuda' if use_cuda else 'cpu')
    model.to(device)

    # 检测多GPU环境
    if use_cuda:
        gpu_count = torch.cuda.device_count()
        if gpu_count > 1:
            model = nn.DataParallel(model)

    return model, device
def multiprocessing_update(task, task_cfg, index, im, dataqueue, resultqueue):
    # build model
    Model = model_builder.build_model(task,
                                      task_cfg.model).to(torch.device("cuda"))
    Model.eval()
    target_pos = []
    target_sz = []
    im_z_crops = []
    lost = []
    features = []
    tracking_index = []
    total_num = 0
    avg_chans = np.mean(im, axis=(0, 1))
    im_h, im_w = im.shape[0], im.shape[1]
    z_size = hyper_params['z_size']
    x_size = hyper_params['x_size']
    context_amount = hyper_params['context_amount']
    phase = hyper_params['phase_init']
    phase_track = hyper_params['phase_track']
    score_size = (
        hyper_params['x_size'] - hyper_params['z_size']
    ) // hyper_params['total_stride'] + 1 - hyper_params['num_conv3x3'] * 2
    if hyper_params['windowing'] == 'cosine':
        window = np.outer(np.hanning(score_size), np.hanning(score_size))
        window = window.reshape(-1)
    elif hyper_params['windowing'] == 'uniform':
        window = np.ones((score_size, score_size))
    else:
        window = np.ones((score_size, score_size))

    def init(state, im_x, total_num):
        for i in range(len(state)):
            target_pos.append(state[i][:2])
            target_sz.append(state[i][2:4])
            tracking_index.append(index * 100 + total_num + i)
            im_z_crop, _ = get_crop(im_x,
                                    target_pos[i],
                                    target_sz[i],
                                    z_size,
                                    avg_chans=avg_chans,
                                    context_amount=context_amount,
                                    func_get_subwindow=get_subwindow_tracking)
            im_z_crops.append(im_z_crop)
            array = torch.from_numpy(
                np.ascontiguousarray(
                    im_z_crops[i].transpose(2, 0, 1)[np.newaxis, ...],
                    np.float32)).to(torch.device("cuda"))
            lost.append(0)
            with torch.no_grad():
                features.append(Model(array, phase=phase))

    def delete_node(j):
        try:
            del target_pos[j]
            del target_sz[j]
            del features[j]
            del tracking_index[j]
            del lost[j]
        except Exception as error:
            print("delete error", error)

    while True:
        try:
            im_x, state, delete = dataqueue.get()
        except Exception as error:
            print(error)
            continue
        else:
            if len(state) > 0:
                init(state, im_x, total_num)
                total_num += len(state)
                continue
            if len(delete) > 0:
                delete_list = []
                for i in delete:
                    if i in tracking_index:
                        # print("delete",i)
                        node = tracking_index.index(i)
                        delete_node(node)

            result = []
            im = im_x.copy()
            del im_x, state, delete
            for i in range(len(features)):
                im_x_crop, scale_x = get_crop(
                    im,
                    target_pos[i],
                    target_sz[i],
                    z_size,
                    x_size=x_size,
                    avg_chans=avg_chans,
                    context_amount=context_amount,
                    func_get_subwindow=get_subwindow_tracking)
                array = torch.from_numpy(
                    np.ascontiguousarray(
                        im_x_crop.transpose(2, 0, 1)[np.newaxis, ...],
                        np.float32)).to(torch.device("cuda"))
                with torch.no_grad():
                    score, box, cls, ctr, *args = Model(array,
                                                        *features[i],
                                                        phase=phase_track)

                box = tensor_to_numpy(box[0])
                score = tensor_to_numpy(score[0])[:, 0]
                cls = tensor_to_numpy(cls[0])
                ctr = tensor_to_numpy(ctr[0])
                box_wh = xyxy2cxywh(box)

                # #lost goal
                if score.max() < 0.2:
                    lost[i] += 1
                    continue
                elif lost[i] > 0:
                    lost[i] -= 1
                best_pscore_id, pscore, penalty = postprocess_score(
                    score, box_wh, target_sz[i], scale_x, window)
                # box post-processing
                new_target_pos, new_target_sz = postprocess_box(
                    best_pscore_id, score, box_wh, target_pos[i], target_sz[i],
                    scale_x, x_size, penalty)
                new_target_pos, new_target_sz = restrict_box(
                    im_h, im_w, new_target_pos, new_target_sz)

                # save underlying state
                target_pos[i], target_sz[i] = new_target_pos, new_target_sz

                # return rect format
                track_rect = cxywh2xywh(
                    np.concatenate([target_pos[i], target_sz[i]], axis=-1))
                result.append(track_rect)

            delete_list = []
            for i in range(len(features)):
                if lost[i] > 10:
                    delete_list = []
                    delete_list.append(i)
            for i in delete_list:
                delete_node(i)
            # print(index, len(features))
            resultqueue.put([result, tracking_index])
Esempio n. 9
0
def train_model():
    num_epochs = 200
    learning_rate = 0.000001
    batch_size = 10
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(device)
    model = ConvNet().to(device)
    """
    dataxy=get_data()
    with open("psd_score.txt","wb") as f:
        pickle.dump(dataxy,f)
    """
    dataxy = []
    with open('psd_score.txt', 'rb') as file:
        dataxy = pickle.load(file)

    x = np.array(dataxy[0])
    y = np.array(dataxy[1])
    train_x_origin, test_x_origin, train_y_origin, test_y_origin = train_test_split(
        x, y)

    traindataset = ds(train_x_origin, train_y_origin, len(train_x_origin))
    testdataset = ds(test_x_origin, test_y_origin, len(test_x_origin))
    loss = nn.MSELoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

    train_loader = DataLoader(dataset=traindataset,
                              batch_size=batch_size,
                              shuffle=True,
                              num_workers=0)
    test_loader = DataLoader(dataset=testdataset,
                             batch_size=batch_size,
                             shuffle=True,
                             num_workers=0)

    train_loss = []
    test_loss = []

    for epoch in tqdm(range(num_epochs)):
        cur_train_loss = []
        cur_test_loss = []
        for i, (data, labels) in enumerate(train_loader):
            data = data.to(device).reshape(-1, 1, 5, 31)
            data = data.type(torch.FloatTensor).to(device)
            labels = labels.to(device)
            y_pred = model(data).type(torch.FloatTensor).to(device)
            labels = labels.type(torch.FloatTensor).to(device)
            l = loss(y_pred, labels)
            l.backward()
            optimizer.step()
            optimizer.zero_grad()
            cur_train_loss.append(l.item())
        train_loss.append(np.array(cur_train_loss).mean())
        print("train loss ", np.array(cur_train_loss).mean())

        for i, (test_x, test_y) in enumerate(test_loader):
            test_x = test_x.reshape(-1, 1, 5,
                                    31).type(torch.FloatTensor).to(device)
            test_y = test_y.to(device)
            out = model(test_x)
            l = loss(out, test_y)
            cur_test_loss.append(l.item())
        test_loss.append(np.array(cur_test_loss).mean())
        print("test loss ", np.array(cur_test_loss).mean())

    plt.plot(train_loss, label="train_loss")
    plt.plot(test_loss, label="test_loss")
    plt.legend()
    plt.show()
    predict = []
    torch.save(model.state_dict(), "model1.pt")
    for i in range(len(test_x_origin)):
        predict.append(
            model(
                torch.from_numpy(test_x_origin[i]).reshape(-1, 1, 5, 31).type(
                    torch.FloatTensor).to(device)).item())
    return test_y_origin, predict
Esempio n. 10
0
import os
import ast
import sys
sys.path.append('BengaliAI-Kaggle/')
from src.data.dataset import BengaliDatasetTrain
from model_dispatcher import MODEL_DISPATCHER
from torch import torch
import torch.nn as nn
from tqdm import tqdm

DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
TRAINING_FOLDS_CSV = os.environ.get('TRAINING_FOLDS_CSV')
IMG_WIDTH = int(os.environ.get('IMG_WIDTH'))
IMG_HEIGHT = int(os.environ.get('IMG_HEIGHT'))
EPOCHS = int(os.environ.get('EPOCHS'))

TRAIN_BATCH_SIZE = int(os.environ.get('TRAIN_BATCH_SIZE'))
TEST_BATCH_SIZE = int(os.environ.get('TEST_BATCH_SIZE'))

MODEL_MEAN = ast.literal_eval(os.environ.get('MODEL_MEAN'))
MODEL_STD = ast.literal_eval(os.environ.get('MODEL_STD'))

TRAINING_FOLDS = ast.literal_eval(os.environ.get('TRAINING_FOLDS'))
VALIDATION_FOLDS = ast.literal_eval(os.environ.get('VALIDATION_FOLDS'))
BASE_MODEL = os.environ.get('BASE_MODEL')


def loss_fn(outputs, targets):
    o1, o2, o3 = outputs
    t1, t2, t3 = targets
    l1 = nn.CrossEntropyLoss()(o1, t1)
Esempio n. 11
0
def train_model():
    num_epochs = 100
    learning_rate = 0.00001
    batch_size = 5
    channel_size = 3
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(device)
    model = ConvNet().to(device)
    #model=AlexNet().to(device)
    """
    dataxy=get_data()
    with open("psd_score.txt","wb") as f:
        pickle.dump(dataxy,f)
    """
    dataxy = []
    with open(
            "D:/code/code/eegemotion/git/model/corr_classify/corr_score_classify.txt",
            "rb") as f:
        dataxy = pickle.load(f)
    x = np.array(dataxy[0])
    y = np.array(dataxy[1])
    train_x_origin, test_x_origin, train_y_origin, test_y_origin = train_test_split(
        x, y)

    traindataset = ds(train_x_origin, train_y_origin, len(train_x_origin))
    testdataset = ds(test_x_origin, test_y_origin, len(test_x_origin))
    loss = nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

    train_loader = DataLoader(dataset=traindataset,
                              batch_size=batch_size,
                              shuffle=True,
                              num_workers=0)
    test_loader = DataLoader(dataset=testdataset,
                             batch_size=batch_size,
                             shuffle=True,
                             num_workers=0)

    train_loss = []
    test_loss = []
    train_accuracy = []
    test_accuracy = []
    for epoch in tqdm(range(num_epochs)):
        cur_train_loss = []
        cur_test_loss = []
        right, total = 0, 0
        for i, (data, labels) in enumerate(train_loader):
            data = data.to(device).reshape(-1, channel_size, 31, 31)
            data = data.type(torch.FloatTensor).to(device)
            y_pred = model(data).to(device)
            labels = labels.to(device).long()
            l = loss(y_pred, labels)
            l.backward()
            optimizer.step()
            optimizer.zero_grad()
            cur_train_loss.append(l.item())
            right += (torch.argmax(y_pred, dim=1) == labels).sum().item()
            total += batch_size

        train_loss.append(np.array(cur_train_loss).mean())
        train_accuracy.append(right / total)
        print("train loss ",
              np.array(cur_train_loss).mean(), " accuracy: ", right / total)

        right, total = 0, 0
        for i, (test_x, test_y) in enumerate(test_loader):
            test_x = test_x.reshape(-1, channel_size, 31,
                                    31).type(torch.FloatTensor).to(device)
            test_y = test_y.to(device).long()
            out = model(test_x)
            l = loss(out, test_y)
            cur_test_loss.append(l.item())

            right += (torch.argmax(out, dim=1) == test_y).sum().item()
            total += batch_size

        test_loss.append(np.array(cur_test_loss).mean())
        test_accuracy.append(right / total)
        print("test loss ",
              np.array(cur_test_loss).mean(), " accuracy: ", right / total)

    plt.plot(train_loss, label="train_loss")
    plt.plot(test_loss, label="test_loss")
    plt.legend()
    plt.show()

    plt.plot(train_accuracy, label="train_accuracy")
    plt.plot(test_accuracy, label="test_accuracy")
    plt.legend()
    plt.show()

    torch.save(model.state_dict(),
               "D:/code/code/eegemotion/git/model/corr_classify/model.pt")
def main():
##########################################################################################################
    #preparation part
    args = arg_parse()
    confidence = float(args.confidence)
    nms_thesh = float(args.nms_thresh)
    start = 0
    CUDA = torch.cuda.is_available()
    
    num_classes = 80
    
    model = Darknet(cfgfile)
    model.load_weights(weightsfile)
    
    model.net_info["height"] = args.reso
    inp_dim = int(model.net_info["height"])
    
    assert inp_dim % 32 == 0                   #assert后面语句为false时触发,中断程序
    assert inp_dim > 32

    if CUDA:
        model.cuda()
            
    model.eval()
    
    global confirm
    global person
    
    fps = 0.0
    count = 0
    frame = 0    
    person = []
    confirm = False
    reconfirm = False
    count_yolo = 0
    model_filename = 'model_data/mars-small128.pb'
    encoder = gdet.create_box_encoder(model_filename,batch_size=1) 
    metric = nn_matching.NearestNeighborDistanceMetric("cosine", max_cosine_distance, nn_budget)
    tracker = Tracker(metric)
    #record the video
    fourcc = cv2.VideoWriter_fourcc(*'XVID')
    #out = cv2.VideoWriter('output/testwrite_normal.avi',fourcc, 15.0, (640,480),True)

    cap = cv2.VideoCapture(0)

    detect_time = []
    recogn_time = []
    kalman_time = []
    aux_time = []
    while True:
        start = time.time()  
        ret, color_image = cap.read()
        '''
        frames = pipeline.wait_for_frames()
        color_frame = frames.get_color_frame()
        color_image = np.asanyarray(color_frame.get_data())
        '''
        if color_image is None:
            break
        img, orig_im, dim = prep_image(color_image, inp_dim)
        
        im_dim = torch.FloatTensor(dim).repeat(1,2)             
##########################################################################################################
        #people detection part                
        if CUDA:
            im_dim = im_dim.cuda()
            img = img.cuda()
        time_a = time.time()
        if count_yolo %3 == 0:                                                               #detect people every 3 frames
            output = model(Variable(img), CUDA)                         #适配后的图像放进yolo网络中,得到检测的结果
            output = write_results(output, confidence, num_classes, nms = True, nms_conf = nms_thesh)         


            if type(output) == int:
                fps  = ( fps + (1./(time.time()-start)) ) / 2
                print("fps= %f"%(fps))
                cv2.imshow("frame", orig_im)
                key = cv2.waitKey(1)
                if key & 0xFF == ord('q'):
                    break
                continue
        
            output[:,1:5] = torch.clamp(output[:,1:5], 0.0, float(inp_dim))/inp_dim                #夹紧张量,限制在一个区间内
        
            #im_dim = im_dim.repeat(output.size(0), 1)
            output[:,[1,3]] *= color_image.shape[1]
            output[:,[2,4]] *= color_image.shape[0]
            output = output.cpu().numpy() 
            output = sellect_person(output)                                       #把标签不是人的output去掉,减少计算量
            output = np.array(output)
            output_update = output
        elif count_yolo %3 != 0:
            output = output_update
        count_yolo += 1
        list(map(lambda x: write(x, orig_im), output))                #把结果加到原来的图像中   
        #output的[0,1:4]分别为框的左上和右下的点的位置
        detect_time.append(time.time() - time_a)
##########################################################################################################
        time_a = time.time()
        #kalman filter part
        outputs_tlwh = to_tlwh(output)                             ##把output数据变成适合kalman更新的类型
        features = encoder(orig_im,outputs_tlwh)
        detections = [Detection(output_tlwh, 1.0, feature) for output_tlwh, feature in zip(outputs_tlwh, features)]

        # Run non-maxima suppression.
        boxes = np.array([d.tlwh for d in detections])
        scores = np.array([d.confidence for d in detections])
        indices = preprocessing.non_max_suppression(boxes, nms_max_overlap, scores)
        detections = [detections[i] for i in indices]

        # Call the tracker
        tracker.predict()
        tracker.update(detections)

        for track in tracker.tracks:
            if not track.is_confirmed() or track.time_since_update > 1:
                continue 
            box = track.to_tlbr()
            cv2.rectangle(orig_im, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])),(255,255,255), 2)
            cv2.putText(orig_im, str(track.track_id),(int(box[0]), int(box[1])),0, 5e-3 * 200, (0,255,0),2)  
        
        kalman_time.append(time.time() - time_a)
##########################################################################################################
        #face recognition part
        time_a = time.time()
        if confirm == False:
            saved_model = './ArcFace/model/068.pth'
            name_list = os.listdir('./users')
            path_list = [os.path.join('./users',i,'%s.txt'%(i)) for i in name_list]
            total_features = np.empty((128,),np.float32)

            for i in path_list:
                temp = np.loadtxt(i)
                total_features = np.vstack((total_features,temp))
            total_features = total_features[1:]

            #threshold = 0.30896     #阈值并不合适,可能是因为训练集和测试集的差异所致!!!
            threshold = 0.5
            model_facenet = mobileFaceNet()
            model_facenet.load_state_dict(torch.load(saved_model)['backbone_net_list'])
            model_facenet.eval()
            #use_cuda = torch.cuda.is_available() and True
            #device = torch.device("cuda" if use_cuda else "cpu")
            device = torch.device("cuda")

            # is_cuda_avilable
            trans = transforms.Compose([
                transforms.Resize((112,112)),
                transforms.ToTensor(),
                transforms.Normalize([0.5,0.5,0.5],[0.5,0.5,0.5])
            ])
            model_facenet.to(device)

            img = Image.fromarray(color_image)
            bboxes, landmark = detect_faces(img)                                                                  #首先检测脸

            if len(bboxes) == 0:
                print('detect no people')
            else:
                for bbox in bboxes:
                    loc_x_y = [bbox[2], bbox[1]]
                    person_img = color_image[int(bbox[1]):int(bbox[3]), int(bbox[0]):int(bbox[2])].copy()              #从图像中截取框
                    feature = np.squeeze(get_feature(person_img, model_facenet, trans, device))                               #框里的图像计算feature
                    cos_distance = cosin_metric(total_features, feature)
                    index = np.argmax(cos_distance)
                    if  cos_distance[index] <= threshold:
                        continue
                    person = name_list[index]  
                    #在这里加框加文字
                    orig_im = draw_ch_zn(orig_im,person,font,loc_x_y)                                                                    #加名字
                    cv2.rectangle(orig_im,(int(bbox[0]),int(bbox[1])),(int(bbox[2]),int(bbox[3])),(0,0,255))           #加box
            #cv2.imshow("frame", orig_im)

##########################################################################################################
            #confirmpart
            print('confirmation rate: {} %'.format(count*10))
            cv2.putText(orig_im, 'confirmation rate: {} %'.format(count*10), (10,30),cv2.FONT_HERSHEY_PLAIN, 2, [0,255,0], 2)
            if len(bboxes)!=0 and len(output)!=0:
                if bboxes[0,0]>output[0,1] and bboxes[0,1]>output[0,2] and bboxes[0,2]<output[0,3] and bboxes[0,3]<output[0,4] and person:
                    count+=1
            frame+=1
            if count>=10 and frame<=30:
                confirm = True
                print('confirm the face is belong to that people')
            elif  frame >= 30:
                print('fail confirm, and start again')
                reconfirm = True
                count = 0
                frame = 0
            if reconfirm == True:
                cv2.putText(orig_im, 'fail confirm, and start again', (10,60),cv2.FONT_HERSHEY_PLAIN, 2, [0,255,0], 2)                   
##########################################################################################################
        recogn_time.append(time.time() - time_a)
        time_a = time.time()
        #show the final output result
        if not confirm:
            cv2.putText(orig_im, 'still not confirm', (output[0,1].astype(np.int32)+100,output[0,2].astype(np.int32)+20),
                                     cv2.FONT_HERSHEY_PLAIN, 2, [0,0,255], 2)
        #把识别的名字加上去
        if confirm:  
            for track in tracker.tracks:
                bbox = track.to_tlbr()
                if track.track_id == 1:
                    cv2.putText(orig_im, person, (int(bbox[0])+100,int(bbox[1])+20),
                                            cv2.FONT_HERSHEY_PLAIN, 2, [0,255,0], 2)
                
                    #rate.sleep()
        cv2.imshow("frame", orig_im)
        #out.write(orig_im)
        key = cv2.waitKey(1)
        if key & 0xFF == ord('q'):
            break
        
        aux_time.append(time.time()-time_a)
        fps  = ( fps + (1./(time.time()-start)) ) / 2
        print("fps= %f"%(fps))
    #calculate how long each part takes
    avg_detect_time = np.mean(detect_time)
    avg_recogn_time = np.mean(recogn_time)
    avg_kalman_time = np.mean(kalman_time)
    avg_aux_time = np.mean(aux_time)
    print("avg detect: {}".format(avg_detect_time))
    print("avg recogn: {}".format(avg_recogn_time))
    print("avg kalman: {}".format(avg_kalman_time))
    print("avg aux: {}".format(avg_aux_time))
    print("avg fps: {}".format(1/(avg_detect_time + avg_recogn_time + avg_kalman_time + avg_aux_time)))
Esempio n. 13
0
#get the other funcions
import dataset
import model
import trainer
#%%
from dataset import Cancer
from trainer import Trainer
from model import get_model
from torch.autograd import Variable
from visuals import plot_losses

import matplotlib.pyplot as plt
#%%
# CULaneDataset = importlib.reload(dataset).CULaneDataset

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

#%%
data_dir = './data'
labels_dir = './data/train_labels.csv'

# WhaleDataset = Whales('train', data_dir, labels_dir, aug=False, preprocess=True, ToTensor=True)
# WhaleDataset_val = Whales('val', data_dir, labels_dir, aug=False, preprocess=True, ToTensor=True)

#images are (96,96,3)
CancerDataset = Cancer('train',
                       data_dir,
                       labels_dir,
                       aug=False,
                       preprocess=True,
                       ToTensor=True)