Skip to content

Commit

Permalink
more refactoring, store/restore DOM instead of using revert functions…
Browse files Browse the repository at this point in the history 
…, colour images, capture image alt text
  • Loading branch information
@seanmcguire12
seanmcguire12 committed Oct 3, 2024
1 parent 917057c commit 4c5edf4
Show file tree
Hide file tree
Showing 2 changed files with 73 additions and 145 deletions.
37 changes: 16 additions & 21 deletions tarsier/core.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -122,6 +122,7 @@ async def page_to_text_colour_tag(
tagless: bool = False,
) -> Tuple[str, dict[int, TagMetadata]]:
adapter = adapter_factory(driver)
stored_dom = await self.store_dom(driver)

coloured_elems, tag_to_xpath, inserted_id_strings = await self._colour_based_tagify(
adapter, tag_text_elements, tagless
Expand All @@ -139,7 +140,7 @@ async def page_to_text_colour_tag(
adapter, all_detected_coloured_elems, inserted_id_strings, tagless
)

await self._cleanup(adapter)
await self.restore_dom(driver, stored_dom)

annotations_formatted = format_text(combined_annotations)
return annotations_formatted, tag_to_xpath
Expand Down Expand Up @@ -369,8 +370,6 @@ async def _create_annotations(
else:
annotations.append(tag_annotation)

await self._remove_text_bounding_boxes(adapter)

fixed_top_annotations = self.sort_annotations(fixed_top_annotations)
annotations = self.sort_annotations(annotations)
fixed_bottom_annotations = self.sort_annotations(fixed_bottom_annotations)
Expand All @@ -381,39 +380,35 @@ async def _create_annotations(

return combined_annotations

async def _cleanup(self, adapter: BrowserAdapter) -> None:
await self._remove_text_bounding_boxes(adapter)
await self._revert_colour_tagging(adapter)

async def _remove_tags(self, adapter: BrowserAdapter) -> None:
async def _store_dom(self, adapter: BrowserAdapter) -> str:
await self._load_tarsier_utils(adapter)
script = "return window.removeTags();"
stored_dom = await adapter.run_js("return window.storeDOM();")
return stored_dom

await adapter.run_js(script)

async def remove_tags(self, driver: AnyDriver) -> None:
async def store_dom(self, driver: AnyDriver) -> str:
adapter = adapter_factory(driver)
await self._remove_tags(adapter)
return await self._store_dom(adapter)

async def _remove_text_bounding_boxes(self, adapter: BrowserAdapter) -> None:
async def _restore_dom(self, adapter: BrowserAdapter, stored_dom: str) -> None:
await self._load_tarsier_utils(adapter)
script = "return window.removeTextBoundingBoxes();"

script = f"return window.restoreDOM({json.dumps(stored_dom)});"
await adapter.run_js(script)

async def remove_text_bounding_boxes(self, driver: AnyDriver) -> None:
async def restore_dom(self, driver: AnyDriver, stored_dom: str) -> None:
adapter = adapter_factory(driver)
await self._remove_text_bounding_boxes(adapter)
await self._restore_dom(adapter, stored_dom)


async def _revert_colour_tagging(self, adapter: BrowserAdapter) -> None:
async def _remove_tags(self, adapter: BrowserAdapter) -> None:
await self._load_tarsier_utils(adapter)
script = "return window.revertColourBasedTagify();"
script = "return window.removeTags();"

await adapter.run_js(script)

async def revert_colour_tagging(self, driver: AnyDriver) -> None:
async def remove_tags(self, driver: AnyDriver) -> None:
adapter = adapter_factory(driver)
await self._revert_colour_tagging(adapter)
await self._remove_tags(adapter)

async def _load_tarsier_utils(self, adapter: BrowserAdapter) -> None:
await adapter.run_js(self._js_utils)
Expand Down
181 changes: 57 additions & 124 deletions tarsier/tag_utils.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,6 @@ interface Window {
insertedIdStrings: string[];
};
hideNonColouredElements: () => void;
getElementHtmlByXPath: (xpath: string) => string;
createTextBoundingBoxes: () => void;
documentDimensions: () => { width: number; height: number };
getElementBoundingBoxes: (xpath: string) => {
Expand All @@ -47,8 +46,8 @@ interface Window {
reColourElements: (colouredElems: ColouredElem[]) => ColouredElem[];
disableTransitionsAndAnimations: () => void;
enableTransitionsAndAnimations: () => void;
removeTextBoundingBoxes: () => void;
revertColourBasedTagify: () => void;
restoreDOM: (storedDOM: string) => void;
storeDOM: () => string;
}

interface TagMetadata {
Expand All @@ -68,6 +67,24 @@ const tarsierSelector = `#${tarsierId}`;
const reworkdVisibilityAttribute = "reworkd-original-visibility";
type TagSymbol = "#" | "$" | "@" | "%" | "";

let originalDOM = document.body.cloneNode(true);

window.storeDOM = () => {
originalDOM = document.body.cloneNode(true);
console.log("DOM state stored.");
return document.body.outerHTML;
};


window.restoreDOM = (storedDOM) => {
console.log("Restoring DOM");
if (storedDOM) {
document.body.innerHTML = storedDOM;
} else {
console.error("No DOM state was provided.");
}
};

const elIsVisible = (el: HTMLElement) => {
const rect = el.getBoundingClientRect();
const computedStyle = window.getComputedStyle(el);
Expand Down Expand Up @@ -870,8 +887,7 @@ window.hideNonColouredElements = () => {

if (
!element.hasAttribute("data-colored") ||
element.getAttribute("data-colored") !== "true" ||
isImageElement(element)
element.getAttribute("data-colored") !== "true"
) {
element.style.visibility = "hidden";
} else {
Expand Down Expand Up @@ -979,7 +995,7 @@ window.colourBasedTagify = (
tagMappingWithTagMeta: { [p: number]: TagMetadata };
insertedIdStrings: string[];
} => {
const { tagMapping, tagMappingWithTagMeta } = createTagMappings(tagLeafTexts);
const tagMappingWithTagMeta = window.tagifyWebpage(tagLeafTexts);

window.removeTags();

Expand All @@ -988,34 +1004,19 @@ window.colourBasedTagify = (
tagless,
);

const elements = collectElementsToColor(tagMapping);
const elements = collectElementsToColor(tagMappingWithTagMeta);

const colorAssignments = getColorsForElements(elements);

const colorMapping = createColorMappingAndApplyStyles(
elements,
colorAssignments,
tagMapping,
tagMappingWithTagMeta,
);

return { colorMapping, tagMappingWithTagMeta, insertedIdStrings };
};

function createTagMappings(tagLeafTexts: boolean): {
tagMapping: { [key: number]: string };
tagMappingWithTagMeta: { [key: number]: TagMetadata };
} {
const tagMappingWithTagMeta = window.tagifyWebpage(tagLeafTexts);
const tagMapping = Object.entries(tagMappingWithTagMeta).reduce(
(acc, [id, meta]) => {
acc[parseInt(id)] = meta.xpath;
return acc;
},
{} as { [key: number]: string },
);
return { tagMapping, tagMappingWithTagMeta };
}

function insertIdStringsIntoTextNodes(
tagMappingWithTagMeta: { [key: number]: TagMetadata },
tagless: boolean,
Expand All @@ -1042,12 +1043,12 @@ function insertIdStringsIntoTextNodes(
}

function collectElementsToColor(
tagMapping: { [key: number]: string },
tagMappingWithTagMeta: { [key: number]: TagMetadata },
): HTMLElement[] {
const elements: HTMLElement[] = [];
const viewportWidth = window.innerWidth;
Object.keys(tagMapping).forEach((id) => {
let xpath = tagMapping[parseInt(id)];
Object.values(tagMappingWithTagMeta).forEach((meta) => {
const { tarsierId: id, xpath } = meta;
const node = document.evaluate(
xpath,
document,
Expand All @@ -1068,7 +1069,7 @@ function collectElementsToColor(
rect.left >= 0 &&
rect.right <= viewportWidth
) {
node.setAttribute("data-id", id);
node.setAttribute("data-id", id.toString());
elements.push(node);
}
}
Expand All @@ -1088,12 +1089,14 @@ function getColorsForElements(
function createColorMappingAndApplyStyles(
elements: HTMLElement[],
colorAssignments: Map<HTMLElement, string>,
tagMapping: { [key: number]: string },
tagMappingWithTagMeta: { [key: number]: TagMetadata },
): ColouredElem[] {
const colorMapping: ColouredElem[] = [];
const bodyRect = document.body.getBoundingClientRect();
const attribute = "data-colored";
const taggedElements = new Set(Object.values(tagMapping));
const taggedElements = new Set(
Object.values(tagMappingWithTagMeta).map((meta) => meta.xpath)
);

elements.forEach((element) => {
const id = parseInt(element.getAttribute("data-id")!);
Expand All @@ -1104,14 +1107,17 @@ function createColorMappingAndApplyStyles(
(midpoint[0] - bodyRect.left) / bodyRect.width,
(midpoint[1] - bodyRect.top) / bodyRect.height,
];
const idSymbol = createIdSymbol(id, element);

const symbol = getTagSymbol(element) || "";
const idSymbol = `[ ${symbol}${symbol ? " " : ""}${id} ]`;

const { isFixed, fixedPosition } = getFixedPosition(element);

colorMapping.push({
id,
idSymbol,
color,
xpath: tagMapping[id],
xpath: tagMappingWithTagMeta[id].xpath,
midpoint,
normalizedMidpoint,
width: rect.width,
Expand Down Expand Up @@ -1139,6 +1145,8 @@ function applyStylesToElement(
(element as HTMLInputElement).type === "checkbox"
) {
applyStylesToCheckbox(element as HTMLInputElement, color, attribute);
} else if (element.tagName.toLowerCase() === "img") {
applyStylesToImage(element as HTMLImageElement, color, attribute);
} else {
element.style.setProperty("background-color", color, "important");
element.style.setProperty("color", color, "important");
Expand Down Expand Up @@ -1205,6 +1213,24 @@ function applyStylesToCheckbox(
});
}

function applyStylesToImage(element: HTMLImageElement, color: string, attribute: string) {
const imageWidth = element.offsetWidth;
const imageHeight = element.offsetHeight;

const rgbToHex = (rgb: string) => {
const result = rgb.match(/\d+/g);
return result
? result.map((x) => parseInt(x).toString(16).padStart(2, "0")).join("")
: "000000";
};

const hexColor = rgbToHex(color);
const newSrc = `https://craftypixels.com/placeholder-image/${imageWidth}x${imageHeight}/${hexColor}/${hexColor}`;

element.setAttribute("src", newSrc);
element.setAttribute(attribute, "true");
}

function applyStylesToLink(
element: HTMLElement,
taggedElements: Set<string>,
Expand Down Expand Up @@ -1256,60 +1282,6 @@ function applyStylesToLink(
}
}

window.revertColourBasedTagify = () => {
document.querySelectorAll("[data-colored]").forEach((element) => {
const htmlElement = element as HTMLElement;
if (htmlElement.tagName.toLowerCase() === "input" && (htmlElement as HTMLInputElement).type === "checkbox") {
const checkboxElement = htmlElement as HTMLInputElement;
checkboxElement.style.removeProperty("width");
checkboxElement.style.removeProperty("height");
checkboxElement.style.removeProperty("background-color");
checkboxElement.style.removeProperty("border");
checkboxElement.style.removeProperty("appearance");
checkboxElement.style.removeProperty("border-radius");
checkboxElement.style.removeProperty("position");
checkboxElement.style.removeProperty("cursor");
checkboxElement.removeAttribute("data-colored");
} else {
htmlElement.style.removeProperty("background-color");
htmlElement.style.removeProperty("color");
htmlElement.style.removeProperty("border-color");
htmlElement.style.removeProperty("opacity");
htmlElement.removeAttribute("data-colored");
}
});

document.querySelectorAll("[data-id]").forEach((element) => {
const htmlElement = element as HTMLElement;
htmlElement.style.removeProperty("display");
htmlElement.removeAttribute("data-id");
});

document.querySelectorAll("[data-id] *").forEach((child) => {
const htmlElement = child as HTMLElement;
htmlElement.style.removeProperty("visibility");
});
};

window.getElementHtmlByXPath = function (xpath: string): string {
try {
const result = document.evaluate(
xpath,
document,
null,
XPathResult.FIRST_ORDERED_NODE_TYPE,
null,
);
const element = result.singleNodeValue as HTMLElement | null;
return element
? element.outerHTML
: "No element matches the provided XPath.";
} catch (error) {
console.error("Error evaluating XPath:", error);
return "";
}
};

function createIdSymbol(idNum: number, el: HTMLElement): string {
let idStr: string;
if (isInteractable(el)) {
Expand Down Expand Up @@ -1413,45 +1385,6 @@ window.createTextBoundingBoxes = () => {
};


window.removeTextBoundingBoxes = () => {
const styleSheets = Array.from(document.styleSheets);
styleSheets.forEach((styleSheet) => {
try {
if (styleSheet && styleSheet.cssRules) {
Array.from(styleSheet.cssRules).forEach((rule, index) => {
if (
rule instanceof CSSStyleRule &&
(rule.selectorText === ".tarsier-highlighted-word" ||
rule.selectorText === ".tarsier-space")
) {
styleSheet.deleteRule(index);
}
});
}
} catch (error) {
console.error("Error accessing style sheet: ", error);
}
});

// Revert the highlighted elements
document.querySelectorAll(".tarsier-highlighted-word, .tarsier-space").forEach((span) => {
if (span.parentNode) {
span.parentNode.replaceChild(
document.createTextNode(span.textContent || ""),
span
);
}
});

document.querySelectorAll("iframe").forEach((iframe) => {
try {
iframe.contentWindow?.postMessage({ action: "removeHighlight" }, "*");
} catch (error) {
console.error("Error accessing iframe content: ", error);
}
});
};

window.documentDimensions = () => {
return {
width: document.documentElement.scrollWidth,
Expand Down

0 comments on commit 4c5edf4

Please sign in to comment.