Update app.py

app.py

@@ -4,8 +4,11 @@ from PIL import Image
 from ultralyticsplus import YOLO, render_result
 import cv2
 import numpy as np
+from transformers import pipeline
 
 model = YOLO('best (1).pt')
+model2 = pipeline('Kaludi/csgo-weapon-classification')
+name = ['grenade','knife','pistol','rifle']
 
  # for i, r in enumerate(results):
       
@@ -17,6 +20,8 @@ def response(image):
   print(image)
   results = model(image)
   text=""
+  name=""
+    
   for r in results:
     conf = np.array(r.boxes.conf)
     cls = np.array(r.boxes.cls)
@@ -30,13 +35,22 @@ def response(image):
         conef = conef * 100
         text += (f"Detected {name[cl]} with confidence {round(conef,1)}% at ({xy[0]},{xy[1]})\n")
         
+        if cl == 0:
+            name += name[cl]
+        elif cl == 1:
+            name += name[cl]
+        elif cl == 2:
+            name += model2(image)
+        elif cl == 3:
+            name += model2(image)
+
         
     # im_rgb = Image.fromarray(im_rgb)
     
       
-    return text
+    return name, text
+
 
-name = ['grenade','knife','pistol','rifle']
 
 def response2(image: gr.Image = None,image_size: gr.Slider = 640, conf_threshold: gr.Slider = 0.3, iou_threshold: gr.Slider = 0.6):
 
@@ -47,7 +61,7 @@ def response2(image: gr.Image = None,image_size: gr.Slider = 640, conf_threshold
     render = render_result(model=model, image=image, result=results[0], rect_th = 1, text_th = 1)
    
    
-    text = response(image)
+    weapon_name, text_detection = response(image)
    
     
     # xywh = int(results.boxes.xywh)
@@ -56,7 +70,7 @@ def response2(image: gr.Image = None,image_size: gr.Slider = 640, conf_threshold
         
    
             
-    return render,text
+    return render, text_detection, weapon_name
 
 
 inputs = [
@@ -71,7 +85,8 @@ inputs = [
 
 
 outputs = [gr.Image( type="filepath", label="Output Image"),
-           gr.Textbox(label="Result")
+           gr.Textbox(label="Result"),
+           gr.Textbox(label="Weapon Name")
           ]
 
 title = "YOLOv8 Custom Object Detection by Uyen Nguyen"
@@ -82,6 +97,7 @@ examples = [['th (11).jpg', 640, 0.3, 0.6],
             ['th (3).jpg', 640, 0.3, 0.6],
             # ['four.jpg', 832, 0.3, 0.3]]
            ]
+           
 
 # yolo_app = gr.Interface(
 #     fn=yoloV8_func,