Update app.py

app.py

@@ -56,20 +56,7 @@ def yolov8_img_inference(
     # count 'car' objects in the results
     # count_result = results[0].boxes.cls[0].item()
     # count_result = results[0].boxes.cls.tolist()
-    object_counts = {x: 0 for x in names}
-    for prediction in predictions:
-        for c in prediction.boxes.cls:
-            c = int(c)
-            if c in names:
-                object_counts[c] += 1
-            elif c not in names:
-                object_counts[c] = 1
-
-    present_objects = object_counts.copy()
 
-    for i in object_counts:
-        if object_counts[i] < 1:
-            present_objects.pop(i)
     
     # clist= results[0].boxes.cls.tolist()
     # cls = set()
@@ -77,7 +64,7 @@ def yolov8_img_inference(
     #     cls.add(model.names[int(cno)])
         
     # count_result = results.pandas().xyxy[0].value_counts('name')
-    return render, present_objects
+    return render
     # results = model.predict(image, imgsz=image_size, return_outputs=True)
     # results = model.predict(image)
     # object_prediction_list = []
@@ -107,7 +94,23 @@ def yolov8_img_inference(
     # output_image = visualize_object_predictions(image=image, object_prediction_list=object_prediction_list)
     # return output_image['image']
     # render = render_result(model=model, image=image, result=results[0])
-    
+def count_objects(predictions, target_classes):
+    object_counts = {x: 0 for x in target_classes}
+    for prediction in predictions:
+        for c in prediction.boxes.cls:
+            c = int(c)
+            if c in target_classes:
+                object_counts[c] += 1
+            elif c not in target_classes:
+                object_counts[c] = 1
+
+    present_objects = object_counts.copy()
+
+    for i in object_counts:
+        if object_counts[i] < 1:
+            present_objects.pop(i)
+
+    return present_objects    
 
         
 
@@ -127,7 +130,7 @@ inputs_image = [
 title = "Tất cả do anh Đạt"
 
 interface_image = gr.Interface(
-    fn=yolov8_img_inference,
+    fn=yolov8_img_inference, count_objects
     inputs=inputs_image,
     outputs=[gr.Image(type="pil"),gr.Textbox(show_label=False)],
     title=title,