LiamKhoaLe commited on
Commit
26110ee
·
1 Parent(s): 125e09d

Flag unreachable as yellow:

Browse files
Files changed (1) hide show
  1. app.py +37 -19
app.py CHANGED
@@ -137,7 +137,13 @@ def build_masks(seg):
137
  return water_mask, garbage_mask, movable_mask
138
 
139
  # Garbage mask can be highlighted in red
140
- def highlight_chunk_masks_on_frame(frame, labels, objs, color_uncollected=(0, 0, 128), color_collected=(0, 128, 0), alpha=0.8):
 
 
 
 
 
 
141
  """
142
  Overlays semi-transparent colored regions for garbage chunks on the frame.
143
  `objs` must have 'pos' and 'col' keys. The collection status changes the overlay color.
@@ -150,8 +156,13 @@ def highlight_chunk_masks_on_frame(frame, labels, objs, color_uncollected=(0, 0,
150
  continue
151
  mask = (labels == lab).astype(np.uint8)
152
  contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
153
- color = color_collected if obj["col"] else color_uncollected
154
- # drawContours on overlay
 
 
 
 
 
155
  cv2.drawContours(overlay, contours, -1, color, thickness=cv2.FILLED)
156
  # Blend overlay with original frame using alpha
157
  return cv2.addWeighted(overlay, alpha, frame, 1 - alpha, 0)
@@ -201,32 +212,29 @@ def knn_path(start, targets, occ):
201
  reachable = []; unreachable = []
202
  while todo:
203
  # KNN follow a Greedy approach, which may not guarantee shortest path, hence only use A*
204
- # nbrs = NearestNeighbors(n_neighbors=1).fit(todo)
205
- # _,idx = nbrs.kneighbors([cur]); nxt=tuple(todo[idx[0][0]])
206
- # seg = astar(cur, nxt, occ)
207
- # if seg:
208
- # if path and seg[0]==path[-1]: seg=seg[1:]
209
- # path.extend(seg)
210
- # cur = nxt; todo.remove(list(nxt))
211
  best = None
212
  best_len = float('inf')
213
  best_seg = []
214
  # Try A* to each target, find shortest actual path
215
  for t in todo:
216
  seg = astar(cur, tuple(t), occ)
217
- if seg: # and len(seg) < best_len: # index error?
218
  best = tuple(t)
219
  best_len = len(seg)
220
  best_seg = seg
221
  if not best:
222
- print("⚠️ Some garbage unreachable")
223
- break # stop if no reachable targets left
 
 
 
224
  if path and path[-1] == best_seg[0]:
225
- best_seg = best_seg[1:] # avoid duplicate point
226
  path.extend(best_seg)
 
227
  cur = best
228
  todo.remove(list(best))
229
- return path
230
 
231
 
232
  # ── Robot sprite/class -──────────────────────────────────────────────────
@@ -488,18 +496,28 @@ def _pipeline(uid,img_path):
488
  robot = Robot(SPRITE)
489
  # Robot will be spawn on the closest movable mask to top-left
490
  robot.pos = [spawn_x, spawn_y]
491
- path = knn_path(robot.pos, centres, movable_mask)
 
 
492
 
493
  # 4- Video synthesis
494
  out_tmp=f"{OUTPUT_DIR}/{uid}_tmp.mp4"
495
  vw=cv2.VideoWriter(out_tmp,cv2.VideoWriter_fourcc(*"mp4v"),10.0,(640,640))
496
- objs=[{"pos":p,"col":False} for p in centres]
 
497
  bg = bgr.copy()
498
  for _ in range(15000): # safety frames
499
  frame=bg.copy()
500
  # Draw garbage chunk masks in red-to-green (semi-transparent)
501
- frame = highlight_chunk_masks_on_frame(frame, labels, objs) # 🔴 garbage overlay
502
- frame = highlight_water_mask_on_frame(frame, water_mask) # 🔵 water overlay
 
 
 
 
 
 
 
503
  # Draw object detections as red (to green) dots
504
  for o in objs:
505
  color = (0, 0, 128) if not o["col"] else (0, 128, 0)
 
137
  return water_mask, garbage_mask, movable_mask
138
 
139
  # Garbage mask can be highlighted in red
140
+ def highlight_chunk_masks_on_frame(
141
+ frame, labels, objs,
142
+ color_uncollected=(0, 0, 128),
143
+ color_collected=(0, 128, 0),
144
+ color_unreachable=(0, 255, 255),
145
+ alpha=0.8
146
+ ):
147
  """
