Update app.py

app.py

@@ -5,7 +5,10 @@ import importlib.util
 import gradio as gr
 from PIL import Image
 
-# Function to load the TensorFlow Lite model and labels
+# Load the TensorFlow Lite models
+MODEL_DIRS = ['model', 'model_2', 'model_3']
+MODEL_NAMES = ['Multi-class model', 'One Empty class', 'Misalignment class']
+
 def load_model(model_dir):
     GRAPH_NAME = 'detect.tflite'
     LABELMAP_NAME = 'labelmap.txt'
@@ -30,18 +33,18 @@ def load_model(model_dir):
     width = input_details[0]['shape'][2]
     floating_model = (input_details[0]['dtype'] == np.float32)
 
-    return interpreter, labels, input_details, output_details, height, width, floating_model
+    outname = output_details[0]['name']
+    if ('StatefulPartitionedCall' in outname):
+        boxes_idx, classes_idx, scores_idx = 1, 3, 0
+    else:
+        boxes_idx, classes_idx, scores_idx = 0, 1, 2
 
-MODEL_DIRS = {
-    "Multi-class model": 'model',
-    "Empty class": 'model_2',
-    "Misalignment class": 'model_3'
-}
+    return interpreter, labels, input_details, output_details, height, width, floating_model, boxes_idx, classes_idx, scores_idx
 
-input_mean = 127.5
-input_std = 127.5
+# Load default model
+interpreter, labels, input_details, output_details, height, width, floating_model, boxes_idx, classes_idx, scores_idx = load_model(MODEL_DIRS[0])
 
-def perform_detection(image, interpreter, labels, input_details, output_details, height, width, floating_model):
+def perform_detection(image, interpreter, labels):
     imH, imW, _ = image.shape
     image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
     image_resized = cv2.resize(image_rgb, (width, height))
@@ -53,9 +56,9 @@ def perform_detection(image, interpreter, labels, input_details, output_details,
     interpreter.set_tensor(input_details[0]['index'], input_data)
     interpreter.invoke()
 
-    boxes = interpreter.get_tensor(output_details[0]['index'])[0]
-    classes = interpreter.get_tensor(output_details[1]['index'])[0]
-    scores = interpreter.get_tensor(output_details[2]['index'])[0]
+    boxes = interpreter.get_tensor(output_details[boxes_idx]['index'])[0]
+    classes = interpreter.get_tensor(output_details[classes_idx]['index'])[0]
+    scores = interpreter.get_tensor(output_details[scores_idx]['index'])[0]
 
     detections = []
     for i in range(len(scores)):
@@ -79,16 +82,19 @@ def perform_detection(image, interpreter, labels, input_details, output_details,
 def resize_image(image, size=640):
     return cv2.resize(image, (size, size))
 
-def detect_image(input_image, model_selection):
-    interpreter, labels, input_details, output_details, height, width, floating_model = load_model(MODEL_DIRS[model_selection])
+def detect_image(input_image, model_index=0):
+    global interpreter, labels, input_details, output_details, height, width, floating_model, boxes_idx, classes_idx, scores_idx
+    interpreter, labels, input_details, output_details, height, width, floating_model, boxes_idx, classes_idx, scores_idx = load_model(MODEL_DIRS[model_index])
+
     image = np.array(input_image)
     resized_image = resize_image(image, size=640)  # Resize input image
-    result_image = perform_detection(resized_image, interpreter, labels, input_details, output_details, height, width, floating_model)
+    result_image = perform_detection(resized_image, interpreter, labels)
     return Image.fromarray(result_image)
 
-def detect_video(input_video, model_selection):
-    interpreter, labels, input_details, output_details, height, width, floating_model = load_model(MODEL_DIRS[model_selection])
-    
+def detect_video(input_video, model_index=0):
+    global interpreter, labels, input_details, output_details, height, width, floating_model, boxes_idx, classes_idx, scores_idx
+    interpreter, labels, input_details, output_details, height, width, floating_model, boxes_idx, classes_idx, scores_idx = load_model(MODEL_DIRS[model_index])
+
     cap = cv2.VideoCapture(input_video)
     frames = []
 
@@ -98,7 +104,7 @@ def detect_video(input_video, model_selection):
             break
 
         resized_frame = resize_image(frame, size=640)  # Resize each frame
-        result_frame = perform_detection(resized_frame, interpreter, labels, input_details, output_details, height, width, floating_model)
+        result_frame = perform_detection(resized_frame, interpreter, labels)
         frames.append(result_frame)
 
     cap.release()
@@ -118,21 +124,29 @@ def detect_video(input_video, model_selection):
 
     return output_video_path
 
-app = gr.Blocks()
+app = gr.Interface(
+    fn=None,
+    inputs=None,
+    outputs=None,
+    title="Object Detection",
+    description="Detect objects in images and videos.",
+    layout="blocks",
+    theme="compact",
+)
 
 with app:
     with gr.Tab("Image Detection"):
         gr.Markdown("Upload an image for object detection")
-        model_selection = gr.Dropdown(choices=["Multi-class model", "Empty class", "Misalignment class"], label="Select Model")
-        image_input = gr.Image(type="pil", label="Upload an image")
-        image_output = gr.Image(type="pil", label="Detection Result")
-        gr.Button("Submit").click(fn=detect_image, inputs=[image_input, model_selection], outputs=image_output)
+        image_input = gr.inputs.Image(type="pil", label="Upload an image")
+        image_output = gr.outputs.Image(type="pil", label="Detection Result")
+        model_dropdown = gr.inputs.Dropdown(choices=MODEL_NAMES, label="Select Model")
+        gr.Button("Submit").on_click(detect_image, inputs=[image_input, model_dropdown], outputs=image_output)
     
     with gr.Tab("Video Detection"):
         gr.Markdown("Upload a video for object detection")
-        model_selection = gr.Dropdown(choices=["Multi-class model", "Empty class", "Misalignment class"], label="Select Model")
-        video_input = gr.Video(label="Upload a video")
-        video_output = gr.Video(label="Detection Result")
-        gr.Button("Submit").click(fn=detect_video, inputs=[video_input, model_selection], outputs=video_output)
+        video_input = gr.inputs.Video(label="Upload a video")
+        video_output = gr.outputs.Video(label="Detection Result")
+        model_dropdown = gr.inputs.Dropdown(choices=MODEL_NAMES, label="Select Model")
+        gr.Button("Submit").on_click(detect_video, inputs=[video_input, model_dropdown], outputs=video_output)
 
 app.launch()