示例#1
0
class PyTorchModel:

    #create model by loading it in from Google Drive path
    def __init__(self, f):
        trainable_backbone_layers = 5
        pretrained = True
        backbone = resnet_fpn_backbone(
            'resnet50', True, trainable_layers=trainable_backbone_layers)
        self.model = FasterRCNN(backbone,
                                num_classes=10,
                                max_size=3840,
                                min_size=2160,
                                rpn_pre_nms_top_n_train=2000,
                                rpn_pre_nms_top_n_test=2000,
                                rpn_post_nms_top_n_train=2000,
                                rpn_post_nms_top_n_test=2000,
                                box_detections_per_img=100,
                                rpn_nms_thresh=0.01,
                                box_nms_thresh=0.01)

        #num_classes = 10
        #self.model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
        #in_features = self.model.roi_heads.box_predictor.cls_score.in_features
        #self.model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

        device = torch.device(
            'cuda') if torch.cuda.is_available() else torch.device('cpu')
        self.model.to(device)
        if (isinstance(f, str)):  #local file
            print("Loading model from local file at {}".format(f))
            self.model.load_state_dict(torch.load(f, map_location=device))
        elif (isinstance(f, io.BytesIO)):  #stream
            print("Loading model from stream")
            pass

    def predict(self, image) -> List[Label]:
        frame = torchvision.transforms.ToTensor()(image)
        frame = frame[None, :, :]
        self.model.eval()
        prediction = self.model(frame)
        print(prediction)
        boxes = prediction[0]["boxes"]
        labels = prediction[0]["labels"]
        scores = prediction[0]["scores"]
        ret = list()
        for i in range(0, len(boxes)):
            score: float = float(scores[i].item())
            xmin: int = int(boxes[i][0].item())
            ymin: int = int(boxes[i][1].item())
            xmax: int = int(boxes[i][2].item())
            ymax: int = int(boxes[i][3].item())
            group: str = classes[str(labels[i].item())]["category"]
            color: str = classes[str(labels[i].item())]["color"]
            ret.append(Label(i, group, xmin, xmax, ymin, ymax, color, score))
        return ret
示例#2
0
import torch
import torchvision
from class_labels import GROCERY_LIST_V0
from PIL import Image
from torch.autograd import Variable
from torchvision import transforms
from torchvision.models.detection.faster_rcnn import FasterRCNN
from torchvision.models.detection.backbone_utils import resnet_fpn_backbone

model = FasterRCNN(resnet_fpn_backbone("resnet50", False), num_classes=64)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

model_path = "synthdet_faster_rcnn.pth"
model_save = torch.load(model_path, map_location=device) # no CUDA on macOS   
model.load_state_dict(model_save["model"])
model.to(device)
model.eval()

# preprocess on test image
image_path = sys.argv[1]
image = Image.open(image_path)
image_to_tensor = transforms.Compose([
    transforms.ToTensor()
])
tensor = image_to_tensor(image)

# inference
threshold = 0.5

# Start timing
time_inference_start = time.time()