Update app.py

app.py

@@ -9,9 +9,9 @@ import requests
 from io import BytesIO
 import os
 
-model = YOLO('80epoch.pt')
+model = YOLO('50epoch-new-weapon.pt')
 model2 = pipeline('image-classification','Kaludi/csgo-weapon-classification')
-name = ['grenade','knife','pistol','rifle']
+name = ['grenade','knife','missile','pistol','rifle']
 image_directory = "/home/user/app/image"
 video_directory = "/home/user/app/video"
 
@@ -57,10 +57,10 @@ def response2(image: gr.Image = None,image_size: gr.Slider = 640, conf_threshold
     
     
     for r in results:
-        conf = np.array(r.boxes.conf)
-        cls = np.array(r.boxes.cls)
+        conf = np.array(r.boxes.conf.cpu())
+        cls = np.array(r.boxes.cls.cpu())
         cls = cls.astype(int)
-        xywh = np.array(r.boxes.xywh)
+        xywh = np.array(r.boxes.xywh.cpu())
         xywh = xywh.astype(int)  
       
         for con, cl, xy in zip(conf, cls, xywh):
@@ -74,9 +74,11 @@ def response2(image: gr.Image = None,image_size: gr.Slider = 640, conf_threshold
             elif cl == 1:
                 name_weap += name[cl] + '\n'
             elif cl == 2:
+                name_weap += name[cl] + '\n'
+            elif cl == 3:
                 out = model2(image)
                 name_weap += out[0]["label"] + '\n'
-            elif cl == 3:
+            elif cl == 4:
                 out = model2(image)
                 name_weap += out[0]["label"] + '\n'
       
@@ -106,7 +108,8 @@ outputs = [gr.Image( type="pil", label="Output Image"),
 examples = [[os.path.join(image_directory, "th (5).jpg"),640, 0.3, 0.6],
             [os.path.join(image_directory, "th (8).jpg"),640, 0.3, 0.6],
             [os.path.join(image_directory, "th (11).jpg"),640, 0.3, 0.6],
-            [os.path.join(image_directory, "th (3).jpg"),640, 0.3, 0.6]
+            [os.path.join(image_directory, "th (3).jpg"),640, 0.3, 0.6],
+            [os.path.join(image_directory, "th.jpg"),640, 0.3, 0.6]
            ]
 title = 'Weapon Detection Finetuned YOLOv8'
 description = 'Image Size: Defines the image size for inference.\nConfidence Treshold: Sets the minimum confidence threshold for detections.\nIOU Treshold: Intersection Over Union (IoU) threshold for Non-Maximum Suppression (NMS). Useful for reducing duplicates.'