Spaces:
Sleeping
Sleeping
File size: 7,082 Bytes
c02c6ce | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 | """
Image preprocessing and postprocessing utilities.
Handles image loading, resizing, and results conversion.
"""
import cv2
import numpy as np
import logging
from pathlib import Path
from typing import Tuple, List, Dict, Optional
logger = logging.getLogger(__name__)
class ImagePreprocessor:
"""Handle image preprocessing for model inference."""
MAX_SIZE = 640 # Standard YOLO input size
@staticmethod
def load_image(image_path: str) -> np.ndarray:
"""
Load image from file.
Args:
image_path: Path to image file
Returns:
Image as numpy array in BGR format
"""
if not Path(image_path).exists():
raise FileNotFoundError(f"Image file not found: {image_path}")
image = cv2.imread(image_path)
if image is None:
raise ValueError(f"Failed to load image: {image_path}")
logger.info(f"Loaded image from {image_path}, shape: {image.shape}")
return image
@staticmethod
def load_image_from_bytes(image_bytes: bytes) -> np.ndarray:
"""
Load image from bytes.
Args:
image_bytes: Image data as bytes
Returns:
Image as numpy array in BGR format
"""
nparr = np.frombuffer(image_bytes, np.uint8)
image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
if image is None:
raise ValueError("Failed to decode image from bytes")
logger.info(f"Loaded image from bytes, shape: {image.shape}")
return image
@staticmethod
def resize_image(image: np.ndarray, max_size: int = MAX_SIZE) -> np.ndarray:
"""
Resize image while maintaining aspect ratio.
Args:
image: Input image
max_size: Maximum size for longest dimension
Returns:
Resized image
"""
height, width = image.shape[:2]
scale = min(max_size / max(height, width), 1.0)
new_width = int(width * scale)
new_height = int(height * scale)
resized = cv2.resize(image, (new_width, new_height), interpolation=cv2.INTER_LINEAR)
logger.info(f"Resized image from {image.shape} to {resized.shape}")
return resized
@staticmethod
def get_image_info(image: np.ndarray) -> Dict:
"""Get basic image information."""
height, width = image.shape[:2]
return {
"width": width,
"height": height,
"channels": image.shape[2] if len(image.shape) > 2 else 1,
"size_mb": (image.nbytes / (1024 * 1024))
}
class ResultPostprocessor:
"""Convert model predictions to clean output format."""
@staticmethod
def process_detections(
predictions,
class_names: List[str],
conf_threshold: float = 0.5
) -> List[Dict]:
"""
Convert YOLO predictions to structured format.
Args:
predictions: YOLOv5 predictions object
class_names: List of class names
conf_threshold: Confidence threshold for filtering
Returns:
List of detection dictionaries
"""
detections = []
# Extract predictions
xyxy = predictions.xyxy[0].cpu().numpy() # Bounding boxes
confs = predictions.conf[0].cpu().numpy() # Confidences
classes = predictions.cls[0].cpu().numpy().astype(int) # Class indices
for bbox, conf, cls_idx in zip(xyxy, confs, classes):
if conf >= conf_threshold:
x1, y1, x2, y2 = bbox
detection = {
"class_id": int(cls_idx),
"class_name": class_names[cls_idx] if cls_idx < len(class_names) else "unknown",
"confidence": float(conf),
"bbox": {
"x1": float(x1),
"y1": float(y1),
"x2": float(x2),
"y2": float(y2)
},
"center": {
"x": float((x1 + x2) / 2),
"y": float((y1 + y2) / 2)
},
"width": float(x2 - x1),
"height": float(y2 - y1)
}
detections.append(detection)
logger.info(f"Processed {len(detections)} detections")
return detections
@staticmethod
def merge_overlapping_detections(
detections: List[Dict],
iou_threshold: float = 0.5
) -> List[Dict]:
"""
Merge nearby detections of the same class using IoU.
Args:
detections: List of detection dictionaries
iou_threshold: IoU threshold for merging
Returns:
Merged detections
"""
if not detections:
return detections
# Sort by confidence (descending)
sorted_dets = sorted(detections, key=lambda x: x["confidence"], reverse=True)
merged = []
used = set()
for i, det in enumerate(sorted_dets):
if i in used:
continue
current_class = det["class_name"]
bbox1 = det["bbox"]
for j, other in enumerate(sorted_dets[i+1:], start=i+1):
if j in used:
continue
# Only merge same class
if other["class_name"] != current_class:
continue
bbox2 = other["bbox"]
iou = ResultPostprocessor._calculate_iou(bbox1, bbox2)
if iou >= iou_threshold:
used.add(j)
merged.append(det)
logger.info(f"Merged overlapping detections: {len(detections)} -> {len(merged)}")
return merged
@staticmethod
def _calculate_iou(bbox1: Dict, bbox2: Dict) -> float:
"""Calculate Intersection over Union of two bboxes."""
x1_min, y1_min = bbox1["x1"], bbox1["y1"]
x1_max, y1_max = bbox1["x2"], bbox1["y2"]
x2_min, y2_min = bbox2["x1"], bbox2["y1"]
x2_max, y2_max = bbox2["x2"], bbox2["y2"]
# Intersection area
xi_min = max(x1_min, x2_min)
yi_min = max(y1_min, y2_min)
xi_max = min(x1_max, x2_max)
yi_max = min(y1_max, y2_max)
if xi_max < xi_min or yi_max < yi_min:
return 0.0
intersection = (xi_max - xi_min) * (yi_max - yi_min)
# Union area
area1 = (x1_max - x1_min) * (y1_max - y1_min)
area2 = (x2_max - x2_min) * (y2_max - y2_min)
union = area1 + area2 - intersection
return intersection / union if union > 0 else 0.0
|