Update app.py

app.py

@@ -201,18 +201,20 @@ def generate_pdf(accuracy, roc_score, report_html, confusion_matrix_plot):
     buffer = BytesIO()
     c = canvas.Canvas(buffer, pagesize=letter)
     
+    # Add content to PDF
     c.drawString(100, 750, f"Model Results")
     c.drawString(100, 730, f"Accuracy: {accuracy:.2f}")
     c.drawString(100, 710, f"ROC Score: {roc_score:.2f}")
     
     y_position = 690
-    for line in report_html.split('<tr>')[2:]:
+    for line in report_html.replace("<br>", "\n").splitlines():
         if y_position < 50:
             c.showPage()
             y_position = 750
         c.drawString(100, y_position, line.strip())
         y_position -= 20
 
+    # Save Confusion Matrix Plot as an Image and Add it to the PDF
     img_buffer = BytesIO()
     confusion_matrix_plot.savefig(img_buffer, format="png")
     img_buffer.seek(0)
@@ -230,47 +232,74 @@ def create_gradio_interface():
 
         with gr.Tabs():
             with gr.Tab("Single Image Detection"):
-                inp = gr.Image(type='pil')
-                in_url = gr.Textbox(label="Image URL")
-                load_btn = gr.Button("Load URL")
-                btn = gr.Button("Detect AI")
-                message = gr.HTML()
-
-                output_html = gr.HTML()
-                output_label = gr.Label(label="Output")
+                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"):
-                zip_file = gr.File(label="Upload Zip", file_types=[".zip"], file_count="single")
-                zip_process_btn = gr.Button("Process Zip")
-
-                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")
-
-                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")
-
-                download_pdf_btn = gr.Button("Download Results as PDF")
-                pdf_output = gr.File(label="Download PDF", visible=False)
+                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")
+
+                    download_pdf_btn = gr.Button("Download Results as PDF")
+                    pdf_output = gr.File(label="Download PDF", visible=False)
 
                 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])
+        btn.click(
+            lambda img: detector.predict(img),
+            inp,
+            [output_html, output_label]
+        )
 
-        def on_download_pdf(accuracy, roc_score, report_html, confusion_matrix_plot):
-            pdf_buffer = generate_pdf(accuracy, roc_score, report_html, confusion_matrix_plot)
-            pdf_buffer.seek(0)
-            return pdf_buffer
+        def enable_zip_btn(file):
+            return gr.Button.update(interactive=file is not None)
 
-        download_pdf_btn.click(
-            on_download_pdf,
-            inputs=[output_acc, output_roc, output_report, output_plots],
-            outputs=pdf_output
-        )
+        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,
@@ -284,8 +313,41 @@ def create_gradio_interface():
             [output_acc, output_roc, output_report, output_plots, output_fp_fn]
         )
 
+        def on_download_pdf(accuracy, roc_score, report_html, confusion_matrix_plot):
+            pdf_buffer = generate_pdf(accuracy, roc_score, report_html, confusion_matrix_plot)
+            pdf_buffer.seek(0)
+            return pdf_buffer
+
+        download_pdf_btn.click(
+            on_download_pdf,
+            inputs=[output_acc, output_roc, output_report, output_plots],
+            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  # Reset outputs
+            ]
+
+        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
+            ]
+        )
+
     return app
 
 if __name__ == "__main__":
     app = create_gradio_interface()
-    app.launch(show_api=False, max_threads=24, show_error=True)
+    app.launch(
+        show_api=False,
+        max_threads=24,
+        show_error=True
+    )