import gradio as gr
from transformers import pipeline

# Initialize the classification pipeline with your private model and token
pipe = pipeline(
    task="image-classification",
    model="scfive/weaponsclass",  # Replace with your model repository identifier
)

# Define custom CSS to style the header, footer, and main container
custom_css = """
/* Header styling */
#header {
    background-color: #003366;
    color: #ffffff;
    padding: 20px;
    text-align: center;
    font-family: 'Arial', sans-serif;
}
#header img {
    max-width: 100px;
    vertical-align: middle;
    margin-right: 15px;
}

/* Footer styling */
#footer {
    background-color: #f0f4f8;
    color: #003366;
    padding: 10px;
    text-align: center;
    font-family: 'Arial', sans-serif;
    font-size: 0.9em;
}

/* Main container styling */
.main-container {
    padding: 20px;
    margin: auto;
    max-width: 800px;
}

/* Button styling */
button {
    font-size: 1em;
    padding: 10px 20px;
}
"""

# Define the image classification function
def classify_image(image_path):
    result = pipe(image_path)
    # Return the label from the first result
    return result[0]['label']

# Build the Gradio Blocks layout
with gr.Blocks(css=custom_css) as demo:
    # Header: includes a logo and the app name
    gr.HTML("""
    <div id="header">
      <img src="https://via.placeholder.com/100" alt="Logo"/>
      <span style="font-size: 2em; font-weight: bold;">Weapon Classifier</span>
    </div>
    """)
    
    # Main content area with an image input and text output
    with gr.Column(elem_classes="main-container"):
        gr.HTML("<h2>Upload an image to classify the weapon</h2>")
        with gr.Row():
            image_input = gr.Image(type="filepath", label="Weapon Image")
            output_text = gr.Textbox(label="Classification Result")
        btn = gr.Button("Classify")
        btn.click(fn=classify_image, inputs=image_input, outputs=output_text)
    
    # Footer: explains the app is for game use only
    gr.HTML("""
    <div id="footer">
      This application is intended for game use only. © 2025 My Company.
    </div>
    """)

# Launch the app
demo.launch()