Rename app_model_2.py to model_2.py

app_model_2.py → model_2.py

@@ -1,11 +1,12 @@
+# model_2.py
 import os
 import cv2
 import numpy as np
 import importlib.util
-import gradio as gr
 from PIL import Image
+import gradio as gr
+from common_detection import perform_detection
 
-# Load the TensorFlow Lite model for Model 1
 MODEL_DIR = 'model_2'
 GRAPH_NAME = 'detect.tflite'
 LABELMAP_NAME = 'labelmap.txt'
@@ -13,22 +14,18 @@ LABELMAP_NAME = 'labelmap.txt'
 pkg = importlib.util.find_spec('tflite_runtime')
 if pkg:
     from tflite_runtime.interpreter import Interpreter
-    from tflite_runtime.interpreter import load_delegate
 else:
     from tensorflow.lite.python.interpreter import Interpreter
-    from tensorflow.lite.python.interpreter import load_delegate
 
 PATH_TO_CKPT = os.path.join(MODEL_DIR, GRAPH_NAME)
 PATH_TO_LABELS = os.path.join(MODEL_DIR, LABELMAP_NAME)
 
-# Load the label map
 with open(PATH_TO_LABELS, 'r') as f:
     labels = [line.strip() for line in f.readlines()]
 
 if labels[0] == '???':
     del(labels[0])
 
-# Load the TensorFlow Lite model
 interpreter = Interpreter(model_path=PATH_TO_CKPT)
 interpreter.allocate_tensors()
 
@@ -38,65 +35,38 @@ height = input_details[0]['shape'][1]
 width = input_details[0]['shape'][2]
 floating_model = (input_details[0]['dtype'] == np.float32)
 
-input_mean = 127.5
-input_std = 127.5
-
-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
-
-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))
-    input_data = np.expand_dims(image_resized, axis=0)
-
-    if floating_model:
-        input_data = (np.float32(input_data) - input_mean) / input_std
-
-    interpreter.set_tensor(input_details[0]['index'], input_data)
-    interpreter.invoke()
-
-    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)):
-        if ((scores[i] > 0.5) and (scores[i] <= 1.0)):
-            ymin = int(max(1, (boxes[i][0] * imH)))
-            xmin = int(max(1, (boxes[i][1] * imW)))
-            ymax = int(min(imH, (boxes[i][2] * imH)))
-            xmax = int(min(imW, (boxes[i][3] * imW)))
-            
-            cv2.rectangle(image, (xmin, ymin), (xmax, ymax), (10, 255, 0), 2)
-            object_name = labels[int(classes[i])]
-            label = '%s: %d%%' % (object_name, int(scores[i] * 100))
-            labelSize, baseLine = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.7, 2)
-            label_ymin = max(ymin, labelSize[1] + 10)
-            cv2.rectangle(image, (xmin, label_ymin - labelSize[1] - 10), (xmin + labelSize[0], label_ymin + baseLine - 10), (255, 255, 255), cv2.FILLED)
-            cv2.putText(image, label, (xmin, label_ymin - 7), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 0), 2)
-
-            detections.append([object_name, scores[i], xmin, ymin, xmax, ymax])
-    return image
-
-def resize_image(image, size=640):
-    return cv2.resize(image, (size, size))
-
 def detect_image(input_image):
     image = np.array(input_image)
-    resized_image = resize_image(image, size=640)  # Resize input image
-    result_image = perform_detection(resized_image, interpreter, labels)
+    resized_image = cv2.resize(image, (640, 640))
+    result_image = perform_detection(resized_image, interpreter, labels, input_details, output_details, height, width, floating_model)
     return Image.fromarray(result_image)
 
-app = gr.Interface(
-    detect_image,
-    inputs=gr.inputs.Image(type="pil", label="Upload an image"),
-    outputs="image",
-    title="Model 2: Empty-class Object Detection",
-    theme="compact"
-)
+def detect_video(input_video):
+    cap = cv2.VideoCapture(input_video)
+    frames = []
+
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+
+        resized_frame = cv2.resize(frame, (640, 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()
 
-app.launch()
+    return output_video_path