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
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()