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

Highlight chunk in red mask not dot

Browse files
Files changed (1) hide show
  1. app.py +17 -5
app.py CHANGED
@@ -116,6 +116,15 @@ def build_masks(seg):
116
  movable_mask = water_mask | garbage_mask
117
  return water_mask, garbage_mask, movable_mask
118
 
 
 
 
 
 
 
 
 
 
119
  # ── A* and KNN over binary water grid ─────────────────────────────────
120
  def astar(start, goal, occ):
121
  h = lambda a,b: abs(a[0]-b[0])+abs(a[1]-b[1])
@@ -386,7 +395,7 @@ def _pipeline(uid,img_path):
386
  centres = []
387
  for x1, y1, x2, y2 in detections: # Define IoU heuristic
388
  '''
389
- We conduct a 30% allowance whether the center
390
  of the detected garbage's bbox lies within the travelable zone
391
  which was segmented earlier to be the water and garbage zone
392
  '''
@@ -396,7 +405,7 @@ def _pipeline(uid,img_path):
396
  box_mask = movable_mask[y1:y2, x1:x2] # ← use MOVABLE mask
397
  if box_mask.size == 0:
398
  continue
399
- if np.count_nonzero(box_mask) / box_mask.size >= 0.3:
400
  centres.append([int((x1 + x2) / 2), int((y1 + y2) / 2)])
401
  # add chunk centres and deduplicate
402
  centres.extend(chunk_centres)
@@ -417,10 +426,13 @@ def _pipeline(uid,img_path):
417
  bg = bgr.copy()
418
  for _ in range(15000): # safety frames
419
  frame=bg.copy()
420
- # draw garbage
 
 
421
  for o in objs:
422
- color=(0,0,255) if not o["col"] else (0,255,0)
423
- x,y=o["pos"]; cv2.circle(frame,(x,y),6,color,-1)
 
424
  # robot
425
  robot.step(path)
426
  rx,ry=robot.pos; sp=robot.png
 
116
  movable_mask = water_mask | garbage_mask
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):
130
  h = lambda a,b: abs(a[0]-b[0])+abs(a[1]-b[1])
 
395
  centres = []
396
  for x1, y1, x2, y2 in detections: # Define IoU heuristic
397
  '''
398
+ We conduct a 20% allowance whether the center
399
  of the detected garbage's bbox lies within the travelable zone
400
  which was segmented earlier to be the water and garbage zone
401
  '''
 
405
  box_mask = movable_mask[y1:y2, x1:x2] # ← use MOVABLE mask
406
  if box_mask.size == 0:
407
  continue
408
+ if np.count_nonzero(box_mask) / box_mask.size >= 0.2:
409
  centres.append([int((x1 + x2) / 2), int((y1 + y2) / 2)])
410
  # add chunk centres and deduplicate
411
  centres.extend(chunk_centres)
 
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
437
  robot.step(path)
438
  rx,ry=robot.pos; sp=robot.png