api/hud_server.py

import io, json
from fastapi import FastAPI, File, UploadFile, Form, HTTPException, Request
from pydantic import BaseModel
from fastapi.responses import StreamingResponse, JSONResponse
from fastapi.middleware.cors import CORSMiddleware
import numpy as np
from PIL import Image
import base64
from typing import List, Optional
import torch
from core.processing import embed_text, get_dino_boxes_from_prompt, get_sam_mask, expand_coords_shape, embed_image_dino_large
from core.storage import query_vector_db_by_image_embedding, query_vector_db_by_text_embedding, get_object_info_from_graph, set_object_primary_location_hierarchy, get_object_location_chain
from core.storage import get_object_owners, add_owner_by_person_id, add_owner_by_person_name, get_all_locations_for_house
from core.image_processing import apply_mask, crop_to_mask_size, encode_image_to_base64

app = FastAPI()
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])


class Point(BaseModel):
    x: float
    y: float

class Point3D(BaseModel):
    x: float
    y: float
    z: float

class MaskRequest(BaseModel):
    image_base64: str  # base64 encoded PNG image
    points: List[Point]
    labels: List[int]
    prompt: str
    return_raw_mask: bool = False
    return_rgb_mask: bool = False
    return_embeddings: bool = False

class BoundingBox(BaseModel):
    x: int
    y: int
    width: int
    height: int

class MaskResponse(BaseModel):
    raw_mask_base64: str
    rgb_mask_base64: str
    embedding: List[float]
    bounding_box: BoundingBox

class ObjectQueryByEmbeddingRequest(BaseModel):
    embedding_type: str  # "image" or "text"
    embedding: List[float]
    k: int = 5  # default to 5 if not specified
    house_id: Optional[str] = None  # Optional house ID to filter results

class ObjectQueryByDescriptionRequest(BaseModel):
    description: str
    k: int = 5
    house_id: str = None  # Optional house ID to filter results

class ObjectQueryResultEntry(BaseModel):
    object_id: str
    aggregated_similarity: float
    probability: float
    descriptions: List[str]

class ObjectInfoRequest(BaseModel):
    house_id: str
    object_id: str

class ObjectInfoResponse(BaseModel):
    object_id: str
    house_id: str
    description: str

class SetPrimaryLocationRequest(BaseModel):
    house_id: str
    object_id: str
    location_hierarchy: List[str]  # Example: ["Kitchen", "Left Upper Cabinet", "Middle Shelf"]

class ObjectLocationRequest(BaseModel):
    house_id: str
    object_id: Optional[str] = None
    include_images: bool = False

class LocationInfo(BaseModel):
    name: str
    image_uri: Optional[str] = None
    image_base64: Optional[str] = None
    location_x: Optional[float] = 0
    location_y: Optional[float] = 0
    location_z: Optional[float] = 0
    shape: Optional[str] = None
    radius: Optional[float] = 0
    height: Optional[float] = 0
    width: Optional[float] = 0
    depth: Optional[float] = 0


class ObjectLocationResponse(BaseModel):
    object_id: Optional[str] = None
    house_id: str
    locations: List[LocationInfo]

class Person(BaseModel):
    person_id: str
    name: Optional[str]
    nickname: Optional [str]
    age: Optional[int]
    type: str = "person"  # e.g., "person", "dog", "robot", etc.
    image_uri: Optional[str] = None

class ObjectOwnersResponse(BaseModel):
    object_id: str
    house_id: str
    owners: List[Person] # Or a more complex model if needed

class AddOwnerByIdRequest(BaseModel):
    house_id: str
    object_id: str
    person_id: str

class AddOwnerByNameRequest(BaseModel):
    house_id: str
    object_id: str
    name: str
    type: Optional[str] = "person"



@app.middleware("http")
async def log_requests(request: Request, call_next):
    print(f"[REQ] {request.method} {request.url}")
    return await call_next(request)

@app.get("/")
async def root():
    return {"message": "Hello, World!"}

@app.post("/object/log_location", response_model=str)
async def log_location(request: Point3D):
    try:
        print(
            f"[LogLocation] "
            f"x:{request.x:.2f} "
            f"y:{request.y:.2f} "
            f"z:{request.z:.2f}"
        )


        response = "log location successful"

        return response

    except Exception as e:
        raise HTTPException(500, f"log location failed: {str(e)}")


@app.post("/object/get_mask", response_model=MaskResponse)
async def mask_endpoint(request: MaskRequest):
    try:
        # Decode base64 image
        image_bytes = base64.b64decode(request.image_base64)
        img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
        img_np = np.array(img)

        # Convert points to numpy array
        point_coords = np.array([[p.x, p.y] for p in request.points], dtype=np.float32)
        point_labels = np.array(request.labels, dtype=np.int32)

        # Optionally get bounding boxes if a prompt is provided
        sam_boxes = None
        if request.prompt:
            sam_boxes = get_dino_boxes_from_prompt(img_np, request.prompt)
            point_coords, point_labels = expand_coords_shape(point_coords, point_labels, sam_boxes.shape[0])

        # Generate the mask
        mask, bbox = get_sam_mask(img_np, None, None, sam_boxes)
        mask_img = apply_mask(img_np, mask, "remove")
        mask_img = crop_to_mask_size(mask_img, mask)

        # Encode images to base64
        mask_raw_base64 = encode_image_to_base64(mask * 255) if request.return_raw_mask else ""
        masked_rgb_base64 = encode_image_to_base64(mask_img) if request.return_rgb_mask else ""
        embedding = embed_image_dino_large(mask_img).tolist() if request.return_embeddings else None  

        response = MaskResponse(
                raw_mask_base64=mask_raw_base64,
                rgb_mask_base64=masked_rgb_base64,
                embedding=embedding,
                bounding_box=BoundingBox(**bbox)
            )

        return response

    except Exception as e:
        raise HTTPException(500, f"Mask generation failed: {str(e)}")
    


@app.post("/object/query_by_embedding", response_model=List[ObjectQueryResultEntry])
async def query_by_embedding(query: ObjectQueryByEmbeddingRequest):
    try:
        k = 5 #query.k
        if query.embedding_type == "text":
            query_vector = np.array(query.embedding, dtype=np.float32)
            results = query_vector_db_by_text_embedding(query_vector, k, query.house_id)
        elif query.embedding_type == "image": 
            query_vector = np.array(query.embedding, dtype=np.float32)
            results = query_vector_db_by_image_embedding(query_vector, k, query.house_id)
        else:
            raise HTTPException(status_code=400, detail="Invalid embedding type. Use 'text' or 'image'.")

        object_scores = {}
        object_views = {}
        for result in results:
            obj_id = result.payload.get("object_id")
            score = result.score
            desc = result.payload.get("description") or "No description available"
            if obj_id in object_scores:
                object_scores[obj_id] = max(object_scores[obj_id], score)
                object_views[obj_id].append(desc)
            else:
                object_scores[obj_id] = score
                object_views[obj_id] = [desc]

        all_scores = np.array(list(object_scores.values()))
        exp_scores = np.exp(all_scores)
        probabilities = exp_scores / np.sum(exp_scores) if np.sum(exp_scores) > 0 else np.zeros_like(exp_scores)

        results = []
        for i, (obj_id, score) in enumerate(object_scores.items()):

            
            results.append(ObjectQueryResultEntry(
                object_id=obj_id,
                aggregated_similarity=float(score),
                probability=float(probabilities[i]),
                descriptions=object_views[obj_id]
            ))

        return results
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/object/query_by_description", response_model=List[ObjectQueryResultEntry])
async def query_by_description(query: ObjectQueryByDescriptionRequest):
    try:
        # Embed the description to get the text embedding
        embedding_vector = embed_text(query.description)

        # Call your existing embedding-based query
        embedding_request = ObjectQueryByEmbeddingRequest(
            embedding_type="text",
            embedding=embedding_vector.tolist(),
            k=query.k,
            house_id=query.house_id
        )

        return await query_by_embedding(embedding_request)

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/object/get_info", response_model=ObjectInfoResponse)
async def get_object_info_endpoint(request: ObjectInfoRequest):
    description = get_object_info_from_graph(request.house_id, request.object_id)
    if description is None:
        raise HTTPException(status_code=404, detail="Object not found in household")

    return ObjectInfoResponse(
        object_id=request.object_id,
        house_id=request.house_id,
        description=description
    )


@app.post("/object/set_primary_location")
async def set_primary_location(request: SetPrimaryLocationRequest):
    try:
        set_object_primary_location_hierarchy(
            object_id=request.object_id,
            house_id=request.house_id,
            location_hierarchy=request.location_hierarchy
        )
        return {"status": "success", "message": f"Primary location set for object {request.object_id}"}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/object/get_primary_location", response_model=ObjectLocationResponse)
async def get_object_location(request: ObjectLocationRequest):
    try:
        locations = get_object_location_chain(
            house_id=request.house_id,
            object_id=request.object_id,
            include_images=request.include_images
        )
        return ObjectLocationResponse(
            object_id=request.object_id,
            house_id=request.house_id,
            locations=locations
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/house/get_all_locations", response_model=ObjectLocationResponse)
async def get_object_location(request: ObjectLocationRequest):
    try:
        locations = get_all_locations_for_house(
            house_id=request.house_id,
            include_images=request.include_images
        )
        return ObjectLocationResponse(
            house_id=request.house_id,
            locations=locations
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/object/get_owners", response_model=ObjectOwnersResponse)
async def get_object_owners_handler(request: ObjectLocationRequest):
    try:
        owners = get_object_owners(
            house_id=request.house_id,
            object_id=request.object_id
        )
        return ObjectOwnersResponse(
            object_id=request.object_id,
            house_id=request.house_id,
            owners=owners
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/object/add_owner_by_id", response_model=Person)
async def api_add_owner_by_id(request: AddOwnerByIdRequest):
    try:
        p = add_owner_by_person_id(
            house_id=request.house_id,
            object_id=request.object_id,
            person_id=request.person_id
        )
        if not p:
            raise HTTPException(status_code=404, detail="Person not found.")
        return p
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
@app.post("/object/add_owner_by_name", response_model=Person)
async def api_add_owner_by_name(request: AddOwnerByNameRequest):
    try:
        p = add_owner_by_person_name(
            house_id=request.house_id,
            object_id=request.object_id,
            name=request.name,
            type=request.type
        )
        if not p:
            raise HTTPException(status_code=500, detail="Failed to create owner.")
        return p
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


if __name__ == "__main__":
    import uvicorn
    uvicorn.run("api.hud_server:app", host="0.0.0.0", port=8000)