|
1 | 1 | import numpy as np
|
2 | 2 | import cv2
|
| 3 | +from skimage.measure import approximate_polygon, find_contours |
3 | 4 |
|
4 | 5 |
|
5 | 6 | MASK_THRESHOLD = .5
|
@@ -43,22 +44,21 @@ def segm_postprocess(box: list, raw_cls_mask, im_h, im_w, threshold):
|
43 | 44 | y = box[4] * height
|
44 | 45 | right = box[5] * width
|
45 | 46 | bottom = box[6] * height
|
46 |
| - mask = masks[index][label] |
| 47 | + mask = masks[index][label - 1] |
47 | 48 |
|
48 | 49 | mask = segm_postprocess((x, y, right, bottom),
|
49 | 50 | mask,
|
50 | 51 | height,
|
51 | 52 | width,
|
52 | 53 | MASK_THRESHOLD)
|
53 | 54 |
|
54 |
| - contour, _ = cv2.findContours(mask, |
55 |
| - cv2.RETR_EXTERNAL, |
56 |
| - cv2.CHAIN_APPROX_TC89_KCOS) |
| 55 | + contours = find_contours(mask, MASK_THRESHOLD) |
| 56 | + contour = contours[0] |
| 57 | + contour = np.flip(contour, axis=1) |
| 58 | + contour = approximate_polygon(contour, tolerance=2.5) |
| 59 | + segmentation = contour.tolist() |
57 | 60 |
|
58 |
| - contour = contour[0] |
59 |
| - contour = contour.tolist() |
60 |
| - contour = [x[0] for x in contour] |
61 | 61 |
|
62 | 62 | # NOTE: if you want to see the boxes, uncomment next line
|
63 | 63 | # results.add_box(x, y, right, bottom, label, frame_number)
|
64 |
| - results.add_polygon(contour, label, frame_number) |
| 64 | + results.add_polygon(segmentation, label, frame_number) |
0 commit comments