LiamKhoaLe commited on
Commit
896eec6
·
1 Parent(s): 264c76c

Update step methodology (stepping faster). Add water mask overlay. Reduce overlay alpha to 0.3

Browse files
Files changed (1) hide show
  1. app.py +25 -9
app.py CHANGED
@@ -134,7 +134,7 @@ def build_masks(seg):
134
  return water_mask, garbage_mask, movable_mask
135
 
136
  # Garbage mask can be highlighted in red
137
- def highlight_chunk_masks_on_frame(frame, labels, objs, color_uncollected=(0, 0, 128), color_collected=(0, 128, 0), alpha=0.6):
138
  """
139
  Overlays semi-transparent colored regions for garbage chunks on the frame.
140
  `objs` must have 'pos' and 'col' keys. The collection status changes the overlay color.
@@ -151,6 +151,16 @@ def highlight_chunk_masks_on_frame(frame, labels, objs, color_uncollected=(0, 0,
151
  cv2.drawContours(overlay, contours, -1, color, thickness=cv2.FILLED)
152
  # Blend overlay with original frame using alpha
153
  return cv2.addWeighted(overlay, alpha, frame, 1 - alpha, 0)
 
 
 
 
 
 
 
 
 
 
154
 
155
  # ── A* and KNN over binary water grid ─────────────────────────────────
156
  def astar(start, goal, occ):
@@ -191,7 +201,7 @@ def knn_path(start, targets, occ):
191
 
192
  # ── Robot sprite/class -──────────────────────────────────────────────────
193
  class Robot:
194
- def __init__(self, sprite, speed=2000): # Declare the robot's physical stats and routing (position, speed, movement, path)
195
  img = Image.open(sprite).convert("RGBA").resize((40, 40))
196
  self.png = np.array(img)
197
  if self.png.shape[-1] != 4:
@@ -199,12 +209,17 @@ class Robot:
199
  self.png = np.array(Image.open(sprite).convert("RGBA").resize((40,40)))
200
  self.pos = [0,0]; self.speed=speed
201
  def step(self, path):
202
- if not path: return
203
- dx,dy = path[0][0]-self.pos[0], path[0][1]-self.pos[1]
204
- dist = (dx*dx+dy*dy)**0.5
205
- if dist<=self.speed:
206
- self.pos=list(path.pop(0)); return
207
- r=self.speed/dist; self.pos=[int(self.pos[0]+dx*r), int(self.pos[1]+dy*r)]
 
 
 
 
 
208
 
209
  # ── FastAPI & HTML content (original styling) ───────────────────────────
210
  # HTML Content for UI (streamed with FastAPI HTML renderer)
@@ -458,7 +473,8 @@ def _pipeline(uid,img_path):
458
  for _ in range(15000): # safety frames
459
  frame=bg.copy()
460
  # Draw garbage chunk masks in red-to-green (semi-transparent)
461
- frame = highlight_chunk_masks_on_frame(frame, labels, objs)
 
462
  # Draw object detections as red (to green) dots
463
  for o in objs:
464
  color = (0, 0, 128) if not o["col"] else (0, 128, 0)
 
134
  return water_mask, garbage_mask, movable_mask
135
 
136
  # Garbage mask can be highlighted in red
137
+ def highlight_chunk_masks_on_frame(frame, labels, objs, color_uncollected=(0, 0, 128), color_collected=(0, 128, 0), alpha=0.3):
138
  """
139
  Overlays semi-transparent colored regions for garbage chunks on the frame.
140
  `objs` must have 'pos' and 'col' keys. The collection status changes the overlay color.
 
151
  cv2.drawContours(overlay, contours, -1, color, thickness=cv2.FILLED)
152
  # Blend overlay with original frame using alpha
153
  return cv2.addWeighted(overlay, alpha, frame, 1 - alpha, 0)
154
+ # Water mask to be blue
155
+ def highlight_water_mask_on_frame(frame, binary_mask, color=(255, 0, 0), alpha=0.3):
156
+ """
157
+ Overlays semi-transparent colored mask (binary) on the frame.
158
+ """
159
+ overlay = frame.copy()
160
+ mask = binary_mask.astype(np.uint8) * 255
161
+ contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
162
+ cv2.drawContours(overlay, contours, -1, color, thickness=cv2.FILLED)
163
+ return cv2.addWeighted(overlay, alpha, frame, 1 - alpha, 0)
164
 
165
  # ── A* and KNN over binary water grid ─────────────────────────────────
166
  def astar(start, goal, occ):
 
201
 
202
  # ── Robot sprite/class -──────────────────────────────────────────────────
203
  class Robot:
204
+ def __init__(self, sprite, speed=5000): # Declare the robot's physical stats and routing (position, speed, movement, path)
205
  img = Image.open(sprite).convert("RGBA").resize((40, 40))
206
  self.png = np.array(img)
207
  if self.png.shape[-1] != 4:
 
209
  self.png = np.array(Image.open(sprite).convert("RGBA").resize((40,40)))
210
  self.pos = [0,0]; self.speed=speed
211
  def step(self, path):
212
+ while path:
213
+ dx, dy = path[0][0] - self.pos[0], path[0][1] - self.pos[1]
214
+ dist = (dx * dx + dy * dy) ** 0.5
215
+ if dist <= self.speed:
216
+ self.pos = list(path.pop(0))
217
+ else:
218
+ r = self.speed / dist
219
+ self.pos = [int(self.pos[0] + dx * r), int(self.pos[1] + dy * r)]
220
+ # Break after one logical move to avoid overshooting
221
+ break
222
+
223
 
224
  # ── FastAPI & HTML content (original styling) ───────────────────────────
225
  # HTML Content for UI (streamed with FastAPI HTML renderer)
 
473
  for _ in range(15000): # safety frames
474
  frame=bg.copy()
475
  # Draw garbage chunk masks in red-to-green (semi-transparent)
476
+ frame = highlight_chunk_masks_on_frame(frame, labels, objs) # 🔴 garbage overlay
477
+ frame = highlight_water_mask_on_frame(frame, water_mask) # 🔵 water overlay
478
  # Draw object detections as red (to green) dots
479
  for o in objs:
480
  color = (0, 0, 128) if not o["col"] else (0, 128, 0)