|
from PIL import Image |
|
import numpy as np |
|
import torch |
|
from torchvision import transforms, models |
|
from onnx import numpy_helper |
|
import os |
|
import onnxruntime as rt |
|
from matplotlib.colors import hsv_to_rgb |
|
import cv2 |
|
import gradio as gr |
|
|
|
import matplotlib.pyplot as plt |
|
import matplotlib.patches as patches |
|
|
|
import pycocotools.mask as mask_util |
|
|
|
def preprocess(image): |
|
|
|
ratio = 800.0 / min(image.size[0], image.size[1]) |
|
image = image.resize((int(ratio * image.size[0]), int(ratio * image.size[1])), Image.BILINEAR) |
|
|
|
|
|
image = np.array(image)[:, :, [2, 1, 0]].astype('float32') |
|
|
|
|
|
image = np.transpose(image, [2, 0, 1]) |
|
|
|
|
|
mean_vec = np.array([102.9801, 115.9465, 122.7717]) |
|
for i in range(image.shape[0]): |
|
image[i, :, :] = image[i, :, :] - mean_vec[i] |
|
|
|
|
|
import math |
|
padded_h = int(math.ceil(image.shape[1] / 32) * 32) |
|
padded_w = int(math.ceil(image.shape[2] / 32) * 32) |
|
|
|
padded_image = np.zeros((3, padded_h, padded_w), dtype=np.float32) |
|
padded_image[:, :image.shape[1], :image.shape[2]] = image |
|
image = padded_image |
|
|
|
return image |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
os.system("wget https://github.com/AK391/models/raw/main/vision/object_detection_segmentation/mask-rcnn/model/MaskRCNN-10.onnx") |
|
sess = rt.InferenceSession("MaskRCNN-10.onnx") |
|
|
|
outputs = sess.get_outputs() |
|
|
|
|
|
classes = [line.rstrip('\n') for line in open('coco_classes.txt')] |
|
|
|
|
|
def display_objdetect_image(image, boxes, labels, scores, masks, score_threshold=0.7): |
|
|
|
ratio = 800.0 / min(image.size[0], image.size[1]) |
|
boxes /= ratio |
|
|
|
_, ax = plt.subplots(1, figsize=(12,9)) |
|
|
|
image = np.array(image) |
|
|
|
for mask, box, label, score in zip(masks, boxes, labels, scores): |
|
|
|
if score <= score_threshold: |
|
continue |
|
|
|
|
|
mask = mask[0, :, :, None] |
|
int_box = [int(i) for i in box] |
|
mask = cv2.resize(mask, (int_box[2]-int_box[0]+1, int_box[3]-int_box[1]+1)) |
|
mask = mask > 0.5 |
|
im_mask = np.zeros((image.shape[0], image.shape[1]), dtype=np.uint8) |
|
x_0 = max(int_box[0], 0) |
|
x_1 = min(int_box[2] + 1, image.shape[1]) |
|
y_0 = max(int_box[1], 0) |
|
y_1 = min(int_box[3] + 1, image.shape[0]) |
|
mask_y_0 = max(y_0 - box[1], 0) |
|
mask_y_1 = mask_y_0 + y_1 - y_0 |
|
mask_x_0 = max(x_0 - box[0], 0) |
|
mask_x_1 = mask_x_0 + x_1 - x_0 |
|
im_mask[y_0:y_1, x_0:x_1] = mask[ |
|
mask_y_0 : mask_y_1, mask_x_0 : mask_x_1 |
|
] |
|
im_mask = im_mask[:, :, None] |
|
|
|
|
|
contours, hierarchy = cv2.findContours( |
|
im_mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE |
|
) |
|
|
|
image = cv2.drawContours(image, contours, -1, 25, 3) |
|
|
|
rect = patches.Rectangle((box[0], box[1]), box[2] - box[0], box[3] - box[1], linewidth=1, edgecolor='b', facecolor='none') |
|
ax.annotate(classes[label] + ':' + str(np.round(score, 2)), (box[0], box[1]), color='w', fontsize=12) |
|
ax.add_patch(rect) |
|
|
|
ax.imshow(image) |
|
plt.axis('off') |
|
plt.savefig('out.png', bbox_inches='tight') |
|
|
|
|
|
def inference(img): |
|
input_image = Image.open(img) |
|
orig_tensor = np.asarray(input_image) |
|
input_tensor = preprocess(input_image) |
|
|
|
output_names = list(map(lambda output: output.name, outputs)) |
|
input_name = sess.get_inputs()[0].name |
|
|
|
boxes, labels, scores, masks = sess.run(output_names, {input_name: input_tensor}) |
|
display_objdetect_image(input_image, boxes, labels, scores, masks) |
|
return 'out.png' |
|
|
|
title="Mask R-CNN" |
|
description="This model is a real-time neural network for object instance segmentation that detects 80 different classes." |
|
examples=[["examplemask-rcnn.jpeg"]] |
|
gr.Interface(inference,gr.inputs.Image(type="filepath"),gr.outputs.Image(type="file"),title=title,description=description,examples=examples).launch(enable_queue=True) |
