sai-Rohan/Dark_Pattern_Detection_API
0
1"""2Optimized inference pipeline for Multimodal Dark Pattern Detection.3Pipeline4--------5Image6 |7 +--> DINOv2 --------------------+8 | |9 +--> YOLO --> UI Graph --> GAT -+--> Fusion --> Classifier10 | |11 +--> EasyOCR --> Structured Text12 |13 +--> ModernBERT14Designed for:15 - FastAPI16 - Hugging Face Docker Space17 - GPU/CPU execution18 - One-time model loading19 - PIL Image input20 - No training code21 - No dataset dependencies22 - No disk-based inference cache23"""24 25from __future__ import annotations26 27import copy28import dataclasses29import logging30import math31import os32import re33import unicodedata34import difflib35import json36 37from dataclasses import dataclass38from typing import Any, Dict, List, Optional, Tuple, Union39 40import numpy as np41from PIL import Image42from huggingface_hub import snapshot_download43import torch44import torch.nn as nn45import torch.nn.functional as F46 47from torch_geometric.data import Data48from torch_geometric.nn import (49 GATv2Conv,50 global_mean_pool,51)52 53import easyocr54 55from transformers import (56 AutoTokenizer,57 AutoModel,58 Dinov2Model,59)60 61from huggingface_hub import hf_hub_download62 63 64# ============================================================65# LOGGING66# ============================================================67 68logging.basicConfig(69 level=logging.INFO,70 format="%(asctime)s | %(levelname)s | %(message)s",71)72 73logger = logging.getLogger("dark-pattern-inference")74 75 76# ============================================================77# DEVICE78# ============================================================79 80DEVICE = torch.device(81 "cuda" if torch.cuda.is_available() else "cpu"82)83 84logger.info("Device: %s", DEVICE)85 86if DEVICE.type == "cuda":87 logger.info(88 "GPU: %s",89 torch.cuda.get_device_name(0)90 )91 92 93# ============================================================94# MODEL REPOSITORY95# ============================================================96 97MODEL_REPO = "sai-Rohan/Dark_pattern_Detection_models"98 99 100# ============================================================101# INFERENCE CONFIG102# ============================================================103 104@dataclass105class Config:106 107 # --------------------------------------------------------108 # DINOv2109 # --------------------------------------------------------110 111 DINO_IMAGE_SIZE: int = 224112 DINO_OUTPUT_DIM: int = 768113 114 # --------------------------------------------------------115 # YOLO116 # --------------------------------------------------------117 118 YOLO_CONF: float = 0.25119 YOLO_IOU: float = 0.45120 YOLO_IMAGE_SIZE: int = 640121 YOLO_MAX_DETECTIONS: int = 300122 123 # --------------------------------------------------------124 # OCR125 # --------------------------------------------------------126 127 OCR_CONF: float = 0.65128 MIN_TEXT_CHARS: int = 2129 130 # --------------------------------------------------------131 # ModernBERT132 # --------------------------------------------------------133 134 TEXT_MAX_LENGTH: int = 256135 TEXT_OUTPUT_DIM: int = 768136 TEXT_POOLING: str = "mean"137 138 # --------------------------------------------------------139 # Fusion140 # --------------------------------------------------------141 142 FUSION_DIM: int = 256143 PROJECTION_DROPOUT: float = 0.2144 145 FUSION_MODE: str = "weighted"146 147 CLASSIFIER_DROPOUT_1: float = 0.3148 CLASSIFIER_DROPOUT_2: float = 0.2149 150 # --------------------------------------------------------151 # Graph152 # --------------------------------------------------------153 154 CONTAINMENT_THRESHOLD: float = 0.80155 CONTAINMENT_MIN_AREA_RATIO: float = 1.15156 157 ALIGNMENT_TOLERANCE: float = 0.02158 159 LARGE_BOX_WIDTH_RATIO: float = 0.85160 LARGE_BOX_HEIGHT_RATIO: float = 0.85161 162 ROW_OVERLAP_THRESHOLD: float = 0.5163 ROW_CENTER_TOLERANCE: float = 0.035164 165 COLUMN_OVERLAP_THRESHOLD: float = 0.5166 COLUMN_CENTER_TOLERANCE: float = 0.035167 168 OVERLAP_IOU_THRESHOLD: float = 0.20169 170 PROXIMITY_DISTANCE_THRESHOLD: float = 0.15171 172 READING_ORDER_ROW_TOLERANCE: float = 0.035173 174 ENABLE_CONNECTIVITY_FALLBACK: bool = True175 176 MIN_VALID_BOX_SIZE: float = 0.001177 178 # --------------------------------------------------------179 # Localization (Leave-One-Component-Out attribution)180 #181 # IMPORTANT:182 # There is no trained localization checkpoint. These183 # settings only control the *post-hoc* perturbation-based184 # attribution procedure that reuses the existing trained185 # classifier. See `leave_one_component_out_attribution`.186 # --------------------------------------------------------187 188 LOCALIZATION_TOP_K: int = 5189 190 # Components with attribution_logit <= this value are not191 # considered "supporting" the dark-pattern prediction.192 LOCALIZATION_MIN_ATTRIBUTION: float = 0.0193 194 195CFG = Config()196 197 198 199def get_device():200 """201 Determine the device at runtime.202 203 This must be called after the ZeroGPU allocation.204 """205 206 if torch.cuda.is_available():207 device = torch.device("cuda")208 else:209 device = torch.device("cpu")210 211 logger.info("Using device: %s", device)212 213 if device.type == "cuda":214 logger.info(215 "GPU: %s",216 torch.cuda.get_device_name(0)217 )218 219 return device220 221 222 223 224# ============================================================225# CHECKPOINT DOWNLOADER226# ============================================================227 228def download_model(filename: str) -> str:229 """230 Download one model file from the Hugging Face model repository.231 huggingface_hub automatically caches downloaded files.232 """233 234 logger.info("Preparing model: %s", filename)235 236 return hf_hub_download(237 repo_id=MODEL_REPO,238 filename=filename,239 )240 241 242# ============================================================243# DINO PREPROCESSOR244# ============================================================245 246IMAGENET_MEAN = np.array(247 [0.485, 0.456, 0.406],248 dtype=np.float32,249)250 251IMAGENET_STD = np.array(252 [0.229, 0.224, 0.225],253 dtype=np.float32,254)255 256 257class DinoPreprocessor:258 259 def __init__(self, image_size: int = 224):260 self.image_size = image_size261 self.resize_size = int(262 image_size * 256 / 224263 )264 265 def __call__(266 self,267 image: Image.Image,268 ) -> torch.Tensor:269 270 image = image.convert("RGB")271 272 w, h = image.size273 274 scale = self.resize_size / min(w, h)275 276 new_w = max(277 1,278 round(w * scale)279 )280 281 new_h = max(282 1,283 round(h * scale)284 )285 286 image = image.resize(287 (new_w, new_h),288 Image.Resampling.BICUBIC,289 )290 291 left = (new_w - self.image_size) // 2292 top = (new_h - self.image_size) // 2293 294 image = image.crop(295 (296 left,297 top,298 left + self.image_size,299 top + self.image_size,300 )301 )302 303 arr = (304 np.asarray(image).astype(305 np.float32306 ) / 255.0307 )308 309 arr = (310 arr - IMAGENET_MEAN311 ) / IMAGENET_STD312 313 tensor = torch.from_numpy(314 arr.transpose(2, 0, 1)315 ).float()316 317 return tensor318 319 320# ============================================================321# DINOv2322# ============================================================323 324class DinoExtractor(nn.Module):325 326 def __init__(327 self,328 checkpoint_path: str,329 output_dim: int = 768,330 device=None,331 ):332 333 super().__init__()334 self.device = (device if device is not None else torch.device("cpu"))335 336 self.output_dim = output_dim337 338 logger.info(339 "Loading DINOv2 checkpoint..."340 )341 342 checkpoint = torch.load(343 checkpoint_path,344 map_location="cpu",345 weights_only=False,346 )347 348 if "teacher" in checkpoint:349 raw_state = checkpoint["teacher"]350 351 elif "student" in checkpoint:352 raw_state = checkpoint["student"]353 354 else:355 raw_state = checkpoint356 357 state_dict = {}358 359 for key, value in raw_state.items():360 361 key = key.replace(362 "module.",363 ""364 )365 366 if key.startswith(367 "backbone."368 ):369 key = key[370 len("backbone.") :371 ]372 373 state_dict[key] = value374 375 logger.info(376 "DINO backbone parameters: %d",377 len(state_dict),378 )379 380 model = Dinov2Model.from_pretrained(381 "facebook/dinov2-base"382 )383 384 missing, unexpected = (385 model.load_state_dict(386 state_dict,387 strict=False,388 )389 )390 391 if len(missing) > 10:392 393 raise RuntimeError(394 "DINOv2 checkpoint mismatch: "395 f"{len(missing)} missing keys"396 )397 398 if model.config.hidden_size != output_dim:399 400 raise RuntimeError(401 f"DINO output dimension is "402 f"{model.config.hidden_size}, "403 f"expected {output_dim}"404 )405 406 self.model = model.to(self.device)407 408 self.model.eval()409 410 for parameter in self.model.parameters():411 parameter.requires_grad_(False)412 413 logger.info(414 "DINOv2 loaded successfully"415 )416 417 @torch.inference_mode()418 def forward(419 self,420 pixel_values: torch.Tensor,421 ) -> torch.Tensor:422 423 output = self.model(424 pixel_values=pixel_values425 )426 427 if (428 getattr(429 output,430 "pooler_output",431 None,432 )433 is not None434 ):435 feature = output.pooler_output436 437 else:438 feature = (439 output440 .last_hidden_state[:, 0]441 )442 443 if feature.shape[-1] != self.output_dim:444 raise RuntimeError(445 f"DINO feature dimension "446 f"{feature.shape[-1]} != "447 f"{self.output_dim}"448 )449 450 return feature451 452 453# ============================================================454# YOLO455# ============================================================456 457class YoloDetector:458 459 def __init__(460 self,461 checkpoint_path: str,462 conf: float,463 iou: float,464 image_size: int,465 max_det: int,466 device=None,467 ):468 469 470 self.device = (device if device is not None else torch.device("cpu"))471 472 from ultralytics.models.yolo.detect import (473 DetectionPredictor,474 )475 476 self.DetectionPredictor = (477 DetectionPredictor478 )479 480 logger.info(481 "Loading YOLO..."482 )483 484 checkpoint = torch.load(485 checkpoint_path,486 map_location=self.device,487 weights_only=False,488 )489 490 if checkpoint.get("ema") is not None:491 492 self.model = checkpoint["ema"]493 494 elif checkpoint.get("model") is not None:495 496 self.model = checkpoint["model"]497 498 else:499 500 raise RuntimeError(501 "YOLO checkpoint contains "502 "neither 'ema' nor 'model'"503 )504 505 self.model = (506 self.model507 .to(self.device)508 .eval()509 )510 511 for parameter in self.model.parameters():512 parameter.requires_grad_(False)513 514 self.conf = conf515 self.iou = iou516 self.image_size = image_size517 self.max_det = max_det518 519 names = self.model.names520 521 if isinstance(names, dict):522 523 self.class_names = dict(524 names525 )526 527 else:528 529 self.class_names = {530 i: name531 for i, name in enumerate(names)532 }533 534 logger.info(535 "YOLO classes: %s",536 self.class_names,537 )538 539 # IMPORTANT:540 # Build predictor ONCE instead of541 # once per request.542 self.predictor = (543 self.DetectionPredictor(544 overrides={545 "conf": self.conf,546 "iou": self.iou,547 "imgsz": self.image_size,548 "max_det": self.max_det,549 "device": self.device,550 "verbose": False,551 }552 )553 )554 555 logger.info(556 "YOLO loaded successfully"557 )558 559 @torch.inference_mode()560 def detect(561 self,562 image: Image.Image,563 ) -> List[Dict[str, Any]]:564 565 image = image.convert("RGB")566 567 results = self.predictor(568 source=image,569 model=self.model,570 )571 572 if not results:573 return []574 575 result = results[0]576 577 if (578 result.boxes is None579 or len(result.boxes) == 0580 ):581 return []582 583 boxes = (584 result.boxes.xyxy585 .detach()586 .cpu()587 .numpy()588 )589 590 confidences = (591 result.boxes.conf592 .detach()593 .cpu()594 .numpy()595 )596 597 class_ids = (598 result.boxes.cls599 .detach()600 .cpu()601 .numpy()602 .astype(int)603 )604 605 img_h, img_w = (606 result.orig_shape607 )608 609 detections = []610 611 for box, conf, class_id in zip(612 boxes,613 confidences,614 class_ids,615 ):616 617 x1, y1, x2, y2 = box618 619 detections.append(620 {621 "class_id": int(class_id),622 623 "class_name":624 self.class_names.get(625 int(class_id),626 "Other",627 ),628 629 "confidence":630 float(conf),631 632 "bbox": [633 float(x1),634 float(y1),635 float(x2),636 float(y2),637 ],638 639 "image_width":640 int(img_w),641 642 "image_height":643 int(img_h),644 }645 )646 647 return detections648 649 650# ============================================================651# GRAPH CONFIG652# ============================================================653 654@dataclass655class GraphConfig:656 657 containment_threshold: float = 0.80658 containment_min_area_ratio: float = 1.15659 660 alignment_tolerance: float = 0.02661 662 large_box_width_ratio: float = 0.85663 large_box_height_ratio: float = 0.85664 665 row_overlap_threshold: float = 0.5666 row_center_tolerance: float = 0.035667 668 column_overlap_threshold: float = 0.5669 column_center_tolerance: float = 0.035670 671 overlap_iou_threshold: float = 0.20672 673 proximity_distance_threshold: float = 0.15674 675 reading_order_row_tolerance: float = 0.035676 677 enable_connectivity_fallback: bool = True678 679 min_valid_box_size: float = 0.001680 681 682# ============================================================683# GRAPH SCHEMA684# ============================================================685 686@dataclass687class GraphSchema:688 689 class_names: List[str]690 691 num_classes: int692 693 node_feature_names: List[str]694 695 node_feature_dim: int696 697 edge_relation_names: List[str]698 699 edge_continuous_names: List[str]700 701 edge_feature_dim: int702 703 graph_config: Dict[str, Any]704 705 schema_version: str = "1.0"706 707 @staticmethod708 def load(path: str):709 710 with open(711 path,712 "r",713 encoding="utf-8",714 ) as f:715 716 data = json.load(f)717 718 return GraphSchema(719 **data720 )721 722 723# ============================================================724# GRAPH DEFINITIONS725# ============================================================726 727EDGE_RELATION_NAMES = [728 729 "contains",730 "inside",731 732 "left_aligned",733 "right_aligned",734 735 "top_aligned",736 "bottom_aligned",737 738 "horizontal_center_aligned",739 "vertical_center_aligned",740 741 "same_row",742 "same_column",743 744 "overlap",745 746 "reading_order",747 748 "controlled_proximity",749 750 "connectivity_fallback",751]752 753NUM_RELATIONS = len(754 EDGE_RELATION_NAMES755)756 757EDGE_CONTINUOUS_NAMES = [758 "norm_center_distance",759 "angle_sin",760 "angle_cos",761]762 763NODE_NUMERIC_NAMES = [764 "norm_center_x",765 "norm_center_y",766 "norm_width",767 "norm_height",768 "norm_area",769 "confidence",770]771 772 773# ============================================================774# BOX775# ============================================================776 777class Box:778 779 __slots__ = (780 "x1",781 "y1",782 "x2",783 "y2",784 "cx",785 "cy",786 "w",787 "h",788 "area",789 )790 791 def __init__(792 self,793 x1,794 y1,795 x2,796 y2,797 ):798 799 self.x1 = x1800 self.y1 = y1801 self.x2 = x2802 self.y2 = y2803 804 self.cx = (805 x1 + x2806 ) / 2.0807 808 self.cy = (809 y1 + y2810 ) / 2.0811 812 self.w = max(813 x2 - x1,814 0.0,815 )816 817 self.h = max(818 y2 - y1,819 0.0,820 )821 822 self.area = (823 self.w * self.h824 )825 826 827def intersection_area(828 a: Box,829 b: Box,830) -> float:831 832 ix1 = max(833 a.x1,834 b.x1,835 )836 837 iy1 = max(838 a.y1,839 b.y1,840 )841 842 ix2 = min(843 a.x2,844 b.x2,845 )846 847 iy2 = min(848 a.y2,849 b.y2,850 )851 852 return (853 max(854 0.0,855 ix2 - ix1,856 )857 *858 max(859 0.0,860 iy2 - iy1,861 )862 )863 864 865def box_iou(866 a: Box,867 b: Box,868) -> float:869 870 intersection = (871 intersection_area(872 a,873 b,874 )875 )876 877 union = (878 a.area879 + b.area880 - intersection881 )882 883 return (884 intersection / union885 if union > 0886 else 0.0887 )888 889 890def directed_containment(891 a: Box,892 b: Box,893 cfg: GraphConfig,894) -> bool:895 896 if b.area <= 0:897 return False898 899 intersection = (900 intersection_area(901 a,902 b,903 )904 )905 906 coverage = (907 intersection / b.area908 )909 910 return (911 coverage912 >= cfg.containment_threshold913 and914 a.area915 >= b.area916 * cfg.containment_min_area_ratio917 )918 919 920# ============================================================921# GRAPH RELATIONSHIPS922# ============================================================923 924def alignment_flags(925 a: Box,926 b: Box,927 cfg: GraphConfig,928):929 930 a_wide = (931 a.w >= cfg.large_box_width_ratio932 )933 934 b_wide = (935 b.w >= cfg.large_box_width_ratio936 )937 938 a_tall = (939 a.h >= cfg.large_box_height_ratio940 )941 942 b_tall = (943 b.h >= cfg.large_box_height_ratio944 )945 946 tolerance = (947 cfg.alignment_tolerance948 )949 950 flags = {951 "left_aligned": False,952 "right_aligned": False,953 "top_aligned": False,954 "bottom_aligned": False,955 "horizontal_center_aligned": False,956 "vertical_center_aligned": False,957 }958 959 if not (a_wide or b_wide):960 961 flags["left_aligned"] = (962 abs(a.x1 - b.x1)963 <= tolerance964 )965 966 flags["right_aligned"] = (967 abs(a.x2 - b.x2)968 <= tolerance969 )970 971 flags[972 "horizontal_center_aligned"973 ] = (974 abs(a.cx - b.cx)975 <= tolerance976 )977 978 if not (a_tall or b_tall):979 980 flags["top_aligned"] = (981 abs(a.y1 - b.y1)982 <= tolerance983 )984 985 flags["bottom_aligned"] = (986 abs(a.y2 - b.y2)987 <= tolerance988 )989 990 flags[991 "vertical_center_aligned"992 ] = (993 abs(a.cy - b.cy)994 <= tolerance995 )996 997 return flags998 999 1000def vertical_overlap_ratio(1001 a: Box,1002 b: Box,1003):1004 1005 iy1 = max(1006 a.y1,1007 b.y1,1008 )1009 1010 iy2 = min(1011 a.y2,1012 b.y2,1013 )1014 1015 overlap = max(1016 0.0,1017 iy2 - iy1,1018 )1019 1020 minimum = max(1021 min(a.h, b.h),1022 1e-6,1023 )1024 1025 return overlap / minimum1026 1027 1028def horizontal_overlap_ratio(1029 a: Box,1030 b: Box,1031):1032 1033 ix1 = max(1034 a.x1,1035 b.x1,1036 )1037 1038 ix2 = min(1039 a.x2,1040 b.x2,1041 )1042 1043 overlap = max(1044 0.0,1045 ix2 - ix1,1046 )1047 1048 minimum = max(1049 min(a.w, b.w),1050 1e-6,1051 )1052 1053 return overlap / minimum1054 1055 1056def same_row(1057 a: Box,1058 b: Box,1059 cfg: GraphConfig,1060):1061 1062 return (1063 vertical_overlap_ratio(1064 a,1065 b,1066 )1067 >= cfg.row_overlap_threshold1068 and1069 abs(a.cy - b.cy)1070 <= cfg.row_center_tolerance1071 )1072 1073 1074def same_column(1075 a: Box,1076 b: Box,1077 cfg: GraphConfig,1078):1079 1080 return (1081 horizontal_overlap_ratio(1082 a,1083 b,1084 )1085 >= cfg.column_overlap_threshold1086 and1087 abs(a.cx - b.cx)1088 <= cfg.column_center_tolerance1089 )1090 1091 1092def reading_order_edges(1093 boxes: List[Box],1094 cfg: GraphConfig,1095):1096 1097 if not boxes:1098 return []1099 1100 ordered = sorted(1101 range(len(boxes)),1102 key=lambda i: boxes[i].cy,1103 )1104 1105 rows = []1106 1107 current = [ordered[0]]1108 1109 current_y = boxes[1110 ordered[0]1111 ].cy1112 1113 for idx in ordered[1:]:1114 1115 y = boxes[idx].cy1116 1117 if (1118 abs(y - current_y)1119 <= cfg.reading_order_row_tolerance1120 ):1121 1122 current.append(idx)1123 1124 else:1125 1126 rows.append(current)1127 1128 current = [idx]1129 1130 current_y = y1131 1132 rows.append(current)1133 1134 final_order = []1135 1136 for row in rows:1137 1138 final_order.extend(1139 sorted(1140 row,1141 key=lambda i:1142 boxes[i].cx,1143 )1144 )1145 1146 return list(1147 zip(1148 final_order[:-1],1149 final_order[1:],1150 )1151 )1152 1153 1154# ============================================================1155# GRAPH BUILDER1156# ============================================================1157 1158def empty_graph(1159 schema: GraphSchema,1160):1161 1162 data = Data(1163 x=torch.zeros(1164 (1165 0,1166 schema.node_feature_dim,1167 ),1168 dtype=torch.float32,1169 ),1170 1171 edge_index=torch.zeros(1172 (2, 0),1173 dtype=torch.long,1174 ),1175 1176 edge_attr=torch.zeros(1177 (1178 0,1179 schema.edge_feature_dim,1180 ),1181 dtype=torch.float32,1182 ),1183 1184 # Explicit (empty) mapping from graph node -> original1185 # YOLO detection index. Kept even for the empty graph so1186 # downstream code can always rely on `detection_indices`1187 # being present.1188 detection_indices=torch.zeros(1189 (0,),1190 dtype=torch.long,1191 ),1192 1193 num_nodes=0,1194 )1195 1196 data.is_empty = True1197 1198 return data1199 1200