PlanetEyeFarm12/construction-deforestation-api
0
1"""2============================================================34-Model Ensemble Detection API — 10x Faster (Frame Skip)4============================================================5"""6 7from fastapi import FastAPI, File, UploadFile, HTTPException8from fastapi.responses import FileResponse, JSONResponse9from fastapi.middleware.cors import CORSMiddleware10import uvicorn11import cv212import os13import shutil14import tempfile15import zipfile16from collections import Counter, defaultdict17from datetime import datetime18from ultralytics import YOLO19 20app = FastAPI(21 title = "4-Model Construction Ensemble API",22 description = "Deforestation + Materials + PPE + Equipment",23 version = "4.0.0"24)25app.add_middleware(26 CORSMiddleware,27 allow_origins=["*"], allow_methods=["*"], allow_headers=["*"]28)29 30# ── Load 4 models ──────────────────────────────────────────31print("Loading models...")32MODELS = {33 "deforestation" : YOLO("deforestation.pt"),34 "materials" : YOLO("construction.pt"),35 "ppe" : YOLO("ppe_safety.pt"),36 "equipment" : YOLO("construction_equipment.pt"),37}38CLASSES = {name: model.names for name, model in MODELS.items()}39for name, cls in CLASSES.items():40 print(f" OK {name:25s}: {list(cls.values())}")41 42MODEL_COLORS = {43 "deforestation" : (0, 50, 220),44 "materials" : (0, 180, 50 ),45 "ppe" : (220, 80, 0 ),46 "equipment" : (0, 200, 255),47}48MODEL_LABELS = {49 "deforestation" : "DEFORESTATION",50 "materials" : "MATERIALS",51 "ppe" : "PPE SAFETY",52 "equipment" : "EQUIPMENT",53}54 55CONF_THRESHOLD = 0.3556SNAPSHOT_EVERY = 2 # log every N seconds57PROCESS_EVERY = 10 # run models every Nth frame (10 = 10x faster)58 59 60# ══════════════════════════════════════════════════════════61# DRAWING62# ══════════════════════════════════════════════════════════63 64def draw_boxes(frame, results, class_names, color):65 for box in results.boxes:66 cls_id = int(box.cls[0])67 conf = float(box.conf[0])68 x1, y1, x2, y2 = map(int, box.xyxy[0])69 name = class_names.get(cls_id, f"cls{cls_id}")70 label = f"{name} {conf:.0%}"71 cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)72 font = cv2.FONT_HERSHEY_SIMPLEX73 (tw, th), bl = cv2.getTextSize(label, font, 0.55, 1)74 cv2.rectangle(frame, (x1, y1-th-bl-6), (x1+tw+6, y1), color, -1)75 cv2.putText(frame, label, (x1+3, y1-bl-3),76 font, 0.55, (255, 255, 255), 1, cv2.LINE_AA)77 return frame78 79 80def draw_legend(frame, counts_per_model):81 font = cv2.FONT_HERSHEY_SIMPLEX82 line_h, pad, w = 20, 8, 31083 lines = []84 for model_name, color in MODEL_COLORS.items():85 counts = counts_per_model.get(model_name, {})86 total = sum(counts.values())87 lines.append((f"[ {MODEL_LABELS[model_name]} ] {total} obj", color, True))88 if counts:89 for cls, n in sorted(counts.items()):90 lines.append((f" {cls}: {n}", (210, 210, 210), False))91 else:92 lines.append((" none", (120, 120, 120), False))93 h = len(lines) * line_h + pad * 294 overlay = frame.copy()95 cv2.rectangle(overlay, (5, 5), (5+w, 5+h), (15, 15, 15), -1)96 cv2.addWeighted(overlay, 0.75, frame, 0.25, 0, frame)97 for i, (text, color, bold) in enumerate(lines):98 y = 5 + pad + (i+1) * line_h99 cv2.putText(frame, text, (12, y),100 font, 0.47, color, 2 if bold else 1, cv2.LINE_AA)101 return frame102 103 104# ══════════════════════════════════════════════════════════105# LOG106# ══════════════════════════════════════════════════════════107 108def snapshot_block(ts, second, counts_per_model):109 W = 65110 lines = [111 "",112 "+" + "-" * W + "+",113 f"| TIMESTAMP: {ts} ({second:.1f}s)".ljust(W+1) + "|",114 "+" + "-" * W + "+",115 ]116 grand = 0117 for model_name, label in [118 ("deforestation", "[DEFORESTATION]"),119 ("materials", "[MATERIALS] "),120 ("ppe", "[PPE SAFETY] "),121 ("equipment", "[EQUIPMENT] "),122 ]:123 counts = counts_per_model.get(model_name, {})124 total = sum(counts.values())125 grand += total126 lines.append(f"| {label} Total detections: {total}".ljust(W+1) + "|")127 if counts:128 for cls, n in sorted(counts.items()):129 bar = "|" * min(n, 20)130 lines.append(f"| {cls:<28s} {n:>3d} {bar}".ljust(W+1) + "|")131 else:132 lines.append("| (no detections)".ljust(W+1) + "|")133 lines.append("|" + " " * W + "|")134 lines.append("+" + "-" * W + "+")135 lines.append(f"| TOTAL ALL MODELS: {grand}".ljust(W+1) + "|")136 lines.append("+" + "-" * W + "+")137 return lines138 139 140# ══════════════════════════════════════════════════════════141# VIDEO PROCESSOR — 10x FASTER WITH FRAME SKIPPING142# ══════════════════════════════════════════════════════════143 144def process_video(input_path, output_dir, snapshot_every=2):145 cap = cv2.VideoCapture(input_path)146 fps = cap.get(cv2.CAP_PROP_FPS) or 25147 width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))148 height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))149 total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))150 dur = total / fps151 152 out_raw = os.path.join(output_dir, "out_raw.mp4")153 out_mp4 = os.path.join(output_dir, "annotated.mp4")154 out_txt = os.path.join(output_dir, "annotation.txt")155 156 writer = cv2.VideoWriter(157 out_raw, cv2.VideoWriter_fourcc(*"mp4v"), fps, (width, height)158 )159 160 W = 65161 log = [162 "=" * (W+2),163 " CONSTRUCTION SITE — 4 MODEL ENSEMBLE DETECTION LOG",164 "=" * (W+2),165 f" Generated : {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}",166 f" Video : {os.path.basename(input_path)}",167 f" Resolution : {width}x{height} | FPS: {fps:.1f}",168 f" Duration : {int(dur//60)}m {int(dur%60)}s ({total} frames)",169 f" Snapshot : every {snapshot_every} seconds",170 f" Speed mode : process every {PROCESS_EVERY} frames (10x faster)",171 "-" * (W+2),172 " MODELS & CLASSES:",173 f" [DEFORESTATION] : {', '.join(CLASSES['deforestation'].values())}",174 f" [MATERIALS] : {', '.join(CLASSES['materials'].values())}",175 f" [PPE SAFETY] : {', '.join(CLASSES['ppe'].values())}",176 f" [EQUIPMENT] : {', '.join(CLASSES['equipment'].values())}",177 "=" * (W+2),178 ]179 180 cumulative = defaultdict(Counter)181 snap_times = list(range(0, int(dur) + 1, snapshot_every))182 snap_idx = 0183 frame_idx = 0184 all_snaps = []185 186 # Cache last results for skipped frames187 last_results = None188 last_frame_counts = {name: Counter() for name in MODELS}189 190 while cap.isOpened():191 ret, frame = cap.read()192 if not ret:193 break194 195 current_sec = frame_idx / fps196 197 # ── Run models only every PROCESS_EVERY frames ──────198 if frame_idx % PROCESS_EVERY == 0:199 last_results = {200 name: model.predict(201 frame, conf=CONF_THRESHOLD, iou=0.45, verbose=False)[0]202 for name, model in MODELS.items()203 }204 last_frame_counts = {205 name: Counter(206 CLASSES[name][int(b.cls[0])] for b in res.boxes207 )208 for name, res in last_results.items()209 }210 211 # ── Always draw last known boxes on every frame ──────212 if last_results is not None:213 for name, res in last_results.items():214 frame = draw_boxes(215 frame, res, CLASSES[name], MODEL_COLORS[name])216 frame = draw_legend(frame, last_frame_counts)217 218 # ── Snapshot log every N seconds ─────────────────────219 if snap_idx < len(snap_times) and current_sec >= snap_times[snap_idx]:220 for name, cnt in last_frame_counts.items():221 cumulative[name].update(cnt)222 223 mins = int(current_sec) // 60224 secs = int(current_sec) % 60225 ts = f"{mins:02d}:{secs:02d}"226 227 log += snapshot_block(ts, current_sec, last_frame_counts)228 229 all_snaps.append({230 "timestamp" : ts,231 "seconds" : round(current_sec, 1),232 "deforestation": dict(last_frame_counts["deforestation"]),233 "materials" : dict(last_frame_counts["materials"]),234 "ppe" : dict(last_frame_counts["ppe"]),235 "equipment" : dict(last_frame_counts["equipment"]),236 "total_objects": sum(237 sum(c.values()) for c in last_frame_counts.values()),238 })239 snap_idx += 1240 241 writer.write(frame)242 frame_idx += 1243 244 # Progress log every 100 frames245 if frame_idx % 100 == 0:246 print(f" Processed {frame_idx}/{total} frames "247 f"({frame_idx/total*100:.1f}%)")248 249 cap.release()250 writer.release()251 252 # ── Overall summary ────────────────────────────────────253 log += ["", "=" * (W+2), " OVERALL SUMMARY (entire video)", "=" * (W+2)]254 grand_total = 0255 for model_name, label in [256 ("deforestation", "[DEFORESTATION]"),257 ("materials", "[MATERIALS] "),258 ("ppe", "[PPE SAFETY] "),259 ("equipment", "[EQUIPMENT] "),260 ]:261 total_count = sum(cumulative[model_name].values())262 grand_total += total_count263 log.append(f" {label} Total: {total_count}")264 for cls, n in sorted(cumulative[model_name].items()):265 log.append(f" {cls:<30s}: {n}")266 log.append("")267 log.append(f" GRAND TOTAL (all models): {grand_total}")268 log.append("=" * (W+2))269 log.append("END OF LOG")270 271 with open(out_txt, "w", encoding="utf-8") as f:272 f.write("\n".join(log))273 274 os.system(f'ffmpeg -i "{out_raw}" -vcodec libx264 -crf 23 '275 f'"{out_mp4}" -y -loglevel quiet')276 277 return out_mp4, out_txt, all_snaps, frame_idx278 279 280# ══════════════════════════════════════════════════════════281# ROUTES282# ══════════════════════════════════════════════════════════283 284@app.get("/")285def root():286 return {287 "status" : "running",288 "version": "4.0.0 — 10x faster",289 "models" : {n: list(c.values()) for n, c in CLASSES.items()},290 }291 292@app.get("/health")293def health():294 return {"status": "ok", "models_loaded": len(MODELS)}295 296@app.get("/classes")297def get_classes():298 return {n: list(c.values()) for n, c in CLASSES.items()}299 300@app.post("/detect/video")301async def detect_video(302 file: UploadFile = File(...),303 snapshot_every: int = 2304):305 if not file.filename.lower().endswith((".mp4",".avi",".mov",".mkv")):306 raise HTTPException(400, "Only video files: mp4, avi, mov, mkv")307 tmp = tempfile.mkdtemp()308 try:309 inp = os.path.join(tmp, "input.mp4")310 with open(inp, "wb") as f:311 shutil.copyfileobj(file.file, f)312 out_mp4, out_txt, _, _ = process_video(inp, tmp, snapshot_every)313 zip_path = os.path.join(tmp, "results.zip")314 with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:315 zf.write(out_mp4, "annotated.mp4")316 zf.write(out_txt, "annotation.txt")317 return FileResponse(zip_path, media_type="application/zip",318 filename="ensemble_results.zip")319 except Exception as e:320 shutil.rmtree(tmp, ignore_errors=True)321 raise HTTPException(500, str(e))322 323@app.post("/detect/json")324async def detect_json(325 file: UploadFile = File(...),326 snapshot_every: int = 2327):328 if not file.filename.lower().endswith((".mp4",".avi",".mov",".mkv")):329 raise HTTPException(400, "Only video files: mp4, avi, mov, mkv")330 tmp = tempfile.mkdtemp()331 try:332 inp = os.path.join(tmp, "input.mp4")333 with open(inp, "wb") as f:334 shutil.copyfileobj(file.file, f)335 _, _, snaps, frames = process_video(inp, tmp, snapshot_every)336 return JSONResponse({337 "status" : "success",338 "total_frames" : frames,339 "snapshot_every": f"{snapshot_every}s",340 "models" : {n: list(c.values()) for n, c in CLASSES.items()},341 "snapshots" : snaps,342 })343 except Exception as e:344 shutil.rmtree(tmp, ignore_errors=True)345 raise HTTPException(500, str(e))346 347if __name__ == "__main__":348 uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=False)