|
|
楼主 |
发表于 2024-8-11 14:43:53
|
显示全部楼层
import cv2
import numpy as np
import os
from samples import coco
from mrcnn import utils
from mrcnn import model as modellib
from datetime import datetime
# Load the pre-trained model data
ROOT_DIR = os.getcwd()
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)
# Change the config infermation
class InferenceConfig(coco.CocoConfig):
GPU_COUNT = 1
IMAGES_PER_GPU = 1
config = InferenceConfig()
# COCO dataset object names
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 drier', 'toothbrush'
]
# ==========================================================
# This function is used to change the colorful background information
def apply_mask_black(image, mask):
image[:, :, 0] = np.where(
mask == 0,
0,
255
)
image[:, :, 1] = np.where(
mask == 0,
0,
255
)
image[:, :, 2] = np.where(
mask == 0,
0,
255
)
return image
# This function is used to show the object detection result in original image.
def display_instances(image, boxes, masks, ids, names, scores):
# max_area will save the largest object for all the detection results
max_area = 0
# n_instances saves the amount of all objects
n_instances = boxes.shape[0]
if not n_instances:
print('NO INSTANCES TO DISPLAY')
else:
assert boxes.shape[0] == masks.shape[-1] == ids.shape[0]
mask = np.array([0])
# mask = False
for i in range(n_instances):
if not np.any(boxes[i]):
continue
# compute the square of each object
y1, x1, y2, x2 = boxes[i]
square = (y2 - y1) * (x2 - x1)
# use label to select person object from all the 80 classes in COCO dataset
label = names[ids[i]]
if label == 'person':
# save the largest object in the image as main character
# other people will be regarded as background
if square > max_area:
max_area = square
mask = masks[:, :, i]
else:
continue
else:
continue
# apply mask for the image
# by mistake you put apply_mask inside for loop or you can write continue in if also
image = apply_mask_black(image, mask)
return image if not mask.sum() ==0 else mask
# return image
def img_black(image_path):
images_list = os.listdir(image_path)
out_path = /home/nvidia/xxx/
for name in sorted(images_list):
original_image = os.path.join(image_path, name)
image = cv2.imread(original_image)
# gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
results = model.detect([image], verbose=0)
r = results[0]
frame = display_instances(image, r['rois'], r['masks'], r['class_ids'], class_names, r['scores'])
if frame.sum() == 0:
continue
else:
# 生成每个图片对应的同名文件夹，里面放有该图片
output_dir = out_path+name+'/'
os.mkdir(output_dir)
cv2.imwrite(output_dir+name,frame)
if __name__ == '__main__':
time = datetime.now()
# images path
image_path = '/home/pcm/Mask-GaitSet/IMG/images'
img_black(image_path)
print('Evaluation complete. Cost:', datetime.now() - time)
这个代码要怎么改
@FishC |
|
[复制链接]
( 粤ICP备18085999号-1 | 粤公网安备 44051102000585号)
狗仔卡
显身卡
发表于 2024-8-11 14:23:31