MLBench/Contour_Detection_Paper
0
1from fastapi import FastAPI, HTTPException, UploadFile, File, Form2from pydantic import BaseModel3import numpy as np4from PIL import Image5import io, uuid, os, shutil, timeit6from datetime import datetime7from fastapi.staticfiles import StaticFiles8from fastapi.middleware.cors import CORSMiddleware9from fastapi.responses import FileResponse10 11# Import your paper-based prediction function12from app import (13 predict_full_paper,14 ReferenceBoxNotDetectedError,15 FingerCutOverlapError,16 MultipleObjectsError,17 NoObjectDetectedError,18 PaperNotDetectedError19)20 21app = FastAPI()22 23# Allow CORS if needed24app.add_middleware(25 CORSMiddleware,26 allow_origins=["*"],27 allow_methods=["*"],28 allow_headers=["*"],29)30 31BASE_URL = "https://app.us-central1.run.app"32 33OUTPUT_DIR = os.path.abspath("./outputs")34os.makedirs(OUTPUT_DIR, exist_ok=True)35 36UPDATES_DIR = os.path.abspath("./updates")37os.makedirs(UPDATES_DIR, exist_ok=True)38 39# Mount static directories with normal StaticFiles40app.mount("/outputs", StaticFiles(directory=OUTPUT_DIR), name="outputs")41app.mount("/updates", StaticFiles(directory=UPDATES_DIR), name="updates")42 43 44def save_and_build_urls(45 session_id: str,46 dxf_path: str,47 output_image: np.ndarray = None,48 outlines: np.ndarray = None,49 mask: np.ndarray = None,50 endpoint_type: str = "predict",51 paper_size: str = None,52 offset_value: float = None,53 offset_unit: str = "mm",54 finger_cut: str = "Off"55):56 """Helper to save all artifacts and return public URLs."""57 request_dir = os.path.join(OUTPUT_DIR, session_id)58 os.makedirs(request_dir, exist_ok=True)59 60 # Get current date61 current_date = datetime.utcnow().strftime("%d-%m-%Y")62 63 # Format offset value with underscore instead of dot64 offset_str = f"{offset_value:.3f}".replace(".", "_") if offset_value is not None else "0_000"65 66 # Create descriptive DXF filename67 if paper_size and offset_value is not None:68 dxf_fn = f"DXF_{current_date}_{paper_size}_{offset_str}{offset_unit}"69 if finger_cut == "On":70 dxf_fn += "_fingercut"71 dxf_fn += ".dxf"72 else:73 dxf_fn = f"DXF_{current_date}.dxf"74 75 # Full path for DXF76 new_dxf_path = os.path.join(request_dir, dxf_fn)77 78 # Copy DXF file79 if os.path.exists(dxf_path):80 shutil.copy(dxf_path, new_dxf_path)81 else:82 # Fallback if your DXF generator returns bytes or string83 with open(new_dxf_path, "wb") as f:84 if isinstance(dxf_path, (bytes, bytearray)):85 f.write(dxf_path)86 else:87 f.write(str(dxf_path).encode("utf-8"))88 89 urls = {90 "dxf_url": f"{BASE_URL}/download/{session_id}/{dxf_fn}",91 }92 93 # Save optional images if provided94 if output_image is not None:95 out_fn = "annotated_image.jpg"96 out_path = os.path.join(request_dir, out_fn)97 Image.fromarray(output_image).save(out_path)98 urls["output_image_url"] = f"{BASE_URL}/outputs/{session_id}/{out_fn}"99 100 if outlines is not None:101 outlines_fn = "outlines.jpg"102 outlines_path = os.path.join(request_dir, outlines_fn)103 Image.fromarray(outlines).save(outlines_path)104 urls["outlines_url"] = f"{BASE_URL}/outputs/{session_id}/{outlines_fn}"105 106 if mask is not None:107 mask_fn = "mask.jpg"108 mask_path = os.path.join(request_dir, mask_fn)109 Image.fromarray(mask).save(mask_path)110 urls["mask_url"] = f"{BASE_URL}/outputs/{session_id}/{mask_fn}"111 112 return urls113 114 115# Add new endpoint for downloading DXF files116@app.get("/download/{session_id}/{filename}")117async def download_file(session_id: str, filename: str):118 file_path = os.path.join(OUTPUT_DIR, session_id, filename)119 if not os.path.exists(file_path):120 raise HTTPException(status_code=404, detail="File not found")121 122 return FileResponse(123 path=file_path,124 filename=filename,125 media_type="application/x-dxf",126 headers={"Content-Disposition": f"attachment; filename={filename}"}127 )128 129 130@app.post("/predict_paper_simple")131async def predict_paper_simple_api(132 file: UploadFile = File(...),133 paper_size: str = Form(..., regex="^(A4|A3|US Letter)$"),134):135 """136 Simple paper-based predict: image + paper size → DXF only137 Default: 0mm offset, no finger cuts138 """139 session_id = str(uuid.uuid4())140 try:141 img_bytes = await file.read()142 image = np.array(Image.open(io.BytesIO(img_bytes)).convert("RGB"))143 except Exception:144 raise HTTPException(400, "Invalid image upload")145 146 try:147 start = timeit.default_timer()148 149 # Call predict_full_paper with default values150 dxf_path, ann_img, outlines_img, mask_img, scale_info = predict_full_paper(151 image=image,152 paper_size=paper_size,153 offset_value_mm=0.0, # No offset154 offset_unit="mm",155 enable_finger_cut="Off", # No finger cuts156 selected_outputs=[] # DXF only157 )158 159 elapsed = timeit.default_timer() - start160 print(f"[{session_id}] predict_paper_simple in {elapsed:.2f}s - {scale_info}")161 162 urls = save_and_build_urls(163 session_id=session_id,164 dxf_path=dxf_path,165 endpoint_type="predict_paper_simple",166 paper_size=paper_size,167 offset_value=0.0,168 offset_unit="mm",169 finger_cut="Off"170 )171 172 # Add scaling info to response173 urls["scale_info"] = scale_info174 return urls175 176 except (ReferenceBoxNotDetectedError, PaperNotDetectedError):177 raise HTTPException(status_code=400, detail="Error detecting paper! Please ensure the paper is clearly visible and try again.")178 except (MultipleObjectsError):179 raise HTTPException(status_code=400, detail="Multiple objects detected! Please place only a single object on the paper.")180 except (NoObjectDetectedError):181 raise HTTPException(status_code=400, detail="No object detected! Please ensure an object is placed on the paper.")182 except FingerCutOverlapError:183 raise HTTPException(status_code=400, detail="There was an overlap with fingercuts! Please try again to generate dxf.")184 except Exception as e:185 print(f"Error in predict_paper_simple: {str(e)}")186 raise HTTPException(status_code=500, detail="Error processing image! Please try again with a clearer image.")187 188 189@app.post("/predict_paper_with_offset")190async def predict_paper_with_offset_api(191 file: UploadFile = File(...),192 paper_size: str = Form(..., regex="^(A4|A3|US Letter)$"),193 offset_value: float = Form(...),194 offset_unit: str = Form(..., regex="^(mm|inches)$"),195 include_images: bool = Form(False) # Optional: include preview images196):197 """198 Paper-based predict with offset: image + paper size + offset → DXF + optional images199 """200 session_id = str(uuid.uuid4())201 try:202 img_bytes = await file.read()203 image = np.array(Image.open(io.BytesIO(img_bytes)).convert("RGB"))204 except Exception:205 raise HTTPException(400, "Invalid image upload")206 207 # Validate offset208 if offset_value < 0:209 raise HTTPException(400, "Offset value cannot be negative")210 if offset_value > 50: # Reasonable upper limit211 raise HTTPException(400, "Offset value too large (max 50)")212 213 try:214 start = timeit.default_timer()215 216 # Determine which outputs to include217 selected_outputs = ["Annotated Image", "Outlines", "Mask"] if include_images else []218 219 dxf_path, ann_img, outlines_img, mask_img, scale_info = predict_full_paper(220 image=image,221 paper_size=paper_size,222 offset_value_mm=offset_value,223 offset_unit=offset_unit,224 enable_finger_cut="Off", # No finger cuts225 selected_outputs=selected_outputs226 )227 228 elapsed = timeit.default_timer() - start229 print(f"[{session_id}] predict_paper_with_offset in {elapsed:.2f}s - {scale_info}")230 231 urls = save_and_build_urls(232 session_id=session_id,233 dxf_path=dxf_path,234 output_image=ann_img if include_images else None,235 outlines=outlines_img if include_images else None,236 mask=mask_img if include_images else None,237 endpoint_type="predict_paper_with_offset",238 paper_size=paper_size,239 offset_value=offset_value,240 offset_unit=offset_unit,241 finger_cut="Off"242 )243 244 urls["scale_info"] = scale_info245 return urls246 247 except (ReferenceBoxNotDetectedError, PaperNotDetectedError):248 raise HTTPException(status_code=400, detail="Error detecting paper! Please ensure the paper is clearly visible and try again.")249 except (MultipleObjectsError):250 raise HTTPException(status_code=400, detail="Multiple objects detected! Please place only a single object on the paper.")251 except (NoObjectDetectedError):252 raise HTTPException(status_code=400, detail="No object detected! Please ensure an object is placed on the paper.")253 except FingerCutOverlapError:254 raise HTTPException(status_code=400, detail="There was an overlap with fingercuts! Please try again to generate dxf.")255 except Exception as e:256 print(f"Error in predict_paper_with_offset: {str(e)}")257 raise HTTPException(status_code=500, detail="Error processing image! Please try again with a clearer image.")258 259 260@app.post("/predict_paper_full")261async def predict_paper_full_api(262 file: UploadFile = File(...),263 paper_size: str = Form(..., regex="^(A4|A3|US Letter)$"),264 offset_value: float = Form(...),265 offset_unit: str = Form(..., regex="^(mm|inches)$"),266 enable_finger_cut: str = Form(..., regex="^(On|Off)$"),267 include_images: bool = Form(False) # Optional: include preview images268):269 """270 Full paper-based predict: image + paper size + offset + finger cuts → DXF + optional images271 """272 session_id = str(uuid.uuid4())273 try:274 img_bytes = await file.read()275 image = np.array(Image.open(io.BytesIO(img_bytes)).convert("RGB"))276 except Exception:277 raise HTTPException(400, "Invalid image upload")278 279 # Validate offset280 if offset_value < 0:281 raise HTTPException(400, "Offset value cannot be negative")282 if offset_value > 50:283 raise HTTPException(400, "Offset value too large (max 50)")284 285 try:286 start = timeit.default_timer()287 288 # Determine which outputs to include289 selected_outputs = ["Annotated Image", "Outlines", "Mask"] if include_images else []290 291 dxf_path, ann_img, outlines_img, mask_img, scale_info = predict_full_paper(292 image=image,293 paper_size=paper_size,294 offset_value_mm=offset_value,295 offset_unit=offset_unit,296 enable_finger_cut=enable_finger_cut,297 selected_outputs=selected_outputs298 )299 300 elapsed = timeit.default_timer() - start301 print(f"[{session_id}] predict_paper_full in {elapsed:.2f}s - {scale_info}")302 303 urls = save_and_build_urls(304 session_id=session_id,305 dxf_path=dxf_path,306 output_image=ann_img if include_images else None,307 outlines=outlines_img if include_images else None,308 mask=mask_img if include_images else None,309 endpoint_type="predict_paper_full",310 paper_size=paper_size,311 offset_value=offset_value,312 offset_unit=offset_unit,313 finger_cut=enable_finger_cut314 )315 316 urls["scale_info"] = scale_info317 return urls318 319 except (ReferenceBoxNotDetectedError, PaperNotDetectedError):320 raise HTTPException(status_code=400, detail="Error detecting paper! Please ensure the paper is clearly visible and try again.")321 except (MultipleObjectsError):322 raise HTTPException(status_code=400, detail="Multiple objects detected! Please place only a single object on the paper.")323 except (NoObjectDetectedError):324 raise HTTPException(status_code=400, detail="No object detected! Please ensure an object is placed on the paper.")325 except FingerCutOverlapError:326 raise HTTPException(status_code=400, detail="There was an overlap with fingercuts! Please try again to generate dxf.")327 except Exception as e:328 print(f"Error in predict_paper_full: {str(e)}")329 raise HTTPException(status_code=500, detail="Error processing image! Please try again with a clearer image.")330 331 332# Keep the legacy endpoints for backward compatibility (optional)333@app.post("/predict1")334async def predict1_api(335 file: UploadFile = File(...)336):337 """338 Legacy endpoint - redirects to simple paper-based prediction with A4 default339 """340 return await predict_paper_simple_api(file=file, paper_size="A4")341 342 343@app.post("/predict2")344async def predict2_api(345 file: UploadFile = File(...),346 enable_fillet: str = Form(..., regex="^(On|Off)$"),347 fillet_value_mm: float = Form(...)348):349 """350 Legacy endpoint - redirects to paper-based prediction with offset351 Note: Fillet functionality mapped to offset for compatibility352 """353 # Map fillet to offset (you might want to adjust this logic)354 offset_value = fillet_value_mm if enable_fillet == "On" else 0.0355 356 return await predict_paper_with_offset_api(357 file=file,358 paper_size="A4", # Default to A4359 offset_value=offset_value,360 offset_unit="mm",361 include_images=True362 )363 364 365@app.post("/predict3")366async def predict3_api(367 file: UploadFile = File(...),368 enable_fillet: str = Form(..., regex="^(On|Off)$"),369 fillet_value_mm: float = Form(...),370 enable_finger_cut: str = Form(..., regex="^(On|Off)$")371):372 """373 Legacy endpoint - redirects to full paper-based prediction374 """375 offset_value = fillet_value_mm if enable_fillet == "On" else 0.0376 377 return await predict_paper_full_api(378 file=file,379 paper_size="A4", # Default to A4380 offset_value=offset_value,381 offset_unit="mm",382 enable_finger_cut=enable_finger_cut,383 include_images=True384 )385 386 387@app.post("/update")388async def update_files(389 output_image: UploadFile = File(...),390 outlines_image: UploadFile = File(...),391 mask_image: UploadFile = File(...),392 dxf_file: UploadFile = File(...)393):394 session_id = str(uuid.uuid4())395 update_dir = os.path.join(UPDATES_DIR, session_id)396 os.makedirs(update_dir, exist_ok=True)397 398 try:399 upload_map = {400 "output_image": output_image,401 "outlines_image": outlines_image,402 "mask_image": mask_image,403 "dxf_file": dxf_file,404 }405 urls = {}406 for key, up in upload_map.items():407 fn = up.filename408 path = os.path.join(update_dir, fn)409 with open(path, "wb") as f:410 shutil.copyfileobj(up.file, f)411 urls[key] = f"{BASE_URL}/updates/{session_id}/{fn}"412 413 return {"session_id": session_id, "uploaded": urls}414 415 except Exception as e:416 raise HTTPException(500, f"Update failed: {e}")417 418 419from fastapi import Response420 421@app.get("/health")422def health():423 return Response(content="OK", status_code=200)424 425 426@app.get("/")427def root():428 return {429 "message": "Paper-based DXF Generator API",430 "endpoints": [431 "/predict_paper_simple - Simple DXF generation with paper reference",432 "/predict_paper_with_offset - DXF generation with contour offset",433 "/predict_paper_full - Full DXF generation with all features",434 "/predict1, /predict2, /predict3 - Legacy endpoints (backward compatibility)"435 ],436 "paper_sizes": ["A4", "A3", "US Letter"],437 "units": ["mm", "inches"]438 }439 440 441if __name__ == "__main__":442 import uvicorn443 port = int(os.environ.get("PORT", 8080))444 print(f"Starting FastAPI server on 0.0.0.0:{port}...")445 uvicorn.run(app, host="0.0.0.0", port=port)