check person class only

dustysys · May 8, 2023 · 6f59a8e · 6f59a8e
1 parent 9238cf8
commit 6f59a8e
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/scripts/ddetailer.py b/scripts/ddetailer.py
@@ -518,7 +518,7 @@ def inference(image, modelname, conf_thres, label):
         results = inference_mmdet_segm(image, modelname, conf_thres, label)
     return results
 
-def inference_mmdet_segm(image, modelname, conf_thres, label):
+def inference_mmdet_segm(image, modelname, conf_thres, label, classname="person"):
     model_checkpoint = modelpath(modelname)
     model_config = os.path.splitext(model_checkpoint)[0] + ".py"
     model_device = get_device()
@@ -563,6 +563,8 @@ def inference_mmdet_segm(image, modelname, conf_thres, label):
         results = [[],[],[],[]]
 
     for i in filter_inds:
+        if classname != "None" and classes[labels[i]] != classname:
+            continue
         results[0].append(label + "-" + classes[labels[i]])
         results[1].append(bboxes[i])
         results[2].append(segms[i])