Here560/custom-model

1import os2import base643import io4from flask import Flask, request, jsonify5from ultralytics import YOLO6from PIL import Image7 8app = Flask(__name__)9 10model = YOLO("best.pt")11 12@app.route("/", methods=["GET"])13def health():14    return jsonify({ "status": "Food Detection API is running" })15 16@app.route("/detect", methods=["POST"])17def detect():18    try:19        data = request.get_json()20        if not data or "image" not in data:21            return jsonify({ "error": "No image provided" }), 40022 23        img_bytes = base64.b64decode(data["image"])24        img = Image.open(io.BytesIO(img_bytes)).convert("RGB")25 26        results = model(img)27        result = results[0]28 29        # Classification model (result.probs exists)30        if result.probs is not None:31            top_index = result.probs.top132            confidence = float(result.probs.top1conf)33            label = result.names[top_index]34            if confidence < 0.3:35                return jsonify({ "label": "Unknown", "confidence": confidence })36            return jsonify({ "label": label, "confidence": confidence })37 38        # Detection model (result.boxes exists — find highest confidence box)39        if result.boxes is not None and len(result.boxes) > 0:40            boxes = result.boxes41            best_idx = int(boxes.conf.argmax())42            confidence = float(boxes.conf[best_idx])43            class_id = int(boxes.cls[best_idx])44            label = result.names[class_id]45            if confidence < 0.3:46                return jsonify({ "label": "Unknown", "confidence": confidence })47            return jsonify({ "label": label, "confidence": confidence })48 49        # Nothing detected at all50        return jsonify({ "label": "Unknown", "confidence": 0.0 })51 52    except Exception as e:53        return jsonify({ "error": str(e) }), 50054 55 56if __name__ == "__main__":57    app.run(host="0.0.0.0", port=7860)