Lookii125/image_editing_agent

1import io2import cv23from PIL import Image4from fastapi import FastAPI, UploadFile, File, Form, Body5from fastapi.responses import StreamingResponse6from fastapi.staticfiles import StaticFiles7from pydantic import BaseModel8from dotenv import load_dotenv9 10from langchain_core.messages import HumanMessage, AIMessage11from detection_image_state import ImageState12from Detection_agent import compiled_agent, register_img_state13 14load_dotenv()15 16app = FastAPI()17 18sessions = {}19 20def get_state(session_key: str) -> ImageState:21    if session_key not in sessions:22        sessions[session_key] = {23            "state": ImageState(),24            "chat_history": []25        }26        register_img_state(session_key, sessions[session_key]["state"])27    return sessions[session_key]["state"]28 29def get_history(session_key: str):30    return sessions.get(session_key, {}).get("chat_history", [])31 32@app.get("/api/state")33def read_state(session_key: str):34    state = get_state(session_key)35    return {36        "has_image": state.current is not None,37        "history_index": state.history_index,38        "has_detections": bool(state.detections),39        "detections_stale": state.detections_are_stale40    }41 42class ActionRequest(BaseModel):43    session_key: str44    action: str45 46@app.post("/api/action")47def perform_action(req: ActionRequest):48    state = get_state(req.session_key)49    msg = ""50    if req.action == "undo":51        if state.undo():52            if hasattr(state, "code_log") and len(state.code_log) > 0:53                state.code_log.pop()54            msg = "Popped latest modification state!"55        else:56            msg = "Workspace history clear."57    elif req.action == "back":58        msg = f"Navigated backward to step {state.history_index}" if state.go_back() else "At original baseline."59    elif req.action == "forward":60        msg = f"Navigating forward to step {state.history_index}" if state.go_forward() else "At latest step."61    elif req.action == "clear":62        state.reset()63        sessions[req.session_key]["chat_history"] = []64        msg = "Canvas Cleared"65    66    return {"message": msg}67 68@app.post("/api/upload")69async def upload_image(session_key: str = Form(...), file: UploadFile = File(...)):70    state = get_state(session_key)71    contents = await file.read()72    pil_img = Image.open(io.BytesIO(contents))73    state.load_image(pil_img)74    sessions[session_key]["chat_history"] = []75    return {"status": "success"}76 77@app.get("/api/image/{img_type}")78def get_image(img_type: str, session_key: str, show_detections: str = "false"):79    state = get_state(session_key)80    if state.current is None:81        return {"error": "No image loaded"}82 83    if img_type == "original":84        img_bgr = state.original85    else:86        if show_detections.lower() == "true" and state.detections and not state.detections_are_stale:87            img_bgr = state.draw_all_detections_op()88        else:89            img_bgr = state.current90 91    out_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)92    out_pil = Image.fromarray(out_rgb)93    buf = io.BytesIO()94    out_pil.save(buf, format="PNG")95    buf.seek(0)96    return StreamingResponse(buf, media_type="image/png")97 98@app.get("/api/code")99def get_code(session_key: str):100    state = get_state(session_key)101    return {"code_log": state.code_log}102 103class ChatRequest(BaseModel):104    session_key: str105    query: str106 107@app.post("/api/chat")108def process_chat(req: ChatRequest):109    state = get_state(req.session_key)110    history = get_history(req.session_key)111    112    if state.current is None:113        return {"response": "Please upload an image first."}114        115    try:116        short_context = []117        if len(history) >= 2:118            for entry in history[-2:]:119                if entry["role"] == "user": short_context.append(HumanMessage(content=entry["content"]))120                else: short_context.append(AIMessage(content=entry["content"]))121 122        graph_inputs = {123            "input": req.query,124            "img_state_ref": req.session_key,125            "chat_history": short_context,126            "messages": [HumanMessage(content=req.query)],127            "parsed_action": {},128            "output": "",129        }130 131        graph_output = compiled_agent.invoke(graph_inputs)132        133        if graph_output.get("output"):134            ai_res = graph_output["output"]135        else:136            final_action = graph_output.get("parsed_action", {})137            ai_res = final_action.get("response", "Processing complete.")138            139    except Exception as e:140        ai_res = f"Pipeline execution fault: {str(e)}"141        142    history.append({"role": "user", "content": req.query})143    history.append({"role": "assistant", "content": ai_res})144    145    return {"response": ai_res}146 147app.mount("/", StaticFiles(directory="static", html=True), name="static")148 149if __name__ == "__main__":150    import uvicorn151    uvicorn.run(app, host="0.0.0.0", port=7860)