model_api.py

import cv2
import base64
import json
import time
from collections import Counter
from fastapi import FastAPI
from fastapi.responses import JSONResponse
from ultralytics import YOLO
import numpy as np

# Initialize FastAPI app and YOLO model
app = FastAPI()
model = YOLO("yolov8n.pt")

# Load your product database (example)
with open("database.json", "r") as f:
    product_db = json.load(f)

# Function to detect products, annotate the image, and save it
def detect_products():
    # Open the webcam feed
    time.sleep(5)
    cap = cv2.VideoCapture(0)

    # Check if the camera opened successfully
    if not cap.isOpened():
        return {"error": "Failed to open camera"}

    # Capture a frame from the camera
    ret, frame = cap.read()
    
    # Check if the frame was successfully captured
    if not ret:
        cap.release()
        return {"error": "Failed to capture frame"}

    # If the frame is blank, report it
    if np.all(frame == 0):
        cap.release()
        return {"error": "Captured frame is blank"}

    # Run YOLO model on the frame
    results = model(frame)
    boxes = results[0].boxes
    names = results[0].names

    product_counts = Counter()

    # Loop through detected boxes and count products
    for box in boxes:
        class_id = int(box.cls[0])
        label = names[class_id]
        product_name = label.split("_")[0]
        if product_name in product_db:
            product_counts[product_name] += 1

    # Annotate the frame with bounding boxes and labels
    annotated_frame = results[0].plot()

    # Save the annotated image to disk
    output_image_path = "output_image.jpg"
    cv2.imwrite(output_image_path, annotated_frame)
    print(f"Annotated image saved to {output_image_path}")

    # Convert annotated frame to base64
    _, encoded_image = cv2.imencode('.jpg', annotated_frame)
    base64_image = base64.b64encode(encoded_image).decode('utf-8')

    # Prepare response data (including the detected products)
    data = []
    for product, count in product_counts.items():
        product_info = product_db.get(product, {})
        data.append({
            "product": product,
            "category": product_info.get("category", "N/A"),
            "price": product_info.get("price", "N/A"),
            "count": count
        })

    # Return both the product data and the base64 image
    return {
        "result": data,
        #"image": base64_image,
        #"image_path": output_image_path  # Include the path to the saved image
    }

@app.get("/predict")
def predict():
    # Detect products and get the annotated image with bounding boxes
    result = detect_products()

    # If the result contains an error, return it
    if "error" in result:
        return JSONResponse(status_code=400, content=result)

    # Return the detected products, base64-encoded image, and image file path
    return JSONResponse(content=result)