| import torch |
| import torchvision.transforms as transforms |
| import wandb |
| import numpy as np |
| from PIL import Image, ImageDraw, ImageFont |
|
|
| DISEASES = ["Atelectasis", "Calcification", "Cardiomegaly", "Consolidation", "Diffuse Nodule", "Effusion", |
| "Emphysema", "Fibrosis", "Fracture", "Mass", "Nodule","Pleural Thickening", "Pneumothorax"] |
|
|
| COLOR_MAP = { |
| "Atelectasis": (255, 0, 0), |
| "Calcification": (0, 255, 0), |
| "Cardiomegaly": (0, 0, 255), |
| "Consolidation": (255, 255, 0), |
| "Diffuse Nodule": (255, 165, 0), |
| "Effusion": (0, 255, 255), |
| "Emphysema": (255, 0, 255), |
| "Fibrosis": (128, 0, 128), |
| "Fracture": (255, 192, 203), |
| "Mass": (173, 255, 47), |
| "Nodule": (0, 128, 255), |
| "Pleural Thickening": (75, 0, 130), |
| "Pneumothorax": (255, 105, 180) |
| } |
|
|
| def visualize_sample(img, target, alpha=0.3): |
|
|
| img_pil = transforms.ToPILImage()(img).convert("RGB") |
| overlay = Image.new("RGBA", img_pil.size, (0, 0, 0, 0)) |
|
|
| if not target or target.get("labels") is None or target["labels"].shape[0] == 0: |
| return img_pil |
|
|
| masks = target.get("masks", None) |
| |
| boxes = target.get("boxes", None) |
| labels = target["labels"].numpy() |
| scores = target.get("scores", None) |
|
|
| if masks is not None and masks.numel() > 0: |
| for i, mask in enumerate(masks): |
| disease_idx = labels[i] - 1 |
| if 0 <= disease_idx < len(DISEASES): |
| disease_name = DISEASES[disease_idx] |
| mask_color = COLOR_MAP[disease_name] |
| |
| mask = mask.squeeze(0) |
| mask_pil = Image.fromarray((mask.cpu().numpy() * 255).astype(np.uint8), mode="L") |
| color_mask = Image.new("RGBA", img_pil.size, mask_color + (int(255 * alpha),)) |
| overlay.paste(color_mask, (0, 0), mask_pil) |
|
|
| img_pil = Image.alpha_composite(img_pil.convert("RGBA"), overlay) |
| draw = ImageDraw.Draw(img_pil) |
| font = ImageFont.load_default() |
|
|
| if boxes is not None and boxes.numel() > 0: |
| boxes = boxes.numpy() |
| for i, box in enumerate(boxes): |
| disease_idx = labels[i] - 1 |
| if 0 <= disease_idx < len(DISEASES): |
| disease_name = DISEASES[disease_idx] |
| color = COLOR_MAP[disease_name] |
| x_min, y_min, x_max, y_max = box |
| |
| text = disease_name |
| if scores is not None: |
| text += f" {scores[i]:.2f}" |
| text_position = (x_min, max(y_min - 12, 0)) |
| draw.text(text_position, text, fill=color, font=font) |
|
|
| return img_pil.convert("RGB") |
|
|
|
|
| def image_log(images, targets, preds, epoch, mode="val"): |
| batch_size = len(images) |
| num_samples = min(8, batch_size) |
|
|
| vis_images = [] |
| for i in range(num_samples): |
| img_sample = images[i].detach().cpu().numpy() |
| target_sample = {k: v.cpu() for k, v in targets[i].items()} |
| pred_sample = {k: v.cpu() for k, v in preds[i].items()} if preds[i] else {} |
| |
| img_sample = (np.transpose(img_sample, (1, 2, 0)) * 255).astype(np.uint8) |
| target_draw = visualize_sample(img_sample, target_sample) |
| output_draw = visualize_sample(img_sample, pred_sample) |
| |
| row_image = np.concatenate([img_sample, np.array(target_draw), np.array(output_draw)], axis=0) |
| vis_images.append(row_image) |
|
|
| final_image = np.concatenate(vis_images, axis=1).astype(np.uint8) |
| pil_image = Image.fromarray(final_image) |
| wandb.log({f"{mode}_batch_visualization": wandb.Image(pil_image, caption=f"Epoch {epoch}")}) |
|
|
|
|
| def image_save(images, preds): |
| batch_size = len(images) |
| num_samples = min(8, batch_size) |
|
|
| vis_images = [] |
| for i in range(num_samples): |
| img_sample = images[i].detach().cpu().numpy() |
| pred_sample = {k: v.cpu() for k, v in preds[i].items()} if preds[i] else {} |
| |
| img_sample = (np.transpose(img_sample, (1, 2, 0)) * 255).astype(np.uint8) |
| output_draw = visualize_sample(img_sample, pred_sample) |
| vis_images.append(output_draw) |
|
|
| return vis_images |
|
|
|
|
