Как я могу обнаружить небо с помощью модели детектора Facebook? - PullRequest
0 голосов
/ 02 апреля 2020

Я намеревался использовать модель детектрона2 Facebook для пользовательских целей. Мне нужно обнаружить небо по изображению и получить только вывод «да» или «нет», основываясь на том, находится ли небо на изображении или нет. Но сейчас я не могу понять, как получить этот конкретный c вывод. Я могу обнаружить объекты на изображении, включая небо, но я не знаю, как изменить вывод таким образом, чтобы он только сообщал, обнаружил ли он класс "sky" или нет.

Это весь мой код:

from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog
import numpy as np
import cv2
import os
im = cv2.imread("ToTest.jpg")

# Create config
cfg = get_cfg()
cfg.merge_from_file("https://github.com/facebookresearch/detectron2/blob/master/configs/COCO-PanopticSegmentation/panoptic_fpn_R_50_1x.yaml")
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5  # set threshold for this model
cfg.MODEL.DEVICE = "cpu"
cfg.MODEL.WEIGHTS = "https://dl.fbaipublicfiles.com/detectron2/COCO-PanopticSegmentation/panoptic_fpn_R_50_1x/139514544/model_final_dbfeb4.pkl"

# Create predictor
predictor = DefaultPredictor(cfg)


# Make prediction
outputs = predictor(im)
#print(outputs)

print(outputs["instances"].pred_classes)
print(outputs["instances"].pred_boxes)

# We can use `Visualizer` to draw the predictions on the image.
panoptic_seg, segments_info = predictor(im)["panoptic_seg"]
v = Visualizer(im[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)
v = v.draw_panoptic_seg_predictions(panoptic_seg.to("cpu"), segments_info)

directory = r"directory"
filename = 'savedOutput.jpg'
cv2.imwrite(filename, v.get_image()[:,:,::-1])

Наконец, я нашел только этот код, который что-то говорил о редактировании выходных данных или что-то еще, но я вообще не мог этого понять. Может быть, кто-нибудь может помочь мне с этим.

images = ImageList.from_tensors(...)  # preprocessed input tensor
model = build_model(cfg)
features = model.backbone(images.tensor)
proposals, _ = model.proposal_generator(images, features)
instances = model.roi_heads._forward_box(features, proposals)
mask_features = [features[f] for f in model.roi_heads.in_features]
mask_features = model.roi_heads.mask_pooler(mask_features, [x.pred_boxes for x in instances])
...