Пример #1
0
def demo():
    args = parse_args()
    print('call with args: {}'.format(args))

    # input images
    images_dir = 'images'
    images_names = ['image1.jpg', 'image2.jpg']

    classes = ('aeroplane', 'bicycle', 'bird', 'boat',
                         'bottle', 'bus', 'car', 'cat', 'chair',
                         'cow', 'diningtable', 'dog', 'horse',
                         'motorbike', 'person', 'pottedplant',
                         'sheep', 'sofa', 'train', 'tvmonitor')

    model = Yolov2()
    weight_loader = WeightLoader()
    weight_loader.load(model, 'yolo-voc.weights')
    print('loaded')

    # model_path = os.path.join(args.output_dir, args.model_name + '.pth')
    # print('loading model from {}'.format(model_path))
    # if torch.cuda.is_available():
    #     checkpoint = torch.load(model_path)
    # else:
    #     checkpoint = torch.load(model_path, map_location='cpu')
    # model.load_state_dict(checkpoint['model'])

    if args.use_cuda:
        model.cuda()

    model.eval()
    print('model loaded')

    for image_name in images_names:
        image_path = os.path.join(images_dir, image_name)
        img = Image.open(image_path)
        im_data, im_info = prepare_im_data(img)

        if args.use_cuda:
            im_data_variable = Variable(im_data).cuda()
        else:
            im_data_variable = Variable(im_data)

        tic = time.time()

        yolo_output = model(im_data_variable)
        yolo_output = [item[0].data for item in yolo_output]
        detections = yolo_eval(yolo_output, im_info, conf_threshold=0.6, nms_threshold=0.4)

        toc = time.time()
        cost_time = toc - tic
        print('im detect, cost time {:4f}, FPS: {}'.format(
            toc-tic, int(1 / cost_time)))

        det_boxes = detections[:, :5].cpu().numpy()
        det_classes = detections[:, -1].long().cpu().numpy()
        im2show = draw_detection_boxes(img, det_boxes, det_classes, class_names=classes)
        plt.figure()
        plt.imshow(im2show)
        plt.show()
