Update app.py

app.py

@@ -1,31 +1,102 @@
 # app.py
+import os
+import cv2
+import numpy as np
+import importlib.util
+from PIL import Image
 import gradio as gr
-import importlib
+from common_detection import perform_detection, resize_image
 
-def load_model(model_name):
-    module = importlib.import_module(model_name)
-    return module
+# Function to load the TensorFlow Lite model and labels
+def load_model_and_labels(model_dir):
+    pkg = importlib.util.find_spec('tflite_runtime')
+    if pkg:
+        from tflite_runtime.interpreter import Interpreter
+    else:
+        from tensorflow.lite.python.interpreter import Interpreter
+
+    PATH_TO_CKPT = os.path.join(model_dir, 'detect.tflite')
+    PATH_TO_LABELS = os.path.join(model_dir, 'labelmap.txt')
+
+    with open(PATH_TO_LABELS, 'r') as f:
+        labels = [line.strip() for line in f.readlines()]
+
+    if labels[0] == '???':
+        del(labels[0])
+
+    interpreter = Interpreter(model_path=PATH_TO_CKPT)
+    interpreter.allocate_tensors()
+
+    input_details = interpreter.get_input_details()
+    output_details = interpreter.get_output_details()
+    height = input_details[0]['shape'][1]
+    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
+
+# Load models
+models = {
+    "Multi-class model": "model",
+    "Empty class": "model_2",
+    "Misalignment class": "model_3"
+}
+
+# Function to perform image detection
+def detect_image(model_choice, input_image):
+    model_dir = models[model_choice]
+    interpreter, labels, input_details, output_details, height, width, floating_model = load_model_and_labels(model_dir)
+    image = np.array(input_image)
+    resized_image = resize_image(image, size=640)
+    result_image = perform_detection(resized_image, interpreter, labels, input_details, output_details, height, width, floating_model)
+    return Image.fromarray(result_image)
+
+# Function to perform video detection
+def detect_video(model_choice, input_video):
+    model_dir = models[model_choice]
+    interpreter, labels, input_details, output_details, height, width, floating_model = load_model_and_labels(model_dir)
+    cap = cv2.VideoCapture(input_video)
+    frames = []
+
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+
+        resized_frame = resize_image(frame, size=640)
+        result_frame = perform_detection(resized_frame, interpreter, labels, input_details, output_details, height, width, floating_model)
+        frames.append(result_frame)
+
+    cap.release()
+
+    if not frames:
+        raise ValueError("No frames were read from the video.")
+
+    height, width, layers = frames[0].shape
+    size = (width, height)
+    output_video_path = "result_" + os.path.basename(input_video)
+    out = cv2.VideoWriter(output_video_path, cv2.VideoWriter_fourcc(*'mp4v'), 15, size)
+
+    for frame in frames:
+        out.write(frame)
+
+    out.release()
+
+    return output_video_path
 
 app = gr.Blocks()
 
 with app:
     gr.Markdown("## Object Detection using TensorFlow Lite Models")
     with gr.Row():
-        model_choice = gr.Dropdown(label="Select Model", choices=["model_1", "model_2", "model_3"])
-        image_input = gr.Image(type="pil", label="Upload an image")
-        image_output = gr.Image(type="pil", label="Detection Result")
-        video_input = gr.Video(label="Upload a video")
-        video_output = gr.Video(label="Detection Result")
-        
-        def image_detection(model_name, input_image):
-            model = load_model(model_name)
-            return model.detect_image(input_image)
-        
-        def video_detection(model_name, input_video):
-            model = load_model(model_name)
-            return model.detect_video(input_video)
-        
-        gr.Button("Submit Image").click(fn=image_detection, inputs=[model_choice, image_input], outputs=image_output)
-        gr.Button("Submit Video").click(fn=video_detection, inputs=[model_choice, video_input], outputs=video_output)
-
-app.launch()
+        model_choice = gr.Dropdown(label="Select Model", choices=["Multi-class model", "Empty class", "Misalignment class"])
+        with gr.Tab("Image Detection"):
+            image_input = gr.Image(type="pil", label="Upload an image")
+            image_output = gr.Image(type="pil", label="Detection Result")
+            gr.Button("Submit Image").click(fn=detect_image, inputs=[model_choice, image_input], outputs=image_output)
+        with gr.Tab("Video Detection"):
+            video_input = gr.Video(label="Upload a video")
+            video_output = gr.Video(label="Detection Result")
+            gr.Button("Submit Video").click(fn=detect_video, inputs=[model_choice, video_input], outputs=video_output)
+
+app.launch()