import gradio as gr
import subprocess
import os

def crowd_counting(image):
    # Save the uploaded image
    image_path = "test/uploaded.jpg"
    image.save(image_path)

    # Run the crowd counting model using subprocess
    
    
    command = "python3 detect.py --weights weights/crowdhuman_yolov5m.pt --source {} --head --project runs/output --exist-ok".format(image_path)
    subprocess.run(command, shell=True)

    # Read the total_boxes from the file
    total_boxes_path = "runs/output/output.txt"
    with open(total_boxes_path, "r") as f:
        total_boxes = f.read()

    # Get the output image
    output_image = "runs/output/output.jpg"

    # Return the output image and total_boxes
    return output_image, total_boxes

# Define the input and output interfaces
inputs = gr.inputs.Image(type="pil", label="Input Image")
outputs = [gr.outputs.Image(type="pil", label="Output Image"), gr.outputs.Textbox(label="Total (Head) Count")]

# Define the title and description
title = "Crowd Counting"
description = "<div style='text-align: center;'>This is a crowd counting application that uses a deep learning model to count the number of heads in an image.<br>Made by HTX (Q3) </div>"

# Create the Gradio interface without the flag button
gradio_interface = gr.Interface(fn=crowd_counting, inputs=inputs, outputs=outputs, title=title, description=description, allow_flagging="never")

# Run the Gradio interface
gradio_interface.launch()