import gradio as gr import os import sys import json from PIL import Image,ImageDraw import tempfile from inference_sdk import InferenceHTTPClient from ultralytics import YOLO import ultralytics 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= API_KEY ) 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 localOutlineBox(image, selected_model="yolo_accurate"): model_paths = { "yolo_accurate": "./vt_dataset_yolov12_v7_weights.pt", "yolo_extra_large": "./VT_dataset_2_Yolov12_Extra_large.pt" } model_path = model_paths.get(selected_model, "./vt_dataset_yolov12_v7_weights.pt") model_path = "./vt_dataset_yolov12_v7_weights.pt" if os.path.exists(model_path): print("model exists") print(f'current model path is: {model_path}') else: print("model is not available") model = YOLO(model_path) yolo_ret = model(image, verbose=False) useful_ret = yolo_ret[0] names = model.names all_box_info = [] for bb in useful_ret.boxes: box_info = {} x, y, w, h = bb.xywh[0].tolist() cls_name = names[int(bb.cls)] box_info['class'] = cls_name box_info['x'] = x box_info['y'] = y box_info['width'] = w box_info['height'] = h box_info['confidence'] = float(bb.conf) box_info['class_id'] = int(bb.cls) all_box_info.append(box_info) print(f"total length of all box_info is: {len(all_box_info)}") return all_box_info 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) kept_other_boxes = [] for i, box in enumerate(other_boxes): should_keep = True for kept_box in kept_other_boxes: if calculate_iou(box, kept_box) > iou_threshold: should_keep = False break if should_keep: kept_other_boxes.append(box) return kept_other_boxes + kept_boxes def getCenterXDis(box1, box2): x0_1, y0_1, x1_1, y1_1 = box1 x0_2, y0_2, x1_2, y1_2 = box2 #centeral corrdinates center_x1 = (x0_1 + x1_1) / 2 center_y1 = (y0_1 + y1_1) / 2 center_x2 = (x0_2 + x1_2) / 2 center_y2 = (y0_2 + y1_2) / 2 return abs(center_x1 - center_x2) def to_ordinal(n): """ Converts an integer to its ordinal string representation. e.g., 1 -> "1st", 2 -> "2nd", 13 -> "13th" """ if not isinstance(n, int): raise TypeError("Input must be an integer.") # Check for 11th, 12th, 13th, which are special cases if 11 <= (n % 100) <= 13: suffix = 'th' else: # Check the last digit for all other cases last_digit = n % 10 if last_digit == 1: suffix = 'st' elif last_digit == 2: suffix = 'nd' elif last_digit == 3: suffix = 'rd' else: suffix = 'th' return f"{n}{suffix}" def packFilterBoxInfo(filter_box_info): # 数字类别映射 digit_classes = { 'one': '1', 'two': '2', 'three': '3', 'four': '4', 'five': '5', 'six': '6', 'seven': '7', 'eight': '8', 'nine': '9', 'zero': '0' } fraction_boxes = [] number_boxes = [] for box in filter_box_info: if box['class'] == 'fraction': fraction_boxes.append(box) elif box['class'] in digit_classes: number_boxes.append(box) fraction_boxes.sort(key=lambda x: x['x'] - x['width']/2) fraction_values = [] for frac_box in fraction_boxes: # fraction框的边界 frac_x = frac_box['x'] frac_y = frac_box['y'] frac_width = frac_box['width'] frac_height = frac_box['height'] # 定义分子分母的区域 numerator_numbers = [] denominator_numbers = [] # 遍历所有数字，判断是否在当前fraction框内 for num_box in number_boxes: # 检查数字是否在fraction框的水平范围内 if (frac_x - frac_width/2 <= num_box['x'] <= frac_x + frac_width/2 and frac_y - frac_height/2 <= num_box['y'] <= frac_y + frac_height/2): # 获取数字值 digit = digit_classes[num_box['class']] # 根据y坐标判断是分子还是分母 if num_box['y'] < frac_y: # 在分数线上方 numerator_numbers.append((num_box['x'], num_box['y'], num_box['width'], num_box['height'], digit)) else: # 在分数线下方 denominator_numbers.append((num_box['x'], num_box['y'], num_box['width'], num_box['height'], digit)) # 按x坐标排序 numerator_numbers.sort(key=lambda x: x[0]-x[2]/2) denominator_numbers.sort(key=lambda x: x[0]-x[2]/2) # 提取排序后的数字 numerator = ''.join(digit for _, _, _, _, digit in numerator_numbers) denominator = ''.join(digit for _, _, _, _, digit in denominator_numbers) if numerator == "": numerator = "?" if denominator == "": denominator = "?" fraction_values.append(f"{numerator}/{denominator}") return fraction_values #Assume its coordinate are top-left, bottom-right def getOverlap(box1, box2): b1_x1, b1_y1, b1_x2, b1_y2 = box1 b2_x1, b2_y1, b2_x2, b2_y2 = box2 inter_x1 = max(b1_x1, b2_x1) inter_y1= max(b1_y1, b2_y1) inter_x2 = min(b1_x2, b2_x2) inter_y2 = min(b1_y2, b2_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 #b1_area = abs(b1_x2 - b1_x1) * abs(b1_y2 - b1_y1) b2_area = abs(b2_x2 - b2_x1) * abs(b2_y2 - b2_y1) return inter_area / b2_area def tick2fraction(ticks, fractions): ret = [] used_ticks = set() for fi, frac in enumerate(fractions): all_dis = [] for ti, tick in enumerate(ticks): if ti not in used_ticks: dis = getCenterXDis(tick, frac) all_dis.append((dis, ti)) if len(all_dis) == 0: #print(f"no tick found for fraction {fi}") #no tick found for this fraction break all_dis.sort(key=lambda x: x[0]) min_dis_index = all_dis[0][1] used_ticks.add(min_dis_index) ret.append(f"T{min_dis_index}-F{fi}") return ret def generate_textual_description(box_info): fraction_values = packFilterBoxInfo(box_info) # Create a dictionary to store information by class ID class_summary = {c: [] for c in classes} for box in box_info: c_name = box['class'] if c_name not in class_summary: continue else: x, y, w, h = box['x'], box['y'], box['width'], box['height'] 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 kept_zero_boxes = [] for zero_box in class_summary['zero']: kept_zero = True for fra_box in class_summary['fraction']: if getOverlap(fra_box, zero_box) >= 0.5: kept_zero = False break for tick_box in class_summary['tick']: if getOverlap(tick_box, zero_box) >= 0.5: kept_zero = False break if kept_zero: kept_zero_boxes.append(zero_box) kept_one_boxes = [] for one_box in class_summary['one']: kept_one = True for fra_box in class_summary['fraction']: if getOverlap(fra_box, one_box) >= 0.5: kept_one = False break for tick_box in class_summary['tick']: if getOverlap(tick_box, one_box) >= 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 len(kept_zero_boxes) >= 1: left_most_zero_cor = kept_zero_boxes[0] textual_description += f"\nThere 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 len(kept_one_boxes) >= 1: right_most_one_cor = kept_one_boxes[-1] textual_description += f"\nThere 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(): class_name = cid if class_name not in present_classes: continue count = len(boxes) boxes.sort(key=lambda x: x[0]) # it has been the x of the top-left corner if count > 0: textual_description += f"\nThere are {count} {class_name}s. Their coordinates are: " for box in boxes: textual_description += f"(({box[0]:.2f}, {box[1]:.2f}), ({box[2]:.2f}, {box[3]:.2f})), " if (class_name == "fraction"): textual_description += f"\nThe fraction numbers from left to right are: {fraction_values}. " tick2fra = tick2fraction(class_summary['tick'], class_summary['fraction']) tick2fraction_des = "" for cor in tick2fra: tick_part, fraction_part = cor.split('-') fraction_idx = int(fraction_part[1:]) tick_idx = int(tick_part[1:]) tick2fraction_des += f"{to_ordinal(fraction_idx + 1)} fraction is associated with {to_ordinal(tick_idx + 1)} tick. " textual_description += tick2fraction_des return textual_description def drawWithAllBox_info(pil_image, box_info): colors = ['red', 'green', 'blue', 'orange', 'purple', 'cyan', 'magenta', 'yellow', 'brown', 'pink', 'gray', 'lime', 'navy'] draw = ImageDraw.Draw(pil_image) for box in box_info: x, y, w, h = box['x'], box['y'], box['width'], box['height'] class_id = box['class_id'] color = 'black' if class_id < len(colors): color = colors[class_id] draw.rectangle([x-w/2, y-h/2, x+w/2, y+h/2], outline=color, width=2) return pil_image def online_process_image(image): if image is None: # Ensure we always return 3 values to prevent errors. return None, "", "" pil_image = image.copy() if hasattr(image, 'copy') else Image.fromarray(image) roboflow_ret = RoboFlowGetOutlineBoxesPIL(pil_image) all_box_info = parse_roboflow_result(roboflow_ret, classes) del roboflow_ret kept_box_info = filter_overlapping_boxes(all_box_info) del all_box_info boxed_img = drawWithAllBox_info(pil_image, kept_box_info) textual = generate_textual_description(kept_box_info) json_str = json.dumps(kept_box_info, indent=2) return boxed_img, textual, json_str def process_image(image, selected_model = "yolo_accurate"): print("start processing image") if image is None: # Ensure we always return 3 values to prevent errors. return None, "", "" #pil_image = image.copy() if hasattr(image, 'copy') else Image.fromarray(image) all_box_info = localOutlineBox(image, selected_model) kept_box_info = filter_overlapping_boxes(all_box_info) del all_box_info boxed_img = drawWithAllBox_info(image, kept_box_info) textual = generate_textual_description(kept_box_info) json_str = json.dumps(kept_box_info, indent=2) return boxed_img, textual, json_str with gr.Blocks( css=r""" /* 精确命中这个 Textbox 的文本区 */ textarea[aria-label="Textual Description"] { font-size: 28px !important; line-height: 1.6 !important; font-family: "Segoe UI", "Helvetica Neue", Arial, sans-serif !important; color: #222 !important; } /* 兜底选择器 */ #desc textarea, .large-font textarea, div[data-testid="textbox"] textarea { font-size: 24px !important; } """ ) as demo: # --- 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") with gr.Row(): model_list = gr.Dropdown( choices=["yolo_accurate", "yolo_extra_large"], value="yolo_accurate", label="Select Model", info="Choose the YOLO model for detection" ) 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, elem_id="desc", elem_classes=["large-font"]) json_state = gr.State("") download_btn = gr.DownloadButton( label="Download Box Info as JSON", interactive=False # Start as disabled. ) # --- Backend Functions --- def create_json_file(json_str): """Creates a temp file with JSON content and returns its path.""" if not json_str: return None with tempfile.NamedTemporaryFile( prefix="detection_info_", mode='w', delete=False, suffix='.json', encoding='utf-8' ) as f: f.write(json_str) return f.name def _process_and_prepare_download(image, selected_model): """Processes the image, creates the JSON file, and updates the UI.""" boxed_img, textual, json_str = process_image(image, selected_model) filepath = create_json_file(json_str) # Use the legacy gr.update() for compatibility with older Gradio versions. download_update = gr.update(value=filepath, interactive=True) return boxed_img, textual, json_str, download_update 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_and_prepare_download, inputs=[img_input, model_list], # The output now includes the download button itself. outputs=[img_out, text_out, json_state, download_btn] ) # The download_btn no longer needs its own click event. demo.launch()