File size: 1,815 Bytes
e803eba | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 |
# Simple inference script for wildlife ensemble detector
import torch
import cv2
import numpy as np
from detectron2 import model_zoo
from detectron2.config import get_cfg
from detectron2.engine import DefaultPredictor
def load_ensemble_models(inception_path, resnet_path, class_names):
"""Load both ensemble models for inference"""
# You'll need to copy the model_architecture.py content here
# and register the backbones before loading
# Setup configs (simplified)
cfg_inception = get_cfg()
cfg_inception.merge_from_file(model_zoo.get_config_file("COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml"))
cfg_inception.MODEL.BACKBONE.NAME = "InceptionBackboneWrapper"
cfg_inception.MODEL.ROI_HEADS.NUM_CLASSES = len(class_names)
cfg_inception.MODEL.WEIGHTS = inception_path
cfg_inception.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
cfg_resnet = get_cfg()
cfg_resnet.merge_from_file(model_zoo.get_config_file("COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml"))
cfg_resnet.MODEL.BACKBONE.NAME = "ResNetBackboneWrapper"
cfg_resnet.MODEL.ROI_HEADS.NUM_CLASSES = len(class_names)
cfg_resnet.MODEL.WEIGHTS = resnet_path
cfg_resnet.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
predictor_inception = DefaultPredictor(cfg_inception)
predictor_resnet = DefaultPredictor(cfg_resnet)
return predictor_inception, predictor_resnet
def predict_ensemble(image_path, predictor_inception, predictor_resnet, class_names):
"""Run ensemble inference on an image"""
img = cv2.imread(image_path)
# Get predictions from both models
outputs_inc = predictor_inception(img)
outputs_res = predictor_resnet(img)
# Combine predictions (simplified)
# Add your ensemble logic here
return outputs_inc, outputs_res
|