Update app.py

app.py

@@ -20,9 +20,25 @@ processor = AutoProcessor.from_pretrained("microsoft/Florence-2-base-ft", trust_
 onnx_model = YOLO("models/best.onnx", task='detect')
 
 
+def filter_detections(detections, target_class_name="mark"):
+    indices_to_keep = [i for i, class_name in enumerate(detections.data['class_name']) if
+                       class_name == target_class_name]
+
+    filtered_xyxy = detections.xyxy[indices_to_keep]
+    filtered_confidence = detections.confidence[indices_to_keep]
+    filtered_class_id = detections.class_id[indices_to_keep]
+    filtered_class_name = detections.data['class_name'][indices_to_keep]
+    detections.xyxy = filtered_xyxy
+    detections.confidence = filtered_confidence
+    detections.class_id = filtered_class_id
+    detections.data['class_name'] = filtered_class_name
+    return detections
+
+    
 def ends_with_number(s):
     return s[-1].isdigit()
 
+
 def ocr(image, prompt="<OCR>"):
     original_height, original_width = image.shape[:2]
     inputs = processor(text=prompt, images=image, return_tensors="pt").to(device)
@@ -45,6 +61,7 @@ def ocr(image, prompt="<OCR>"):
 
     return parsed_answer
 
+
 def parse_detection(detections):
     parsed_rows = []
     for i in range(len(detections.xyxy)):
@@ -103,6 +120,7 @@ def cut_and_save_image(image, parsed_detections, output_dir):
             output_path_list.append(output_path)
     return output_path_list
 
+
 def analysis(progress=gr.Progress()):
     progress(0, desc="Analyzing...")
     list_files = glob.glob("output/*.png")
@@ -123,6 +141,7 @@ def analysis(progress=gr.Progress()):
         print("Time taken:", time.time() - start_time)
     return pd.DataFrame(results.items(), columns=['Mark', 'Total']).reset_index(drop=False).rename(columns={'index': 'No.'})
 
+
 def inference(
     image_path,
     conf_threshold,
@@ -143,6 +162,7 @@ def inference(
 
     results = onnx_model(image, conf=conf_threshold, iou=iou_threshold)[0]
     detections = sv.Detections.from_ultralytics(results)
+    detections = filter_detections(detections)
     parsed_detections = parse_detection(detections)
     output_dir = "output"
     # Check if the output directory exists, clear all the files inside