BIBEK8108/POTHOLE_DETECTION
0
1"""2======================================================================3FILE: app.py4PURPOSE: Flask web app — user uploads a photo or video,5 YOLOv8 detects potholes, and returns the annotated result.6======================================================================7"""8 9import os10import sys11import time12from flask import (13 Flask, render_template, request,14 redirect, url_for, send_file, flash, jsonify, Response15)16from werkzeug.utils import secure_filename17 18# Add src/ to Python path19sys.path.insert(0, os.path.join(os.path.dirname(__file__), "src"))20from detector import detect_image, detect_video21 22# ──────────────────────────────────────────────────────────────────23# APP SETUP24# ──────────────────────────────────────────────────────────────────25app = Flask(__name__)26app.secret_key = "pothole_detection_v2_2024"27 28BASE_DIR = os.path.dirname(os.path.abspath(__file__))29UPLOAD_FOLDER = os.path.join(BASE_DIR, "uploads")30OUTPUT_FOLDER = os.path.join(BASE_DIR, "outputs")31MODEL_PATH = os.path.join(BASE_DIR, "models", "best.pt")32 33app.config["MAX_CONTENT_LENGTH"] = 500 * 1024 * 1024 # 500 MB max34 35 36ALLOWED_IMAGE = {"jpg", "jpeg", "png", "bmp", "webp"}37ALLOWED_VIDEO = {"mp4", "avi", "mov", "mkv"}38 39os.makedirs(UPLOAD_FOLDER, exist_ok=True)40os.makedirs(OUTPUT_FOLDER, exist_ok=True)41 42 43def allowed_file(filename: str, allowed_set: set) -> bool:44 return "." in filename and filename.rsplit(".", 1)[1].lower() in allowed_set45 46 47def is_image(filename: str) -> bool:48 return allowed_file(filename, ALLOWED_IMAGE)49 50 51def is_video(filename: str) -> bool:52 return allowed_file(filename, ALLOWED_VIDEO)53 54 55# ──────────────────────────────────────────────────────────────────56# ROUTE: Home / Upload Page57# ──────────────────────────────────────────────────────────────────58@app.route("/")59def index():60 """Show the main upload page."""61 return render_template("index.html")62 63 64# ──────────────────────────────────────────────────────────────────65# ROUTE: Process Upload66# ──────────────────────────────────────────────────────────────────67@app.route("/process", methods=["POST"])68def process():69 """70 Receive uploaded file → run YOLOv8 → return annotated result.71 Supports both images and videos.72 """73 if "file" not in request.files:74 flash("No file uploaded!", "error")75 return redirect(url_for("index"))76 77 uploaded = request.files["file"]78 if uploaded.filename == "":79 flash("Please select a file!", "error")80 return redirect(url_for("index"))81 82 filename = secure_filename(uploaded.filename)83 confidence = float(request.form.get("confidence", 0.5))84 85 # Determine file type86 if is_image(filename):87 file_type = "image"88 elif is_video(filename):89 file_type = "video"90 else:91 flash("Unsupported file type! Use JPG, PNG, MP4, AVI, MOV, or MKV.", "error")92 return redirect(url_for("index"))93 94 # Check model exists95 if not os.path.exists(MODEL_PATH):96 flash("Model not found! Place best.pt in the models/ folder.", "error")97 return redirect(url_for("index"))98 99 # Save uploaded file100 input_path = os.path.join(UPLOAD_FOLDER, filename)101 uploaded.save(input_path)102 103 try:104 start_time = time.time()105 106 if file_type == "image":107 output_path, stats = detect_image(108 image_path=input_path,109 model_path=MODEL_PATH,110 output_dir=OUTPUT_FOLDER,111 confidence=confidence112 )113 output_filename = os.path.basename(output_path)114 else:115 skip_frames = int(request.form.get("skip_frames", 3))116 output_path, stats = detect_video(117 video_path=input_path,118 model_path=MODEL_PATH,119 output_dir=OUTPUT_FOLDER,120 confidence=confidence,121 skip_frames=skip_frames122 )123 output_filename = os.path.basename(output_path)124 125 elapsed = round(time.time() - start_time, 1)126 stats["elapsed_seconds"] = elapsed127 stats["file_type"] = file_type128 stats["output_filename"] = output_filename129 stats["original_filename"] = uploaded.filename130 # Images always preview fine in browser; videos depend on codec131 if file_type == "image":132 stats["browser_compatible"] = True133 134 return render_template("result.html", stats=stats, file_type=file_type,135 output_filename=output_filename)136 137 except Exception as e:138 import traceback139 traceback.print_exc()140 flash(f"Error during detection: {str(e)}", "error")141 return redirect(url_for("index"))142 143 144# ──────────────────────────────────────────────────────────────────145# ROUTE: Download Output File146# ──────────────────────────────────────────────────────────────────147@app.route("/download/<filename>")148def download(filename):149 """Send annotated output file to browser for download."""150 file_path = os.path.join(OUTPUT_FOLDER, filename)151 if not os.path.exists(file_path):152 return "File not found", 404153 return send_file(file_path, as_attachment=True)154 155 156# ──────────────────────────────────────────────────────────────────157# ROUTE: Preview / View Output File158# ──────────────────────────────────────────────────────────────────159@app.route("/preview/<filename>")160def preview(filename):161 """162 Serve the output file for inline viewing.163 For VIDEOS: implements HTTP Range request support (206 Partial Content),164 which is required by ALL browsers to play video inline. Without this,165 browsers silently refuse to stream/play the video.166 """167 file_path = os.path.join(OUTPUT_FOLDER, filename)168 if not os.path.exists(file_path):169 return "File not found", 404170 171 ext = filename.rsplit(".", 1)[-1].lower()172 mime_map = {173 "mp4": "video/mp4",174 "jpg": "image/jpeg",175 "jpeg": "image/jpeg",176 "png": "image/png",177 }178 mimetype = mime_map.get(ext, "application/octet-stream")179 180 # Images: simple send_file is fine (no range requests needed)181 if ext in ("jpg", "jpeg", "png", "bmp", "webp"):182 return send_file(file_path, mimetype=mimetype)183 184 # ── Videos: MUST support HTTP Range requests for browser inline playback ──185 file_size = os.path.getsize(file_path)186 range_header = request.headers.get("Range", None)187 188 if range_header:189 # Parse "bytes=start-end"190 byte_range = range_header.strip().replace("bytes=", "")191 parts = byte_range.split("-")192 start = int(parts[0]) if parts[0] else 0193 end = int(parts[1]) if len(parts) > 1 and parts[1] else file_size - 1194 end = min(end, file_size - 1)195 length = end - start + 1196 197 def generate_chunk(s, l):198 with open(file_path, "rb") as f:199 f.seek(s)200 remaining = l201 while remaining > 0:202 chunk = f.read(min(65536, remaining)) # 64 KB chunks203 if not chunk:204 break205 remaining -= len(chunk)206 yield chunk207 208 resp = Response(209 generate_chunk(start, length),210 status=206, # Partial Content211 mimetype=mimetype,212 direct_passthrough=True213 )214 resp.headers["Content-Range"] = f"bytes {start}-{end}/{file_size}"215 resp.headers["Accept-Ranges"] = "bytes"216 resp.headers["Content-Length"] = str(length)217 return resp218 219 # No Range header — serve whole file but advertise range support220 def stream_full():221 with open(file_path, "rb") as f:222 while True:223 chunk = f.read(65536)224 if not chunk:225 break226 yield chunk227 228 resp = Response(stream_full(), status=200, mimetype=mimetype,229 direct_passthrough=True)230 resp.headers["Accept-Ranges"] = "bytes"231 resp.headers["Content-Length"] = str(file_size)232 return resp233 234 235# ──────────────────────────────────────────────────────────────────236# START SERVER237# ──────────────────────────────────────────────────────────────────238if __name__ == "__main__":239 print("\n" + "=" * 55)240 print(" POTHOLE DETECTION v2")241 print("=" * 55)242 print(f" Model : {MODEL_PATH}")243 print(f" Uploads : {UPLOAD_FOLDER}")244 print(f" Outputs : {OUTPUT_FOLDER}")245 print("=" * 55)246 print(" Open: http://localhost:5001")247 print("=" * 55 + "\n")248 app.run(debug=True, host="0.0.0.0", port=5001)249