| 97 | |
| 98 | |
| 99 | class YoloDetector: |
| 100 | # fmt: off |
| 101 | yolo_classes = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', |
| 102 | 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', |
| 103 | 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', |
| 104 | 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', |
| 105 | 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush'] |
| 106 | # fmt: on |
| 107 | |
| 108 | yolo_alias = { |
| 109 | "bicycle": ["bicycle"], |
| 110 | "car": ["car", "truck"], |
| 111 | "bus": ["bus", "truck"], |
| 112 | "motorcycle": ["motorcycle"], |
| 113 | "boat": ["boat"], |
| 114 | "fire hydrant": ["fire hydrant", "parking meter"], |
| 115 | "parking meter": ["fire hydrant", "parking meter"], |
| 116 | "traffic light": ["traffic light"], |
| 117 | } |
| 118 | |
| 119 | def __init__(self) -> None: |
| 120 | pass |
| 121 | |
| 122 | def detect_image(self, image: cv2.typing.MatLike, tile_amount: int, task_type: str) -> List[bool]: |
| 123 | response = [False for _ in range(tile_amount)] |
| 124 | height, width, _ = image.shape |
| 125 | tiles_per_row = int(math.sqrt(tile_amount)) |
| 126 | tile_width, tile_height = width // tiles_per_row, height // tiles_per_row |
| 127 | |
| 128 | outputs = detection_models.yolo_model.predict(image, verbose=False, conf=0.2, iou=0.3) # , save=True |
| 129 | results = outputs[0] |
| 130 | |
| 131 | for result in results: |
| 132 | assert result |
| 133 | # Check if correct task type |
| 134 | class_index = int(result.boxes.cls[0]) |
| 135 | if self.yolo_classes[class_index] not in self.yolo_alias[task_type]: |
| 136 | continue |
| 137 | |
| 138 | masks = result.masks |
| 139 | mask = masks.xy[0] |
| 140 | response = calculate_segmentation_response(mask, response, tile_width, tile_height, tiles_per_row) |
| 141 | |
| 142 | # In AreaCaptcha Mode, Calculate Tiles inside of boundary, which arent covered by mask point |
| 143 | if tile_amount == 16: |
| 144 | coords = result.boxes.xyxy.flatten().tolist() |
| 145 | points_start, point_end = coords[:2], coords[2:] |
| 146 | tiles_in_bbox = get_tiles_in_bounding_box(image, tile_amount, tuple(points_start), tuple(point_end)) |
| 147 | # Appending True Tiles to Response but not making True ones False again |
| 148 | response = [x or y for x, y in zip(response, tiles_in_bbox)] |
| 149 | |
| 150 | return response |
| 151 | |
| 152 | |
| 153 | class ClipDetector: |