Pytorch/torchvision - 如何增加可检测对象的限制

如何解决Pytorch/torchvision - 如何增加可检测对象的限制

我是 Pytorch 的新手，到目前为止，它令人难以置信。我用它来计算图像中的药丸数量。我发现在我的大多数图像中，它检测到的最大对象数是 100。对于下面的图片，它达到了 100 的最大计数，置信度约为 0.6。之后它不再增加，甚至降低到 0.1 置信度。我无法在文档或在线任何其他地方找到任何内容。我正在使用 fastrcnn_resnet50_fpn 模型。下面是加载训练模型和评估图像的代码。任何能够计算所有对象的技巧或什至不同的包都会非常有用。

## Loading the trained module
loaded_model = get_model(num_classes = 2)
loaded_model.load_state_dict(torch.load('Pillcount/model'))

os.chdir('../pytorchobjdet/vision')

class CountDataset(torch.utils.data.Dataset):
    def __init__(self,root,data_file,transforms=None):
        self.root = root
        self.transforms = transforms
        self.imgs = sorted(os.listdir(os.path.join(root,"count")))
        self.path_to_data_file = data_file
    def __getitem__(self,idx):
        # load images and bounding boxes
        img_path = os.path.join(self.root,"count",self.imgs[idx])
        img = Image.open(img_path).convert("RGB")
        box_list = parse_one_annot(self.path_to_data_file,self.imgs[idx])
        boxes = torch.as_tensor(box_list,dtype=torch.float32)
        num_objs = len(box_list)
        # there is only one class
        labels = torch.ones((num_objs,),dtype=torch.int64)
        image_id = torch.tensor([idx])
        area = (boxes[:,3] - boxes[:,1]) * (boxes[:,2] - boxes[:,0])
        # suppose all instances are not crowd
        iscrowd = torch.zeros((num_objs,dtype=torch.int64)
        target = {}
        target["boxes"] = boxes
        target["labels"] = labels
        target["image_id"] = image_id
        target["area"] = area
        target["iscrowd"] = iscrowd
        if self.transforms is not None:
                img,target = self.transforms(img,target)
        return img,target
    def __len__(self):
            return len(self.imgs)


dataset_count = CountDataset(root='../../Pill_Object_Detection',data_file = "../../Pill_Object_Detection/count_labels.csv",transforms = get_transform(train=False))

idx = 1
img,_ = dataset_count[idx]
#put the model in evaluation mode
loaded_model.eval()
with torch.no_grad():
   prediction = loaded_model([img])
image = Image.fromarray(img.mul(255).permute(1,2,0).byte().numpy())
draw = ImageDraw.Draw(image)
# draw groundtruth
count = 0
for element in range(len(prediction[0]["boxes"])):
   boxes = prediction[0]["boxes"][element].cpu().numpy()
   score = np.round(prediction[0]["scores"][element].cpu().numpy(),decimals= 4)
   if score > 0.6:
      draw.rectangle([(boxes[0],boxes[1]),(boxes[2],boxes[3])],outline ="red",width =3)
      draw.text((boxes[0],text = str(score))
      count +=1
print(f'count = {count}')
image

解决方法

上述评论中的建议非常有帮助。我使用了 YOLO5vs 模型，它完成了令人难以置信的工作。这个 tutorial 有一个超级简单的设置，您可以将带注释的图像上传到 roboflow，然后它为几乎所有当前的对象检测器设置了一些 google colab 教程。这是结果。我只需要提供更好质量的训练数据，但对于我提供的几张图片来说，它做得非常好。它可以在同一图像中计算超过 150 个对象，没问题。