Пример #2
0
def test():
    args = parse_args()
    args.conf_thresh = 0.005
    args.nms_thresh = 0.45
    if args.vis:
        args.conf_thresh = 0.5
    print('Called with args:')
    print(args)

    # prepare dataset

    val_imdb = get_imdb(args.dataset)

    val_dataset = RoiDataset(val_imdb, train=False)
    val_dataloader = DataLoader(val_dataset,
                                batch_size=args.batch_size,
                                shuffle=False)

    # load model
    model = Yolov2(arch=args.arch)
    # weight_loader = WeightLoader()
    # weight_loader.load(model, 'yolo-voc.weights')
    # print('loaded')

    model_path = os.path.join(args.output_dir, args.model_name + '.pth')
    print('loading model from {}'.format(model_path))
    if torch.cuda.is_available():
        checkpoint = torch.load(model_path)
    else:
        checkpoint = torch.load(model_path, map_location='cpu')
    model.load_state_dict(checkpoint['model'])

    if args.use_cuda:
        model.cuda()

    model.eval()
    print('model loaded')

    dataset_size = len(val_imdb.image_index)
    print('classes: ', val_imdb.num_classes)

    all_boxes = [[[] for _ in range(dataset_size)]
                 for _ in range(val_imdb.num_classes)]

    det_file = os.path.join(args.output_dir, 'detections.pkl')

    results = []

    img_id = -1
    with torch.no_grad():
        for batch, (im_data, im_infos) in enumerate(val_dataloader):
            if args.use_cuda:
                im_data_variable = Variable(im_data).cuda()
            else:
                im_data_variable = Variable(im_data)

            yolo_outputs = model(im_data_variable)
            for i in range(im_data.size(0)):
                img_id += 1
                output = [item[i].data for item in yolo_outputs]
                im_info = {'width': im_infos[i][0], 'height': im_infos[i][1]}
                detections = yolo_eval(output,
                                       im_info,
                                       conf_threshold=args.conf_thresh,
                                       nms_threshold=args.nms_thresh)

                if img_id % 100 == 0:
                    print('im detect [{}/{}]'.format(img_id + 1,
                                                     len(val_dataset)))
                if len(detections) > 0:
                    for cls in range(val_imdb.num_classes):
                        inds = torch.nonzero(detections[:, -1] == cls).view(-1)
                        if inds.numel() > 0:
                            cls_det = torch.zeros((inds.numel(), 5))
                            cls_det[:, :4] = detections[inds, :4]
                            cls_det[:, 4] = detections[inds,
                                                       4] * detections[inds, 5]
                            all_boxes[cls][img_id] = cls_det.cpu().numpy()

                img = Image.open(val_imdb.image_path_at(img_id))
                if len(detections) > 0:
                    detect_result = {}

                    boxes = detections[:, :5].cpu().numpy()
                    classes = detections[:, -1].long().cpu().numpy()
                    class_names = val_imdb.classes

                    num_boxes = boxes.shape[0]

                    labels = []

                    for i in range(num_boxes):
                        det_bbox = tuple(
                            np.round(boxes[i, :4]).astype(np.int64))
                        score = boxes[i, 4]
                        gt_class_ind = classes[i]
                        class_name = class_names[gt_class_ind]
                        disp_str = '{}: {:.2f}'.format(class_name, score)

                        bbox = tuple(np.round(boxes[i, :4]).astype(np.int64))

                        xmin = bbox[0]
                        ymin = bbox[1]
                        xmax = bbox[2]
                        ymax = bbox[3]

                        box2d = {}
                        box2d["x1"] = str(xmin)
                        box2d["y1"] = str(ymin)
                        box2d["x2"] = str(xmax)
                        box2d["y2"] = str(ymax)

                        bbox = {}
                        bbox["box2d"] = box2d
                        bbox["category"] = class_name

                        labels.append(bbox)

                    detect_result["ImageID"] = os.path.basename(
                        val_imdb.image_path_at(img_id))
                    detect_result["labels"] = labels

                    results.append(detect_result)

                if args.vis:
                    img = Image.open(val_imdb.image_path_at(img_id))
                    if len(detections) == 0:
                        continue
                    det_boxes = detections[:, :5].cpu().numpy()
                    det_classes = detections[:, -1].long().cpu().numpy()
                    im2show = draw_detection_boxes(
                        img,
                        det_boxes,
                        det_classes,
                        class_names=val_imdb.classes)
                    plt.figure()
                    plt.imshow(im2show)
                    plt.show()

            #if img_id > 10:
            #    break

    print(results)
    results_file = os.path.join(args.output_dir, 'detections.json')
    with open(results_file, 'w') as f:
        json.dump(results,
                  f,
                  ensure_ascii=False,
                  indent=4,
                  sort_keys=True,
                  separators=(',', ': '))

    with open(det_file, 'wb') as f:
        pickle.dump(all_boxes, f, pickle.HIGHEST_PROTOCOL)

    val_imdb.evaluate_detections(all_boxes, output_dir=args.output_dir)
