Update app.py

app.py

@@ -37,9 +37,7 @@ def yolov8_inference(
         anchor = h if h > w else w
         
         # Create the dictionary by filtering list1 and list2 based on list3
-        selected_classes = {key: value for key, value in zip(["anus", "make_love", "nipple", "penis", "vagina"],\
-                                                             [0, 1, 2, 3, 4]) \
-                            if key in selected_labels_list}
+        selected_classes = [[0, 1, 2, 3, 4][["anus", "make_love", "nipple", "penis", "vagina"].index(item)] for item in selected_labels_list]
         
         # print(filtered_mapping)
         # selected_classes = [0, 1, 2, 3, 4] # all classes