LiamKhoaLe commited on
Commit
5478d4e
·
0 Parent(s):

Restore HF Space backend

Browse files
Files changed (12) hide show
  1. .gitignore +2 -0
  2. Dockerfile +72 -0
  3. README.md +12 -0
  4. app.py +447 -0
  5. download_models.py +24 -0
  6. entrypoint.sh +3 -0
  7. icon.png +0 -0
  8. model/donothing.txt +1 -0
  9. requirements.txt +28 -0
  10. setup.py +28 -0
  11. sprite.png +0 -0
  12. uploads/donothing.txt +1 -0
.gitignore ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ bot.py
2
+ bin
Dockerfile ADDED
@@ -0,0 +1,72 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ FROM python:3.9
2
+
3
+ # Set working directory inside the container
4
+ WORKDIR /app
5
+
6
+ # Reduce Docker image size by cleaning up unused files
7
+ RUN apt-get clean && rm -rf /var/lib/apt/lists/* /tmp/* ~/.cache/
8
+
9
+ # Install system dependencies
10
+ RUN apt-get update && apt-get install -y \
11
+ libgl1-mesa-glx libglib2.0-0 fontconfig \
12
+ git wget ffmpeg && \
13
+ rm -rf /var/lib/apt/lists/*
14
+
15
+ # Copy entrypoint script and make it executable
16
+ COPY entrypoint.sh /entrypoint.sh
17
+ RUN chmod +x /entrypoint.sh
18
+
19
+ # Create non-root user
20
+ RUN useradd -m -u 1000 user
21
+ USER root
22
+ ENV HOME=/home/user
23
+
24
+ # Install Python dependencies
25
+ COPY requirements.txt ./
26
+ RUN pip install --no-cache-dir -r requirements.txt
27
+
28
+ # Install Hugging Face CLI + required ML dependencies (hf_hub is at 0.14.1 and may cause a crash)
29
+ RUN pip install --no-cache-dir "huggingface_hub>=0.16.0,<0.21.0" --extra-index-url "https://download.pytorch.org/whl/cpu"
30
+
31
+ # Install LoveDA repository dependencies (last line ignore error flag on)
32
+ RUN git clone https://github.com/Junjue-Wang/LoveDA.git /home/user/LoveDA
33
+ WORKDIR /home/user/LoveDA
34
+ RUN pip install --no-cache-dir -r requirements.txt || true
35
+
36
+ # Return to working dir
37
+ WORKDIR $HOME/app
38
+
39
+ # Ensure Hugging Face CLI is accessible
40
+ ENV PATH="/home/user/.local/bin:$PATH"
41
+
42
+ # Create model/cache folders
43
+ RUN mkdir -p /home/user/app/model /home/user/app/cache/uploads && \
44
+ chown -R user:user /home/user/app && chmod -R 777 /home/user/app
45
+
46
+ # Switch back to non-root user
47
+ USER user
48
+ WORKDIR $HOME/app
49
+
50
+ # Copy application source code
51
+ COPY --chown=user . $HOME/app
52
+
53
+ # Pre-download all Hugging Face models
54
+ RUN python -c "from huggingface_hub import snapshot_download; snapshot_download(repo_id='facebook/detr-resnet-50', local_dir='/home/user/app/model/detr', local_dir_use_symlinks=False)"
55
+ RUN wget -O $HOME/app/model/garbage_detector.pt https://huggingface.co/BinKhoaLe1812/Garbage_Detection/resolve/main/garbage_detector.pt
56
+ RUN wget -O $HOME/app/model/yolov5-detect-trash-classification.pt https://huggingface.co/turhancan97/yolov5-detect-trash-classification/resolve/main/yolov5s.pt
57
+
58
+ # Verify model setup
59
+ RUN python setup.py
60
+
61
+ # Clean outputs folder
62
+ RUN rm -rf /home/user/app/outputs/*.mp4
63
+
64
+ # Copy static images
65
+ COPY sprite.png /home/user/app/sprite.png
66
+ COPY icon.png /home/user/app/icon.png
67
+
68
+ # Expose FastAPI port
69
+ EXPOSE 7860
70
+
71
+ # Start app
72
+ ENTRYPOINT ["/entrypoint.sh"]
README.md ADDED
@@ -0,0 +1,12 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ title: 'Sall-e Garbage Detection'
3
+ emoji: ♻️
4
+ colorFrom: green
5
+ colorTo: blue
6
+ sdk: docker
7
+ pinned: false
8
+ license: apache-2.0
9
+ short_description: 'Incorporating multi-modal garbage detection images'
10
+ ---
11
+
12
+ Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
app.py ADDED
@@ -0,0 +1,447 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Access: https://BinKhoaLe1812-Sall-eGarbageDetection.hf.space/ui
2
+
3
+ # ───────────────────────── app.py (Sall-e demo) ─────────────────────────
4
+ # FastAPI ▸ upload image ▸ multi-model garbage detection ▸ ADE-20K
5
+ # semantic segmentation (Water / Garbage) ▸ A* + KNN navigation ▸ H.264 video
6
+ # =======================================================================
7
+
8
+ import os, uuid, threading, shutil, time, heapq, cv2, numpy as np
9
+ from PIL import Image
10
+ import uvicorn
11
+ from fastapi import FastAPI, File, UploadFile, Request
12
+ from fastapi.responses import HTMLResponse, StreamingResponse, Response
13
+ from fastapi.staticfiles import StaticFiles
14
+
15
+ # ── Vision libs ─────────────────────────────────────────────────────────
16
+ import torch, yolov5, ffmpeg
17
+ from ultralytics import YOLO
18
+ from transformers import (
19
+ DetrImageProcessor, DetrForObjectDetection,
20
+ SegformerFeatureExtractor, SegformerForSemanticSegmentation
21
+ )
22
+ from sklearn.neighbors import NearestNeighbors
23
+
24
+ # ── Folders / files ─────────────────────────────────────────────────────
25
+ BASE = "/home/user/app"
26
+ CACHE = f"{BASE}/cache"
27
+ UPLOAD_DIR = f"{CACHE}/uploads"
28
+ OUTPUT_DIR = f"{BASE}/outputs"
29
+ MODEL_DIR = f"{BASE}/model"
30
+ SPRITE = f"{BASE}/sprite.png"
31
+
32
+ os.makedirs(UPLOAD_DIR, exist_ok=True)
33
+ os.makedirs(OUTPUT_DIR, exist_ok=True)
34
+ os.makedirs(CACHE , exist_ok=True)
35
+ os.environ["TRANSFORMERS_CACHE"] = CACHE
36
+ os.environ["HF_HOME"] = CACHE
37
+
38
+ # ── Load models once ───────────────────────────────────────────────────
39
+ print("🔄 Loading models …")
40
+ model_self = YOLO(f"{MODEL_DIR}/garbage_detector.pt") # YOLOv11(l)
41
+ model_yolo5 = yolov5.load(f"{MODEL_DIR}/yolov5-detect-trash-classification.pt")
42
+ processor_detr = DetrImageProcessor.from_pretrained(f"{MODEL_DIR}/detr")
43
+ model_detr = DetrForObjectDetection.from_pretrained(f"{MODEL_DIR}/detr")
44
+ feat_extractor = SegformerFeatureExtractor.from_pretrained(
45
+ "nvidia/segformer-b4-finetuned-ade-512-512")
46
+ segformer = SegformerForSemanticSegmentation.from_pretrained(
47
+ "nvidia/segformer-b4-finetuned-ade-512-512")
48
+ print("✅ Models ready\n")
49
+
50
+ # ── ADE-20K palette + custom mapping (verbatim) ─────────────────────────
51
+ # ADE20K palette
52
+ ade_palette = np.array([
53
+ [0, 0, 0], [120, 120, 120], [180, 120, 120], [6, 230, 230], [80, 50, 50],
54
+ [4, 200, 3], [120, 120, 80], [140, 140, 140], [204, 5, 255], [230, 230, 230],
55
+ [4, 250, 7], [224, 5, 255], [235, 255, 7], [150, 5, 61], [120, 120, 70],
56
+ [8, 255, 51], [255, 6, 82], [143, 255, 140], [204, 255, 4], [255, 51, 7],
57
+ [204, 70, 3], [0, 102, 200], [61, 230, 250], [255, 6, 51], [11, 102, 255],
58
+ [255, 7, 71], [255, 9, 224], [9, 7, 230], [220, 220, 220], [255, 9, 92],
59
+ [112, 9, 255], [8, 255, 214], [7, 255, 224], [255, 184, 6], [10, 255, 71],
60
+ [255, 41, 10], [7, 255, 255], [224, 255, 8], [102, 8, 255], [255, 61, 6],
61
+ [255, 194, 7], [255, 122, 8], [0, 255, 20], [255, 8, 41], [255, 5, 153],
62
+ [6, 51, 255], [235, 12, 255], [160, 150, 20], [0, 163, 255], [140, 140, 140],
63
+ [250, 10, 15], [20, 255, 0], [31, 255, 0], [255, 31, 0], [255, 224, 0],
64
+ [153, 255, 0], [0, 0, 255], [255, 71, 0], [0, 235, 255], [0, 173, 255],
65
+ [31, 0, 255], [11, 200, 200], [255, 82, 0], [0, 255, 245], [0, 61, 255],
66
+ [0, 255, 112], [0, 255, 133], [255, 0, 0], [255, 163, 0], [255, 102, 0],
67
+ [194, 255, 0], [0, 143, 255], [51, 255, 0], [0, 82, 255], [0, 255, 41],
68
+ [0, 255, 173], [10, 0, 255], [173, 255, 0], [0, 255, 153], [255, 92, 0],
69
+ [255, 0, 255], [255, 0, 245], [255, 0, 102], [255, 173, 0], [255, 0, 20],
70
+ [255, 184, 184], [0, 31, 255], [0, 255, 61], [0, 71, 255], [255, 0, 204],
71
+ [0, 255, 194], [0, 255, 82], [0, 10, 255], [0, 112, 255], [51, 0, 255],
72
+ [0, 194, 255], [0, 122, 255], [0, 255, 163], [255, 153, 0], [0, 255, 10],
73
+ [255, 112, 0], [143, 255, 0], [82, 0, 255], [163, 255, 0], [255, 235, 0],
74
+ [8, 184, 170], [133, 0, 255], [0, 255, 92], [184, 0, 255], [255, 0, 31],
75
+ [0, 184, 255], [0, 214, 255], [255, 0, 112], [92, 255, 0], [0, 224, 255],
76
+ [112, 224, 255], [70, 184, 160], [163, 0, 255], [153, 0, 255], [71, 255, 0],
77
+ [255, 0, 163], [255, 204, 0], [255, 0, 143], [0, 255, 235], [133, 255, 0],
78
+ [255, 0, 235], [245, 0, 255], [255, 0, 122], [255, 245, 0], [10, 190, 212],
79
+ [214, 255, 0], [0, 204, 255], [20, 0, 255], [255, 255, 0], [0, 153, 255],
80
+ [0, 41, 255], [0, 255, 204], [41, 0, 255], [41, 255, 0], [173, 0, 255],
81
+ [0, 245, 255], [71, 0, 255], [122, 0, 255], [0, 255, 184], [0, 92, 255],
82
+ [184, 255, 0], [0, 133, 255], [255, 214, 0], [25, 194, 194], [102, 255, 0],
83
+ [92, 0, 255]
84
+ ], dtype=np.uint8)
85
+
86
+ custom_class_map = {
87
+ "Garbage": [(150, 5, 61)],
88
+ "Water": [(0, 102, 200), (11, 102, 255), (31, 0, 255)],
89
+ "Grass / Vegetation": [(10, 255, 71), (143, 255, 140)],
90
+ "Tree / Natural Obstacle": [(4, 200, 3), (235, 12, 255), (255, 6, 82), (255, 163, 0)],
91
+ "Sand / Soil / Ground": [(80, 50, 50), (230, 230, 230)],
92
+ "Buildings / Structures": [(255, 0, 255), (184, 0, 255), (120, 120, 120), (7, 255, 224)],
93
+ "Sky / Background": [(180, 120, 120)],
94
+ "Undetecable": [(0, 0, 0)],
95
+ "Unknown Class": []
96
+ }
97
+ TOL = 30 # RGB tolerance
98
+
99
+ # Masking zones (Garbage and Water zone to be travelable)
100
+ def build_masks(seg):
101
+ """
102
+ Returns three binary masks at (H,W):
103
+ water_mask – 1 = water
104
+ garbage_mask – 1 = semantic “Garbage” pixels
105
+ movable_mask – union of water & garbage (robot can travel here)
106
+ """
107
+ decoded = ade_palette[seg]
108
+ water_mask = np.zeros(seg.shape, np.uint8)
109
+ garbage_mask = np.zeros_like(water_mask)
110
+ # Append water pixels to water_mask
111
+ for rgb in custom_class_map["Water"]:
112
+ water_mask |= (np.abs(decoded - rgb).max(axis=-1) <= TOL)
113
+ # Append gb pixels to garbage_mask
114
+ for rgb in custom_class_map["Garbage"]:
115
+ garbage_mask |= (np.abs(decoded - rgb).max(axis=-1) <= TOL)
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])
122
+ N8 = [(-1,-1),(-1,0),(-1,1),(0,-1),(0,1),(1,-1),(1,0),(1,1)]
123
+ openq=[(0,start)]; g={start:0}; came={}
124
+ while openq:
125
+ _,cur=heapq.heappop(openq)
126
+ if cur==goal:
127
+ p=[cur]; # reconstruct
128
+ while cur in came: cur=came[cur]; p.append(cur)
129
+ return p[::-1]
130
+ for dx,dy in N8:
131
+ nx,ny=cur[0]+dx,cur[1]+dy
132
+ if not (0<=nx<640 and 0<=ny<640): continue
133
+ if occ[ny,nx]==0: continue
134
+ ng=g[cur]+1
135
+ if (nx,ny) not in g or ng<g[(nx,ny)]:
136
+ g[(nx,ny)]=ng
137
+ f=ng+h((nx,ny),goal)
138
+ heapq.heappush(openq,(f,(nx,ny)))
139
+ came[(nx,ny)]=cur
140
+ return []
141
+
142
+ # KNN fit
143
+ def knn_path(start, targets, occ):
144
+ todo = targets[:]; path=[]
145
+ cur = tuple(start)
146
+ while todo:
147
+ nbrs = NearestNeighbors(n_neighbors=1).fit(todo)
148
+ _,idx = nbrs.kneighbors([cur]); nxt=tuple(todo[idx[0][0]])
149
+ seg = astar(cur, nxt, occ)
150
+ if seg:
151
+ if path and seg[0]==path[-1]: seg=seg[1:]
152
+ path.extend(seg)
153
+ cur = nxt; todo.remove(list(nxt))
154
+ return path
155
+
156
+ # ── Robot sprite/class -──────────────────────────────────────────────────
157
+ class Robot:
158
+ def __init__(self, sprite, speed=200): # Declare the robot's physical stats and routing (position, speed, movement, path)
159
+ self.png = np.array(Image.open(sprite).convert("RGBA").resize((40,40)))
160
+ self.pos = [0,0]; self.speed=speed
161
+ def step(self, path):
162
+ if not path: return
163
+ dx,dy = path[0][0]-self.pos[0], path[0][1]-self.pos[1]
164
+ dist = (dx*dx+dy*dy)**0.5
165
+ if dist<=self.speed:
166
+ self.pos=list(path.pop(0)); return
167
+ r=self.speed/dist; self.pos=[int(self.pos[0]+dx*r), int(self.pos[1]+dy*r)]
168
+
169
+ # ── FastAPI & HTML content (original styling) ───────────────────────────
170
+ # HTML Content for UI (streamed with FastAPI HTML renderer)
171
+ HTML_CONTENT = """
172
+ <!DOCTYPE html>
173
+ <html>
174
+ <head>
175
+ <title>Sall-e Garbage Detection</title>
176
+ <link rel="website icon" type="png" href="/static/icon.png" >
177
+ <style>
178
+ body {
179
+ font-family: 'Roboto', sans-serif; background: linear-gradient(270deg, rgb(44, 13, 58), rgb(13, 58, 56)); color: white; text-align: center; margin: 0; padding: 50px;
180
+ }
181
+ h1 {
182
+ font-size: 40px;
183
+ background: linear-gradient(to right, #f32170, #ff6b08, #cf23cf, #eedd44);
184
+ -webkit-text-fill-color: transparent;
185
+ -webkit-background-clip: text;
186
+ font-weight: bold;
187
+ }
188
+ #upload-container {
189
+ background: rgba(255, 255, 255, 0.2); padding: 20px; width: 70%; border-radius: 10px; display: inline-block; box-shadow: 0px 0px 10px rgba(255, 255, 255, 0.3);
190
+ }
191
+ #upload {
192
+ font-size: 18px; padding: 10px; border-radius: 5px; border: none; background: #fff; cursor: pointer;
193
+ }
194
+ #loader {
195
+ margin-top: 10px; margin-left: auto; margin-right: auto; width: 60px; height: 60px; font-size: 12px; text-align: center;
196
+ }
197
+ p {
198
+ margin-top: 10px; font-size: 12px; color: #3498db;
199
+ }
200
+ #spinner {
201
+ border: 8px solid #f3f3f3; border-top: 8px solid rgb(117 7 7); border-radius: 50%; animation: spin 1s linear infinite; width: 40px; height: 40px; margin: auto;
202
+ }
203
+ @keyframes spin {
204
+ 0% { transform: rotate(0deg); }
205
+ 100% { transform: rotate(360deg); }
206
+ }
207
+ #outputVideo {
208
+ margin-top: 20px; width: 70%; margin-left: auto; margin-right: auto; max-width: 640px; border-radius: 10px; box-shadow: 0px 0px 10px rgba(255, 255, 255, 0.3);
209
+ }
210
+ #downloadBtn {
211
+ display: block; visibility: hidden; width: 20%; margin-top: 20px; margin-left: auto; margin-right: auto; padding: 10px 15px; font-size: 16px; background: #27ae60; color: white; border: none; border-radius: 5px; cursor: pointer; text-decoration: none;
212
+ }
213
+ #downloadBtn:hover {
214
+ background: #950606;
215
+ }
216
+ .hidden {
217
+ display: none;
218
+ }
219
+ @media (max-width: 860px) {
220
+ h1 { font-size: 30px; }
221
+ }
222
+ @media (max-width: 720px) {
223
+ h1 { font-size: 25px; }
224
+ #upload { font-size: 15px; }
225
+ #downloadBtn { font-size: 13px; }
226
+ }
227
+ @media (max-width: 580px) {
228
+ h1 { font-size: 20px; }
229
+ #upload { font-size: 10px; }
230
+ #downloadBtn { font-size: 10px; }
231
+ }
232
+ @media (max-width: 580px) {
233
+ h1 { font-size: 10px; }
234
+ }
235
+ @media (max-width: 460px) {
236
+ #upload { font-size: 7px; }
237
+ }
238
+ @media (max-width: 400px) {
239
+ h1 { font-size: 14px; }
240
+ }
241
+ @media (max-width: 370px) {
242
+ h1 { font-size: 11px; }
243
+ #upload { font-size: 5px; }
244
+ #downloadBtn { font-size: 7px; }
245
+ }
246
+ @media (max-width: 330px) {
247
+ h1 { font-size: 8px; }
248
+ #upload { font-size: 3px; }
249
+ #downloadBtn { font-size: 5px; }
250
+ }
251
+ </style>
252
+ </head>
253
+ <body>
254
+ <h1>Upload an Image for Garbage Detection</h1>
255
+ <div id="upload-container">
256
+ <input type="file" id="upload" accept="image/*">
257
+ </div>
258
+ <div id="loader" class="loader hidden">
259
+ <div id="spinner"></div>
260
+ <!-- <p>Garbage detection model processing...</p> -->
261
+ </div>
262
+ <video id="outputVideo" class="outputVideo" controls></video>
263
+ <a id="downloadBtn" class="downloadBtn">Download Video</a>
264
+ <script>
265
+ document.addEventListener("DOMContentLoaded", function() {
266
+ document.getElementById("outputVideo").classList.add("hidden");
267
+ document.getElementById("downloadBtn").style.visibility = "hidden";
268
+ });
269
+ document.getElementById('upload').addEventListener('change', async function(event) {
270
+ event.preventDefault();
271
+ const loader = document.getElementById("loader");
272
+ const outputVideo = document.getElementById("outputVideo");
273
+ const downloadBtn = document.getElementById("downloadBtn");
274
+ let file = event.target.files[0];
275
+ if (file) {
276
+ let formData = new FormData();
277
+ formData.append("file", file);
278
+ loader.classList.remove("hidden");
279
+ outputVideo.classList.add("hidden");
280
+ document.getElementById("downloadBtn").style.visibility = "hidden";
281
+ let response = await fetch('/upload/', { method: 'POST', body: formData });
282
+ let result = await response.json();
283
+ let user_id = result.user_id;
284
+ while (true) {
285
+ let checkResponse = await fetch(`/check_video/${user_id}`);
286
+ let checkResult = await checkResponse.json();
287
+ if (checkResult.ready) break;
288
+ await new Promise(resolve => setTimeout(resolve, 3000)); // Wait 3s before checking again
289
+ }
290
+ loader.classList.add("hidden");
291
+ let videoUrl = `/video/${user_id}?t=${new Date().getTime()}`;
292
+ outputVideo.src = videoUrl;
293
+ outputVideo.load();
294
+ outputVideo.play();
295
+ outputVideo.setAttribute("crossOrigin", "anonymous");
296
+ outputVideo.classList.remove("hidden");
297
+ downloadBtn.href = videoUrl;
298
+ document.getElementById("downloadBtn").style.visibility = "visible";
299
+ }
300
+ });
301
+ document.getElementById('outputVideo').addEventListener('error', function() {
302
+ console.log("⚠️ Video could not be played, showing download button instead.");
303
+ document.getElementById('outputVideo').classList.add("hidden");
304
+ document.getElementById("downloadBtn").style.visibility = "visible";
305
+ });
306
+ </script>
307
+ </body>
308
+ </html>
309
+ """
310
+
311
+ # ── Static-web ────────────────────────────��─────────────────────────────
312
+ app = FastAPI()
313
+ app.mount("/static", StaticFiles(directory=BASE), name="static")
314
+ video_ready={}
315
+ @app.get("/ui", response_class=HTMLResponse)
316
+ def ui(): return HTML_CONTENT
317
+ def _uid(): return uuid.uuid4().hex[:8]
318
+
319
+ # ── End-points ──────────────────────────────────────────────────────────
320
+ # User upload environment img here
321
+ @app.post("/upload/")
322
+ async def upload(file:UploadFile=File(...)):
323
+ uid=_uid(); dest=f"{UPLOAD_DIR}/{uid}_{file.filename}"
324
+ with open(dest,"wb") as bf: shutil.copyfileobj(file.file,bf)
325
+ threading.Thread(target=_pipeline, args=(uid,dest)).start()
326
+ return {"user_id":uid}
327
+
328
+ # Health check, make sure the video generator is alive and debug which video id is processed (multiple video can be processed at 1 worker)
329
+ @app.get("/check_video/{uid}")
330
+ def chk(uid:str): return {"ready":video_ready.get(uid,False)}
331
+
332
+ # Where the final video being saved
333
+ @app.get("/video/{uid}")
334
+ def stream(uid:str):
335
+ vid=f"{OUTPUT_DIR}/{uid}.mp4"
336
+ if not os.path.exists(vid): return Response(status_code=404)
337
+ return StreamingResponse(open(vid,"rb"), media_type="video/mp4")
338
+
339
+ # ── Core pipeline (runs in background thread) ───────────────────────────
340
+ def _pipeline(uid,img_path):
341
+ print(f"▶️ [{uid}] processing")
342
+ bgr=cv2.resize(cv2.imread(img_path),(640,640)); rgb=cv2.cvtColor(bgr,cv2.COLOR_BGR2RGB)
343
+ pil=Image.fromarray(rgb)
344
+
345
+ # 1- Segmentation → masking each segmented zone with pytorch
346
+ with torch.no_grad():
347
+ inputs = feat_extractor(pil, return_tensors="pt")
348
+ seg_logits = segformer(**inputs).logits
349
+ # Tensor run by CPU
350
+ seg_tensor = seg_logits.argmax(1)[0].cpu()
351
+ if seg_tensor.numel() == 0:
352
+ print(f"❌ [{uid}] segmentation failed (empty tensor)")
353
+ video_ready[uid] = True
354
+ return
355
+ # Resize the tensor to 640x640
356
+ seg = cv2.resize(seg_tensor.numpy(), (640, 640), interpolation=cv2.INTER_NEAREST)
357
+ print(f"🧪 [{uid}] segmentation input shape: {inputs['pixel_values'].shape}")
358
+ water_mask, garbage_mask, movable_mask = build_masks(seg) # movable zone = water and garbage masks
359
+
360
+ # 2- Garbage detection (3 models) → keep centres on water
361
+ detections=[]
362
+ # Detect garbage chunks (from segmentation)
363
+ num_cc, labels = cv2.connectedComponents(garbage_mask.astype(np.uint8))
364
+ chunk_centres = []
365
+ for lab in range(1, num_cc):
366
+ ys, xs = np.where(labels == lab)
367
+ if xs.size == 0: # safety
368
+ continue
369
+ chunk_centres.append([int(xs.mean()), int(ys.mean())])
370
+ print(f"🧠 {len(chunk_centres)} garbage chunk detected")
371
+ # Detect garbage object by within travelable zones
372
+ for r in model_self(bgr): # YOLOv11 (self-trained)
373
+ detections += [b.xyxy[0].tolist() for b in r.boxes]
374
+ for r in model_yolo5(bgr): # YOLOv5
375
+ if hasattr(r, 'pred') and len(r.pred) > 0:
376
+ detections += [p[:4].tolist() for p in r.pred[0]]
377
+ inp=processor_detr(images=pil,return_tensors="pt")
378
+ with torch.no_grad(): out=model_detr(**inp) # DETR
379
+ post = processor_detr.post_process_object_detection(
380
+ outputs=out,
381
+ target_sizes=torch.tensor([pil.size[::-1]]),
382
+ threshold=0.5
383
+ )[0]
384
+ detections += [b.tolist() for b in post["boxes"]]
385
+ # centre & mask filter (the garbage lies within travelable zone are collectable)
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
+ '''
393
+ x1, y1, x2, y2 = map(int, [x1, y1, x2, y2])
394
+ x1 = max(0, min(x1, 639)); y1 = max(0, min(y1, 639))
395
+ x2 = max(0, min(x2, 639)); y2 = max(0, min(y2, 639))
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)
403
+ centres = [list(c) for c in {tuple(c) for c in centres}]
404
+ if not centres: # No garbages within travelable zone
405
+ print(f"🛑 [{uid}] no reachable garbage"); video_ready[uid]=True; return
406
+ else: # Garbage within valid travelable zone
407
+ print(f"🧠 {len(centres)} garbage objects on water selected from {len(detections)} detections")
408
+
409
+ # 3- Global route
410
+ robot = Robot(SPRITE)
411
+ path = knn_path(robot.pos, centres, movable_mask)
412
+
413
+ # 4- Video synthesis
414
+ out_tmp=f"{OUTPUT_DIR}/{uid}_tmp.mp4"
415
+ vw=cv2.VideoWriter(out_tmp,cv2.VideoWriter_fourcc(*"mp4v"),10.0,(640,640))
416
+ objs=[{"pos":p,"col":False} for p in centres]
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
427
+ a=sp[:,:,3]/255.; bgroi=frame[ry:ry+40,rx:rx+40]
428
+ for c in range(3): bgroi[:,:,c]=a*sp[:,:,c]+(1-a)*bgroi[:,:,c]
429
+ frame[ry:ry+40,rx:rx+40]=bgroi
430
+ # collection check
431
+ for o in objs:
432
+ if not o["col"] and np.hypot(o["pos"][0]-rx,o["pos"][1]-ry)<=20:
433
+ o["col"]=True
434
+ vw.write(frame)
435
+ if all(o["col"] for o in objs): break
436
+ if not path: break
437
+ vw.release()
438
+
439
+ # 5- Convert to H.264
440
+ final=f"{OUTPUT_DIR}/{uid}.mp4"
441
+ ffmpeg.input(out_tmp).output(final,vcodec="libx264",pix_fmt="yuv420p").run(overwrite_output=True,quiet=True)
442
+ os.remove(out_tmp); video_ready[uid]=True
443
+ print(f"✅ [{uid}] video ready → {final}")
444
+
445
+ # ── Run locally (HF Space ignores since built with Docker image) ────────
446
+ if __name__=="__main__":
447
+ uvicorn.run(app,host="0.0.0.0",port=7860)
download_models.py ADDED
@@ -0,0 +1,24 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from huggingface_hub import hf_hub_download
2
+ import os
3
+ import shutil
4
+
5
+ # Define model download directory
6
+ model_dir = "/home/user/app/model"
7
+ cache_dir = "/home/user/app/cache"
8
+
9
+ max_cache_size = 500 * 1024 * 1024 # 500MB
10
+
11
+ # Ensure the directory exists
12
+ os.makedirs(model_dir, exist_ok=True)
13
+ os.makedirs(cache_dir, exist_ok=True)
14
+
15
+ # Download DETR model (only the model weights (167MB)) and save to local model directory
16
+ print("🚀 Downloading DETR model from Hugging Face...")
17
+ hf_hub_download(repo_id="facebook/detr-resnet-50", filename="pytorch_model.bin", local_dir=f"{model_dir}/detr", cache_dir=cache_dir)
18
+
19
+ print("✅ DETR model downloaded successfully!")
20
+
21
+ if os.path.exists(cache_dir) and shutil.disk_usage(cache_dir).used > max_cache_size:
22
+ print("🗑️ Clearing Hugging Face cache to free up space...")
23
+ shutil.rmtree(cache_dir)
24
+ os.makedirs(cache_dir, exist_ok=True)
entrypoint.sh ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ #!/bin/bash
2
+ echo "🚀 Starting FastAPI server with Uvicorn..."
3
+ exec python -m uvicorn app:app --host 0.0.0.0 --port 7860
icon.png ADDED
model/donothing.txt ADDED
@@ -0,0 +1 @@
 
 
1
+ this directory serves as temporary folder placing tmp files
requirements.txt ADDED
@@ -0,0 +1,28 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Image Processing
2
+ opencv-python-headless
3
+ pillow
4
+ numpy
5
+ pillow-avif-plugin
6
+
7
+ # Vision Models
8
+ torch
9
+ torchvision
10
+ torchaudio
11
+ scikit-learn
12
+ ultralytics
13
+ timm
14
+ yolov5
15
+ huggingface_hub>=0.20.3
16
+ transformers==4.37.2
17
+ accelerate==0.27.2
18
+
19
+ # Video Processing
20
+ ffmpeg-python
21
+
22
+ # Server
23
+ fastapi
24
+ uvicorn
25
+ python-multipart
26
+
27
+ # Git (needed for clone/setup LoveDA repo)
28
+ git-lfs
setup.py ADDED
@@ -0,0 +1,28 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+
3
+ # Dir path
4
+ CACHE_DIR = "/home/user/app/cache"
5
+ MODEL_DIR = "/home/user/app/model"
6
+
7
+ # Verify structures:
8
+
9
+ # Model Dir
10
+ def print_model():
11
+ print("\n📂 Model Structure (Build Level):")
12
+ for root, dirs, files in os.walk(MODEL_DIR):
13
+ print(f"📁 {root}/")
14
+ for file in files:
15
+ print(f" 📄 {file}")
16
+
17
+ # Cache Dir
18
+ def print_cache():
19
+ print("\n📂 Cache Structure (Build Level):")
20
+ for root, dirs, files in os.walk(CACHE_DIR):
21
+ print(f"📁 {root}/")
22
+ for file in files:
23
+ print(f" 📄 {file}")
24
+
25
+
26
+ # Show
27
+ print_model()
28
+ print_cache()
sprite.png ADDED
uploads/donothing.txt ADDED
@@ -0,0 +1 @@
 
 
1
+ this directory serves as temporary folder placing tmp files