owlting / ai_basketball_games_video_editor Goto Github PK

  def detect(self, model, img, image_size):
        model.eval()
        
        IN_IMAGE_H, IN_IMAGE_W = image_size
        
        sized = cv2.resize(img, (IN_IMAGE_W, IN_IMAGE_H))
        sized = cv2.cvtColor(sized, cv2.COLOR_BGR2RGB)        
        
        t0 = time.time()

        if type(sized) == np.ndarray and len(sized.shape) == 3:  # cv2 image
            sized = torch.from_numpy(sized.transpose(2, 0, 1)).float().div(255.0).unsqueeze(0)
        elif type(sized) == np.ndarray and len(sized.shape) == 4:
            sized = torch.from_numpy(sized.transpose(0, 3, 1, 2)).float().div(255.0)
        else:
            print("unknow image type")
            exit(-1)
        
        use_cuda = 1
        if use_cuda:
            sized = sized.cuda()
        sized = torch.autograd.Variable(sized)

        t1 = time.time()
        
        with torch.no_grad():
            output = model(sized)

        t2 = time.time()

#         print('-----------------------------------')
#         print('           Preprocess : %f' % (t1 - t0))
#         print('      Model Inference : %f' % (t2 - t1))
#         print('-----------------------------------')

        boxes = post_processing(img, 0.4, 0.6, output)

        return boxes

        m = Darknet(cfg_path)
#         m.print_network()
        m.load_weights(weight_path)
        print('Loading weights from %s... Done!' % (weight_path))

        if use_cuda:
            m.cuda()

        self.num_classes = m.num_classes
        self.class_names = load_class_names(namesfile_path)
        self.engine = m
        self.image_size = inference_size

  with torch.no_grad():
            output = model(sized)