Update handler.py

handler.py

@@ -30,21 +30,35 @@ class EndpointHandler:
             A :obj:`list` | `dict`: will be serialized and returned
         """
 
+        # inputs = data.get("inputs")
+        # videos = read_video(inputs)
+        # with torch.no_grad():
+        #     outputs = self.model(videos)
+        # logits = outputs.logits
+        # _, predicted = torch.max(logits, 1)
+        # return predicted.tolist()
+
         inputs = data.get("inputs")
         videos = read_video(inputs)
+        
         with torch.no_grad():
             outputs = self.model(videos)
+        
         logits = outputs.logits
-        _, predicted = torch.max(logits, 1)
-        return predicted.tolist()
+        probabilities = torch.softmax(logits, dim=1)
+        
+        # Top 3
+        top_probs, top_indices = torch.topk(probabilities, 3, dim=1)
+        
+        top_probs_list = top_probs.tolist()
+        top_indices_list = top_indices.tolist()
+        
+        top_results = []
+        for i in range(len(top_indices_list)):
+            top_results.append({
+                "class_indices": top_indices_list[i],
+                "probabilities": top_probs_list[i]
+            })
+        
+        return top_results
         
-        # inputs = data.get("inputs")
-        # if not inputs:
-        #     return {"error": "No video input provided"}
-
-        # # 비디오 파일 경로
-        # video_path = inputs.get("video_path")
-        # if not video_path or not os.path.exists(video_path):
-        #     return {"error": "Invalid or missing video file"}
-
-        # return {"predicted_class": 1}