Update app.py

app.py

@@ -1,46 +1,46 @@
-import gradio as gr
-from pathlib import Path
-from PIL import Image
-
-# Import your model classes (adjust import paths as needed)
-from app.src.vit_load import VITDocumentClassifier
-from app.src.vgg16_load import VGGDocumentClassifier
-from app.src.constant import vit_model_path, vit_mlb_path, vgg_model_path, vgg_mlb_path
-
-# Load models once at startup
-vit_model = VITDocumentClassifier(vit_model_path, vit_mlb_path)
-vgg_model = VGGDocumentClassifier(vgg_model_path, vgg_mlb_path)
-
-def predict_vit(image, cut_off):
-    if image is None:
-        return "Please upload an image."
-    temp_path = "temp_vit_image.png"
-    image.save(temp_path)
-    result = vit_model.predict(Path(temp_path), cut_off)
-    return f"ViT Prediction: {result}"
-
-def predict_vgg(image):
-    if image is None:
-        return "Please upload an image."
-    temp_path = "temp_vgg_image.png"
-    image.save(temp_path)
-    result = vgg_model.predict(Path(temp_path))
-    return f"VGG16 Prediction: {result}"
-
-with gr.Blocks() as demo:
-    gr.Markdown("# Document Classification Demo\nUpload an image and choose a model to classify it.")
-    with gr.Row():
-        with gr.Column():
-            image_input = gr.Image(type="pil", label="Upload Image")
-            cut_off = gr.Slider(0, 1, value=0.5, label="ViT Cutoff Threshold")
-        with gr.Column():
-            result_output = gr.Textbox(label="Prediction Result", interactive=False)
-    with gr.Row():
-        vit_btn = gr.Button("Predict with ViT Model")
-        vgg_btn = gr.Button("Predict with VGG16 Model")
-
-    vit_btn.click(fn=predict_vit, inputs=[image_input, cut_off], outputs=result_output)
-    vgg_btn.click(fn=predict_vgg, inputs=image_input, outputs=result_output)
-
-if __name__ == "__main__":
-    demo.launch() 
+import gradio as gr
+from pathlib import Path
+from PIL import Image
+
+# Import your model classes (adjust import paths as needed)
+from app.src.vit_load import VITDocumentClassifier
+from app.src.vgg16_load import VGGDocumentClassifier
+from app.src.constant import vit_model_path, vit_mlb_path, vgg_model_path, vgg_mlb_path
+
+# Load models once at startup
+vit_model = VITDocumentClassifier(vit_model_path, vit_mlb_path)
+vgg_model = VGGDocumentClassifier(vgg_model_path, vgg_mlb_path)
+
+def predict_vit(image, cut_off):
+    if image is None:
+        return "Please upload an image."
+    temp_path = "temp_vit_image.png"
+    image.save(temp_path)
+    result = vit_model.predict(Path(temp_path), cut_off)
+    return f"ViT Prediction: {result}"
+
+def predict_vgg(image):
+    if image is None:
+        return "Please upload an image."
+    temp_path = "temp_vgg_image.png"
+    image.save(temp_path)
+    result = vgg_model.predict(Path(temp_path))
+    return f"VGG16 Prediction: {result}"
+
+with gr.Blocks() as demo:
+    gr.Markdown("# Document Classification Demo\nUpload an image and choose a model to classify it.")
+    with gr.Row():
+        with gr.Column():
+            image_input = gr.Image(type="pil", label="Upload Image")
+            cut_off = gr.Slider(0, 1, value=0.5, label="ViT Cutoff Threshold")
+        with gr.Column():
+            result_output = gr.Textbox(label="Prediction Result", interactive=False)
+    with gr.Row():
+        vit_btn = gr.Button("Predict with ViT Model")
+        vgg_btn = gr.Button("Predict with VGG16 Model")
+
+    vit_btn.click(fn=predict_vit, inputs=[image_input, cut_off], outputs=result_output)
+    vgg_btn.click(fn=predict_vgg, inputs=image_input, outputs=result_output)
+
+if __name__ == "__main__":
+    demo.launch(ssr_mode=False, share=True)