Update app.py

app.py

@@ -14,16 +14,26 @@ def detect_objects_on_image(image_path):
     results = model.predict(image_path)
     result = results[0]
     output = []
+
+    class_names_mapping = {
+        "DPHF": "Double Person Helmet",
+        "DPNH": "Double Person No Helmet",
+        "SPHF": "Single Person Helmet",
+        "SPNH": "Single Person No Helmet",
+        "NP": "Number Plate"
+    }
+
     for box in result.boxes:
-        x1, y1, x2, y2 = [
-            round(x) for x in box.xyxy[0].tolist()
-            
-        ]
+        x1, y1, x2, y2 = [round(x) for x in box.xyxy[0].tolist()]
         class_id = box.cls[0].item()
         prob = round(box.conf[0].item(), 2)
+
+        class_name = class_names_mapping.get(result.names[class_id], result.names[class_id])
+
         output.append([
-            x1, y1, x2, y2, result.names[class_id], prob
+            x1, y1, x2, y2, class_name, prob
         ])
+
         cv2.rectangle(
             image,
             (x1, y1),
@@ -32,8 +42,9 @@ def detect_objects_on_image(image_path):
             thickness=2,
             lineType=cv2.LINE_AA
         )
-        
-        cv2.putText(image,result.names[class_id],  (x1, y1), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (36,255,12), 2)
+
+        cv2.putText(image, class_name, (x1, y1), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (36, 255, 12), 2)
+
     return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)