app.py

html += "<h3>False Negatives (Real images classified as AI):</h3>"
    html += '<div class="image-grid">'
    for img_name, img_data in false_negatives:
        html += f'''
        <div class="image-item">
            <img src="data:image/jpeg;base64,{img_data}" alt="{img_name}">
            <p>{img_name}</p>
        </div>
        '''
    html += '</div>'

    return html

def load_url(url):
    try:
        urllib.request.urlretrieve(url, "temp_image.png")
        image = Image.open("temp_image.png")
        message = "Image Loaded"
    except Exception as e:
        image = None
        message = f"Image not Found<br>Error: {e}"
    return image, message

detector = AIDetector()

def create_gradio_interface():
    with gr.Blocks() as app:
        gr.Markdown("""<center><h1>AI Image Detector</h1></center>""")

        with gr.Tabs():
            with gr.Tab("Single Image Detection"):
                with gr.Column():
                    inp = gr.Image(type='pil')
                    in_url = gr.Textbox(label="Image URL")
                    with gr.Row():
                        load_btn = gr.Button("Load URL")
                        btn = gr.Button("Detect AI")
                    message = gr.HTML()

                with gr.Group():
                    with gr.Box():
                        gr.HTML(f"""<b>Testing on Model: <a href='https://huggingface.co/{MODEL_NAME}'>{MODEL_NAME}</a></b>""")
                        output_html = gr.HTML()
                        output_label = gr.Label(label="Output")

            with gr.Tab("Batch Image Processing"):
                with gr.Accordion("Upload Zip File (max 100MB)", open=False):
                    zip_file = gr.File(
                        label="Upload Zip (must contain 'real' and 'ai' folders)",
                        file_types=[".zip"],
                        file_count="single",
                        max_file_size=100  # 100 MB limit
                    )
                    zip_process_btn = gr.Button("Process Zip", interactive=False)

                with gr.Accordion("Upload Individual Files (for datasets over 100MB)", open=False):
                    with gr.Row():
                        ai_files = gr.File(
                            label="Upload AI Images",
                            file_types=["image"],
                            file_count="multiple"
                        )
                        real_files = gr.File(
                            label="Upload Real Images",
                            file_types=["image"],
                            file_count="multiple"
                        )
                    individual_process_btn = gr.Button("Process Individual Files", interactive=False)

                with gr.Group():
                    gr.Markdown(f"### Results for {MODEL_NAME}")
                    output_acc = gr.Label(label="Accuracy")
                    output_roc = gr.Label(label="ROC Score")
                    output_report = gr.HTML(label="Classification Report")
                    output_plots = gr.Plot(label="Confusion Matrix and ROC Curve")
                    output_fp_fn = gr.HTML(label="False Positives and Negatives")
                    
                    # Add export button and PDF output
                    export_btn = gr.Button("Export Results to PDF", variant="primary")
                    pdf_output = gr.File(label="Downloaded PDF")

                reset_btn = gr.Button("Reset")

        load_btn.click(load_url, in_url, [inp, message])
        btn.click(
            lambda img: detector.predict(img),
            inp,
            [output_html, output_label]
        )

        def enable_zip_btn(file):
            return gr.Button.update(interactive=file is not None)

        def enable_individual_btn(ai_files, real_files):
            return gr.Button.update(interactive=(ai_files is not None and real_files is not None))

        zip_file.upload(enable_zip_btn, zip_file, zip_process_btn)
        
        ai_files.upload(enable_individual_btn, [ai_files, real_files], individual_process_btn)
        real_files.upload(enable_individual_btn, [ai_files, real_files], individual_process_btn)

        zip_process_btn.click(
            process_zip,
            zip_file,
            [output_acc, output_roc, output_report, output_plots, output_fp_fn]
        )

        individual_process_btn.click(
            process_files,
            [ai_files, real_files],
            [output_acc, output_roc, output_report, output_plots, output_fp_fn]
        )

        # Add export button click handler
        export_btn.click(
            export_to_pdf,
            inputs=[output_acc, output_roc, output_report, output_plots, output_fp_fn],
            outputs=pdf_output
        )

        def reset_interface():
            return [
                None, None, None, None, None,  # Reset inputs
                gr.Button.update(interactive=False),  # Reset zip process button
                gr.Button.update(interactive=False),  # Reset individual process button
                None, None, None, None, None, None  # Reset outputs (including PDF)
            ]

        reset_btn.click(
            reset_interface,
            inputs=None,
            outputs=[
                zip_file, ai_files, real_files,
                output_acc, output_roc, output_report, output_plots, output_fp_fn,
                zip_process_btn, individual_process_btn, pdf_output
            ]
        )

    return app

if __name__ == "__main__":
    app = create_gradio_interface()
    app.launch(
        show_api=False,
        max_threads=24,
        show_error=True
    )