Update app.py

app.py

@@ -5,11 +5,18 @@ from PIL import Image
 import traceback
 import re
 import torch
+import numpy as np
+import tensorflow as tf
+from tensorflow.keras.models import load_model # type: ignore
 import argparse
 from transformers import AutoModel, AutoTokenizer
 
+# Configuration for image classification model
+class_names = ['Calculus', 'Dental Caries', 'Gingivitis', 'Hypodontia', 'Tooth Discoloration']
+cnn_model = load_model('new_model2.h5')
+
 # Argparser
-parser = argparse.ArgumentParser(description='demo')
+parser = argparse.ArgumentParser(description='app')
 parser.add_argument('--device', type=str, default='cpu', help='cpu')
 parser.add_argument('--dtype', type=str, default='fp32', help='fp32')
 args = parser.parse_args()
@@ -50,6 +57,27 @@ top_p_slider = {'minimum': 0, 'maximum': 1, 'value': 0.8, 'step': 0.05, 'interac
 top_k_slider = {'minimum': 0, 'maximum': 200, 'value': 100, 'step': 1, 'interactive': True, 'label': 'Top K'}
 temperature_slider = {'minimum': 0, 'maximum': 2, 'value': 0.7, 'step': 0.05, 'interactive': True, 'label': 'Temperature'}
 
+def classify_images(image):
+    # Check if the image is None
+    if image is None:
+        return "No image uploaded. Please upload a dental image."
+
+    # Resize and preprocess the image
+    try:
+        input_image = tf.image.resize(image, (180, 180))  # Resize to expected input size
+        input_image_array = tf.keras.utils.img_to_array(input_image)
+        input_image_exp_dim = tf.expand_dims(input_image_array, axis=0)
+
+        # Make predictions
+        predictions = cnn_model.predict(input_image_exp_dim)
+        result = tf.nn.softmax(predictions[0])
+        
+        # Prepare the outcome message
+        outcome = f'The image belongs to {class_names[np.argmax(result)]} with a score of {np.max(result) * 100:.2f}%'
+        return outcome
+    except Exception as e:
+        return f"Error processing the image: {str(e)}"
+    
 def create_component(params, comp='Slider'):
     if comp == 'Slider':
         return gr.Slider(**params)
@@ -116,26 +144,33 @@ def clear(chat_bot, app_session):
     chat_bot.clear()
     return chat_bot
 
-with gr.Blocks() as demo:
+with gr.Blocks() as app:
     gr.Markdown("<h1 style='text-align: center;'>Medical Assistant</h1>")
 
-    with gr.Row():
-        with gr.Column(scale=2, min_width=300):
-            app_session = gr.State({'sts': None, 'ctx': None, 'img': None})
-            bt_pic = gr.Image(label="Upload an image to start")
-            txt_message = gr.Textbox(label="Ask your question...")
-        
-        with gr.Column(scale=2, min_width=300):
-            chat_bot = gr.Chatbot(label=f"Chatbot")
-            clear_button = gr.Button(value='Clear')
-            txt_message.submit(
-                respond, 
-                [txt_message, chat_bot, app_session], 
-                [txt_message, chat_bot, app_session]
-            )
-
-            bt_pic.upload(lambda: None, None, chat_bot, queue=False).then(upload_img, inputs=[bt_pic, chat_bot, app_session], outputs=[chat_bot, app_session])
-            clear_button.click(clear, [chat_bot, app_session], chat_bot)
+    with gr.Tab("Image Classification"):
+        with gr.Row():
+            image_input = gr.Image(type='numpy', label="Upload Dental Image")
+            classification_output = gr.Label(num_top_classes=5, label="Classification Results")
+        image_input.change(fn=classify_images, inputs=image_input, outputs=classification_output)
+    
+    with gr.Tab("Medical Chatbot"):
+        with gr.Row():
+            with gr.Column(scale=2, min_width=300):
+                app_session = gr.State({'sts': None, 'ctx': None, 'img': None})
+                bt_pic = gr.Image(label="Upload an image to start")
+                txt_message = gr.Textbox(label="Ask your question...")
+            
+            with gr.Column(scale=2, min_width=300):
+                chat_bot = gr.Chatbot(label=f"Chatbot")
+                clear_button = gr.Button(value='Clear')
+                txt_message.submit(
+                    respond, 
+                    [txt_message, chat_bot, app_session], 
+                    [txt_message, chat_bot, app_session]
+                )
+
+                bt_pic.upload(lambda: None, None, chat_bot, queue=False).then(upload_img, inputs=[bt_pic, chat_bot, app_session], outputs=[chat_bot, app_session])
+                clear_button.click(clear, [chat_bot, app_session], chat_bot)
 
 # Launch
-demo.launch(share=True, debug=True, show_api=False)
+app.launch(share=True, debug=True, show_api=False)