router fix

backend/app.py

@@ -1,19 +1,18 @@
-from fastapi import FastAPI, File, UploadFile
+from fastapi import FastAPI, File, UploadFile, Query
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.staticfiles import StaticFiles
-from PIL import Image
+from fastapi.responses import JSONResponse
+from fastapi.routing import APIRouter
+from typing import List
+import base64
+import gdown
 import io
 import os
+import pickle
 import time
 import numpy as np
-import pickle
-import base64
+from PIL import Image
 from io import BytesIO
-from fastapi.responses import JSONResponse
-from fastapi import Query
-from typing import List
-import gdown
-
 
 segmentationColors = [
     (255, 0, 0),
@@ -32,7 +31,6 @@ segmentationColors = [
 
 data_path = os.getenv('DATA_PATH')
 if not os.path.exists(data_path):
-    # raise FileNotFoundError(f"The data file at {data_path} was not found.")
     url = os.getenv('DATA_URL')
     file_id = url.split('/')[-2]
     direct_link = f"https://drive.google.com/uc?id={file_id}"
@@ -54,59 +52,33 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# app.mount("/", StaticFiles(directory="/app/frontend/build", html=True), name="frontend")
-# Serve the React frontend
+# Serve the React frontend if available
 frontend_path = "/app/frontend/build"
 if os.path.exists(frontend_path):
     app.mount("/", StaticFiles(directory=frontend_path, html=True), name="frontend")
 else:
     print(f"Warning: Frontend build directory '{frontend_path}' does not exist.")
 
-api = FastAPI()
+# Create a router for the API endpoints
+router = APIRouter()
 
 def overlay_mask(base_image, mask_image, color_idx):
-    """
-    Given the base_image and the 0/1 mask, overlay the mask with the color indexed by the color_idx.
-    """
-    # Convert inputs to NumPy arrays
     overlay = np.array(base_image, dtype=np.uint8)
     mask = np.array(mask_image).astype(bool)
-
     if overlay.shape[:2] != mask.shape:
         raise ValueError("Base image and mask must have the same dimensions.")
-
-    # Ensure color index is valid
     if not (0 <= color_idx < len(segmentationColors)):
         raise ValueError(f"Color index {color_idx} is out of bounds.")
-    
-    # Retrieve color
     color = np.array(segmentationColors[color_idx], dtype=np.uint8)
-    
-    # Print debugging info (optional)
-    print(f'Overlay shape: {overlay.shape}')
-    print(f'Mask shape: {mask.shape}')
-    print(f'Color idx: {color_idx}')
-    print(f'Color: {color}')
-    
-    # Apply color blending
     overlay[mask] = (overlay[mask] * 0.4 + color * 0.6).astype(np.uint8)
-    
-    # Convert back to Image
     return Image.fromarray(overlay)
 
-
 def convert_to_pil(image):
-    """
-    Ensure the image is a PIL Image.
-    """
     if isinstance(image, np.ndarray):
         return Image.fromarray(image)
     return image
 
 async def return_thumbnails():
-    """
-    Return a list of thumbnail images.
-    """
     thumbnails = []
     for item in data:
         pil_image = convert_to_pil(item['image'])
@@ -114,67 +86,56 @@ async def return_thumbnails():
     return thumbnails
 
 def rgb_to_hex(rgb):
-    """Convert an RGB tuple to a HEX string."""
     return "#{:02x}{:02x}{:02x}".format(rgb[0], rgb[1], rgb[2])
 
 async def return_state_data(state):
-    print(state)
-    """
-    Return state-specific data including overlays and object validity.
-    """
     image_data = data[state['image_index']]
     base_image = convert_to_pil(image_data['image'])
-
     response = {
         'mask_overlayed_image': base_image,
         'valid_object_color_tuples': [],
         'invalid_objects': []
     }
-
     mask_data = image_data['mask_data'].get(state['detail_level'], {})
-    # print(mask_data)
-
     for object_type, mask_info in mask_data.items():
         if mask_info['valid']:
             idx = len(response['valid_object_color_tuples'])
             if idx in state['object_list']:
-                print(f"Overlaying mask for {object_type} with color {segmentationColors[idx]} as {idx} is in {state['object_list']}")
-                response['mask_overlayed_image'] = overlay_mask(response['mask_overlayed_image'], mask_info['mask'], idx)
-            # append the color at idx in hex
+                response['mask_overlayed_image'] = overlay_mask(
+                    response['mask_overlayed_image'],
+                    mask_info['mask'], idx
+                )
             color = segmentationColors[idx]
             response['valid_object_color_tuples'].append((object_type, rgb_to_hex(color)))
         else:
             response['invalid_objects'].append(object_type)
-
     buffer = BytesIO()
-    response['mask_overlayed_image'].save(buffer, format="PNG")  # Save image to a buffer in PNG format
-    base64_str = base64.b64encode(buffer.getvalue()).decode("utf-8")  # Encode to Base64
+    response['mask_overlayed_image'].save(buffer, format="PNG")
+    base64_str = base64.b64encode(buffer.getvalue()).decode("utf-8")
     response['mask_overlayed_image'] = base64_str
     return response
 
-@api.get("/return_thumbnails")
+@router.get("/return_thumbnails")
 async def return_thumbnails_endpoint():
     thumbnails = await return_thumbnails()
     encoded_images = []
     for thumbnail in thumbnails:
         buffer = BytesIO()
-        thumbnail.save(buffer, format="PNG")  # Save image to a buffer in PNG format
-        base64_str = base64.b64encode(buffer.getvalue()).decode("utf-8")  # Encode to Base64
+        thumbnail.save(buffer, format="PNG")
+        base64_str = base64.b64encode(buffer.getvalue()).decode("utf-8")
         encoded_images.append(base64_str)
     return JSONResponse(content={"thumbnails": encoded_images})
 
-@api.get("/return_state_data")
+@router.get("/return_state_data")
 async def return_state_data_endpoint(
     image_index: int = Query(...),
     detail_level: int = Query(...),
     object_list: str = Query(...)
 ):
-    print(object_list)
     if object_list == 'None':
         object_list = []
     else:
         object_list = [int(x) for x in object_list.split(",")]
-        print(object_list)
     state = {
         "image_index": image_index,
         "detail_level": detail_level,
@@ -183,8 +144,9 @@ async def return_state_data_endpoint(
     response = await return_state_data(state)
     return response
 
-app.mount("/api", api)
+# Include the router with a prefix, making endpoints accessible under /api
+app.include_router(router, prefix="/api")
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=7860)