Spaces:
Sleeping
Sleeping
| """ | |
| Screenshot Annotator | |
| Generates annotated screenshots with red circles highlighting visual differences | |
| """ | |
| from typing import List, Dict, Any | |
| from state_schema import VisualDifference | |
| import os | |
| class ScreenshotAnnotator: | |
| """Annotates screenshots with visual difference indicators.""" | |
| def annotate_screenshot( | |
| screenshot_path: str, | |
| differences: List[VisualDifference], | |
| output_path: str, | |
| viewport_name: str = "desktop" | |
| ) -> bool: | |
| """ | |
| Annotate screenshot with red circles for differences. | |
| Args: | |
| screenshot_path: Path to original screenshot | |
| differences: List of visual differences | |
| output_path: Path to save annotated screenshot | |
| viewport_name: Viewport name (desktop/mobile) | |
| Returns: | |
| True if successful | |
| """ | |
| try: | |
| from PIL import Image, ImageDraw | |
| if not os.path.exists(screenshot_path): | |
| return False | |
| # Load image | |
| img = Image.open(screenshot_path).convert('RGB') | |
| draw = ImageDraw.Draw(img, 'RGBA') | |
| # Filter differences for this viewport | |
| viewport_diffs = [d for d in differences if d.viewport == viewport_name] | |
| # Draw circles and labels for each difference | |
| circle_radius = 30 | |
| label_offset = 40 | |
| for idx, diff in enumerate(viewport_diffs): | |
| if not diff.location: | |
| # Generate random location if not provided | |
| x = (idx % 3) * 300 + 150 | |
| y = (idx // 3) * 300 + 150 | |
| else: | |
| x = diff.location.get("x", 150) | |
| y = diff.location.get("y", 150) | |
| # Draw circle | |
| circle_color = ScreenshotAnnotator._get_color_by_severity(diff.severity) | |
| draw.ellipse( | |
| [(x - circle_radius, y - circle_radius), | |
| (x + circle_radius, y + circle_radius)], | |
| outline=circle_color, | |
| width=3 | |
| ) | |
| # Draw number label | |
| label_number = str(idx + 1) | |
| draw.text( | |
| (x - 8, y - 8), | |
| label_number, | |
| fill=(255, 255, 255), | |
| font=None | |
| ) | |
| # Create output directory | |
| os.makedirs(os.path.dirname(output_path), exist_ok=True) | |
| # Save annotated image | |
| img.save(output_path) | |
| return True | |
| except Exception as e: | |
| print(f"Error annotating screenshot: {str(e)}") | |
| return False | |
| def _get_color_by_severity(severity: str) -> tuple: | |
| """Get color based on severity level.""" | |
| if severity == "High": | |
| return (255, 0, 0, 200) # Red | |
| elif severity == "Medium": | |
| return (255, 165, 0, 200) # Orange | |
| else: | |
| return (0, 255, 0, 200) # Green | |
| def create_comparison_image( | |
| figma_screenshot: str, | |
| website_screenshot: str, | |
| figma_annotated: str, | |
| website_annotated: str, | |
| output_path: str | |
| ) -> bool: | |
| """ | |
| Create side-by-side comparison image. | |
| Args: | |
| figma_screenshot: Original Figma screenshot | |
| website_screenshot: Original website screenshot | |
| figma_annotated: Annotated Figma screenshot | |
| website_annotated: Annotated website screenshot | |
| output_path: Path to save comparison | |
| Returns: | |
| True if successful | |
| """ | |
| try: | |
| from PIL import Image | |
| # Load annotated images | |
| figma_img = Image.open(figma_annotated) | |
| website_img = Image.open(website_annotated) | |
| # Resize to same height | |
| max_height = max(figma_img.height, website_img.height) | |
| figma_img = figma_img.resize( | |
| (int(figma_img.width * max_height / figma_img.height), max_height) | |
| ) | |
| website_img = website_img.resize( | |
| (int(website_img.width * max_height / website_img.height), max_height) | |
| ) | |
| # Create side-by-side image | |
| total_width = figma_img.width + website_img.width + 20 | |
| comparison = Image.new('RGB', (total_width, max_height), (255, 255, 255)) | |
| comparison.paste(figma_img, (0, 0)) | |
| comparison.paste(website_img, (figma_img.width + 20, 0)) | |
| # Create output directory | |
| os.makedirs(os.path.dirname(output_path), exist_ok=True) | |
| # Save comparison | |
| comparison.save(output_path) | |
| return True | |
| except Exception as e: | |
| print(f"Error creating comparison image: {str(e)}") | |
| return False | |
| def annotate_all_screenshots(state: Dict[str, Any]) -> Dict[str, Any]: | |
| """ | |
| Annotate all screenshots with visual differences. | |
| Args: | |
| state: Current workflow state | |
| Returns: | |
| Updated state with annotated screenshots | |
| """ | |
| print("\n📸 Generating Annotated Screenshots...") | |
| annotator = ScreenshotAnnotator() | |
| figma_screenshots = state.get("figma_screenshots", {}) | |
| website_screenshots = state.get("website_screenshots", {}) | |
| visual_differences = state.get("visual_differences", []) | |
| execution_id = state.get("execution_id", "unknown") | |
| # Convert dict differences back to VisualDifference objects for the annotator | |
| vd_objects = [] | |
| for d in visual_differences: | |
| if isinstance(d, dict): | |
| vd_objects.append(VisualDifference( | |
| category=d.get("category", "visual"), | |
| severity=d.get("severity", "Medium"), | |
| issue_id=d.get("issue_id", "unknown"), | |
| title=d.get("title", "Difference"), | |
| description=d.get("description", ""), | |
| design_value=d.get("design_value", ""), | |
| website_value=d.get("website_value", ""), | |
| viewport=d.get("viewport", "desktop"), | |
| location=d.get("location") | |
| )) | |
| else: | |
| vd_objects.append(d) | |
| annotated_screenshots = state.get("annotated_screenshots", {}) | |
| for viewport_name, figma_path in figma_screenshots.items(): | |
| if viewport_name not in website_screenshots: | |
| continue | |
| print(f" 📝 Annotating {viewport_name} screenshots...") | |
| website_path = website_screenshots[viewport_name] | |
| figma_annotated = f"data/annotated/{execution_id}_figma_{viewport_name}.png" | |
| website_annotated = f"data/annotated/{execution_id}_website_{viewport_name}.png" | |
| # Annotate Figma screenshot | |
| if annotator.annotate_screenshot(figma_path, vd_objects, figma_annotated, viewport_name): | |
| print(f" ✓ Figma annotated: {figma_annotated}") | |
| # Annotate website screenshot | |
| if annotator.annotate_screenshot(website_path, vd_objects, website_annotated, viewport_name): | |
| print(f" ✓ Website annotated: {website_annotated}") | |
| # Create comparison | |
| comparison_path = f"data/comparisons/{execution_id}_{viewport_name}_comparison.png" | |
| if annotator.create_comparison_image(figma_path, website_path, figma_annotated, website_annotated, comparison_path): | |
| print(f" ✓ Comparison created: {comparison_path}") | |
| annotated_screenshots[f"{viewport_name}_comparison"] = comparison_path | |
| state["annotated_screenshots"] = annotated_screenshots | |
| return state | |