如何解决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 个对象,没问题。
版权声明:本文内容由互联网用户自发贡献,该文观点与技术仅代表作者本人。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如发现本站有涉嫌侵权/违法违规的内容, 请发送邮件至 dio@foxmail.com 举报,一经查实,本站将立刻删除。