Пример #3
0
def demo():
    args = parse_args()
    print('call with args: {}'.format(args))

    # input images
    images_dir = 'images'
    images_names = ['trainval1.jpg', 'trainval2.jpg', 'test1.jpg', 'test2.jpg']


    classes = ("car", "bus", "truck", "svehicle", "pedestrian", "motorbike", "bicycle", "train", "signal", "signs")

    model = Yolov2(arch=args.arch)
    #weight_loader = WeightLoader()
    #weight_loader.load(model, 'yolo-voc.weights')
    #print('loaded')

    model_path = os.path.join(args.output_dir, args.model_name + '.pth')
    print('loading model from {}'.format(model_path))
    if torch.cuda.is_available():
        checkpoint = torch.load(model_path)
    else:
        checkpoint = torch.load(model_path, map_location='cpu')
    model.load_state_dict(checkpoint['model'])

    if args.use_cuda:
        model.cuda()

    model.eval()
    print('model loaded')
    print(model)

    ## generate weight
    idx = 0
    for ii, module in enumerate(model.trunk.features):
        #print("ii",ii)
        #print("module",module)
        #print(type(module))
        #print(module.__dict__)

        if isinstance( module, torch.nn.modules.conv.Conv2d):
            print("conv2d layer_%d" % idx)
            #print("weight",module.__dict__['_parameters']['weight'])
            weight = module.__dict__['_parameters']['weight']
            weight = weight.detach().numpy() # nn.tensor -> numpy
            #print(weight)
            #print(weight.shape)
            
            if ii == 0:
                header_w = float_dtype+' w_%s[%d][%d][%d][%d]=\n' % ((str(ii),)+(weight.shape)) + arr2header(weight)
                save_header('./weight_l0.h', header_w)

            #print("bias",module.__dict__['_parameters']['bias'])
            bias = module.__dict__['_parameters']['bias']
            bias = bias.detach().numpy() # nn.tensor -> numpy
            #print(bias)

            if ii == 0:
                coef = bias.reshape(-1,).astype(float_np_dtype)
                header = float_dtype+(' b_%s[%d]=\n' % ((str(ii)),len(coef))) + arr2header(coef)
                save_header('./bias_l0.h', header)

            idx += 1
    exit()


    for image_name in images_names:
        image_path = os.path.join(images_dir, image_name)
        img = Image.open(image_path)
        im_data, im_info = prepare_im_data(img)

        if args.use_cuda:
            im_data_variable = Variable(im_data).cuda()
        else:
            im_data_variable = Variable(im_data)

        tic = time.time()

        yolo_output = model(im_data_variable)
        yolo_output = [item[0].data for item in yolo_output]
        detections = yolo_eval(yolo_output, im_info, conf_threshold=0.2, nms_threshold=0.4)
        ##print(detections)

        toc = time.time()
        cost_time = toc - tic
        print('im detect, cost time {:4f}, FPS: {}'.format(
            toc-tic, int(1 / cost_time)))

        det_boxes = detections[:, :5].cpu().numpy()
        det_classes = detections[:, -1].long().cpu().numpy()
        im2show = draw_detection_boxes(img, det_boxes, det_classes, class_names=classes)
        plt.figure()
        plt.imshow(im2show)
        #plt.show()

        save_image_path = os.path.join(images_dir, image_name + "_detect.jpg")
        print("save -> " + save_image_path)
        plt.savefig(save_image_path)
