Update app.py

Multiple frames and data stats

app.py

@@ -7,20 +7,26 @@ from pydantic import BaseModel
 import os
 import io
 import base64
+from typing import List
+from fastapi.openapi.docs import get_swagger_ui_html
 
 # Set JAX to use CPU platform (adjust if GPU is needed)
 os.environ['JAX_PLATFORMS'] = 'cpu'
 
-# Load the model once globally (assumes it's cached locally)
+# Load the model once globally
 model = OctoModel.load_pretrained("hf://rail-berkeley/octo-small-1.5")
 
 # Initialize FastAPI app
-app = FastAPI(title="Octo Model Inference API")
+app = FastAPI(
+    title="Octo Model Inference API",
+    docs_url="/"  # Swagger UI at root
+)
 
 # Define request body model
 class InferenceRequest(BaseModel):
-    image_base64: str  # Base64-encoded image string
+    image_base64: List[str]  # List of base64-encoded images in time sequence
     task: str = "pick up the fork"  # Default task
+    window_size: int = 2  # Default window size, configurable
 
 # Health check endpoint
 @app.get("/health")
@@ -29,37 +35,54 @@ async def health_check():
 
 # Inference endpoint
 @app.post("/predict")
-async def predict(request: InferenceRequest):
+async def predict(request: InferenceRequest, dataset_name: str = "bridge_dataset"):
     try:
-        # Decode base64 image
-        img_base64 = request.image_base64
-        if img_base64.startswith("data:image"):
-            img_base64 = img_base64.split(",")[1]
-        
-        img_data = base64.b64decode(img_base64)
-        img = Image.open(io.BytesIO(img_data)).resize((256, 256))
-        img = np.array(img)
+        # Validate input
+        if len(request.image_base64) < request.window_size:
+            raise HTTPException(
+                status_code=400,
+                detail=f"At least {request.window_size} images required for the specified window size"
+            )
 
-        # Add batch and time horizon dimensions
-        img = img[np.newaxis, np.newaxis, ...]  # Shape: (1, 1, 256, 256, 3)
+        # Process images
+        images = []
+        for img_base64 in request.image_base64:
+            if img_base64.startswith("data:image"):
+                img_base64 = img_base64.split(",")[1]
+            img_data = base64.b64decode(img_base64)
+            img = Image.open(io.BytesIO(img_data)).resize((256, 256))
+            img = np.array(img)
+            images.append(img)
+
+        # Stack all images and add batch dimension
+        img_array = np.stack(images)[np.newaxis, ...]  # Shape: (1, T, 256, 256, 3)
         observation = {
-            "image_primary": img,
-            "timestep_pad_mask": np.array([[True]])
+            "image_primary": img_array,
+            "timestep_pad_mask": np.full((1, len(images)), True, dtype=bool)  # Shape: (1, T)
         }
 
         # Create task and predict actions
         task_obj = model.create_tasks(texts=[request.task])
         actions = model.sample_actions(
-            observation, 
-            task_obj, 
-            unnormalization_statistics=model.dataset_statistics["bridge_dataset"]["action"], 
+            observation,
+            task_obj,
+            unnormalization_statistics=model.dataset_statistics[dataset_name]["action"],
             rng=jax.random.PRNGKey(0)
         )
-        actions = actions[0]
+        actions = actions[0]  # Remove batch dimension, Shape: (horizon, action_dim)
 
-        # Convert NumPy array to list for JSON response
+        # Convert to list for JSON response
         actions_list = actions.tolist()
 
         return {"actions": actions_list}
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error processing request: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Error processing request: {str(e)}")
+
+# Custom Swagger UI route (optional)
+@app.get("/docs", include_in_schema=False)
+async def custom_swagger_ui_html():
+    return get_swagger_ui_html(
+        openapi_url=app.openapi_url,
+        title=app.title + " - Swagger UI",
+        oauth2_redirect_url=app.swagger_ui_oauth2_redirect_url,
+    )