예제 #1
0
 def detect_object(self, img, conf_thres, nms_thres):
     img = transforms.ToTensor()(img)
     img, _ = pad_to_square(img, 0)
     # Resize
     img = resize(img, self.img_size)
     img = img.to(self.device)
     img = img.unsqueeze(0)
     detection = self.model_yolov3(img)
     # From [x,y,w,h] to [x_l,y_l,x_r,y_r]
     detection = non_max_suppression(detection, conf_thres, nms_thres)
     if detection[0] is not None:
         detection = rescale_boxes(detection[0], self.img_size, (600, 600))
     return detection
예제 #2
0
        im_dim_list = im_dim_list.cuda()

    # 5 resim için 10sn sürüyor
    start_outputs_loop = time.time()

    # bu for içinde tek tek resim işlemeye çalış
    for i, batch in enumerate(im_batches):
        # load the image

        if CUDA:
            batch = batch.cuda()
        with torch.no_grad():
            prediction = model(Variable(batch))  # bu işlem uzun sürüyor

        prediction = non_max_suppression(prediction,
                                         confidence,
                                         num_classes,
                                         nms_conf=nms_thesh)
        try:
            prediction[:,
                       0] += i * batch_size  # transform the atribute from index in batch to index in imlist
        except TypeError as e:
            # yazdir(img[0])
            outputs_names = pd.Series(imlist).apply(
                lambda x: "{}/{}".format(args.outputs,
                                         x.split("/")[-1]))
            list(map(cv2.imwrite, outputs_names, loaded_ims))
            os.remove(os.path.join(images, img[0]))
            print("\nTYPE ERROR\n")
            continue

        if not write:  # If we have't initialised output
while True:
    ret, frame = cap.read()

    if ret:
        img = prep_image(frame, inp_dim)
        im_dim = frame.shape[1], frame.shape[0]
        im_dim = torch.FloatTensor(im_dim).repeat(1, 2)

        if CUDA:
            im_dim = im_dim.cuda()
            img = img.cuda()

        with torch.no_grad():
            output = model(Variable(img, volatile=True))
        output = non_max_suppression(output,
                                     confidence,
                                     num_classes,
                                     nms_conf=nms_thesh)

        if type(output) == int:
            frames += 1
            print("FPS of the video is {:5.4f}".format(frames /
                                                       (time.time() - start)))
            cv2.imshow("frame", frame)
            key = cv2.waitKey(1)
            if key & 0xFF == ord('q'):
                break
            continue

        obj_counter = output[:, -1]
        obj_counter = Counter(obj_counter.numpy().astype(int).tolist())