Пример #4
0
def test():
    args = parse_args()
    args.conf_thresh = 0.005
    args.nms_thresh = 0.45
    if args.vis:
        args.conf_thresh = 0.5
    print('Called with args:')
    print(args)

    # prepare dataset

    if args.dataset == 'voc07trainval':
        args.imdbval_name = 'voc_2007_trainval'

    elif args.dataset == 'voc07test':
        args.imdbval_name = 'voc_2007_test'

    else:
        raise NotImplementedError

    val_imdb = get_imdb(args.imdbval_name)

    val_dataset = RoiDataset(val_imdb, train=False)
    val_dataloader = DataLoader(val_dataset,
                                batch_size=args.batch_size,
                                shuffle=False)

    # load model
    model = Yolov2()
    # weight_loader = WeightLoader()
    # weight_loader.load(model, 'yolo-voc.weights')
    # print('loaded')

    model_path = os.path.join(args.output_dir, args.model_name + '.pth')
    print('loading model from {}'.format(model_path))
    if torch.cuda.is_available():
        checkpoint = torch.load(model_path)
    else:
        checkpoint = torch.load(model_path, map_location='cpu')
    model.load_state_dict(checkpoint['model'])

    if args.use_cuda:
        model.cuda()

    model.eval()
    print('model loaded')

    dataset_size = len(val_imdb.image_index)

    all_boxes = [[[] for _ in range(dataset_size)]
                 for _ in range(val_imdb.num_classes)]

    det_file = os.path.join(args.output_dir, 'detections.pkl')

    img_id = -1
    with torch.no_grad():
        for batch, (im_data, im_infos) in enumerate(val_dataloader):
            if args.use_cuda:
                im_data_variable = Variable(im_data).cuda()
            else:
                im_data_variable = Variable(im_data)

            yolo_outputs = model(im_data_variable)
            for i in range(im_data.size(0)):
                img_id += 1
                output = [item[i].data for item in yolo_outputs]
                im_info = {'width': im_infos[i][0], 'height': im_infos[i][1]}
                detections = yolo_eval(output,
                                       im_info,
                                       conf_threshold=args.conf_thresh,
                                       nms_threshold=args.nms_thresh)
                print('im detect [{}/{}]'.format(img_id + 1, len(val_dataset)))
                if len(detections) > 0:
                    for cls in range(val_imdb.num_classes):
                        inds = torch.nonzero(detections[:, -1] == cls).view(-1)
                        if inds.numel() > 0:
                            cls_det = torch.zeros((inds.numel(), 5))
                            cls_det[:, :4] = detections[inds, :4]
                            cls_det[:, 4] = detections[inds,
                                                       4] * detections[inds, 5]
                            all_boxes[cls][img_id] = cls_det.cpu().numpy()

                if args.vis:
                    img = Image.open(val_imdb.image_path_at(img_id))
                    if len(detections) == 0:
                        continue
                    det_boxes = detections[:, :5].cpu().numpy()
                    det_classes = detections[:, -1].long().cpu().numpy()
                    im2show = draw_detection_boxes(
                        img,
                        det_boxes,
                        det_classes,
                        class_names=val_imdb.classes)
                    plt.figure()
                    plt.imshow(im2show)
                    plt.show()

    with open(det_file, 'wb') as f:
        pickle.dump(all_boxes, f, pickle.HIGHEST_PROTOCOL)

    val_imdb.evaluate_detections(all_boxes, output_dir=args.output_dir)
Пример #5
0
def demo():
    args = parse_args()
    print('call with args: {}'.format(args))

    # input images
    images_dir = 'images'
    images_names = ['image1.jpg', 'image2.jpg']

    classes = ('aeroplane', 'bicycle', 'bird', 'boat',
               'bottle', 'bus', 'car', 'cat', 'chair',
                         'cow', 'diningtable', 'dog', 'horse',
                         'motorbike', 'person', 'pottedplant',
                         'sheep', 'sofa', 'train', 'tvmonitor')

    #model = Yolov2()
    #state = torch.load('output/yolov2_epoch_160.pth')['model']
    # model.load_state_dict(state)

    model = torch.load(args.model)
    print('loaded')

    # model_path = os.path.join(args.output_dir, args.model_name + '.pth')
    # print('loading model from {}'.format(model_path))
    # if torch.cuda.is_available():
    #     checkpoint = torch.load(model_path)
    # else:
    #     checkpoint = torch.load(model_path, map_location='cpu')
    # model.load_state_dict(checkpoint['model'])

    # if args.use_cuda:
    model.cuda()

    model.eval()
    print('model loaded')

    for image_name in images_names:
        image_path = os.path.join(images_dir, image_name)
        img = Image.open(image_path)
        im_data, im_info = prepare_im_data(img)
        im_data_variable = Variable(im_data).cuda()

        tic = time.time()

        #delta_pred, conf_pred, class_pred = model(im_data_variable)
        output = model(im_data_variable)
        B, C, H, W = output.size()
        out = output.permute(0, 2, 3,
                             1).contiguous().view(B, H * W * 5, 5 + 20)

        xy_pred = torch.sigmoid(out[:, :, 0:2])
        conf_pred = torch.sigmoid(out[:, :, 4:5])
        hw_pred = torch.exp(out[:, :, 2:4])
        class_score = out[:, :, 5:]
        class_pred = F.softmax(class_score, dim=-1)
        delta_pred = torch.cat([xy_pred, hw_pred], dim=-1)
        yolo_output = [delta_pred, conf_pred, class_pred]
        yolo_output = [item[0].data for item in yolo_output]
        detections = yolo_eval(yolo_output, im_info,
                               conf_threshold=0.5, nms_threshold=0.4)

        toc = time.time()
        cost_time = toc - tic
        print('im detect, cost time {:4f}, FPS: {}'.format(
            toc-tic, int(1 / cost_time)))

        det_boxes = detections[:, :5].cpu().numpy()
        det_classes = detections[:, -1].long().cpu().numpy()
        im2show = draw_detection_boxes(
            img, det_boxes, det_classes, class_names=classes)
        plt.figure()
        plt.imshow(im2show)
        plt.show()
