Update app.py

app.py

@@ -9,6 +9,10 @@ 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):
     # Resize
@@ -51,61 +55,67 @@ outputs = sess.get_outputs()
 
 
 classes = [line.rstrip('\n') for line in open('coco_classes.txt')]
-num_classes = len(classes)
-
-def get_palette():
-    # prepare and return palette
-    palette = [0] * num_classes * 3
-    
-    for hue in range(num_classes):
-        if hue == 0: # Background color
-            colors = (0, 0, 0)
-        else:
-            colors = hsv_to_rgb((hue / num_classes, 0.75, 0.75))
-            
-        for i in range(3):
-            palette[hue * 3 + i] = int(colors[i] * 255)
-            
-    return palette
-
-def colorize(labels):
-    # generate colorized image from output labels and color palette
-    result_img = Image.fromarray(labels).convert('P', colors=num_classes)
-    result_img.putpalette(get_palette())
-    return np.array(result_img.convert('RGB'))
-
-def visualize_output(image, output):
-    assert(image.shape[0] == output.shape[1] and \
-           image.shape[1] == output.shape[2]) # Same height and width
-    assert(output.shape[0] == num_classes)
-    
-    # get classification labels
-    raw_labels = np.argmax(output, axis=0).astype(np.uint8)
-
-    # comput confidence score
-    confidence = float(np.max(output, axis=0).mean())
-
-    # generate segmented image
-    result_img = colorize(raw_labels)
-    
-    # generate blended image
-    blended_img = cv2.addWeighted(image[:, :, ::-1], 0.5, result_img, 0.5, 0)
-    
-    result_img = Image.fromarray(result_img)
-    blended_img = Image.fromarray(blended_img)
-
-    return confidence, result_img, blended_img, raw_labels
+
+
+def display_objdetect_image(image, boxes, labels, scores, masks, 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)
+
+    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)
+
+    ax.imshow(image)
+    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
-  detections = sess.run(output_names, {input_name: input_tensor})
-  output, aux = detections
-  conf, result_img, blended_img, _ = visualize_output(orig_tensor, output[0])
-  return blended_img
+  
+  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."