Update handler.py

handler.py

@@ -1,41 +1,73 @@
 import cv2
 import numpy as np
 import tensorflow as tf
+import time
 import random
 from typing import Any, Dict, List
 
-class EndpointHandler():
-    def __init__(self, path=""):
-        self.model = tf.saved_model.load(f'{path}/my_model')
-        self.classes_1 = ["RoadAccidents", "Fighting", "NormalVideos"]
-        self.locations = ['Miami', 'Smouha', 'Mandara', 'Sporting', 'Montazah']
+class EndpointHandler:
 
+  def __init__(self, path=""):
 
-    def get_top_k(self, probs, k=1, label_map=None):
-        if label_map is None:
-            label_map = self.classes_1
-        top_predictions = tf.argsort(probs, axis=-1, direction='DESCENDING')[:k]
-        top_labels = tf.gather(label_map, top_predictions, axis=-1)
-        top_labels = [label.decode('utf8') for label in top_labels.numpy()]
-        top_probs = tf.gather(probs, top_predictions, axis=-1).numpy()
-        return top_labels[0]
+    self.model = tf.saved_model.load(f'{path}/my_model')
+    self.classes_1 = ["RoadAccidents", "Fighting", "NormalVideos"]
+    self.locations = ['Miami', 'Smouha', 'Mandara', 'Sporting', 'Montazah']
 
-    def perform_action_recognition(self, frame, k=1):
-        outputs = self.model.signatures['serving_default'](image=frame)
-        probs = tf.nn.softmax(outputs['classifier_head_1'])
-        return self.get_top_k(probs[0], k=k)
+  def preprocess_frame(self, frame):
+    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+    frame = cv2.resize(frame, (224, 224))
+    frame = frame.astype('float32') / 255.0
+    frame = np.expand_dims(frame, axis=0)
+    return frame
 
-    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
-        if isinstance(data, str):
-            data = json.loads(data)  # Parse JSON string to dictionary
+  def get_top_k(self, probs, k=1, label_map=None):
+    if label_map is None:
+      label_map = self.classes_1
+    top_predictions = tf.argsort(probs, axis=-1, direction='DESCENDING')[:k]
+    top_labels = tf.gather(label_map, top_predictions, axis=-1)
+    top_labels = [label.decode('utf8') for label in top_labels.numpy()]
+    top_probs = tf.gather(probs, top_predictions, axis=-1).numpy()
+    return top_labels[0]
 
-        frame = np.array(data.get("frame"))
+  def perform_action_recognition(self, url, k=1):
 
-        if frame is None:
-            raise ValueError("'frame' is missing from the request body")
+    cap = cv2.VideoCapture(url)
+    start_time = time.time()
 
-        if not isinstance(frame, np.ndarray):
-            raise ValueError(f"Expected 'frame' to be a np.ndarray, but found {type(frame)}")
+    while True:
+      ret, frame = cap.read()
+      if not ret:
+        break
 
-        prediction = self.perform_action_recognition(frame)
-        return {"prediction": prediction}
+      preprocessed_frame = self.preprocess_frame(frame)
+      outputs = self.model.signatures['serving_default'](image=preprocessed_frame[tf.newaxis])
+      probs = tf.nn.softmax(outputs['classifier_head_1'])
+      current_time = time.time() - start_time
+      m, s = divmod(current_time, 60)
+      h, m = divmod(m, 60)
+      ip_address = url.split("/")[-1]
+      output = {
+          "class": self.get_top_k(probs[0], k=k),
+          "elapsed_time": f"{int(h):02d}:{int(m):02d}:{int(s):02d}",
+          "location": random.choice(self.locations),
+          "ip_address": ip_address
+      }
+      yield output
+
+    cap.release()
+
+  def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+
+    url = data.get("url")
+
+    if url is None:
+      raise ValueError("'url' is missing from the request body")
+
+    if not isinstance(url, str):
+      raise ValueError(f"Expected 'url' to be a str, but found {type(url)}")
+
+    outputs = []
+    for output in self.perform_action_recognition(url):
+      outputs.append(output)
+
+    return outputs