Spaces:
Runtime error
Runtime error
File size: 1,688 Bytes
f1dd031 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 |
import numpy as np
from pycocotools import mask as mask_utils
SHAPE = [720, 1280]
def mask_prepare(track_dict):
scores, masks = [], []
labels = track_dict["labels"]
for instance in labels:
masks.append(mask_utils.decode(instance["rle"]))
scores.append(instance["score"])
return scores, masks, track_dict
def mask_postprocess(scores, masks, track_dict):
sorted_idxs = np.argsort(scores)[::-1] # Sort indices in descending order of scores
processed_area = np.zeros(
SHAPE, dtype=np.uint8
) # Empty mask to record processed areas
for idx in sorted_idxs:
current_mask = masks[idx]
# Remove overlapping parts with already processed areas
current_mask = np.where(processed_area, 0, current_mask)
if current_mask.sum() > 0: # Only keep non-empty masks
# Update processed area
processed_area = np.maximum(processed_area, current_mask)
masks[idx] = current_mask
valid_rle_masks = [
mask_utils.encode(np.asfortranarray(masks[idx]))
for idx in sorted_idxs
if masks[idx].sum() > 0
]
valid_idxs = [idx for idx in sorted_idxs if masks[idx].sum() > 0]
valid_track_dicts = track_dict.copy()
valid_labels = []
for i in range(len(valid_idxs)):
vidx = valid_idxs[i]
if isinstance(valid_rle_masks[i]["counts"], bytes):
valid_rle_masks[i]["counts"] = valid_rle_masks[i]["counts"].decode()
valid_track_dicts["labels"][vidx]["rle"] = valid_rle_masks[i]
valid_labels.append(valid_track_dicts["labels"][vidx])
valid_track_dicts["labels"] = valid_labels
return valid_track_dicts
|