clean up app setup

app.ipynb

@@ -9,7 +9,7 @@
     "import gradio as gr\n",
     "import numpy as np\n",
     "from PIL import Image\n",
-    "from transformers import MaskFormerFeatureExtractor, MaskFormerForInstanceSegmentation"
+    "from transformers import MaskFormerForInstanceSegmentation, MaskFormerImageProcessor"
    ]
   },
   {
@@ -18,15 +18,14 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# feature_extractor = MaskFormerFeatureExtractor.from_pretrained(\"facebook/maskformer-swin-tiny-coco\")\n",
-    "# model = MaskFormerForInstanceSegmentation.from_pretrained(\"facebook/maskformer-swin-tiny-coco\")\n",
-    "feature_extractor = MaskFormerFeatureExtractor.from_pretrained(\"facebook/maskformer-swin-large-coco\")\n",
-    "model = MaskFormerForInstanceSegmentation.from_pretrained(\"facebook/maskformer-swin-large-coco\")\n",
+    "model_id = f\"facebook/maskformer-swin-large-coco\"\n",
+    "\n",
+    "feature_extractor = MaskFormerImageProcessor.from_pretrained(model_id)\n",
+    "model = MaskFormerForInstanceSegmentation.from_pretrained(model_id)\n",
     "\n",
     "with Image.open(\"../color-filter-calculator/assets/Artshack_screen.jpg\") as img:\n",
     "    img_size = (img.height, img.width)\n",
-    "    inputs = feature_extractor(images=img, return_tensors=\"pt\")\n",
-    "    "
+    "    inputs = feature_extractor(images=img, return_tensors=\"pt\")"
    ]
   },
   {

app.py

@@ -2,25 +2,26 @@ import glob
 import gradio as gr
 import numpy as np
 from PIL import Image
-from transformers import MaskFormerFeatureExtractor, MaskFormerForInstanceSegmentation
+from transformers import MaskFormerForInstanceSegmentation, MaskFormerImageProcessor
 
 
-# feature_extractor = MaskFormerFeatureExtractor.from_pretrained("facebook/maskformer-swin-tiny-coco")
-# model = MaskFormerForInstanceSegmentation.from_pretrained("facebook/maskformer-swin-tiny-coco")
-feature_extractor = MaskFormerFeatureExtractor.from_pretrained("facebook/maskformer-swin-large-coco")
-model = MaskFormerForInstanceSegmentation.from_pretrained("facebook/maskformer-swin-large-coco")
-
 example_images = sorted(glob.glob('examples/map*.jpg'))
 
-def visualize_instance_seg_mask(img_in, mask, id2label):
+model_id = f"facebook/maskformer-swin-large-coco"
+vegetation_labels = ["tree-merged", "grass-merged"]
+
+preprocessor = MaskFormerImageProcessor.from_pretrained(model_id)
+model = MaskFormerForInstanceSegmentation.from_pretrained(model_id)
+
+
+def visualize_instance_seg_mask(img_in, mask, id2label, included_labels):
     img_out = np.zeros((mask.shape[0], mask.shape[1], 3))
     image_total_pixels = mask.shape[0] * mask.shape[1]
     label_ids = np.unique(mask)
-    vegetation_labels = ["tree-merged", "grass-merged"]
 
     def get_color(id):
         id_color = (np.random.randint(0, 2), np.random.randint(0, 4), np.random.randint(0, 256))
-        if id2label[id] in vegetation_labels:
+        if id2label[id] in included_labels:
             id_color = (0, 140, 0)
         return id_color
 
@@ -34,13 +35,13 @@ def visualize_instance_seg_mask(img_in, mask, id2label):
 
     image_res = (0.5 * img_in + 0.5 * img_out).astype(np.uint8)
 
-    vegetation_count = sum([id2count[id] for id in label_ids if id2label[id] in vegetation_labels])
+    vegetation_count = sum([id2count[id] for id in label_ids if id2label[id] in included_labels])
 
     dataframe_vegetation_items = [[
         f"{id2label[id]}",
         f"{(100 * id2count[id] / image_total_pixels):.2f} %",
         f"{np.sqrt(id2count[id] / image_total_pixels):.2f} m"
-        ] for id in label_ids if id2label[id] in vegetation_labels]
+        ] for id in label_ids if id2label[id] in included_labels]
     dataframe_all_items = [[
         f"{id2label[id]}",
         f"{(100 * id2count[id] / image_total_pixels):.2f} %",
@@ -65,10 +66,10 @@ def visualize_instance_seg_mask(img_in, mask, id2label):
 def query_image(image_path):
     img = np.array(Image.open(image_path))
     img_size = (img.shape[0], img.shape[1])
-    inputs = feature_extractor(images=img, return_tensors="pt")
+    inputs = preprocessor(images=img, return_tensors="pt")
     outputs = model(**inputs)
-    results = feature_extractor.post_process_semantic_segmentation(outputs=outputs, target_sizes=[img_size])[0]
-    mask_img, dataframe = visualize_instance_seg_mask(img, results.numpy(), model.config.id2label)
+    results = preprocessor.post_process_semantic_segmentation(outputs=outputs, target_sizes=[img_size])[0]
+    mask_img, dataframe = visualize_instance_seg_mask(img, results.numpy(), model.config.id2label, vegetation_labels)
     return mask_img, dataframe
 
 
@@ -83,8 +84,7 @@ demo = gr.Interface(
     allow_flagging="never",
     analytics_enabled=None,
     examples=example_images,
-    # cache_examples=True
+    cache_examples=True
 )
 
-demo.queue(concurrency_count=8, max_size=8)
 demo.launch(show_api=False)