""" 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.""" @staticmethod 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 @staticmethod 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 @staticmethod 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