from
google.colab
import
drive
drive.mount(
'/content/drive'
)
import
os, sys
import
random
import
math
import
numpy as np
import
skimage.io
import
matplotlib
import
matplotlib.pyplot as plt
import
cv2
from
matplotlib.patches
import
Polygon
os.chdir("
/
content
/
drive
/
My Drive
/
Colab Notebooks
/
MRCNN_pure")
sys.path.append("
/
content
/
drive
/
My Drive
/
Colab Notebooks
/
MRCNN_pure")
VIDEO_STREAM
=
"
/
content
/
drive
/
My Drive
/
Colab Notebooks
/
Millery.avi"
VIDEO_STREAM_OUT
=
"
/
content
/
drive
/
My Drive
/
Colab Notebooks
/
Result.avi"
ROOT_DIR
=
os.path.abspath(".")
sys.path.append(ROOT_DIR)
from
mrcnn
import
utils
import
mrcnn.model as modellib
from
mrcnn
import
visualize
sys.path.append(os.path.join(ROOT_DIR, "samples
/
coco
/
"))
import
coco
MODEL_DIR
=
os.path.join(ROOT_DIR, "logs")
COCO_MODEL_PATH
=
os.path.join(ROOT_DIR, "mask_rcnn_coco.h5")
if
not
os.path.exists(COCO_MODEL_PATH):
utils.download_trained_weights(COCO_MODEL_PATH)
IMAGE_DIR
=
os.path.join(ROOT_DIR, "images")
class
InferenceConfig(coco.CocoConfig):
GPU_COUNT
=
1
IMAGES_PER_GPU
=
1
def
display_instances(image, boxes, masks, ids, names, scores):
"""
take the image
and
results
and
apply
the mask, box,
and
Label
"""
n_instances
=
boxes.shape[
0
]
colors
=
visualize.random_colors(n_instances)
if
not
n_instances:
print
(
'NO INSTANCES TO DISPLAY'
)
else
:
assert
boxes.shape[
0
]
=
=
masks.shape[
-
1
]
=
=
ids.shape[
0
]
for
i, color
in
enumerate
(colors):
if
not
np.
any
(boxes[i]):
continue
y1, x1, y2, x2
=
boxes[i]
label
=
names[ids[i]]
score
=
scores[i]
if
scores
is
not
None
else
None
caption
=
'{} {:.2f}'
.
format
(label, score)
if
score
else
label
mask
=
masks[:, :, i]
image
=
visualize.apply_mask(image, mask, color)
image
=
cv2.rectangle(image, (x1, y1), (x2, y2), color,
2
)
image
=
cv2.putText(
image, caption, (x1, y1), cv2.FONT_HERSHEY_COMPLEX,
0.7
, color,
2
)
return
image
config
=
InferenceConfig()
config.display()
model
=
modellib.MaskRCNN(mode
=
"inference", model_dir
=
MODEL_DIR, config
=
config)
model.load_weights(COCO_MODEL_PATH, by_name
=
True
)
class_names
=
[
'BG'
,
'person'
,
'bicycle'
,
'car'
,
'motorcycle'
,
'airplane'
,
'bus'
,
'train'
,
'truck'
,
'boat'
,
'traffic light'
,
'fire hydrant'
,
'stop sign'
,
'parking meter'
,
'bench'
,
'bird'
,
'cat'
,
'dog'
,
'horse'
,
'sheep'
,
'cow'
,
'elephant'
,
'bear'
,
'zebra'
,
'giraffe'
,
'backpack'
,
'umbrella'
,
'handbag'
,
'tie'
,
'suitcase'
,
'frisbee'
,
'skis'
,
'snowboard'
,
'sports ball'
,
'kite'
,
'baseball bat'
,
'baseball glove'
,
'skateboard'
,
'surfboard'
,
'tennis racket'
,
'bottle'
,
'wine glass'
,
'cup'
,
'fork'
,
'knife'
,
'spoon'
,
'bowl'
,
'banana'
,
'apple'
,
'sandwich'
,
'orange'
,
'broccoli'
,
'carrot'
,
'hot dog'
,
'pizza'
,
'donut'
,
'cake'
,
'chair'
,
'couch'
,
'potted plant'
,
'bed'
,
'dining table'
,
'toilet'
,
'tv'
,
'laptop'
,
'mouse'
,
'remote'
,
'keyboard'
,
'cell phone'
,
'microwave'
,
'oven'
,
'toaster'
,
'sink'
,
'refrigerator'
,
'book'
,
'clock'
,
'vase'
,
'scissors'
,
'teddy bear'
,
'hair dryer'
,
'toothbrush'
]
vs
=
cv2.VideoCapture(VIDEO_STREAM)
writer
=
None
vs.
set
(cv2.CAP_PROP_POS_FRAMES,
1000
);
i
=
0
while
i <
20000
:
(grabbed, frame)
=
vs.read()
i
+
=
1
if
not
grabbed:
print
("Not grabbed.")
break
;
results
=
model.detect([frame], verbose
=
1
)
r
=
results[
0
]
masked_frame
=
display_instances(frame, r[
'rois'
], r[
'masks'
], r[
'class_ids'
],
class_names, r[
'scores'
])
if
writer
is
None
:
fourcc
=
cv2.VideoWriter_fourcc(
*
"XVID")
writer
=
cv2.VideoWriter(VIDEO_STREAM_OUT, fourcc,
30
,
(masked_frame.shape[
1
], masked_frame.shape[
0
]),
True
)
writer.write(masked_frame)
print
("[INFO] cleaning up...")
writer.release()