Update app.py

app.py

@@ -1,31 +1,27 @@
 from fastapi import FastAPI, UploadFile, File
-from fastapi.responses import JSONResponse
+from fastapi.middleware.cors import CORSMiddleware
 from PIL import Image
-import torch, torchvision.transforms as T
-from transformers import MobileNetV2ForSemanticSegmentation
-import io
-
-# Load the model
-model = MobileNetV2ForSemanticSegmentation.from_pretrained("seg_model")
-model.eval()
-
-preprocess = T.Compose([
-    T.Resize(513),
-    T.ToTensor(),
-    T.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225])
-])
+from io import BytesIO
+import numpy as np
+import tensorflow as tf
 
 app = FastAPI()
 
-@app.get("/")
-def root():
-    return {"status": "API up for segmentation"}
+# CORS config (optional)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
 
 @app.post("/predict")
 async def predict(file: UploadFile = File(...)):
-    img = Image.open(await file.read()).convert("RGB")
-    x = preprocess(img).unsqueeze(0)
-    with torch.no_grad():
-        outputs = model(x).logits
-        seg = outputs.argmax(1)[0].tolist()
-    return JSONResponse(content={"segmentation_mask": seg})
+    contents = await file.read()
+    img = Image.open(BytesIO(contents)).convert("RGB")
+    img = img.resize((256, 256))  # or whatever your model expects
+    arr = np.array(img) / 255.0
+    arr = np.expand_dims(arr, 0)
+    # prediction = model.predict(arr)
+    # result = do_something_with_prediction(prediction)
+    return {"success": True}