File size: 5,348 Bytes
9946dd2 | 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 | import pandas as pd
import os
import re
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import timm
from sklearn.metrics import classification_report
from sklearn.model_selection import StratifiedGroupKFold
from sklearn.utils.class_weight import compute_class_weight
from submission.utils.utils import ImageData
import torchvision.transforms as transforms
import numpy as np
from tqdm import tqdm
# --- CONFIGURATION ---
BASE_PATH = "/Users/yusufbardolia/Documents/Intelligent System In Medicine/phase_1a"
PATH_TO_IMAGES = os.path.join(BASE_PATH, "images")
PATH_TO_GT = os.path.join(BASE_PATH, "gt_for_classification_multiclass_from_filenames_0_index.csv")
PATH_TO_SPLIT_GT = os.path.join(os.getcwd(), "honest_split_gt.csv")
MODEL_SAVE_PATH = os.path.join("submission", "multiclass_model.pth")
# --- UPGRADES ---
MODEL_NAME = 'efficientnet_b3' # Larger, more powerful model
IMAGE_SIZE = (300, 300) # EfficientNet-B3 native resolution
MAX_EPOCHS = 15
BATCH_SIZE = 16 # Smaller batch for larger model
NUM_CLASSES = 3
LEARNING_RATE = 0.0003
if torch.backends.mps.is_available():
DEVICE = "mps"
print(f"✅ Using Apple M-Series GPU (MPS)")
else:
DEVICE = "cpu"
def create_honest_split():
print("Creating honest, stratified data split...")
df = pd.read_csv(PATH_TO_GT)
surgery_dates = []
for fname in df["file_name"]:
match = re.search(r'(202\d{5})', fname)
surgery_dates.append(match.group(1) if match else "unknown")
groups = np.array(surgery_dates)
y = df["category_id"].values
sgkf = StratifiedGroupKFold(n_splits=5, shuffle=True, random_state=42)
train_idx, val_idx = next(sgkf.split(df, y, groups=groups))
df["validation_set"] = 0
df.loc[val_idx, "validation_set"] = 1
df.to_csv(PATH_TO_SPLIT_GT, index=False)
classes = np.unique(y)
weights = compute_class_weight(class_weight='balanced', classes=classes, y=y[train_idx])
return PATH_TO_SPLIT_GT, torch.tensor(weights, dtype=torch.float32).to(DEVICE)
def main():
split_csv_path, class_weights = create_honest_split()
# 2. Transforms (Heavy Augmentation)
train_transforms = transforms.Compose([
transforms.Resize((320, 320)), # Resize larger first
transforms.RandomCrop(IMAGE_SIZE), # Then random crop (better data aug)
transforms.RandomHorizontalFlip(p=0.5),
transforms.RandomVerticalFlip(p=0.5),
transforms.RandomRotation(degrees=45),
transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
val_transforms = transforms.Compose([
transforms.Resize(IMAGE_SIZE),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
train_dataset = ImageData(PATH_TO_IMAGES, split_csv_path, validation_set=0, transform=train_transforms)
val_dataset = ImageData(PATH_TO_IMAGES, split_csv_path, validation_set=1, transform=val_transforms)
train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=0)
val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=0)
print(f"Loading {MODEL_NAME}...")
model = timm.create_model(MODEL_NAME, pretrained=True, num_classes=NUM_CLASSES)
model = model.to(DEVICE)
criterion = nn.CrossEntropyLoss(weight=class_weights, label_smoothing=0.1)
optimizer = optim.AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=0.01)
scheduler = optim.lr_scheduler.OneCycleLR(optimizer, max_lr=LEARNING_RATE, steps_per_epoch=len(train_loader), epochs=MAX_EPOCHS)
print(f"Starting training...")
best_f1 = 0.0
for epoch in range(MAX_EPOCHS):
model.train()
pbar = tqdm(train_loader, desc=f"Epoch {epoch+1}")
for img, label in pbar:
img, label = img.to(DEVICE), label.to(DEVICE)
optimizer.zero_grad()
output = model(img)
loss = criterion(output, label)
loss.backward()
optimizer.step()
scheduler.step()
pbar.set_postfix({"Loss": f"{loss.item():.4f}"})
# Validation
model.eval()
all_preds = []
all_labels = []
with torch.no_grad():
for img, label in val_loader:
img, label = img.to(DEVICE), label.to(DEVICE)
output = model(img)
preds = torch.argmax(output, dim=1).cpu().numpy()
all_preds.extend(preds)
all_labels.extend(label.cpu().numpy())
report = classification_report(all_labels, all_preds, output_dict=True, zero_division=0)
curr_f1 = report['macro avg']['f1-score']
print(f"Val F1: {curr_f1:.4f}")
if curr_f1 > best_f1:
best_f1 = curr_f1
torch.save(model.state_dict(), MODEL_SAVE_PATH)
print(f"🚀 Saved {MODEL_SAVE_PATH}")
print(f"Done. Best F1: {best_f1:.4f}")
if __name__ == "__main__":
if not os.path.exists("submission"): os.makedirs("submission")
main() |