461 lines
19 KiB
Python
461 lines
19 KiB
Python
"""
|
|
Region detection and processing module for HTDetectionBatchProcessor
|
|
"""
|
|
import torch
|
|
import math
|
|
from enum import Enum
|
|
import numpy as np
|
|
import time
|
|
from PIL import Image, ImageFilter
|
|
|
|
from ..modules.debug_utils import save_debug_image
|
|
|
|
# Re-export needed enums and constants
|
|
class ScaleMode(Enum):
|
|
MAX = "max" # Scale up to 1024 regardless
|
|
UP = "up" # Scale up to next closest bucket
|
|
DOWN = "down" # Scale down to next closest bucket
|
|
CLOSEST = "closest" # Scale to closest bucket
|
|
|
|
class ShortEdgeDivisibility(Enum):
|
|
DIV_BY_8 = 8
|
|
DIV_BY_64 = 64
|
|
|
|
# Target bucket sizes for standardized dimensions
|
|
BUCKET_SIZES = [1024, 768, 512]
|
|
MAX_BUCKET_SIZE = 1024
|
|
MAX_RESOLUTION = 8192
|
|
|
|
def force_crops_output(original_image, working_segs):
|
|
"""
|
|
Create direct visualization of crops from SEGS
|
|
"""
|
|
print("[DEBUG] Creating visualization of detected crops")
|
|
|
|
if working_segs is None or len(working_segs[1]) == 0:
|
|
print("[DEBUG] No segments found for visualization")
|
|
return original_image
|
|
|
|
# Create a new empty batch to hold the crops
|
|
crops_list = []
|
|
|
|
# Extract each segment as a separate image
|
|
for i, seg in enumerate(working_segs[1]):
|
|
x1, y1, x2, y2 = seg.crop_region
|
|
print(f"[DEBUG] Extracting region {i}: {x1},{y1} to {x2},{y2}")
|
|
|
|
try:
|
|
# Extract crop - handling potential dimension issues
|
|
if len(original_image.shape) == 4: # With batch dimension
|
|
crop = original_image[0, y1:y2, x1:x2, :]
|
|
else: # Without batch dimension
|
|
crop = original_image[y1:y2, x1:x2, :]
|
|
|
|
# Add batch dimension if needed
|
|
if len(crop.shape) == 3:
|
|
crop = crop.unsqueeze(0)
|
|
|
|
print(f"[DEBUG] Crop shape: {crop.shape}")
|
|
crops_list.append(crop)
|
|
except Exception as e:
|
|
print(f"[DEBUG] Crop extraction error: {e}")
|
|
|
|
if len(crops_list) == 0:
|
|
print("[DEBUG] No valid crops created")
|
|
return original_image
|
|
|
|
# Stack crops horizontally if possible
|
|
try:
|
|
# Resize all crops to same height
|
|
from torchvision.transforms.functional import resize
|
|
|
|
# Use first crop's height as reference
|
|
target_height = crops_list[0].shape[1]
|
|
|
|
# Resize each crop to match reference height
|
|
resized_crops = []
|
|
for crop in crops_list:
|
|
# Move channels to position expected by resize
|
|
crop_ch_first = crop.permute(0, 3, 1, 2)
|
|
# Calculate new width maintaining aspect ratio
|
|
new_width = int(crop.shape[2] * (target_height / crop.shape[1]))
|
|
# Resize
|
|
resized = resize(crop_ch_first, [target_height, new_width])
|
|
# Move channels back
|
|
resized = resized.permute(0, 2, 3, 1)
|
|
resized_crops.append(resized)
|
|
|
|
# Pad each crop with a border
|
|
padded_crops = []
|
|
for crop in resized_crops:
|
|
# Create slightly larger tensor with black border
|
|
padded = torch.zeros(1, crop.shape[1]+4, crop.shape[2]+4, 3)
|
|
# Place crop in center
|
|
padded[0, 2:-2, 2:-2, :] = crop[0]
|
|
padded_crops.append(padded)
|
|
|
|
# Concatenate horizontally - we need same height for this
|
|
concat_width = sum(crop.shape[2] for crop in padded_crops)
|
|
concat_crop = torch.zeros(1, padded_crops[0].shape[1], concat_width, 3)
|
|
|
|
current_x = 0
|
|
for crop in padded_crops:
|
|
width = crop.shape[2]
|
|
concat_crop[0, :, current_x:current_x+width, :] = crop[0]
|
|
current_x += width
|
|
|
|
print(f"[DEBUG] Concatenated visualization shape: {concat_crop.shape}")
|
|
save_debug_image(concat_crop, "concatenated_crops_visualization")
|
|
return concat_crop
|
|
|
|
except Exception as e:
|
|
print(f"[DEBUG] Concatenation error: {e}")
|
|
# Return first crop as fallback
|
|
return crops_list[0]
|
|
|
|
def prepare_crop_batch(image_batch, crop_regions, original_image):
|
|
"""
|
|
Properly prepare cropped batch for output, ensuring correct dimensions.
|
|
"""
|
|
# If no crops were found, return a modified version of the original image
|
|
if len(crop_regions) == 0:
|
|
print("[DEBUG] No crop regions found, returning marker image")
|
|
# Create a red-tinted version of the original image to indicate no crops
|
|
marker = original_image.clone()
|
|
if marker.shape[-1] == 3: # If channels are in the last dimension
|
|
marker[..., 0] = 1.0 # Set red channel to max
|
|
marker[..., 1] = 0.3 # Reduce green
|
|
marker[..., 2] = 0.3 # Reduce blue
|
|
return marker
|
|
|
|
# Ensure the image_batch has the right dimensions
|
|
if len(image_batch.shape) != 4:
|
|
print(f"[DEBUG] Unusual image batch shape: {image_batch.shape}")
|
|
|
|
# If we have a single image with [H, W, C] format, add batch dimension
|
|
if len(image_batch.shape) == 3 and image_batch.shape[-1] == 3:
|
|
image_batch = image_batch.unsqueeze(0)
|
|
|
|
# If we have [C, H, W] format, convert to [1, H, W, C]
|
|
elif len(image_batch.shape) == 3 and image_batch.shape[0] == 3:
|
|
image_batch = image_batch.permute(1, 2, 0).unsqueeze(0)
|
|
|
|
# If we have [B, C, H, W] format, convert to [B, H, W, C]
|
|
elif len(image_batch.shape) == 4 and image_batch.shape[1] == 3:
|
|
image_batch = image_batch.permute(0, 2, 3, 1)
|
|
|
|
print(f"[DEBUG] Reshaped to: {image_batch.shape}")
|
|
|
|
# Ensure we have proper [B, H, W, C] format with batch size matching crop_regions
|
|
if image_batch.shape[0] != len(crop_regions):
|
|
print(f"[DEBUG] Warning: Batch size {image_batch.shape[0]} doesn't match crop count {len(crop_regions)}")
|
|
|
|
# If we have more crops than batch items, repeat the batch
|
|
if image_batch.shape[0] < len(crop_regions):
|
|
repeats = math.ceil(len(crop_regions) / image_batch.shape[0])
|
|
image_batch = image_batch.repeat(repeats, 1, 1, 1)[:len(crop_regions)]
|
|
|
|
# If we have more batch items than crops, take only what we need
|
|
elif image_batch.shape[0] > len(crop_regions):
|
|
image_batch = image_batch[:len(crop_regions)]
|
|
|
|
save_debug_image(image_batch, "prepared_crop_batch")
|
|
return image_batch
|
|
|
|
def calculate_bucket_dimensions(width, height, scale_mode, short_edge_div, mask_dilation=1.0, bucket_scale_factor=1.0):
|
|
"""
|
|
Calculate the target dimensions based on bucket constraints.
|
|
"""
|
|
# Apply dilation factor to dimensions
|
|
width = int(width * mask_dilation)
|
|
height = int(height * mask_dilation)
|
|
|
|
# Determine long and short edges
|
|
is_landscape = width >= height
|
|
long_edge = width if is_landscape else height
|
|
short_edge = height if is_landscape else width
|
|
|
|
# Cap at MAX_BUCKET_SIZE if larger
|
|
if long_edge > MAX_BUCKET_SIZE:
|
|
target_long_edge = MAX_BUCKET_SIZE
|
|
scale_type = "down"
|
|
else:
|
|
# Select target bucket based on scale mode
|
|
if scale_mode == ScaleMode.MAX:
|
|
# Always scale up to MAX_BUCKET_SIZE
|
|
target_long_edge = MAX_BUCKET_SIZE
|
|
scale_type = "up" if long_edge < MAX_BUCKET_SIZE else "same"
|
|
elif scale_mode == ScaleMode.UP:
|
|
# Find the next bucket size up
|
|
larger_buckets = [b for b in BUCKET_SIZES if b > long_edge]
|
|
if larger_buckets:
|
|
target_long_edge = min(larger_buckets)
|
|
scale_type = "up"
|
|
else:
|
|
target_long_edge = MAX_BUCKET_SIZE
|
|
scale_type = "max"
|
|
elif scale_mode == ScaleMode.DOWN:
|
|
# Find the next bucket size down
|
|
smaller_buckets = [b for b in BUCKET_SIZES if b < long_edge]
|
|
if smaller_buckets:
|
|
target_long_edge = max(smaller_buckets)
|
|
scale_type = "down"
|
|
else:
|
|
target_long_edge = min(BUCKET_SIZES)
|
|
scale_type = "min"
|
|
else: # CLOSEST
|
|
# Find the closest bucket size
|
|
closest_bucket = min(BUCKET_SIZES, key=lambda b: abs(b - long_edge))
|
|
target_long_edge = closest_bucket
|
|
scale_type = "up" if closest_bucket > long_edge else "down" if closest_bucket < long_edge else "same"
|
|
|
|
# Calculate scaling factor
|
|
target_long_edge = int(target_long_edge * bucket_scale_factor)
|
|
scale_factor = target_long_edge / long_edge
|
|
|
|
# Calculate raw short edge size after scaling
|
|
raw_short_edge = short_edge * scale_factor
|
|
|
|
# Adjust short edge to be divisible by the specified divisor
|
|
divisor = short_edge_div.value
|
|
adjusted_short_edge = math.ceil(raw_short_edge / divisor) * divisor
|
|
|
|
# Calculate final dimensions
|
|
if is_landscape:
|
|
target_width = target_long_edge
|
|
target_height = adjusted_short_edge
|
|
else:
|
|
target_width = adjusted_short_edge
|
|
target_height = target_long_edge
|
|
|
|
return (target_width, target_height, scale_type, target_long_edge)
|
|
|
|
def adjust_crop_to_bucket(crop_region, original_img_size, scale_mode, short_edge_div, mask_dilation=1.0, bucket_scale_factor=1.0):
|
|
"""
|
|
Adjust a crop region to fit into the target bucket dimensions.
|
|
"""
|
|
x1, y1, x2, y2 = crop_region
|
|
img_width, img_height = original_img_size
|
|
|
|
# Current crop dimensions
|
|
crop_width = x2 - x1
|
|
crop_height = y2 - y1
|
|
|
|
# Apply mask dilation
|
|
if mask_dilation != 1.0:
|
|
# Calculate expansion needed
|
|
width_expansion = int(crop_width * (mask_dilation - 1))
|
|
height_expansion = int(crop_height * (mask_dilation - 1))
|
|
|
|
# Apply expansion evenly on all sides
|
|
x1 = max(0, x1 - width_expansion // 2)
|
|
y1 = max(0, y1 - height_expansion // 2)
|
|
x2 = min(img_width, x2 + width_expansion // 2)
|
|
y2 = min(img_height, y2 + height_expansion // 2)
|
|
|
|
# Update dimensions
|
|
crop_width = x2 - x1
|
|
crop_height = y2 - y1
|
|
|
|
# Calculate target dimensions for this crop
|
|
target_width, target_height, scale_type, bucket_size = calculate_bucket_dimensions(
|
|
crop_width, crop_height, scale_mode, short_edge_div, mask_dilation, bucket_scale_factor
|
|
)
|
|
|
|
# Track if we need to downscale
|
|
needs_downscale = False
|
|
if target_width < crop_width or target_height < crop_height:
|
|
needs_downscale = True
|
|
|
|
# Calculate required expansion to achieve target ratio
|
|
target_ratio = target_width / target_height
|
|
current_ratio = crop_width / crop_height
|
|
|
|
# Adjust crop to match target ratio
|
|
if current_ratio > target_ratio:
|
|
# Width is the limiting factor, adjust height
|
|
new_height = crop_width / target_ratio
|
|
height_diff = new_height - crop_height
|
|
y1 = max(0, y1 - height_diff / 2)
|
|
y2 = min(img_height, y2 + height_diff / 2)
|
|
else:
|
|
# Height is the limiting factor, adjust width
|
|
new_width = crop_height * target_ratio
|
|
width_diff = new_width - crop_width
|
|
x1 = max(0, x1 - width_diff / 2)
|
|
x2 = min(img_width, x2 + width_diff / 2)
|
|
|
|
# Ensure the crop values are integers
|
|
adjusted_crop = (int(x1), int(y1), int(x2), int(y2))
|
|
target_dims = (int(target_width), int(target_height))
|
|
|
|
return adjusted_crop, target_dims, scale_type, bucket_size, needs_downscale
|
|
|
|
def process_segs_with_buckets(segs, original_image, scale_mode, short_edge_div, mask_dilation=1.0, bucket_scale_factor=1.0):
|
|
"""
|
|
Process SEGS to fit into specified buckets.
|
|
Returns individual crop tensors (not stacked) to handle different dimensions.
|
|
"""
|
|
|
|
start_time = time.time()
|
|
print(f"[USDU_DEBUG] Starting bucket processing for {len(segs[1]) if segs and len(segs) > 1 else 0} segments")
|
|
|
|
if segs is None or len(segs[1]) == 0:
|
|
# Create a 1x1 black pixel image tensor for empty SEGS
|
|
empty_tensor = torch.zeros(1, 1, 1, 3)
|
|
return [empty_tensor], [], [], [], []
|
|
|
|
image_crops = []
|
|
crop_regions = []
|
|
target_dimensions = []
|
|
scale_types = []
|
|
bucket_sizes = []
|
|
|
|
img_height, img_width = original_image.shape[1:3]
|
|
original_img_size = (img_width, img_height)
|
|
|
|
for seg in segs[1]:
|
|
# Get crop region
|
|
x1, y1, x2, y2 = seg.crop_region
|
|
|
|
# Adjust crop region to fit target bucket
|
|
adjusted_crop, target_dims, scale_type, bucket_size, needs_downscale = adjust_crop_to_bucket(
|
|
(x1, y1, x2, y2), original_img_size, scale_mode, short_edge_div, mask_dilation, bucket_scale_factor
|
|
)
|
|
|
|
# Updated logic
|
|
ax1, ay1, ax2, ay2 = adjusted_crop
|
|
|
|
# Crop the image with adjusted coordinates
|
|
cropped_image = original_image[:, ay1:ay2, ax1:ax2, :]
|
|
|
|
# If needs downscaling, do it immediately with Lanczos
|
|
if needs_downscale and scale_type == "down":
|
|
target_w, target_h = target_dims
|
|
from torchvision.transforms.functional import resize
|
|
# Reshape for resize ([B,H,W,C] -> [B,C,H,W])
|
|
reshaped = cropped_image.permute(0, 3, 1, 2)
|
|
# Resize
|
|
resized = resize(reshaped, [target_h, target_w], antialias=True)
|
|
# Back to original shape
|
|
cropped_image = resized.permute(0, 2, 3, 1)
|
|
print(f"[USDU] Pre-downscaled region from {ay2-ay1}x{ax2-ax1} to {target_h}x{target_w}")
|
|
|
|
# Store information
|
|
crop_regions.append(adjusted_crop)
|
|
target_dimensions.append(target_dims)
|
|
scale_types.append(scale_type)
|
|
bucket_sizes.append(bucket_size)
|
|
image_crops.append(cropped_image)
|
|
|
|
# Return the list of crops - DO NOT stack them since they have different dimensions
|
|
if len(image_crops) > 0:
|
|
# Save a sample crop for debugging
|
|
save_debug_image(image_crops[0], "bucket_processed_crop_sample")
|
|
print(f"[USDU_DEBUG] Bucket processing completed in {time.time() - start_time:.2f}s, found {len(crop_regions)} regions")
|
|
return image_crops, crop_regions, target_dimensions, scale_types, bucket_sizes
|
|
else:
|
|
# Fallback 1x1 black pixel for safety
|
|
empty_tensor = torch.zeros(1, 1, 1, 3)
|
|
print(f"[USDU_DEBUG] Bucket processing failed in {time.time() - start_time:.2f}s, found no regions")
|
|
return [empty_tensor], [], [], [], []
|
|
|
|
def segs_to_image_batch(segs, original_image):
|
|
"""
|
|
Convert SEGS to a list of individual image crops for processing.
|
|
Returns a list of image tensors rather than a batched tensor.
|
|
"""
|
|
print(f"[DEBUG] Original image shape: {original_image.shape}")
|
|
|
|
if segs is None or len(segs[1]) == 0:
|
|
print("[DEBUG] No segments found in segs!")
|
|
empty_tensor = torch.zeros(1, 1, 1, 3)
|
|
return [empty_tensor], []
|
|
|
|
image_crops = []
|
|
crop_regions = []
|
|
|
|
print(f"[DEBUG] Creating crops from {len(segs[1])} segments")
|
|
for i, seg in enumerate(segs[1]):
|
|
# Get crop region
|
|
x1, y1, x2, y2 = seg.crop_region
|
|
print(f"[DEBUG] Segment {i} - Crop region: {x1},{y1} to {x2},{y2} size: {x2-x1}x{y2-y1}")
|
|
|
|
# Validate crop region
|
|
if x1 >= x2 or y1 >= y2:
|
|
print(f"[DEBUG] WARNING: Invalid crop region {x1},{y1} to {x2},{y2}")
|
|
continue
|
|
|
|
if x2 > original_image.shape[2] or y2 > original_image.shape[1]:
|
|
print(f"[DEBUG] WARNING: Crop region outside image bounds. Image: {original_image.shape[1:3]}, Crop: {y2}x{x2}")
|
|
# Clamp values
|
|
x2 = min(x2, original_image.shape[2])
|
|
y2 = min(y2, original_image.shape[1])
|
|
|
|
crop_regions.append((x1, y1, x2, y2)) # Store validated region
|
|
|
|
# Explicitly crop the image and ensure proper dimensions
|
|
try:
|
|
cropped_image = original_image[:, y1:y2, x1:x2, :]
|
|
print(f"[DEBUG] Cropped shape: {cropped_image.shape}")
|
|
|
|
# Additional validation
|
|
if cropped_image.numel() == 0:
|
|
print(f"[DEBUG] WARNING: Empty crop for segment {i}")
|
|
continue
|
|
|
|
image_crops.append(cropped_image)
|
|
except Exception as e:
|
|
print(f"[DEBUG] Error cropping segment {i}: {e}")
|
|
|
|
# Return the list of crops without stacking
|
|
if len(image_crops) > 0:
|
|
print(f"[DEBUG] Created {len(image_crops)} individual crops")
|
|
save_debug_image(image_crops[0], "segs_extracted_crop_example")
|
|
return image_crops, crop_regions
|
|
else:
|
|
print("[DEBUG] No valid crops created!")
|
|
empty_tensor = torch.zeros(1, 1, 1, 3)
|
|
return [empty_tensor], []
|
|
|
|
# Export the SimpleDetectorForEach class for segmentation
|
|
class SimpleDetectorForEach:
|
|
"""
|
|
Class providing detection and segmentation functionality.
|
|
This embedded implementation avoids external dependencies.
|
|
"""
|
|
|
|
@classmethod
|
|
def detect(cls, detector, image, threshold, dilation, crop_factor, drop_size,
|
|
crop_min_size=0, crop_square=False, detailer_hook=None, bbox_fill=0.5, segment_threshold=0.5,
|
|
sam_model_opt=None, segm_detector_opt=None):
|
|
"""
|
|
Run detection on an image.
|
|
"""
|
|
# Run basic detection
|
|
segs = detector.detect(image, threshold, dilation, crop_factor, drop_size,
|
|
crop_min_size, crop_square, detailer_hook, bbox_fill)
|
|
|
|
# Run segmentation refinement if available
|
|
if segm_detector_opt is not None and sam_model_opt is not None:
|
|
refined_segs = cls.segment_refined(segs, sam_model_opt, segm_detector_opt, segment_threshold)
|
|
return refined_segs
|
|
|
|
return segs
|
|
|
|
@classmethod
|
|
def segment_refined(cls, segs, sam_model, segm_detector, threshold):
|
|
"""
|
|
Refine detections with segmentation.
|
|
"""
|
|
# Basic segmentation refinement logic
|
|
if segs is None or segs[1] is None or len(segs[1]) == 0:
|
|
return segs
|
|
|
|
try:
|
|
# Use the segmentation detector to refine the bounding boxes
|
|
refined_segs = segm_detector.detect(segs, sam_model=sam_model, threshold=threshold)
|
|
return refined_segs
|
|
except Exception as e:
|
|
print(f"[DEBUG] Segmentation refinement error: {e}")
|
|
return segs |