Update app.py

app.py

@@ -7,6 +7,98 @@ import tempfile
 
 classes = ['zero', 'one', 'two', 'three', 'four', 'five', 'six', 'seven', 'eight', 'nine', 'tick', 'fraction']
 API_KEY = os.environ.get("ROBOFLOW_API_KEY")
+ACCESS_TOKEN = os.environ.get("ACCESS_TOKEN")
+
+
+def RoboFlowGetOutlineBoxesPIL(pil_img):
+    client = InferenceHTTPClient(
+        api_url="https://detect.roboflow.com",
+        api_key= roboflow_api
+    )
+    result = client.run_workflow(
+        workspace_name="mathnet-mmpuo",
+        workflow_id="custom-workflow-2",
+        images={
+            "image": pil_img
+        },
+        use_cache=True # cache workflow definition for 15 minutes
+    )
+    return result
+def calculate_iou(box1, box2):
+    """计算两个框的IoU
+    box格式: (x, y, width, height)
+    """
+    # 计算每个框的左上角和右下角坐标
+    box1_x1 = box1['x'] - box1['width']/2
+    box1_y1 = box1['y'] - box1['height']/2
+    box1_x2 = box1['x'] + box1['width']/2
+    box1_y2 = box1['y'] + box1['height']/2
+    
+    box2_x1 = box2['x'] - box2['width']/2
+    box2_y1 = box2['y'] - box2['height']/2
+    box2_x2 = box2['x'] + box2['width']/2
+    box2_y2 = box2['y'] + box2['height']/2
+    
+    # 计算交集区域的坐标
+    inter_x1 = max(box1_x1, box2_x1)
+    inter_y1 = max(box1_y1, box2_y1)
+    inter_x2 = min(box1_x2, box2_x2)
+    inter_y2 = min(box1_y2, box2_y2)
+    
+    # 计算交集面积
+    if inter_x1 < inter_x2 and inter_y1 < inter_y2:
+        inter_area = (inter_x2 - inter_x1) * (inter_y2 - inter_y1)
+    else:
+        return 0.0
+    
+    # 计算两个框的面积
+    box1_area = box1['width'] * box1['height']
+    box2_area = box2['width'] * box2['height']
+    
+    # 计算并集面积
+    union_area = box1_area + box2_area - inter_area
+    
+    # 返回IoU
+    return inter_area / union_area
+
+def parse_roboflow_result(result, kept_classes):
+    all_box_info = []
+    for box_info in result[0]['predictions']['predictions']['predictions']:
+        if box_info['class'] in kept_classes:
+            all_box_info.append(box_info)
+    return all_box_info
+
+def filter_overlapping_boxes(filter_box_info, iou_threshold=0.5):
+
+    digit_classes = {'zero', 'one', 'two', 'three', 'four', 
+                    'five', 'six', 'seven', 'eight', 'nine'}
+    
+    # 分离数字框和其他框
+    digit_boxes = []
+    other_boxes = []
+    
+    for box in filter_box_info:
+        if box['class'] in digit_classes:
+            digit_boxes.append(box)
+        else:
+            other_boxes.append(box)
+    
+
+    digit_boxes.sort(key=lambda x: x['confidence'], reverse=True)
+    kept_boxes = []
+
+    for i, box in enumerate(digit_boxes):
+        should_keep = True
+        
+        for kept_box in kept_boxes:
+            if calculate_iou(box, kept_box) > iou_threshold:
+                should_keep = False
+                break
+        if should_keep:
+            kept_boxes.append(box)
+    
+    return other_boxes + kept_boxes
+
 def packFilterBoxInfo(filter_box_info):
     # 数字类别映射
     digit_classes = {
@@ -95,23 +187,34 @@ def generate_textual_description(box_info):
             class_summary[c_name].append([x-w/2, y-h/2, x+w/2, y+h/2])
     # Generate a summary for each class
     #the index of the left one
-    class_summary['zero'].sort()
-    left_most_zero_cor = class_summary['zero'][0]
-    left_zero = True 
-    class_summary['one'].sort()
-    right_most_one_cor = class_summary['one'][-1]
-    right_one = True
-    for fra_box in class_summary['fraction']:
-        if getOverlap(fra_box, left_most_zero_cor) >= 0.5:
-            left_zero = False
-        if getOverlap(fra_box, right_most_one_cor) >= 0.5:
-            left_one = False
-    textual_description = ""
+    kept_zero_boxes = []
+    for zero_box in class_summary['zero']:
+        kept_zero = True
+        for fra_box in class_summary['fraction']:
+            if getOverlap(fra_box, left_most_zero_cor) >= 0.5:
+                kept_zero = False
+                break
+        if kept_zero:
+            kept_zero_boxes.append(zero_box)
+    kept_one_boxes = []
+    for one_box in class_summary['fraction']:
+        kept_one = True
+        for fra_box in class_summary['fraction']:
+            if getOverlap(fra_box, right_most_one_cor) >= 0.5:
+                kept_one = False
+                break
+        if kept_one:
+            kept_one_boxes.append(one_box)
+    kept_zero_boxes.sort(key = lambda x: x[0])
+    kept_one_boxes.sort(key = lambda x: x[0])
+    textual_description = "" #final output
     textual_description += "The key elements are interpreted via visual translator. Their coordinates are represented as outlined boxes (top-left, bottom-right)"
     #print(f"The key elements are interpreted via visual translator. Their coordinates are represented as outlined boxes (top-left, bottom-right)")
-    if left_zero:
+    if len(kept_zero_boxes) >= 1:
+        left_most_zero_cor = kept_zero_boxes[0]
         textual_description += f"There is a zero on the left side of the number line. Its coordinate is (({left_most_zero_cor[0]:.2f}, {left_most_zero_cor[1]:.2f}), ({left_most_zero_cor[2]:.2f}, {left_most_zero_cor[3]:.2f}))"
-    if right_one:
+    if len(kept_one_boxes) >= 1:
+        right_most_one_cor = kept_one_boxes[-1]
         textual_description += f"There is a one on the right side of the number line. Its coordinate is (({right_most_one_cor[0]:.2f}, {right_most_one_cor[1]:.2f}), ({right_most_one_cor[2]:.2f}, {right_most_one_cor[3]:.2f}))"
     present_classes = ['fraction', 'tick']
     for cid, boxes in class_summary.items():
@@ -147,41 +250,76 @@ def greet(name):
 
 def process_image(image):
     if image is None:
-        return None, "", ""
+        return None, "", "", None
     pil_image = image.copy() if hasattr(image, 'copy') else Image.fromarray(image)
     boxed_img = drawWithAllBox_info(pil_image, test_box_info)
     textual = generate_textual_description(test_box_info)
     json_str = json.dumps(test_box_info, indent=2)
     return boxed_img, textual, json_str
-'''
-def download_json(json_str):
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".json", mode="w") as f:
-        f.write(json_str)
-        temp_path = f.name
-    return gr.File.update(value=temp_path, visible=True)
-'''
+
+
 with gr.Blocks() as demo:
-    img_input = gr.Image(type="pil", label="Upload Image")
-    run_btn = gr.Button("Run Detection")
-    img_out = gr.Image(type="pil", label="Image with Boxes")
-    text_out = gr.Textbox(label="Textual Description", lines=8)
-    json_state = gr.State("")
-    download_btn = gr.DownloadButton(
-        label="Download Box Info as JSON"
-    )
+    # --- Authentication Layer ---
+    with gr.Row():
+        token_input = gr.Textbox(
+            label="Invite Token",
+            type="password",
+            placeholder="Enter your invite token to unlock the app"
+        )
+        unlock_btn = gr.Button("Unlock")
+
+    status_text = gr.Markdown()
+
+    # --- Main Application (initially hidden) ---
+    with gr.Column(visible=False) as main_app:
+        img_input = gr.Image(type="pil", label="Upload Image")
+        run_btn = gr.Button("Run Detection")
+        img_out = gr.Image(type="pil", label="Image with Boxes")
+        text_out = gr.Textbox(label="Textual Description", lines=8)
+        json_state = gr.State("")
+        download_btn = gr.DownloadButton(
+            label="Download Box Info as JSON"
+        )
+
+    # --- Backend Functions ---
     def _process(image):
         boxed_img, textual, json_str = process_image(image)
         return boxed_img, textual, json_str
+
     def create_and_download_json(json_str):
         if not json_str:
             return None
         with tempfile.NamedTemporaryFile(
-            mode='w', delete=False, prefix="detection_box_", suffix='.json', encoding='utf-8'
+            prefix="detection_info_",
+            mode='w', delete=False, suffix='.json', encoding='utf-8'
         ) as f:
             f.write(json_str)
             return f.name
 
-    run_btn.click(_process, inputs=img_input, outputs=[img_out, text_out, json_state])
-    download_btn.click(create_and_download_json, inputs=json_state, outputs=download_btn)
+    def check_token(token):
+        # Securely check if the token is correct
+        if ACCESS_TOKEN and token == ACCESS_TOKEN:
+            return gr.update(visible=True), "Token accepted. You can now use the application."
+        else:
+            return gr.update(visible=False), "Invalid token. Please try again."
+
+    # --- Event Listeners ---
+    unlock_btn.click(
+        check_token,
+        inputs=token_input,
+        outputs=[main_app, status_text]
+    )
+
+    run_btn.click(
+        _process,
+        inputs=img_input,
+        outputs=[img_out, text_out, json_state]
+    )
+
+    download_btn.click(
+        create_and_download_json,
+        inputs=json_state,
+        outputs=download_btn
+    )
 
 demo.launch()