Hi @Dwight_Foster ,
thank you for your reply. But how can specify that I am only interested in class 1 (person) & class 18 (dog)? I do not want to train again the fasterrcnn_resnet50_fpn model. I want that the model only detects my wished classes and ingnores other 88 classes.
I tried to set the num_classes to 3 but get an error.
My code:
coco_names = ['__background__', 'person', 'dog']

model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True, num_classes=len(coco_names))


device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.eval().to(device)

path='Downloads/input2/test.jpg'
image = Image.open(path)

boxes, classes, labels = detect_utils.predict(image, model, device, 0.8)
image =detect_utils.draw_boxes(boxes, classes, labels, image)
#cv2.imshow('Image', image)
save_name = f"{path.split('/')[-1].split('.')[0]}"
cv2.imwrite(f"outputs/{save_name}.jpg", image)

the error that I get:
RuntimeError: Error(s) in loading state_dict for FasterRCNN:
	size mismatch for roi_heads.box_predictor.cls_score.weight: copying a param with shape torch.Size([91, 1024]) from checkpoint, the shape in current model is torch.Size([3, 1024]).
	size mismatch for roi_heads.box_predictor.cls_score.bias: copying a param with shape torch.Size([91]) from checkpoint, the shape in current model is torch.Size([3]).
	size mismatch for roi_heads.box_predictor.bbox_pred.weight: copying a param with shape torch.Size([364, 1024]) from checkpoint, the shape in current model is torch.Size([12, 1024]).
	size mismatch for roi_heads.box_predictor.bbox_pred.bias: copying a param with shape torch.Size([364]) from checkpoint, the shape in current model is torch.Size([12]).