Update app.py

app.py

@@ -1,7 +1,386 @@
-import gradio as gr
+# ============================================================================
+# FACE AGE & GENDER PREDICTION - COMPLETE TRAINING WITH TRACKIO
+# ============================================================================
+# Generates graphs like your screenshot: train/val loss curves for age, gender, total
+# Logs metrics per step and epoch to TrackIO for real-time visualization
 
-def greet(name):
-    return "Hello " + name + "!!"
+!pip install -q trackio
+
+import os
+import gc
+import numpy as np
+import pandas as pd
+from PIL import Image
+import torch
+from torch import nn
+from torch.utils.data import Dataset, DataLoader
+from torchvision import transforms, models
+import pytorch_lightning as pl
+from pytorch_lightning.callbacks import ModelCheckpoint
+from sklearn.model_selection import train_test_split
+from kaggle_secrets import UserSecretsClient
+import trackio
+
+# ============================================================================
+# GLOBAL SETTINGS
+# ============================================================================
+
+class PipelineSettings:
+    def __init__(self):
+        self.DATA_ROOT_DIR = "/kaggle/input/sep-25-dl-gen-ai-nppe-1/face_dataset"
+        self.TRAIN_CSV_PATH = f"{self.DATA_ROOT_DIR}/train.csv"
+        self.TEST_CSV_PATH = f"{self.DATA_ROOT_DIR}/test.csv"
+        
+        self.INPUT_IMAGE_SIZE = 128
+        self.BATCH_SIZE = 128
+        self.LEARNING_RATE = 1e-3
+        self.NUM_EPOCHS = 10
+        self.AGE_LOSS_WEIGHT = 0.01
+        
+        self.NUM_DATALOADER_WORKERS = os.cpu_count()
+
+settings = PipelineSettings()
+
+# ============================================================================
+# IMAGE AUGMENTATION
+# ============================================================================
+
+class ImageAugmentor:
+    def __init__(self, image_size):
+        self.image_size = image_size
+        self.norm_params = {'mean': [0.485, 0.456, 0.406], 'std': [0.229, 0.224, 0.225]}
+
+    def get_training_transforms(self):
+        return transforms.Compose([
+            transforms.Resize((self.image_size, self.image_size)),
+            transforms.RandomHorizontalFlip(p=0.5),
+            transforms.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1),
+            transforms.ToTensor(),
+            transforms.Normalize(**self.norm_params),
+        ])
+    
+    def get_inference_transforms(self):
+        return transforms.Compose([
+            transforms.Resize((self.image_size, self.image_size)),
+            transforms.ToTensor(),
+            transforms.Normalize(**self.norm_params),
+        ])
+
+# ============================================================================
+# DATASET
+# ============================================================================
+
+class FaceImageDataset(Dataset):
+    def __init__(self, metadata_df, image_dir, image_transform=None):
+        self.metadata = metadata_df
+        self.image_dir = image_dir
+        self.transform = image_transform
+    
+    def __len__(self):
+        return len(self.metadata)
+    
+    def __getitem__(self, idx):
+        row = self.metadata.iloc[idx]
+        image_path = os.path.join(self.image_dir, row['full_path'])
+        image = Image.open(image_path).convert("RGB")
+        
+        if self.transform:
+            image = self.transform(image)
+        
+        gender_target = torch.tensor(row['gender'], dtype=torch.float32)
+        age_target = torch.tensor(row['age'], dtype=torch.float32)
+        return image, gender_target, age_target
+
+# ============================================================================
+# DATA MODULE
+# ============================================================================
+
+class FaceDataModule(pl.LightningDataModule):
+    def __init__(self, config: PipelineSettings):
+        super().__init__()
+        self.cfg = config
+        self.augmentor = ImageAugmentor(self.cfg.INPUT_IMAGE_SIZE)
+        self.train_df, self.val_df = None, None
+
+    def prepare_data(self):
+        pass
+
+    def setup(self, stage=None):
+        if stage == 'fit' or stage is None:
+            full_train = pd.read_csv(self.cfg.TRAIN_CSV_PATH)
+            self.train_df, self.val_df = train_test_split(
+                full_train, test_size=0.15, random_state=42, stratify=full_train['gender']
+            )
+            
+            self.train_dataset = FaceImageDataset(
+                self.train_df, self.cfg.DATA_ROOT_DIR, self.augmentor.get_training_transforms()
+            )
+            self.val_dataset = FaceImageDataset(
+                self.val_df, self.cfg.DATA_ROOT_DIR, self.augmentor.get_inference_transforms()
+            )
+    
+    def train_dataloader(self):
+        return DataLoader(self.train_dataset, batch_size=self.cfg.BATCH_SIZE, 
+                         shuffle=True, num_workers=self.cfg.NUM_DATALOADER_WORKERS)
+    
+    def val_dataloader(self):
+        return DataLoader(self.val_dataset, batch_size=self.cfg.BATCH_SIZE, 
+                         num_workers=self.cfg.NUM_DATALOADER_WORKERS)
+
+# ============================================================================
+# BASE MODEL WITH TRACKIO LOGGING (MATCHES YOUR SCREENSHOT)
+# ============================================================================
+
+class AbstractFaceModel(pl.LightningModule):
+    def __init__(self, learning_rate, age_loss_weight):
+        super().__init__()
+        self.save_hyperparameters()
+        self.lr = learning_rate
+        self.age_weight = age_loss_weight
+        self.gender_loss_fn = nn.BCEWithLogitsLoss()
+        self.age_loss_fn = nn.MSELoss()
+        
+        self.training_step_outputs = []
+        self.validation_step_outputs = []
+
+    def _calculate_losses(self, gender_preds, age_preds, gender_labels, age_labels):
+        gender_loss = self.gender_loss_fn(gender_preds.squeeze(), gender_labels)
+        age_loss = self.age_loss_fn(age_preds.squeeze(), age_labels)
+        total_loss = gender_loss + (age_loss * self.age_weight)
+        return total_loss, gender_loss, age_loss
+
+    def training_step(self, batch, batch_idx):
+        images, gender_labels, age_labels = batch
+        gender_preds, age_preds = self(images)
+        total_loss, gender_loss, age_loss = self._calculate_losses(
+            gender_preds, age_preds, gender_labels, age_labels
+        )
+        
+        # Log to Lightning (progress bar)
+        self.log('train_loss', total_loss, on_step=True, on_epoch=True, prog_bar=True)
+        
+        # Store for TrackIO logging
+        self.training_step_outputs.append({
+            'loss_total': total_loss.detach(),
+            'loss_gender': gender_loss.detach(),
+            'loss_age': age_loss.detach()
+        })
+        
+        # Log to TrackIO per step (creates step-by-step graphs like your screenshot)
+        try:
+            trackio.log({
+                'train/loss_total': total_loss.item(),
+                'train/loss_gender': gender_loss.item(),
+                'train/loss_age': age_loss.item(),
+                'step': self.global_step
+            })
+        except: pass
+        
+        return total_loss
+    
+    def on_train_epoch_end(self):
+        if len(self.training_step_outputs) > 0:
+            # Calculate epoch averages
+            avg_total = torch.stack([x['loss_total'] for x in self.training_step_outputs]).mean()
+            avg_gender = torch.stack([x['loss_gender'] for x in self.training_step_outputs]).mean()
+            avg_age = torch.stack([x['loss_age'] for x in self.training_step_outputs]).mean()
+            
+            # Log epoch summary to TrackIO
+            try:
+                trackio.log({
+                    'train/epoch_loss_total': avg_total.item(),
+                    'train/epoch_loss_gender': avg_gender.item(),
+                    'train/epoch_loss_age': avg_age.item(),
+                    'epoch': self.current_epoch
+                })
+            except: pass
+            
+            self.training_step_outputs.clear()
+    
+    def validation_step(self, batch, batch_idx):
+        images, gender_labels, age_labels = batch
+        gender_preds, age_preds = self(images)
+        total_loss, gender_loss, age_loss = self._calculate_losses(
+            gender_preds, age_preds, gender_labels, age_labels
+        )
+        
+        # Log to Lightning
+        self.log('val_loss', total_loss, on_epoch=True, prog_bar=True)
+        
+        # Store for TrackIO
+        self.validation_step_outputs.append({
+            'loss_total': total_loss.detach(),
+            'loss_gender': gender_loss.detach(),
+            'loss_age': age_loss.detach()
+        })
+    
+    def on_validation_epoch_end(self):
+        if len(self.validation_step_outputs) > 0:
+            # Calculate validation averages
+            avg_total = torch.stack([x['loss_total'] for x in self.validation_step_outputs]).mean()
+            avg_gender = torch.stack([x['loss_gender'] for x in self.validation_step_outputs]).mean()
+            avg_age = torch.stack([x['loss_age'] for x in self.validation_step_outputs]).mean()
+            
+            # Log to TrackIO (creates val graphs like your screenshot)
+            try:
+                trackio.log({
+                    'val/loss_total': avg_total.item(),
+                    'val/loss_gender': avg_gender.item(),
+                    'val/loss_age': avg_age.item(),
+                    'epoch': self.current_epoch
+                })
+            except: pass
+            
+            self.validation_step_outputs.clear()
+    
+    def configure_optimizers(self):
+        return torch.optim.Adam(self.parameters(), lr=self.lr)
+
+# ============================================================================
+# SCRATCH CNN MODEL
+# ============================================================================
+
+class ScratchCNNModel(AbstractFaceModel):
+    def __init__(self, learning_rate, age_loss_weight):
+        super().__init__(learning_rate, age_loss_weight)
+        
+        def conv_block(in_f, out_f):
+            return nn.Sequential(
+                nn.Conv2d(in_f, out_f, 3, padding=1, bias=False),
+                nn.BatchNorm2d(out_f),
+                nn.ReLU(inplace=True),
+                nn.MaxPool2d(2, 2)
+            )
+
+        self.feature_extractor = nn.Sequential(
+            conv_block(3, 32), conv_block(32, 64),
+            conv_block(64, 128), conv_block(128, 256)
+        )
+        
+        probe = torch.randn(1, 3, settings.INPUT_IMAGE_SIZE, settings.INPUT_IMAGE_SIZE)
+        flat_size = self.feature_extractor(probe).view(1, -1).size(1)
+        
+        self.gender_head = nn.Linear(flat_size, 1)
+        self.age_head = nn.Linear(flat_size, 1)
+
+    def forward(self, x):
+        features = torch.flatten(self.feature_extractor(x), 1)
+        return self.gender_head(features), self.age_head(features)
+
+# ============================================================================
+# FINE-TUNED RESNET MODEL
+# ============================================================================
+
+class FineTunedResNetModel(AbstractFaceModel):
+    def __init__(self, learning_rate, age_loss_weight):
+        super().__init__(learning_rate, age_loss_weight)
+        resnet = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
+        num_features = resnet.fc.in_features
+        
+        self.backbone = nn.Sequential(*list(resnet.children())[:-1])
+        self.gender_head = nn.Linear(num_features, 1)
+        self.age_head = nn.Linear(num_features, 1)
+
+    def forward(self, x):
+        features = torch.flatten(self.backbone(x), 1)
+        return self.gender_head(features), self.age_head(features)
+
+# ============================================================================
+# PIPELINE RUNNER
+# ============================================================================
+
+class PipelineRunner:
+    def __init__(self, cfg: PipelineSettings):
+        self.cfg = cfg
+        self.data_module = FaceDataModule(cfg)
+        self._setup_trackio()
+
+    def _setup_trackio(self):
+        try:
+            secrets = UserSecretsClient()
+            hf_token = secrets.get_secret("HUGGINGFACE_TOKEN")
+            os.environ["HF_TOKEN"] = hf_token
+            print("✅ TrackIO auth configured")
+        except Exception as e:
+            print(f"⚠️ TrackIO auth failed: {e}")
+
+    def _train_model(self, model, model_name, run_name):
+        print(f"\n{'='*70}\n🚀 Training: {model_name}\n{'='*70}")
+        
+        # Initialize TrackIO with your space
+        try:
+            trackio.init(
+                space_id="muhammad-bilal1/dlgenai-nppe",  # UPDATE: Your HF space from screenshot
+                project="25-t3-nppe1",
+                group=run_name,
+                config={
+                    "lr": self.cfg.LEARNING_RATE,
+                    "epochs": self.cfg.NUM_EPOCHS,
+                    "batch_size": self.cfg.BATCH_SIZE,
+                    "model": model_name,
+                    "image_size": self.cfg.INPUT_IMAGE_SIZE,
+                    "age_weight": self.cfg.AGE_LOSS_WEIGHT
+                }
+            )
+            print(f"✅ TrackIO initialized: {run_name}")
+        except Exception as e:
+            print(f"⚠️ TrackIO init failed: {e}")
+        
+        # Setup checkpoint callback
+        checkpoint_cb = ModelCheckpoint(
+            monitor='val_loss',
+            dirpath='/kaggle/working/',
+            filename=f'{model_name}-best-model',
+            save_top_k=1,
+            mode='min'
+        )
+        
+        # Train
+        trainer = pl.Trainer(
+            max_epochs=self.cfg.NUM_EPOCHS,
+            accelerator='gpu',
+            devices='auto',
+            strategy="ddp_notebook",
+            callbacks=[checkpoint_cb],
+            log_every_n_steps=10  # Log frequently for smooth graphs
+        )
+        
+        trainer.fit(model, self.data_module)
+        print(f"✅ Checkpoint: {checkpoint_cb.best_model_path}")
+        
+        # Finish TrackIO run
+        try:
+            final_val = trainer.callback_metrics.get('val_loss', torch.tensor(0.0)).item()
+            trackio.log({"final_val_loss": final_val})
+            trackio.finish()
+            print("✅ TrackIO run finished")
+        except Exception as e:
+            print(f"⚠️ TrackIO finish failed: {e}")
+        
+        del model, trainer, checkpoint_cb
+        gc.collect()
+        torch.cuda.empty_cache()
+
+    def execute(self):
+        print("\n🔥 TRAINING PIPELINE STARTED\n")
+        
+        # Train Scratch CNN
+        scratch = ScratchCNNModel(self.cfg.LEARNING_RATE, self.cfg.AGE_LOSS_WEIGHT)
+        self._train_model(scratch, "scratch", "scratch-cnn-run")
+        
+        # Train Fine-Tuned ResNet
+        finetuned = FineTunedResNetModel(self.cfg.LEARNING_RATE, self.cfg.AGE_LOSS_WEIGHT)
+        self._train_model(finetuned, "finetuned", "resnet-finetuned-run")
+        
+        print("\n🎉 TRAINING COMPLETE!")
+        print("📂 Checkpoints: /kaggle/working/")
+        print("📊 TrackIO Dashboard: https://huggingface.co/spaces/muhammad-bilal1/dlgenai-nppe")
+
+# ============================================================================
+# RUN TRAINING
+# ============================================================================
+
+if __name__ == "__main__":
+    pipeline = PipelineRunner(settings)
+    pipeline.execute()
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()