two interfaces

app.py

@@ -15,7 +15,9 @@ 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")
+#model = tf.keras.models.load_model("densenet")
+binary_model = tf.keras.models.load_model("densenet")
+binary_labels = {0: 'healthy', 1: 'patient'}
 
 
 #model = tf.keras.models.load_model('/content/drive/MyDrive/project_image_2023_NO/saved_models/saved_model/densenet')
@@ -26,29 +28,38 @@ def classify_image(inp):
   inp = inp.reshape((-1, 224, 224, 3))
   inp = tf.keras.applications.densenet.preprocess_input(inp)
   prediction = model.predict(inp)
-  confidences = {labels[i]: float(prediction[0][i]) for i in range(2)}
+  confidences = {binary_labels[i]: float(prediction[0][i]) for i in range(2)}
   return confidences
 
 
-image_interface = gr.Interface(fn=classify_image, 
-             inputs=gr.inputs.Image(shape=(224, 224)),
-             outputs=gr.outputs.Label(num_top_classes = 2),
-             title="Demo",
+binary_interface = gr.Interface(fn=classify_image, 
+             inputs=gr.Image(shape=(224, 224)),
+             outputs=gr.Label(num_top_classes = 2),
+             title="Binary Image Classification",
              description="Here's a sample image classification. Enjoy!",
              examples=[['300104.png']])
 
-image_interface.launch()
+# Multi-Class Image Classification
 
-text_interface = pipeline("text-generation", model="densenet")
+multi_model = pipeline('image-classification', model='densenet')
+multi_labels = multi_model.model.config.id2label
 
-def generate_text(prompt):
-  text = text_interface(prompt, max_length=50)[0]['generated_text']
-  return text
+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"
+                      )
 
-text_generator_interface = gr.Interface(fn=generate_text,
-                                         inputs="text",
-                                         outputs="text",
-                                         title="Text Generator",
-                                         description="Generate text using GPT-2",
-                                         examples=[["300104.png"]])
-text_generator_interface.launch()
+iface.launch()