implement image alignment api endpoint

CDCgov · Oct 31, 2024 · 9d0e2bb · 9d0e2bb
1 parent 09c1a98
commit 9d0e2bb
Showing 1 changed file with 13 additions and 0 deletions.
diff --git a/OCR/ocr/api.py b/OCR/ocr/api.py
@@ -9,6 +9,7 @@
 from fastapi.middleware.cors import CORSMiddleware
 
 from ocr.services.image_ocr import ImageOCR
+from ocr.services.alignment import ImageAligner
 from ocr.services.image_segmenter import ImageSegmenter, segment_by_color_bounding_box
 
 app = FastAPI()
@@ -36,6 +37,18 @@ async def health_check():
     return {"status": "UP"}
 
 
+@app.post("/image_alignment/")
+async def image_alignment(source_image: UploadFile, segmentation_template: UploadFile):
+    source_image_np = np.frombuffer(await source_image.read(), np.uint8)
+    source_image_img = cv.imdecode(source_image_np, cv.IMREAD_COLOR)
+
+    segmentation_template_np = np.frombuffer(await segmentation_template.read(), np.uint8)
+    segmentation_template_img = cv.imdecode(segmentation_template_np, cv.IMREAD_COLOR)
+
+    aligner = ImageAligner()
+    return aligner.align(source_image_img, segmentation_template_img)
+
+
 @app.post("/image_file_to_text/")
 async def image_file_to_text(source_image: UploadFile, segmentation_template: UploadFile, labels: str = Form()):
     source_image_np = np.frombuffer(await source_image.read(), np.uint8)