Пример #6
0
def importance_based_compression():
    # ====================load model=====================
    # model1:encoder to get intermediate deep feature
    model1 = HalfMirror()
    pretrained = torch.load('./models/Reconstruction.pkl')
    model1.load_state_dict(pretrained['model'])
    model1.cuda()
    model1.eval()

    # model2:decoder to reconstruct input image
    model2 = Mirror()
    pretrained_dict = model1.state_dict()
    model_dict = model2.state_dict()
    # 将pretrained_dict里不属于model_dict的键剔除掉
    pretrained_dict = {
        k: v
        for k, v in pretrained_dict.items() if k in model_dict
    }
    # 更新现有的model_dict
    model_dict.update(pretrained_dict)
    # 加载我们真正需要的state_dict
    model2.load_state_dict(model_dict)
    model2.cuda()
    model2.eval()

    # model3:decoder to detect objection
    model3 = Detect()
    pretrained = torch.load('./models/Detection.pth')
    model3.load_state_dict(pretrained['model'])
    model3.cuda()
    model3.eval()

    # ====================input image and extract feature====================
    # input the picture
    img = Image.open('./data/23.tif')
    im_data, im_info = prepare_im_data(img)
    im_data_variable = Variable(im_data).cuda()

    # get the feature tensor from split point max_7
    handle_feat = model1.max_7.register_forward_hook(
        get_features_hook)  # conv1
    a = model1(im_data_variable)
    handle_feat.remove()
    feat1 = feat_result[0]
    V_featensor = feat1[0, ...]

    # get the importance fator
    handle_weight = model1.channel_select.register_forward_hook(
        get_weight_hook)
    a = model1(im_data_variable)
    handle_weight.remove()
    weight1 = weight_result[0]
    weight = np.transpose(weight1[0])[0][0]

    # # ====================intermediate deep feature compression====================
    C = V_featensor.shape[0]
    height = V_featensor.shape[1]
    width = V_featensor.shape[2]
    quant_bit = 4
    parameters1, parameters2, pp, indexcode, Qstream, quant_feature = compression(
        V_featensor, quant_bit, height, width, weight)

    #
    # ====================intermediate deep feature decompression====================
    inputpath = './output/compressed.bin'
    outputpath = './output/decompressed.txt'
    decode_feature, featuremap_d, unpacked_feature = decompression(
        inputpath, outputpath, parameters1, parameters2, height, width,
        quant_bit)
    compressedsize = os.path.getsize(inputpath)
    print('compressed_size:{}'.format(compressedsize))
    # ====================complete multi_tasks====================
    inputfeature = decode_feature.reshape(1, C, height, width)
    inputfeature = torch.from_numpy(inputfeature)
    inputfeature = inputfeature.type(torch.FloatTensor)  # 不变
    data_variable = Variable(inputfeature).cuda()

    handle_feat = model2.channel_select.register_forward_hook(get_weight_hook1)
    output = model2(data_variable)
    handle_feat.remove()
    detection_input = feat_result1[0]
    detection_input = torch.from_numpy(detection_input)
    detection_input = detection_input.type(torch.FloatTensor)  # 不变
    data_variable2 = Variable(detection_input).cuda()

    #       ==============input reconstruction===========

    ssim_loss = tools.pytorch_ssim.ssim(im_data_variable, output)
    psnr_loss = psnr(im_data_variable, output)
    pic = to_img(output.cpu().data)
    pic_in = to_img(im_data_variable.cpu().data)
    save_image(pic_in, './output/imagein.png')
    save_image(pic, './output/imageout.png')
    print('SSIM {}  PSNR {}'.format(ssim_loss.data.item(), psnr_loss))

    #      ==============object detection================
    classes = ('aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car',
               'cat', 'chair', 'cow', 'diningtable', 'dog', 'horse',
               'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train',
               'tvmonitor')
    tic = time.time()
    yolo_output = model3(data_variable2)

    yolo_output = [item[0].data for item in yolo_output]

    detections = yolo_eval(yolo_output,
                           im_info,
                           conf_threshold=0.6,
                           nms_threshold=0.4)
    print(detections)
    if len(detections) == 0:
        plt.figure()
        plt.imshow(img)
        plt.show()
    else:
        toc = time.time()
        cost_time = toc - tic
        print('im detect, cost time {:4f}, FPS: {}'.format(
            toc - tic, int(1 / cost_time)))

        det_boxes = detections[:, :5].cpu().numpy()
        det_classes = detections[:, -1].long().cpu().numpy()

        im2show = draw_detection_boxes(img,
                                       det_boxes,
                                       det_classes,
                                       class_names=classes)
        plt.figure()
        plt.imshow(im2show)
        plt.gca().xaxis.set_major_locator(plt.NullLocator())
        plt.gca().yaxis.set_major_locator(plt.NullLocator())
        plt.subplots_adjust(top=1,
                            bottom=0,
                            right=1,
                            left=0,
                            hspace=0,
                            wspace=0)
        plt.margins(0, 0)
        plt.savefig('./output/detect_result.jpg', dpi=150)
        plt.show()
