Spaces:
Running
Running
| import argparse | |
| import json | |
| from pathlib import Path | |
| from bs4 import BeautifulSoup | |
| import cv2 | |
| import re | |
| def main(args): | |
| # --- Phase 1: Crop and Save All Images First --- | |
| # 1. Load data | |
| mapping_data = json.loads(args.mapping.read_text()) | |
| uied_data = json.loads(args.uied.read_text()) | |
| original_image = cv2.imread(str(args.original_image)) | |
| if original_image is None: | |
| raise ValueError(f"Could not load the original image from {args.original_image}") | |
| # Get image shapes to calculate a simple, global scaling factor | |
| H_proc, W_proc, _ = uied_data['img_shape'] | |
| H_orig, W_orig, _ = original_image.shape | |
| scale_x = W_orig / W_proc | |
| scale_y = H_orig / H_proc | |
| print(f"Using global scaling for cropping: scale_x={scale_x:.3f}, scale_y={scale_y:.3f}") | |
| uied_boxes = { | |
| comp['id']: (comp['column_min'], comp['row_min'], comp['width'], comp['height']) | |
| for comp in uied_data['compos'] | |
| } | |
| # 2. Create a directory for cropped images | |
| crop_dir = args.output_html.parent / "cropped_images" | |
| crop_dir.mkdir(exist_ok=True) | |
| print(f"Saving cropped images to: {crop_dir.resolve()}") | |
| # 3. Iterate through mappings and save cropped images to files | |
| for region_id, region_data in mapping_data.items(): | |
| for placeholder_id, uied_id in region_data['mapping'].items(): | |
| if uied_id not in uied_boxes: | |
| print(f"Warning: UIED ID {uied_id} from mapping not found. Skipping placeholder {placeholder_id}.") | |
| continue | |
| uied_bbox = uied_boxes[uied_id] | |
| x_proc, y_proc, w_proc, h_proc = uied_bbox | |
| x_tf = x_proc * scale_x | |
| y_tf = y_proc * scale_y | |
| w_tf = w_proc * scale_x | |
| h_tf = h_proc * scale_y | |
| x1, y1 = int(x_tf), int(y_tf) | |
| x2, y2 = int(x_tf + w_tf), int(y_tf + h_tf) | |
| h_img, w_img, _ = original_image.shape | |
| x1, y1 = max(0, x1), max(0, y1) | |
| x2, y2 = min(w_img, x2), min(h_img, y2) | |
| cropped_img = original_image[y1:y2, x1:x2] | |
| if cropped_img.size == 0: | |
| print(f"Warning: Cropped image for {placeholder_id} is empty. Skipping.") | |
| continue | |
| output_path = crop_dir / f"{placeholder_id}.png" | |
| cv2.imwrite(str(output_path), cropped_img) | |
| # --- Phase 2: Use BeautifulSoup to Replace Placeholders by Order --- | |
| print("\nStarting offline HTML processing with BeautifulSoup...") | |
| html_content = args.gray_html.read_text() | |
| soup = BeautifulSoup(html_content, 'html.parser') | |
| # 1. Find all placeholder elements by their class, in document order. | |
| placeholder_elements = soup.find_all(class_="bg-gray-400") | |
| # 2. Get the placeholder IDs from the mapping file in the correct, sorted order. | |
| def natural_sort_key(s): | |
| return [int(text) if text.isdigit() else text.lower() for text in re.split('([0-9]+)', s)] | |
| ordered_placeholder_ids = [] | |
| # Sort region IDs numerically to process them in order | |
| for region_id in sorted(mapping_data.keys(), key=int): | |
| region_mapping = mapping_data[region_id]['mapping'] | |
| # Sort the placeholder IDs within each region naturally (e.g., ph1, ph2, ph10) | |
| sorted_ph_ids = sorted(region_mapping.keys(), key=natural_sort_key) | |
| ordered_placeholder_ids.extend(sorted_ph_ids) | |
| # 3. Check for count mismatches | |
| if len(placeholder_elements) != len(ordered_placeholder_ids): | |
| print(f"Warning: Mismatch in counts! Found {len(placeholder_elements)} gray boxes in HTML, but {len(ordered_placeholder_ids)} mappings.") | |
| else: | |
| print(f"Found {len(placeholder_elements)} gray boxes to replace.") | |
| # 4. Iterate through both lists, create a proper <img> tag, and replace the placeholder. | |
| for i, ph_element in enumerate(placeholder_elements): | |
| if i >= len(ordered_placeholder_ids): | |
| print(f"Warning: More gray boxes in HTML than mappings. Stopping at box {i+1}.") | |
| break | |
| ph_id = ordered_placeholder_ids[i] | |
| relative_img_path = f"{crop_dir.name}/{ph_id}.png" | |
| # --- Create a new <img> tag and replace the placeholder --- | |
| # a. Get all classes from the original placeholder to preserve styling. | |
| original_classes = ph_element.get('class', []) | |
| if 'bg-gray-400' in original_classes: | |
| original_classes.remove('bg-gray-400') # Remove the placeholder background | |
| # b. Create the new <img> tag | |
| img_tag = soup.new_tag("img", src=relative_img_path) | |
| img_tag['class'] = original_classes | |
| # c. Replace the placeholder with the new image tag. | |
| ph_element.replace_with(img_tag) | |
| # Save the modified HTML | |
| args.output_html.write_text(str(soup)) | |
| print(f"\nSuccessfully replaced {min(len(placeholder_elements), len(ordered_placeholder_ids))} placeholders.") | |
| print(f"Final HTML generated at {args.output_html.resolve()}") | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser(description="Replace placeholder divs in an HTML file with cropped images based on UIED mappings.") | |
| parser.add_argument("--mapping", type=Path, required=False, help="Path to the mapping JSON file from mapping.py.") | |
| parser.add_argument("--uied", type=Path, required=False, help="Path to the UIED JSON file.") | |
| parser.add_argument("--original-image", type=Path, required=False, help="Path to the original screenshot image.") | |
| parser.add_argument("--gray-html", type=Path, required=False, help="Path to the input HTML file with gray placeholders.") | |
| parser.add_argument("--output-html", type=Path, required=False, help="Path to save the final, modified HTML file.") | |
| parser.set_defaults( | |
| mapping=Path('data/tmp/mapping_full_test1.json'), | |
| uied=Path('data/tmp/ip/test1.json'), | |
| original_image=Path('data/input/test1.png'), | |
| gray_html=Path('data/output/test1_layout.html'), | |
| output_html=Path('data/output/test1_layout_final.html') | |
| ) | |
| args = parser.parse_args() | |
| main(args) | |