app.py

import gradio as gr
from dotenv import load_dotenv
from roboflow import Roboflow
import tempfile
import os

# Load environment variables from .env file
load_dotenv()
api_key = os.getenv("ROBOFLOW_API_KEY")
workspace = os.getenv("ROBOFLOW_WORKSPACE")
project_name = os.getenv("ROBOFLOW_PROJECT")
model_version = int(os.getenv("ROBOFLOW_MODEL_VERSION"))

# Initialize Roboflow using the loaded environment variables
rf = Roboflow(api_key=api_key)
project = rf.workspace(workspace).project(project_name)
model = project.version(model_version).model

# Function to handle image input and output
def detect_objects(image):
    # Save the uploaded image as a temporary file
    with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file:
        image.save(temp_file, format="JPEG")
        temp_file_path = temp_file.name

    # Perform prediction on the image
    predictions = model.predict(temp_file_path, confidence=60, overlap=80).json()
    
    # Count the number of objects per class
    class_count = {}
    total_count = 0  # Store the total number of objects

    for prediction in predictions['predictions']:
        class_name = prediction['class']
        if class_name in class_count:
            class_count[class_name] += 1
        else:
            class_count[class_name] = 1
        total_count += 1  # Increment total object count for each prediction

    # Prepare the result text
    result_text = "Product Nestle\n\n"
    
    for class_name, count in class_count.items():
        result_text += f"{class_name}: {count} \n"
    
    result_text += f"\nTotal Product Nestle: {total_count}"
    
    # Save the image with predictions
    output_image_path = "/tmp/prediction.jpg"
    model.predict(temp_file_path, confidence=60, overlap=80).save(output_image_path)
    
    # Remove the temporary file after prediction
    os.remove(temp_file_path)
    
    return output_image_path, result_text

# Create the Gradio interface
with gr.Blocks() as iface:
    with gr.Row():
        image_input = gr.Image(type="pil", label="Input Image")
    with gr.Row():
        with gr.Column():
            image_output = gr.Image(label="Detect Object")
        with gr.Column():
            text_output = gr.Textbox(label="Counting Object")
    gr.Interface(
        fn=detect_objects,
        inputs=image_input,
        outputs=[image_output, text_output],
    )

# Launch the interface
iface.launch()