Пример #7
0
def demo():
    args = parse_args()
    print('call with args: {}'.format(args))

    # input images
    images_dir = 'images'
    images_names = ['trainval1.jpg', 'trainval2.jpg', 'test1.jpg', 'test2.jpg']

    classes = ("car", "bus", "truck", "svehicle", "pedestrian", "motorbike",
               "bicycle", "train", "signal", "signs")

    model = Yolov2(arch=args.arch)
    #weight_loader = WeightLoader()
    #weight_loader.load(model, 'yolo-voc.weights')
    #print('loaded')

    model_path = os.path.join(args.output_dir, args.model_name + '.pth')
    print('loading model from {}'.format(model_path))
    if torch.cuda.is_available():
        checkpoint = torch.load(model_path)
    else:
        checkpoint = torch.load(model_path, map_location='cpu')
    model.load_state_dict(checkpoint['model'])

    if args.use_cuda:
        model.cuda()

    model.eval()
    print('model loaded')

    for image_name in images_names:
        image_path = os.path.join(images_dir, image_name)
        img = Image.open(image_path)
        im_data, im_info = prepare_im_data(img)

        if args.use_cuda:
            im_data_variable = Variable(im_data).cuda()
        else:
            im_data_variable = Variable(im_data)

        tic = time.time()

        yolo_output = model(im_data_variable)
        yolo_output = [item[0].data for item in yolo_output]
        detections = yolo_eval(yolo_output,
                               im_info,
                               conf_threshold=0.2,
                               nms_threshold=0.4)
        ##print(detections)

        toc = time.time()
        cost_time = toc - tic
        print('im detect, cost time {:4f}, FPS: {}'.format(
            toc - tic, int(1 / cost_time)))

        det_boxes = detections[:, :5].cpu().numpy()
        det_classes = detections[:, -1].long().cpu().numpy()
        im2show = draw_detection_boxes(img,
                                       det_boxes,
                                       det_classes,
                                       class_names=classes)
        plt.figure()
        plt.imshow(im2show)
        #plt.show()

        save_image_path = os.path.join(images_dir, image_name + "_detect.jpg")
        print("save -> " + save_image_path)
        plt.savefig(save_image_path)