two interfaces

app.py

@@ -16,50 +16,54 @@ from transformers import pipeline
 # load the model from the Hugging Face Model Hub
 #model = pipeline('image-classification', model='image_classification/densenet')
 #model = tf.keras.models.load_model("densenet")
-binary_model = tf.keras.models.load_model("densenet")
-binary_labels = {0: 'healthy', 1: 'patient'}
+#binary_model = tf.keras.models.load_model("densenet")
+#binary_labels = {0: 'healthy', 1: 'patient'}
+
+# load the binary classification model
+model_binary = tf.keras.models.load_model("densenet_binary")
+
+# load the multi-label classification model
+model_multi = tf.keras.models.load_model("densenet_multi")
+
+# define the labels for the multi-label classification model
+labels_multi = {0: 'healthy', 1: 'mild', 2: 'moderate'}
 
 
 #model = tf.keras.models.load_model('/content/drive/MyDrive/project_image_2023_NO/saved_models/saved_model/densenet')
 #labels = ['Healthy', 'Patient']
-labels = {0: 'healthy', 1: 'patient'}
+#labels = {0: 'healthy', 1: 'patient'}
 
-def classify_image(inp):
+def classify_image_binary(inp):
   inp = inp.reshape((-1, 224, 224, 3))
   inp = tf.keras.applications.densenet.preprocess_input(inp)
-  prediction = model.predict(inp)
-  confidences = {binary_labels[i]: float(prediction[0][i]) for i in range(2)}
-  return confidences
+  prediction = model_binary.predict(inp)
+  confidence = float(prediction[0])
+  label = "healthy" if confidence < 0.5 else "patient"
+  return {label: confidence}
 
 
-binary_interface = gr.Interface(fn=classify_image, 
+def classify_image_multi(inp):
+  inp = inp.reshape((-1, 224, 224, 3))
+  inp = tf.keras.applications.densenet.preprocess_input(inp)
+  prediction = model_multi.predict(inp)
+  confidences = {labels_multi[i]: float(prediction[0][i]) for i in range(3)}
+  return confidences
+    
+iface_binary = gr.Interface(fn=classify_image_binary, 
              inputs=gr.Image(shape=(224, 224)),
-             outputs=gr.Label(num_top_classes = 2),
+             outputs=gr.Label(num_top_classes=2),
              title="Binary Image Classification",
-             description="Here's a sample image classification. Enjoy!",
-             examples=[['300104.png']])
-
-# Multi-Class Image Classification
-
-multi_model = pipeline('image-classification', model='densenet')
-multi_labels = multi_model.model.config.id2label
+             description="Classify an image as healthy or patient.",
+             examples=[['300104.png']]
+             )
 
-def classify_multi_image(inp):
-  prediction = multi_model(inp)
-  confidences = {multi_labels[i]: float(prediction[0][i]) for i in range(len(multi_labels))}
-  return confidences
-
-multi_interface = gr.Interface(fn=classify_multi_image, 
-                               inputs=gr.Image(shape=(224, 224)),
-                               outputs=gr.Label(num_top_classes=len(multi_labels)),
-                               title="Multi-Class Image Classification",
-                               description="Classify an image into multiple classes",
-                               examples=[['300104.png']]
-                               )
-iface = gr.Interface([binary_interface, multi_interface], 
-                      title="Image Classification App",
-                      layout="vertical",
-                      description="Demo for binary and multi-class image classification"
-                      )
+iface_multi = gr.Interface(fn=classify_image_multi, 
+             inputs=gr.Image(shape=(224, 224)),
+             outputs=gr.Label(num_top_classes=3),
+             title="Multi-Label Image Classification",
+             description="Classify an image as healthy, mild, or moderate.",
+             examples=[['300104.png']]
+             )
 
-iface.launch()
+iface_binary.launch()
+iface_multi.launch()