vop

app.py

@@ -71,16 +71,23 @@ def segment_endpoint():
         return {"error": "image file and prompt are required."}, 400
 
     prompt = request.form["prompt"]
-    image = Image.open(request.files["image"].stream).convert("RGB")
+    image = request.files["image"]
+
+    # Check for unsupported file types
+    if not image.content_type.startswith("image/"):
+        return {"error": "unsupported file type. Only image files are allowed."}, 415
 
     try:
-        mask = segment(image, prompt)
+        image_pil = Image.open(image.stream).convert("RGB")
+        mask = segment(image_pil, prompt)
     except ValueError as e:
         return {"error": str(e)}, 422
+    except Exception as e:
+        return {"error": "not supported"}, 500
 
     overlay = request.args.get("overlay", "false").lower() == "true"
     if overlay:
-        colored = np.array(image).copy()
+        colored = np.array(image_pil).copy()
         colored[mask] = [255, 0, 0]  # red overlay
         out_img = Image.fromarray(colored)
     else: