Update app.py

app.py

@@ -33,21 +33,34 @@ def process_img(image):
             for element in boxes:
                 box=np.array(element.xyxy.cpu()).flatten()
                 if element.cls[0].cpu().numpy()==2.0:
-                        lbel='car'
+                    lbel='car'
+                    clr=(0,255,0)
+                    
                 if element.cls[0].cpu().numpy()==0.0:
-                        lbel='bicycle'
+                    lbel='bicycle'
+                    clr=(255,0,0)
+                    
                 if element.cls[0].cpu().numpy()==1.0:
-                        lbel='bus'
+                    lbel='bus'
+                    clr=(0,0,255)
+                    
                 if element.cls[0].cpu().numpy()==3.0:
-                        lbel='motorcycle'
+                    lbel='motorcycle'
+                    clr=(255,0,255)
+                    
                 if element.cls[0].cpu().numpy()==4.0:
-                        lbel='person'
+                    lbel='person'
+                    clr=(255,128,0)
+                    
                 if element.cls[0].cpu().numpy()==5.0:
-                        lbel='train'
+                    lbel='train'
+                    clr=(255,0,128)
+                    
                 if element.cls[0].cpu().numpy()==6.0:
-                        lbel='truck'
+                    lbel='truck'
+                    clr=(0,255,255)
                     
-                ann.box_label(box=box, label=lbel, color=(0,128,0))
+                ann.box_label(box=box, label=lbel, color=clr)
             vis=ann.result()
         else:
             vis = image