gatesla commited on
Commit
3f88c96
1 Parent(s): a2d536f

Updated APP to see if I can get it to work on the website.

Browse files
Files changed (2) hide show
  1. app.py +18 -1
  2. understand.py +1 -1
app.py CHANGED
@@ -5,8 +5,10 @@ import requests, validators
5
  import torch
6
  import pathlib
7
  from PIL import Image
 
 
8
 
9
- from transformers import DetrFeatureExtractor, DetrForSegmentation, MaskFormerImageProcessor, MaskFormerForInstanceSegmentation
10
  from transformers.models.detr.feature_extraction_detr import rgb_to_id
11
 
12
 
@@ -60,6 +62,12 @@ def visualize_prediction(pil_img, output_dict, threshold=0.7, id2label=None):
60
  plt.axis("off")
61
  return fig2img(plt.gcf())
62
 
 
 
 
 
 
 
63
  def segment_images(model_name,url_input,image_input,threshold):
64
  ####
65
  # Get Image Object
@@ -81,6 +89,15 @@ def segment_images(model_name,url_input,image_input,threshold):
81
 
82
  outputs = model(**inputs)
83
  results = processor.post_process_panoptic_segmentation(outputs, target_sizes=[image.size[::-1]])[0]
 
 
 
 
 
 
 
 
 
84
 
85
  pass
86
  else:
 
5
  import torch
6
  import pathlib
7
  from PIL import Image
8
+ import cv2 as cv
9
+ import numpy as np
10
 
11
+ from transformers import DetrImageProcessor, DetrForSegmentation, MaskFormerImageProcessor, MaskFormerForInstanceSegmentation
12
  from transformers.models.detr.feature_extraction_detr import rgb_to_id
13
 
14
 
 
62
  plt.axis("off")
63
  return fig2img(plt.gcf())
64
 
65
+ def contour_map(map_to_use, label_id):
66
+ mask = (map_to_use.cpu().numpy() == label_id)
67
+ visual_mask = (mask * 255).astype(np.uint8)
68
+ contours, hierarchy = cv.findContours(visual_mask, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE)
69
+ return contours, hierarchy
70
+
71
  def segment_images(model_name,url_input,image_input,threshold):
72
  ####
73
  # Get Image Object
 
89
 
90
  outputs = model(**inputs)
91
  results = processor.post_process_panoptic_segmentation(outputs, target_sizes=[image.size[::-1]])[0]
92
+
93
+ return_string = ""
94
+
95
+ for r in results["segments_info"]:
96
+ contour_list, hierarchy = contour_map(results["segmentation"], r["current_id"])
97
+ label_name = model.config.id2label[r["label_id"]]
98
+ return_string += f"ID: {r['id']}, Contour Length: {len(contour_list)}, Label Name: {label_name}, Score: {r['score']}\n"
99
+
100
+ return image, return_string
101
 
102
  pass
103
  else:
understand.py CHANGED
@@ -6,7 +6,7 @@ import numpy as np
6
  from PIL import Image
7
  import cv2 as cv
8
 
9
- from transformers import DetrFeatureExtractor, DetrForSegmentation, MaskFormerImageProcessor, MaskFormerForInstanceSegmentation
10
  # from transformers.models.detr.feature_extraction_detr import rgb_to_id
11
  from transformers.image_transforms import rgb_to_id
12
 
 
6
  from PIL import Image
7
  import cv2 as cv
8
 
9
+ from transformers import DetrImageProcessor, DetrForSegmentation, MaskFormerImageProcessor, MaskFormerForInstanceSegmentation
10
  # from transformers.models.detr.feature_extraction_detr import rgb_to_id
11
  from transformers.image_transforms import rgb_to_id
12