Tobias Cornille commited on
Commit
d197a83
1 Parent(s): dd286da

Fix thing category names

Browse files
Files changed (1) hide show
  1. app.py +8 -3
app.py CHANGED
@@ -290,7 +290,7 @@ def generate_panoptic_mask(
290
  image_array = np.asarray(image)
291
 
292
  # detect boxes for "thing" categories using Grounding DINO
293
- thing_boxes, _ = dino_detection(
294
  dino_model,
295
  image,
296
  image_array,
@@ -357,9 +357,14 @@ def generate_panoptic_mask(
357
  .numpy()
358
  .astype(int)
359
  )
360
- category_names = ["background"] + stuff_category_names + thing_category_names
 
 
 
 
361
  subsection_label_pairs = [
362
- (panoptic_bool_masks[i], category_names[i]) for i in range(len(category_names))
 
363
  ]
364
 
365
  return (image_array, subsection_label_pairs)
 
290
  image_array = np.asarray(image)
291
 
292
  # detect boxes for "thing" categories using Grounding DINO
293
+ thing_boxes, category_ids = dino_detection(
294
  dino_model,
295
  image,
296
  image_array,
 
357
  .numpy()
358
  .astype(int)
359
  )
360
+ panoptic_names = (
361
+ ["background"]
362
+ + stuff_category_names
363
+ + [category_names[category_id] for category_id in category_ids]
364
+ )
365
  subsection_label_pairs = [
366
+ (panoptic_bool_masks[i], panoptic_name)
367
+ for i, panoptic_name in enumerate(panoptic_names)
368
  ]
369
 
370
  return (image_array, subsection_label_pairs)