LiamKhoaLe commited on
Commit
70b6774
·
1 Parent(s): 7c0fd7a

Upd chunk mask to be transparrent

Browse files
Files changed (1) hide show
  1. app.py +19 -8
app.py CHANGED
@@ -117,13 +117,24 @@ def build_masks(seg):
117
  return water_mask, garbage_mask, movable_mask
118
 
119
  # Garbage mask can be highlighted in red
120
- def highlight_chunk_masks_on_frame(frame, labels, num_cc, color=(128, 0, 0)):
121
- """Draw filled red areas over the garbage chunk regions"""
122
- for lab in range(1, num_cc):
 
 
 
 
 
 
 
 
123
  mask = (labels == lab).astype(np.uint8)
124
  contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
125
- cv2.drawContours(frame, contours, -1, color, thickness=cv2.FILLED)
126
- return frame
 
 
 
127
 
128
  # ── A* and KNN over binary water grid ─────────────────────────────────
129
  def astar(start, goal, occ):
@@ -426,11 +437,11 @@ def _pipeline(uid,img_path):
426
  bg = bgr.copy()
427
  for _ in range(15000): # safety frames
428
  frame=bg.copy()
429
- # Draw garbage chunk masks in red
430
- frame = highlight_chunk_masks_on_frame(frame, labels, num_cc, color=(0, 0, 255))
431
  # Draw object detections as red (to green) dots
432
  for o in objs:
433
- color = (0, 0, 255) if not o["col"] else (0, 255, 0)
434
  x, y = o["pos"]
435
  cv2.circle(frame, (x, y), 6, color, -1)
436
  # robot
 
117
  return water_mask, garbage_mask, movable_mask
118
 
119
  # Garbage mask can be highlighted in red
120
+ def highlight_chunk_masks_on_frame(frame, labels, objs, num_cc, color_uncollected=(0, 0, 128), color_collected=(0, 128, 0), alpha=0.3):
121
+ """
122
+ Overlays semi-transparent colored regions for garbage chunks on the frame.
123
+ `objs` must have 'pos' and 'col' keys. The collection status changes the overlay color.
124
+ """
125
+ overlay = frame.copy()
126
+ for i, obj in enumerate(objs):
127
+ x, y = obj["pos"]
128
+ lab = labels[y, x]
129
+ if lab == 0:
130
+ continue
131
  mask = (labels == lab).astype(np.uint8)
132
  contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
133
+ color = color_collected if obj["col"] else color_uncollected
134
+ cv2.drawContours(overlay, contours, -1, color, thickness=cv2.FILLED)
135
+
136
+ # Blend overlay with original frame using alpha
137
+ return cv2.addWeighted(overlay, alpha, frame, 1 - alpha, 0)
138
 
139
  # ── A* and KNN over binary water grid ─────────────────────────────────
140
  def astar(start, goal, occ):
 
437
  bg = bgr.copy()
438
  for _ in range(15000): # safety frames
439
  frame=bg.copy()
440
+ # Draw garbage chunk masks in red-to-green (semi-transparent)
441
+ frame = highlight_chunk_masks_on_frame(frame, labels, objs)
442
  # Draw object detections as red (to green) dots
443
  for o in objs:
444
+ color = (0, 0, 128) if not o["col"] else (0, 128, 0)
445
  x, y = o["pos"]
446
  cv2.circle(frame, (x, y), 6, color, -1)
447
  # robot