MLBench/Contours_Extraction
0
1# from fastapi import FastAPI, HTTPException, UploadFile, File, Form2# from pydantic import BaseModel3# import numpy as np4# from PIL import Image5# import io, uuid, os, shutil, timeit6# from datetime import datetime7# from fastapi.staticfiles import StaticFiles8# from fastapi.middleware.cors import CORSMiddleware9 10# # import your three wrappers11# from app import predict_simple, predict_middle, predict_full12 13# app = FastAPI()14 15# # allow CORS if needed16# app.add_middleware(17# CORSMiddleware,18# allow_origins=["*"],19# allow_methods=["*"],20# allow_headers=["*"],21# )22 23# BASE_URL = "https://snapanddtraceapp-988917236820.us-central1.run.app"24# OUTPUT_DIR = os.path.abspath("./outputs")25# os.makedirs(OUTPUT_DIR, exist_ok=True)26# app.mount("/outputs", StaticFiles(directory=OUTPUT_DIR), name="outputs")27 28# UPDATES_DIR = os.path.abspath("./updates")29# os.makedirs(UPDATES_DIR, exist_ok=True)30# app.mount("/updates", StaticFiles(directory=UPDATES_DIR), name="updates")31 32 33# def save_and_build_urls(34# session_id: str,35# output_image: np.ndarray,36# outlines: np.ndarray,37# dxf_path: str,38# mask: np.ndarray39# ):40# """Helper to save all four artifacts and return public URLs."""41# request_dir = os.path.join(OUTPUT_DIR, session_id)42# os.makedirs(request_dir, exist_ok=True)43 44# # filenames45# out_fn = "overlay.jpg"46# outlines_fn = "outlines.jpg"47# mask_fn = "mask.jpg"48# current_date = datetime.now().strftime("%d-%m-%Y")49# dxf_fn = f"out_{current_date}_{session_id}.dxf"50 51# # full paths52# out_path = os.path.join(request_dir, out_fn)53# outlines_path = os.path.join(request_dir, outlines_fn)54# mask_path = os.path.join(request_dir, mask_fn)55# new_dxf_path = os.path.join(request_dir, dxf_fn)56 57# # save images58# Image.fromarray(output_image).save(out_path)59# Image.fromarray(outlines).save(outlines_path)60# Image.fromarray(mask).save(mask_path)61 62# # copy dx file63# if os.path.exists(dxf_path):64# shutil.copy(dxf_path, new_dxf_path)65# else:66# # fallback if your DXF generator returns bytes or string67# with open(new_dxf_path, "wb") as f:68# if isinstance(dxf_path, (bytes, bytearray)):69# f.write(dxf_path)70# else:71# f.write(str(dxf_path).encode("utf-8"))72 73# # build URLs74# return {75# "output_image_url": f"{BASE_URL}/outputs/{session_id}/{out_fn}",76# "outlines_url": f"{BASE_URL}/outputs/{session_id}/{outlines_fn}",77# "mask_url": f"{BASE_URL}/outputs/{session_id}/{mask_fn}",78# "dxf_url": f"{BASE_URL}/outputs/{session_id}/{dxf_fn}",79# }80 81 82# @app.post("/predict1")83# async def predict1_api(84# file: UploadFile = File(...)85# ):86# """87# Simple predict: only image → overlay, outlines, mask, DXF88# """89# session_id = str(uuid.uuid4())90# try:91# img_bytes = await file.read()92# image = np.array(Image.open(io.BytesIO(img_bytes)).convert("RGB"))93# except Exception:94# raise HTTPException(400, "Invalid image upload")95 96# try:97# start = timeit.default_timer()98# out_img, outlines, dxf_path, mask = predict_simple(image)99# elapsed = timeit.default_timer() - start100# print(f"[{session_id}] predict1 in {elapsed:.2f}s")101 102# return save_and_build_urls(session_id, out_img, outlines, dxf_path, mask)103 104# except Exception as e:105# raise HTTPException(500, f"predict1 failed: {e}")106# except ReferenceBoxNotDetectedError:107# raise HTTPException(status_code=400, detail="Error detecting reference battery! Please try again with a clearer image.")108# except FingerCutOverlapError:109# raise HTTPException(status_code=400, detail="There was an overlap with fingercuts!s Please try again to generate dxf.")110 111 112# @app.post("/predict2")113# async def predict2_api(114# file: UploadFile = File(...),115# enable_fillet: str = Form(..., regex="^(On|Off)$"),116# fillet_value_mm: float = Form(...)117# ):118# """119# Middle predict: image + fillet toggle + fillet value → overlay, outlines, mask, DXF120# """121# session_id = str(uuid.uuid4())122# try:123# img_bytes = await file.read()124# image = np.array(Image.open(io.BytesIO(img_bytes)).convert("RGB"))125# except Exception:126# raise HTTPException(400, "Invalid image upload")127 128# try:129# start = timeit.default_timer()130# out_img, outlines, dxf_path, mask = predict_middle(131# image, enable_fillet, fillet_value_mm132# )133# elapsed = timeit.default_timer() - start134# print(f"[{session_id}] predict2 in {elapsed:.2f}s")135 136# return save_and_build_urls(session_id, out_img, outlines, dxf_path, mask)137 138# except Exception as e:139# raise HTTPException(500, f"predict2 failed: {e}")140# except ReferenceBoxNotDetectedError:141# raise HTTPException(status_code=400, detail="Error detecting reference battery! Please try again with a clearer image.")142# except FingerCutOverlapError:143# raise HTTPException(status_code=400, detail="There was an overlap with fingercuts!s Please try again to generate dxf.")144 145# @app.post("/predict3")146# async def predict3_api(147# file: UploadFile = File(...),148# enable_fillet: str = Form(..., regex="^(On|Off)$"),149# fillet_value_mm: float = Form(...),150# enable_finger_cut: str = Form(..., regex="^(On|Off)$")151# ):152# """153# Full predict: image + fillet toggle/value + finger-cut toggle → overlay, outlines, mask, DXF154# """155# session_id = str(uuid.uuid4())156# try:157# img_bytes = await file.read()158# image = np.array(Image.open(io.BytesIO(img_bytes)).convert("RGB"))159# except Exception:160# raise HTTPException(400, "Invalid image upload")161 162# try:163# start = timeit.default_timer()164# out_img, outlines, dxf_path, mask = predict_full(165# image, enable_fillet, fillet_value_mm, enable_finger_cut166# )167# elapsed = timeit.default_timer() - start168# print(f"[{session_id}] predict3 in {elapsed:.2f}s")169 170# return save_and_build_urls(session_id, out_img, outlines, dxf_path, mask)171 172# except Exception as e:173# raise HTTPException(500, f"predict3 failed: {e}")174# except ReferenceBoxNotDetectedError:175# raise HTTPException(status_code=400, detail="Error detecting reference battery! Please try again with a clearer image.")176# except FingerCutOverlapError:177# raise HTTPException(status_code=400, detail="There was an overlap with fingercuts!s Please try again to generate dxf.")178 179# @app.post("/update")180# async def update_files(181# output_image: UploadFile = File(...),182# outlines_image: UploadFile = File(...),183# mask_image: UploadFile = File(...),184# dxf_file: UploadFile = File(...)185# ):186# session_id = str(uuid.uuid4())187# update_dir = os.path.join(UPDATES_DIR, session_id)188# os.makedirs(update_dir, exist_ok=True)189 190# try:191# upload_map = {192# "output_image": output_image,193# "outlines_image": outlines_image,194# "mask_image": mask_image,195# "dxf_file": dxf_file,196# }197# urls = {}198# for key, up in upload_map.items():199# fn = up.filename200# path = os.path.join(update_dir, fn)201# with open(path, "wb") as f:202# shutil.copyfileobj(up.file, f)203# urls[key] = f"{BASE_URL}/updates/{session_id}/{fn}"204 205# return {"session_id": session_id, "uploaded": urls}206 207# except Exception as e:208# raise HTTPException(500, f"Update failed: {e}")209 210 211# if __name__ == "__main__":212# import uvicorn213# port = int(os.environ.get("PORT", 8082))214# print(f"Starting FastAPI server on 0.0.0.0:{port}...")215# uvicorn.run(app, host="0.0.0.0", port=port)216 217 218 219 220 221 222 223 224 225 226from fastapi import FastAPI, HTTPException, UploadFile, File, Form227from pydantic import BaseModel228import numpy as np229from PIL import Image230import io, uuid, os, shutil, timeit231from datetime import datetime232from fastapi.staticfiles import StaticFiles233from fastapi.middleware.cors import CORSMiddleware234from fastapi.responses import FileResponse235 236# import your three wrappers237from app import predict_simple, predict_middle, predict_full238 239from app import (240 predict_simple, predict_middle, predict_full,241 ReferenceBoxNotDetectedError,242 FingerCutOverlapError243)244 245 246app = FastAPI()247 248# allow CORS if needed249app.add_middleware(250 CORSMiddleware,251 allow_origins=["*"],252 allow_methods=["*"],253 allow_headers=["*"],254)255 256BASE_URL = "https://snapanddtraceapp-988917236820.us-central1.run.app"257 258OUTPUT_DIR = os.path.abspath("./outputs")259os.makedirs(OUTPUT_DIR, exist_ok=True)260 261UPDATES_DIR = os.path.abspath("./updates")262os.makedirs(UPDATES_DIR, exist_ok=True)263 264# Mount static directories with normal StaticFiles265app.mount("/outputs", StaticFiles(directory=OUTPUT_DIR), name="outputs")266app.mount("/updates", StaticFiles(directory=UPDATES_DIR), name="updates")267 268 269def save_and_build_urls(270 session_id: str,271 output_image: np.ndarray,272 outlines: np.ndarray,273 dxf_path: str,274 mask: np.ndarray,275 endpoint_type: str,276 fillet_value: float = None,277 finger_cut: str = None278):279 """Helper to save all four artifacts and return public URLs."""280 request_dir = os.path.join(OUTPUT_DIR, session_id)281 os.makedirs(request_dir, exist_ok=True)282 283 # filenames284 out_fn = "overlay.jpg"285 outlines_fn = "outlines.jpg"286 mask_fn = "mask.jpg"287 288 # Get current date289 current_date = datetime.utcnow().strftime("%d-%m-%Y")290 291 292 # Format fillet value with underscore instead of dot293 fillet_str = f"{fillet_value:.2f}".replace(".", "_") if fillet_value is not None else None294 295 # Determine DXF filename based on endpoint type296 if endpoint_type == "predict1":297 dxf_fn = f"DXF_{current_date}.dxf"298 elif endpoint_type == "predict2":299 dxf_fn = f"DXF_{current_date}.dxf"300 elif endpoint_type == "predict3":301 dxf_fn = f"DXF_{current_date}.dxf"302 303 # full paths304 out_path = os.path.join(request_dir, out_fn)305 outlines_path = os.path.join(request_dir, outlines_fn)306 mask_path = os.path.join(request_dir, mask_fn)307 new_dxf_path = os.path.join(request_dir, dxf_fn)308 309 # save images310 Image.fromarray(output_image).save(out_path)311 Image.fromarray(outlines).save(outlines_path)312 Image.fromarray(mask).save(mask_path)313 314 # copy dxf file315 if os.path.exists(dxf_path):316 shutil.copy(dxf_path, new_dxf_path)317 else:318 # fallback if your DXF generator returns bytes or string319 with open(new_dxf_path, "wb") as f:320 if isinstance(dxf_path, (bytes, bytearray)):321 f.write(dxf_path)322 else:323 f.write(str(dxf_path).encode("utf-8"))324 325 # build URLs with /download prefix for DXF326 return {327 "output_image_url": f"{BASE_URL}/outputs/{session_id}/{out_fn}",328 "outlines_url": f"{BASE_URL}/outputs/{session_id}/{outlines_fn}",329 "mask_url": f"{BASE_URL}/outputs/{session_id}/{mask_fn}",330 "dxf_url": f"{BASE_URL}/download/{session_id}/{dxf_fn}", # Changed to use download endpoint331 }332 333# Add new endpoint for downloading DXF files334@app.get("/download/{session_id}/{filename}")335async def download_file(session_id: str, filename: str):336 file_path = os.path.join(OUTPUT_DIR, session_id, filename)337 if not os.path.exists(file_path):338 raise HTTPException(status_code=404, detail="File not found")339 340 return FileResponse(341 path=file_path,342 filename=filename,343 media_type="application/x-dxf",344 headers={"Content-Disposition": f"attachment; filename={filename}"}345 )346 347 348@app.post("/predict1")349async def predict1_api(350 file: UploadFile = File(...)351):352 """353 Simple predict: only image → overlay, outlines, mask, DXF354 DXF naming format: DXF_DD-MM-YYYY.dxf355 """356 session_id = str(uuid.uuid4())357 try:358 img_bytes = await file.read()359 image = np.array(Image.open(io.BytesIO(img_bytes)).convert("RGB"))360 except Exception:361 raise HTTPException(400, "Invalid image upload")362 363 try:364 start = timeit.default_timer()365 out_img, outlines, dxf_path, mask = predict_simple(image)366 elapsed = timeit.default_timer() - start367 print(f"[{session_id}] predict1 in {elapsed:.2f}s")368 369 return save_and_build_urls(370 session_id=session_id,371 output_image=out_img,372 outlines=outlines,373 dxf_path=dxf_path,374 mask=mask,375 endpoint_type="predict1"376 )377 378 except ReferenceBoxNotDetectedError:379 raise HTTPException(status_code=400, detail="Error detecting reference battery! Please try again with a clearer image.")380 except FingerCutOverlapError:381 raise HTTPException(status_code=400, detail="There was an overlap with fingercuts! Please try again to generate dxf.")382 except HTTPException as e:383 raise e384 except Exception as e:385 raise HTTPException(status_code=500, detail="Error detecting reference battery! Please try again with a clearer image.")386 387@app.post("/predict2")388async def predict2_api(389 file: UploadFile = File(...),390 enable_fillet: str = Form(..., regex="^(On|Off)$"),391 fillet_value_mm: float = Form(...)392):393 """394 Middle predict: image + fillet toggle + fillet value → overlay, outlines, mask, DXF395 DXF naming format: DXF_DD-MM-YYYY_fillet-value_mm.dxf396 """397 session_id = str(uuid.uuid4())398 try:399 img_bytes = await file.read()400 image = np.array(Image.open(io.BytesIO(img_bytes)).convert("RGB"))401 except Exception:402 raise HTTPException(400, "Invalid image upload")403 404 try:405 start = timeit.default_timer()406 out_img, outlines, dxf_path, mask = predict_middle(407 image, enable_fillet, fillet_value_mm408 )409 elapsed = timeit.default_timer() - start410 print(f"[{session_id}] predict2 in {elapsed:.2f}s")411 412 return save_and_build_urls(413 session_id=session_id,414 output_image=out_img,415 outlines=outlines,416 dxf_path=dxf_path,417 mask=mask,418 endpoint_type="predict2",419 fillet_value=fillet_value_mm420 )421 422 except ReferenceBoxNotDetectedError:423 raise HTTPException(status_code=400, detail="Error detecting reference battery! Please try again with a clearer image.")424 except FingerCutOverlapError:425 raise HTTPException(status_code=400, detail="There was an overlap with fingercuts! Please try again to generate dxf.")426 except HTTPException as e:427 raise e428 except Exception as e:429 raise HTTPException(status_code=500, detail="Error detecting reference battery! Please try again with a clearer image.")430 431 432@app.post("/predict3")433async def predict3_api(434 file: UploadFile = File(...),435 enable_fillet: str = Form(..., regex="^(On|Off)$"),436 fillet_value_mm: float = Form(...),437 enable_finger_cut: str = Form(..., regex="^(On|Off)$")438):439 """440 Full predict: image + fillet toggle/value + finger-cut toggle → overlay, outlines, mask, DXF441 DXF naming format: DXF_DD-MM-YYYY_fillet-value_mm_fingercut-On|Off.dxf442 """443 session_id = str(uuid.uuid4())444 try:445 img_bytes = await file.read()446 image = np.array(Image.open(io.BytesIO(img_bytes)).convert("RGB"))447 except Exception:448 raise HTTPException(400, "Invalid image upload")449 450 try:451 start = timeit.default_timer()452 out_img, outlines, dxf_path, mask = predict_full(453 image, enable_fillet, fillet_value_mm, enable_finger_cut454 )455 elapsed = timeit.default_timer() - start456 print(f"[{session_id}] predict3 in {elapsed:.2f}s")457 458 return save_and_build_urls(459 session_id=session_id,460 output_image=out_img,461 outlines=outlines,462 dxf_path=dxf_path,463 mask=mask,464 endpoint_type="predict3",465 fillet_value=fillet_value_mm,466 finger_cut=enable_finger_cut467 )468 469 except ReferenceBoxNotDetectedError:470 raise HTTPException(status_code=400, detail="Error detecting reference battery! Please try again with a clearer image.")471 except FingerCutOverlapError:472 raise HTTPException(status_code=400, detail="There was an overlap with fingercuts! Please try again to generate dxf.")473 except HTTPException as e:474 raise e475 except Exception as e:476 raise HTTPException(status_code=500, detail="Error detecting reference battery! Please try again with a clearer image.")477 478 479@app.post("/update")480async def update_files(481 output_image: UploadFile = File(...),482 outlines_image: UploadFile = File(...),483 mask_image: UploadFile = File(...),484 dxf_file: UploadFile = File(...)485):486 session_id = str(uuid.uuid4())487 update_dir = os.path.join(UPDATES_DIR, session_id)488 os.makedirs(update_dir, exist_ok=True)489 490 try:491 upload_map = {492 "output_image": output_image,493 "outlines_image": outlines_image,494 "mask_image": mask_image,495 "dxf_file": dxf_file,496 }497 urls = {}498 for key, up in upload_map.items():499 fn = up.filename500 path = os.path.join(update_dir, fn)501 with open(path, "wb") as f:502 shutil.copyfileobj(up.file, f)503 urls[key] = f"{BASE_URL}/updates/{session_id}/{fn}"504 505 return {"session_id": session_id, "uploaded": urls}506 507 except Exception as e:508 raise HTTPException(500, f"Update failed: {e}")509 510 511from fastapi import Response512 513@app.get("/health")514def health():515 return Response(content="OK", status_code=200)516 517 518if __name__ == "__main__":519 import uvicorn520 port = int(os.environ.get("PORT", 8080))521 print(f"Starting FastAPI server on 0.0.0.0:{port}...")522 uvicorn.run(app, host="0.0.0.0", port=port)523 524 525 526