Update app.py

app.py

@@ -7,14 +7,10 @@ import gradio as gr
 import torch, torchvision
 print(torch.__version__, torch.cuda.is_available())
 assert torch.__version__.startswith("1.9")   # please manually install torch 1.9 if Colab changes its default version
-# Some basic setup:
-# Setup detectron2 logger
 import detectron2
 from detectron2.utils.logger import setup_logger
-# import some common libraries
 import numpy as np
 import os, json, random
-# import some common detectron2 utilities
 from detectron2 import model_zoo
 from detectron2.engine import DefaultPredictor
 from detectron2.config import get_cfg
@@ -25,35 +21,39 @@ from matplotlib import pyplot as plt
 
 cfg = get_cfg()
 cfg.MODEL.DEVICE='cpu'
-# add project-specific config (e.g., TensorMask) here if you're not running a model in detectron2's core library
 cfg.INPUT.MASK_FORMAT='bitmask'
 cfg.MODEL.ROI_HEADS.NUM_CLASSES = 3 
 cfg.TEST.DETECTIONS_PER_IMAGE = 1000
 cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
 cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5  # set threshold for this model
-# Find a model from detectron2's model zoo. You can use the https://dl.fbaipublicfiles... url as well
 cfg.MODEL.WEIGHTS = "model_final.pth"
 
 predictor = DefaultPredictor(cfg)
 
 
 def inference(img):
+    class_names = ['astro', 'cort', 'sh-sy5y']
     im = np.asarray(Image.open(img).convert('RGB'))
     outputs = predictor(im)
-
+    pred_classes = outputs['instances'].pred_classes.cpu().numpy().tolist()
     take = outputs['instances'].scores >= 0.5 #Threshold
     pred_masks = outputs['instances'].pred_masks[take].cpu().numpy()
+    pred_class = max(set(pred_classes), key=pred_classes.count)
 
     mask = np.stack(pred_masks)
     mask = np.any(mask == 1, axis=0)
 
     p = plt.imshow(im,cmap='gray')
-    p1 = plt.imshow(mask, alpha=0.4)
+    p = plt.imshow(mask, alpha=0.4)
+    p = plt.xticks(fontsize=8)
+    p = plt.yticks(fontsize=8)
+    p = plt.title("cell type: " + class_names[pred_class])
 
     return plt
 
 
 
+
 title = "Sartorius Cell Instance Segmentation"
 description = "Sartorius Cell Instance Segmentation Demo: Current Kaggle competition - kaggle.com/c/sartorius-cell-instance-segmentation"
 article = "<p style='text-align: center'><a href='https://ai.facebook.com/blog/-detectron2-a-pytorch-based-modular-object-detection-library-/' target='_blank'>Detectron2: A PyTorch-based modular object detection library</a> | <a href='https://github.com/facebookresearch/detectron2' target='_blank'>Github Repo</a></p>"