Update app.py

app.py

@@ -48,8 +48,8 @@ def preprocess(image):
 # based on the build flags) when instantiating InferenceSession.
 # For example, if NVIDIA GPU is available and ORT Python package is built with CUDA, then call API as following:
 # onnxruntime.InferenceSession(path/to/model, providers=['CUDAExecutionProvider'])
-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")
+os.system("wget https://github.com/AK391/models/raw/main/vision/object_detection_segmentation/faster-rcnn/model/FasterRCNN-10.onnx")
+sess = rt.InferenceSession("FasterRCNN-10.onnx")
 
 outputs = sess.get_outputs()
 
@@ -57,51 +57,22 @@ 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):
+def display_objdetect_image(image, boxes, labels, scores, score_threshold=0.7):
     # Resize boxes
     ratio = 800.0 / min(image.size[0], image.size[1])
     boxes /= ratio
 
     _, ax = plt.subplots(1, figsize=(12,9))
-
     image = np.array(image)
+    ax.imshow(image)
 
-    for mask, box, label, score in zip(masks, boxes, labels, scores):
-        # Showing boxes with score > 0.7
-        if score <= score_threshold:
-            continue
-
-        # Finding contour based on mask
-        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]
-
-        # OpenCV version 4.x
-        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)
+    # Showing boxes with score > 0.7
+    for box, label, score in zip(boxes, labels, scores):
+        if score > score_threshold:
+            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')
 
@@ -114,11 +85,12 @@ def inference(img):
   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)
+  boxes, labels, scores = sess.run(output_names, {input_name: input_tensor})
+  display_objdetect_image(input_image, boxes, labels, scores)
+
   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."
+title="Faster R-CNN"
+description="This model is a real-time neural network for object detection 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)