Update app.py

app.py

@@ -1,7 +1,8 @@
 import gradio as gr
 import tensorflow as tf
 import gdown
-from PIL import Image
+import numpy as np
+from PIL import Image, ImageDraw
 
 input_shape = (32, 32, 3)
 resized_shape = (224, 224, 3)
@@ -41,15 +42,34 @@ def predict_class(image):
     predicted_class = labels[class_index]
     return predicted_class
 
+# Perform object detection
+def detect_objects(image):
+    img = image.copy()
+    img = tf.image.resize(img, [input_shape[0], input_shape[1]])
+    img = tf.expand_dims(img, axis=0)
+    prediction = model.predict(img)
+    boxes, scores, classes = prediction[0]['detection_boxes'], prediction[0]['detection_scores'], prediction[0]['detection_classes']
+    height, width, _ = img.shape
+    draw = ImageDraw.Draw(image)
+    for box, score, _class in zip(boxes, scores, classes):
+        if score > 0.5:
+            ymin, xmin, ymax, xmax = box
+            left, right, top, bottom = xmin * width, xmax * width, ymin * height, ymax * height
+            draw.rectangle([(left, top), (right, bottom)], outline='red', width=2)
+            draw.text((left, top - 10), labels[int(_class)], fill='red')
+
+    return image
+
 # UI Design
 def classify_image(image):
     predicted_class = predict_class(image)
-    return image, predicted_class
+    image_with_box = detect_objects(image)
+    return image_with_box, predicted_class
 
 inputs = gr.inputs.Image(label="Upload an image")
-outputs = gr.outputs.Image(label="Uploaded Image"), gr.outputs.Textbox(label="Predicted Class", live=True)
+outputs = gr.outputs.Image(label="Output Image"), gr.outputs.Textbox(label="Predicted Class", live=True)
 
-title = "Image Classifier"
-description = "Upload an image and get the predicted class."
+title = "Image Classifier with Object Detection"
+description = "Upload an image and get the predicted class with object detection."
 
 gr.Interface(fn=classify_image, inputs=inputs, outputs=outputs, title=title, description=description).launch(inline=True)