This repository was archived by the owner on Sep 16, 2021. It is now read-only.
forked from cjekal/detectron2_web_app
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathObjectDetector.py
More file actions
86 lines (64 loc) · 2.33 KB
/
ObjectDetector.py
File metadata and controls
86 lines (64 loc) · 2.33 KB
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
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
import cv2 as cv
import json
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
from detectron2.utils.visualizer import ColorMode
from detectron2 import model_zoo
from detectron2.data import MetadataCatalog, DatasetCatalog
from detectron2.modeling import build_model
import torch
import numpy as np
from PIL import Image
class Detector:
def __init__(self):
# set model and test set
self.model = 'mask_rcnn_R_101_FPN_3x.yaml'
# obtain detectron2's default config
self.cfg = get_cfg()
# load values from a file
# self.cfg.merge_from_file("test.yaml")
self.cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_101_FPN_3x.yaml"))
# set device to cpu
self.cfg.MODEL.DEVICE = "cpu"
# set number of classes to detect
self.cfg.MODEL.ROI_HEADS.NUM_CLASSES = 2
# get weights
self.cfg.MODEL.WEIGHTS = "./model.pth"
# self.cfg.MODEL.WEIGHTS = "model_final_f10217.pkl"
# set the testing threshold for this model
self.cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7
# build model from weights
# self.cfg.MODEL.WEIGHTS = self.convert_model_for_inference()
def get_dataset_catalog():
return []
DatasetCatalog.register("my_dataset", get_dataset_catalog)
MetadataCatalog.get("my_dataset").thing_classes = ["Edge", "Weld"]
# build model and convert for inference
def convert_model_for_inference(self):
# build model
model = build_model(self.cfg)
# save as checkpoint
torch.save(model.state_dict(), 'checkpoint.pth')
# return path to inference model
return 'checkpoint.pth'
# detectron model
# adapted from detectron2 colab notebook: https://colab.research.google.com/drive/16jcaJoc6bCFAQ96jDe2HwtXj7BMD_-m5
def inference(self, file):
predictor = DefaultPredictor(self.cfg)
im = cv.imread(file)
outputs = predictor(im)
# with open(self.curr_dir+'/data.txt', 'w') as fp:
# json.dump(outputs['instances'], fp)
# # json.dump(cfg.dump(), fp)
# get metadata
metadata = MetadataCatalog.get("my_dataset")
# visualise
v = Visualizer(im[:, :, ::-1], metadata=metadata, scale=1.0)
v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
# get image
img = v.get_image()
img = Image.fromarray(img)
# write to jpg
# cv.imwrite('img.jpg',v.get_image())
return img