148
  Overlays semi-transparent colored regions for garbage chunks on the frame.
149
  `objs` must have 'pos' and 'col' keys. The collection status changes the overlay color.
 
156
  continue
157
  mask = (labels == lab).astype(np.uint8)
158
  contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
159
+ # Choose color based on status
160
+ if obj.get("unreachable"):
161
+ color = color_unreachable
162
+ elif obj["col"]:
163
+ color = color_collected
164
+ else:
165
+ color = color_uncollected # drawContours on overlay
166
  cv2.drawContours(overlay, contours, -1, color, thickness=cv2.FILLED)
167
  # Blend overlay with original frame using alpha
168
  return cv2.addWeighted(overlay, alpha, frame, 1 - alpha, 0)
 
212
  reachable = []; unreachable = []
213
  while todo:
214
  # KNN follow a Greedy approach, which may not guarantee shortest path, hence only use A*
 
 
 
 
 
 
 
215
  best = None
216
  best_len = float('inf')
217
  best_seg = []
218
  # Try A* to each target, find shortest actual path
219
  for t in todo:
220
  seg = astar(cur, tuple(t), occ)
221
+ if seg and len(seg) < best_len: # index error?
222
  best = tuple(t)
223
  best_len = len(seg)
224
  best_seg = seg
225
  if not best:
226
+ # All remaining in `todo` are unreachable
227
+ for u in todo:
228
+ print(f"⚠️ Garbage unreachable at {u}")
229
+ unreachable.append(u)
230
+ break # no more reachable targets
231
  if path and path[-1] == best_seg[0]:
232
+ best_seg = best_seg[1:] # skip duplicate
233
  path.extend(best_seg)
234
+ reachable.append(list(best))
235
  cur = best
236
  todo.remove(list(best))
237
+ return path, unreachable
238
 
239
 
240
  # ── Robot sprite/class -──────────────────────────────────────────────────
 
496
  robot = Robot(SPRITE)
497
  # Robot will be spawn on the closest movable mask to top-left
498
  robot.pos = [spawn_x, spawn_y]
499
+ path, unreachable = knn_path(robot.pos, centres, movable_mask)
500
+ if unreachable:
501
+ print(f"⚠️ Unreachable garbage chunks at: {unreachable}")
502
 
503
  # 4- Video synthesis
504
  out_tmp=f"{OUTPUT_DIR}/{uid}_tmp.mp4"
505
  vw=cv2.VideoWriter(out_tmp,cv2.VideoWriter_fourcc(*"mp4v"),10.0,(640,640))
506
+ objs = [{"pos": p, "col": False, "unreachable": False} for p in centres if p not in unreachable]
507
+ objs += [{"pos": p, "col": False, "unreachable": True} for p in unreachable]
508
  bg = bgr.copy()
509
  for _ in range(15000): # safety frames
510
  frame=bg.copy()
511
  # Draw garbage chunk masks in red-to-green (semi-transparent)
512
+ frame = highlight_chunk_masks_on_frame(
513
+ frame,
514
+ labels,
515
+ objs,
516
+ color_uncollected=(0, 0, 128), # 🔴
517
+ color_collected=(0, 128, 0), # 🟢
518
+ color_unreachable=(0, 255, 255) # 🟡
519
+ ) # 🔴 garbage overlay
520
+ frame = highlight_water_mask_on_frame(frame, water_mask) # 🔵 water overlay
521
  # Draw object detections as red (to green) dots
522
  for o in objs:
523
  color = (0, 0, 128) if not o["col"] else (0, 128, 0)