File size: 5,649 Bytes
15cffc9
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
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
    )