Erfuuun/Keyhole_AI
0
1import os2import io3import zipfile4import torch5import torch.nn as nn6import numpy as np7from PIL import Image8import albumentations as A9from albumentations.pytorch import ToTensorV210import segmentation_models_pytorch.encoders as smp_encoders11from celery import Celery12 13# 1. INITIALIZE CELERY14celery_app = Celery(15 "ai_tasks",16 broker="redis://localhost:6379/0",17 backend="redis://localhost:6379/0"18)19 20# 2. MODEL ARCHITECTURE21class ConvBlock(nn.Module):22 def __init__(self, in_channels, out_channels, padding=1, batch_norm=True):23 super().__init__()24 self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=padding)25 self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=padding)26 self.relu = nn.ReLU(inplace=True)27 if batch_norm:28 self.bn1 = nn.BatchNorm2d(out_channels)29 self.bn2 = nn.BatchNorm2d(out_channels)30 else:31 self.bn1 = nn.Identity()32 self.bn2 = nn.Identity()33 34 def forward(self, x):35 x = self.conv1(x)36 x = self.bn1(x)37 x = self.relu(x)38 x = self.conv2(x)39 x = self.bn2(x)40 x = self.relu(x)41 return x42 43class CustomUnetPlusPlus(nn.Module):44 def __init__(self, encoder_name="se_resnext101_32x4d", encoder_weights=None, in_channels=3, num_classes=1, decoder_channels=(256, 128, 64, 32), batch_norm=True):45 super().__init__()46 self.encoder = smp_encoders.get_encoder(47 name=encoder_name,48 in_channels=in_channels,49 weights=encoder_weights,50 depth=551 )52 enc_channels = self.encoder.out_channels53 54 self.bottleneck = ConvBlock(enc_channels[5], enc_channels[5], batch_norm=batch_norm)55 self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)56 57 self.x_3_1 = ConvBlock(enc_channels[5] + enc_channels[4], decoder_channels[0], batch_norm=batch_norm)58 self.x_2_1 = ConvBlock(enc_channels[4] + enc_channels[3], decoder_channels[1], batch_norm=batch_norm)59 self.x_1_1 = ConvBlock(enc_channels[3] + enc_channels[2], decoder_channels[2], batch_norm=batch_norm)60 self.x_0_1 = ConvBlock(enc_channels[2] + enc_channels[1], decoder_channels[3], batch_norm=batch_norm)61 62 self.x_2_2 = ConvBlock(decoder_channels[0] + enc_channels[3] + decoder_channels[1], decoder_channels[1], batch_norm=batch_norm)63 self.x_1_2 = ConvBlock(decoder_channels[1] + enc_channels[2] + decoder_channels[2], decoder_channels[2], batch_norm=batch_norm)64 self.x_0_2 = ConvBlock(decoder_channels[2] + enc_channels[1] + decoder_channels[3], decoder_channels[3], batch_norm=batch_norm)65 66 self.x_1_3 = ConvBlock(decoder_channels[1] + enc_channels[2] + decoder_channels[2]*2, decoder_channels[2], batch_norm=batch_norm)67 self.x_0_3 = ConvBlock(decoder_channels[2] + enc_channels[1] + decoder_channels[3]*2, decoder_channels[3], batch_norm=batch_norm)68 69 self.x_0_4 = ConvBlock(decoder_channels[2] + enc_channels[1] + decoder_channels[3]*3, decoder_channels[3], batch_norm=batch_norm)70 71 self.final_up = nn.ConvTranspose2d(decoder_channels[3], 32, kernel_size=2, stride=2)72 self.final_conv = nn.Sequential(73 ConvBlock(32, 32, padding=1, batch_norm=batch_norm),74 nn.Conv2d(32, num_classes, kernel_size=1)75 )76 77 def forward(self, x):78 features = self.encoder(x)79 skip1, skip2, skip3, skip4 = features[1], features[2], features[3], features[4]80 b_out = self.bottleneck(features[5])81 82 x_3_1 = self.x_3_1(torch.cat([self.up(b_out), skip4], 1))83 x_2_1 = self.x_2_1(torch.cat([self.up(skip4), skip3], 1))84 x_1_1 = self.x_1_1(torch.cat([self.up(skip3), skip2], 1))85 x_0_1 = self.x_0_1(torch.cat([self.up(skip2), skip1], 1))86 87 x_2_2 = self.x_2_2(torch.cat([self.up(x_3_1), skip3, x_2_1], 1))88 x_1_2 = self.x_1_2(torch.cat([self.up(x_2_1), skip2, x_1_1], 1))89 x_0_2 = self.x_0_2(torch.cat([self.up(x_1_1), skip1, x_0_1], 1))90 91 x_1_3 = self.x_1_3(torch.cat([self.up(x_2_2), skip2, x_1_1, x_1_2], 1))92 x_0_3 = self.x_0_3(torch.cat([self.up(x_1_2), skip1, x_0_1, x_0_2], 1))93 94 x_0_4 = self.x_0_4(torch.cat([self.up(x_1_3), skip1, x_0_1, x_0_2, x_0_3], 1))95 96 out = self.final_up(x_0_4)97 out = self.final_conv(out)98 return out99 100# 3. LOAD MODEL101MODEL_WEIGHTS = "best_model_se_resnext101_32x4d test 2.pth"102ENCODER_NAME = "se_resnext101_32x4d"103DEVICE = "cuda" if torch.cuda.is_available() else "cpu"104 105print("Worker starting: Loading model...")106model = CustomUnetPlusPlus(encoder_name=ENCODER_NAME, encoder_weights=None).to(DEVICE)107model.load_state_dict(torch.load(MODEL_WEIGHTS, map_location=DEVICE), strict=True)108model.eval()109print("Model loaded successfully!")110 111def get_transforms():112 return A.Compose([113 A.Resize(512, 512),114 A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),115 ToTensorV2()116 ])117 118def create_comparison(original, mask):119 w, h = original.size120 side_by_side = Image.new('RGB', (w * 2, h))121 side_by_side.paste(original, (0, 0))122 side_by_side.paste(mask.convert('RGB'), (w, 0))123 return side_by_side124 125# 4. CELERY TASKS126@celery_app.task(name="process_single")127def process_single_task(input_path, output_path):128 """Processes one image and saves it to the hard drive."""129 try:130 image = Image.open(input_path).convert("RGB")131 w, h = image.size132 image_np = np.array(image)133 134 transform = get_transforms()135 augmented = transform(image=image_np)136 img_tensor = augmented['image'].unsqueeze(0).to(DEVICE)137 138 with torch.no_grad():139 logits = model(img_tensor)140 probs = torch.sigmoid(logits)141 pred_mask = (probs > 0.5).float()142 143 pred_mask = pred_mask.squeeze().cpu().numpy()144 mask_uint8 = ((1 - pred_mask) * 255).astype(np.uint8)145 mask_pil = Image.fromarray(mask_uint8).resize((w, h), resample=Image.NEAREST)146 147 combined_img = create_comparison(image, mask_pil)148 combined_img.save(output_path, format='PNG')149 150 return {"status": "success", "result_path": output_path}151 except Exception as e:152 return {"status": "error", "message": str(e)}153 154@celery_app.task(name="process_folder")155def process_folder_task(input_folder, output_zip_path):156 """Processes an entire folder of images (including subfolders) and creates a zip file."""157 try:158 with zipfile.ZipFile(output_zip_path, "w", zipfile.ZIP_DEFLATED) as out_zip:159 # Walk through all directories and subdirectories160 for root, dirs, files in os.walk(input_folder):161 for filename in files:162 if not filename.lower().endswith(('.png', '.jpg', '.jpeg', '.tif', '.tiff', '.bmp')):163 continue164 165 filepath = os.path.join(root, filename)166 image = Image.open(filepath).convert("RGB")167 w, h = image.size168 image_np = np.array(image)169 170 transform = get_transforms()171 augmented = transform(image=image_np)172 img_tensor = augmented['image'].unsqueeze(0).to(DEVICE)173 174 with torch.no_grad():175 logits = model(img_tensor)176 probs = torch.sigmoid(logits)177 pred_mask = (probs > 0.5).float()178 179 pred_mask = pred_mask.squeeze().cpu().numpy()180 mask_uint8 = ((1 - pred_mask) * 255).astype(np.uint8)181 mask_pil = Image.fromarray(mask_uint8).resize((w, h), resample=Image.NEAREST)182 183 name_part = os.path.splitext(filename)[0]184 185 mask_byte_arr = io.BytesIO()186 mask_pil.save(mask_byte_arr, format='PNG')187 out_zip.writestr(f"only_masks/{name_part}_mask.png", mask_byte_arr.getvalue())188 189 comparison_pil = create_comparison(image, mask_pil)190 comp_byte_arr = io.BytesIO()191 comparison_pil.save(comp_byte_arr, format='PNG')192 out_zip.writestr(f"comparison/{name_part}_comparison.png", comp_byte_arr.getvalue())193 194 return {"status": "success", "result_path": output_zip_path}195 except Exception as e:196 return {"status": "error", "message": str(e)}