Update app.py

app.py

@@ -5,8 +5,7 @@ import torch
 import torch.nn.functional as F
 import torchvision.transforms as transforms
 import torchvision
-from torchkeras.data import get_url_img
-from torchkeras import plots 
+from ultralytics.utils.plotting import Annotator, colors 
 import numpy as np
 import yaml
 from huggingface_hub import hf_hub_download
@@ -28,8 +27,26 @@ def process_img(image):
     with torch.no_grad():
         result = model(source=image)
         if len(result[0].boxes)>0:
-            vis = plots.plot_detection(image,boxes=result[0].boxes,
-                     class_names=list(result[0].names.values()), min_score=0.2)
+            ann=Annotator(im=img)
+            boxes=result[0].boxes.xyxy
+            for element in boxes:
+                box=np.array(element.cpu()).flatten()
+                if result[0].boxes.cls[0].cpu().numpy()==2:
+                        label='car'
+                if result[0].boxes.cls[0].cpu().numpy()==0:
+                        label='bicycle'
+                if result[0].boxes.cls[0].cpu().numpy()==1:
+                        label='bus'
+                if result[0].boxes.cls[0].cpu().numpy()==3:
+                        label='motorcycle'
+                if result[0].boxes.cls[0].cpu().numpy()==4:
+                        label='person'
+                if result[0].boxes.cls[0].cpu().numpy()==5:
+                        label='train'
+                if result[0].boxes.cls[0].cpu().numpy()==6:
+                        label='truck'
+                ann.box_label(box=box, label=label, color=(0,128,0))
+            vis=ann.result()
         else:
             vis = img
     return vis