Update src/inference/lr_model.py

src/inference/lr_model.py

@@ -43,21 +43,38 @@ class LRModel:
         {
           "class_id": int,
           "class_name": str,
-          "probabilities": {class_name: prob, ...}
+          "probabilities": {class_name: prob},
+          "top_k": [
+            {"class_id": int, "class_name": str, "probability": float},
+            ...
+          ]
         }
         """
         x = self.preprocess(image)
         probs = self.model.predict_proba(x)[0]
-        class_id = int(np.argmax(probs))
-        class_name = self.labels[class_id]
+    
+        pred_id = int(np.argmax(probs))
+        pred_name = self.labels[pred_id]
 
-        # Build probability dict (optional)
         prob_dict = {
             self.labels[i]: float(probs[i]) for i in range(len(probs))
         }
-
+    
+        # Top-k (sorted)
+        sorted_indices = np.argsort(probs)[::-1]
+        top_k = min(top_k, len(sorted_indices))
+        top_k_list = []
+        for i in range(top_k):
+            cid = int(sorted_indices[i])
+            top_k_list.append({
+                "class_id": cid,
+                "class_name": self.labels[cid],
+                "probability": float(probs[cid]),
+            })
+    
         return {
-            "class_id": class_id,
-            "class_name": class_name,
-            "probabilities": prob_dict
+            "class_id": pred_id,
+            "class_name": pred_name,
+            "probabilities": prob_dict,
+            "top_k": top_k_list,
         }