Migration with slight changes

app.py

@@ -0,0 +1,207 @@
+import torch
+from pathlib import Path
+from huggingface_hub import hf_hub_download
+from PIL import Image
+from torchvision import transforms
+from medmnist import INFO
+import gradio as gr
+import os
+import base64
+from io import BytesIO
+from huggingface_hub import HfApi
+from datetime import datetime
+import io
+
+from model import resnet18, resnet50
+
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu")
+AUTH_TOKEN = os.getenv("APP_TOKEN")#to acces the app
+DATASET_REPO = os.getenv("Dataset_repo") #"G44mlops/API_received"  
+HF_TOKEN = os.getenv("HF_TOKEN")  #to acces dataset repo
+MODEL = os.getenv("Model_repo")#"G44mlops/ResNet-medmnist"
+
+#taken from Mikolaj code with closed PR
+def load_model_from_hf(
+    repo_id: str,
+    filename: str,
+    model_type: str,
+    num_classes: int,
+    in_channels: int,
+    device: str,
+) -> torch.nn.Module:
+    """Load trained model from Hugging Face Hub.
+    
+    Args:
+        repo_id: Hugging Face repository ID
+        filename: Model checkpoint filename
+        model_type: Type of model ('resnet18' or 'resnet50')
+        num_classes: Number of output classes
+        in_channels: Number of input channels
+        device: Device to load model on
+        
+    Returns:
+        Loaded model in eval mode
+    """
+    print(f"Downloading model from Hugging Face: {repo_id}/{filename}")
+    checkpoint_path = hf_hub_download(repo_id=repo_id, filename=filename)
+
+    # Create model
+    if model_type == "resnet18":
+        model = resnet18(num_classes=num_classes, in_channels=in_channels)
+    else:
+        model = resnet50(num_classes=num_classes, in_channels=in_channels)
+
+    # Load checkpoint
+    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=True)
+    model.load_state_dict(checkpoint["model_state_dict"])
+    model.to(device)
+    model.eval()
+
+    return model
+#taken from Mikolaj code with closed PR
+
+
+# Image preprocessing pipeline (basic so far, can be improved)
+def get_preprocessing_pipeline() -> transforms.Compose:
+    """Get preprocessing pipeline for images."""
+    #getting information on number of image channels (RGB or Grayscale) for trained model
+    info = INFO["organamnist"]  # Using organamnist as reference
+    output_channels = info["n_channels"] # RGB or Grayscale
+    #chosing 'standard' mean and std values for normalization if dataset statistics are not available
+    mean = (0.5,) * output_channels
+    std = (0.5,) * output_channels
+    #preparing transformation pipeline
+    trans = transforms.Compose([
+        transforms.Resize(256),
+        transforms.CenterCrop(224),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=mean, std=std),
+    ])
+    #returning the transformation pipeline
+    return trans
+def get_class_labels(data_flag: str = "organamnist") -> list[str]:
+    """Get class labels for MedMNIST dataset."""
+    #retrieving dataset info
+    info = INFO[data_flag]
+    labels = info["label"]
+    #returning class labels
+    return labels
+
+def save_image_to_hf_folder(image_path, prediction_label):
+    """Upload image to HF dataset folder."""
+    api = HfApi()
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    
+    # Create a text file with metadata
+    metadata = f"prediction: {prediction_label}\ntimestamp: {timestamp}"
+    metadata_path = f"{Path(image_path).stem}_metadata.txt"
+    
+    # Upload image
+    api.upload_file(
+        path_or_fileobj=image_path,
+        path_in_repo=f"uploads/{timestamp}_{Path(image_path).name}",
+        repo_id=DATASET_REPO,
+        repo_type="dataset",
+        token=HF_TOKEN
+    )
+     # Upload metadata as separate file
+    api.upload_file(
+        path_or_fileobj=io.BytesIO(metadata.encode()),
+        path_in_repo=f"uploads/{timestamp}_{Path(image_path).stem}_metadata.txt",
+        repo_id=DATASET_REPO,
+        repo_type="dataset",
+        token=HF_TOKEN
+    )
+
+def classify_images(images) -> str:
+    """Classify images and return formatted HTML with embedded images."""
+    # Handle case with no images
+    if images is None:
+        return "<p>No images uploaded</p>"
+    # Ensure images is a list if(case when only one image is uploaded is problematic without it)
+    if isinstance(images, str):
+        images = [images]
+    #creating HTML structure for results
+    html = "<div style='display: flex; flex-wrap: wrap; gap: 30px; padding: 20px; justify-content: center;'>"
+    #loop over images and classify them
+    for image_path in images:
+        #preparing image for classification
+        img = Image.open(image_path).convert("L")  # Convert to grayscale (as project uses grayscale images)
+        input_tensor = preprocess(img).unsqueeze(0)
+        #forward pass + softmax to get probabilities
+        with torch.no_grad():
+            output = model(input_tensor)
+            probs = torch.nn.functional.softmax(output[0], dim=0)
+            top_class = probs.argmax().item()
+        #getting class label
+        label = class_labels[str(top_class)]
+        #getting image filename
+        filename = Path(image_path).name
+        #Preparing image for embedding in HTML (base64 encoding)
+        buffered = BytesIO()
+        img.save(buffered, format="JPEG")
+        img_str = base64.b64encode(buffered.getvalue()).decode()
+        #adding current image block to HTML
+        html += f"""
+        <div style='border: 2px solid #ddd; padding: 15px; border-radius: 8px; background: #f9f9f9; width: 280px;'>
+            <p style='font-size: 14px; color: #666; margin: 0 0 10px 0; text-align: center; font-weight: bold;'>{filename}</p>
+            <img src='data:image/jpeg;base64,{img_str}' style='width: 250px; height: 250px; object-fit: contain; display: block; margin: 0 auto 10px;'>
+            <p style='font-size: 18px; color: #0066cc; margin: 10px 0 0 0; text-align: center; font-weight: bold;'>{label}</p>
+        </div>
+        """
+        # Save image and metadata to HF dataset folder
+        save_image_to_hf_folder(image_path, label)
+    #closing HTML container
+    html += "</div>"
+    #returning results
+    return html
+
+###main code to launch Gradio app###
+
+#prepare model and preprocessing pipeline (kind of backend)
+model =  load_model_from_hf(#taken from Mikolaj code with closed PR
+            repo_id=MODEL, 
+            filename="resnet18_best.pth",
+            model_type="resnet18",
+            num_classes=11,
+            in_channels=1,
+            device=DEVICE,
+        )
+preprocess = get_preprocessing_pipeline()
+class_labels = get_class_labels()
+#preparing Gradio interface (frontend)
+with gr.Blocks() as demo:
+    #app "title"
+    gr.Markdown("<h1 style='text-align: center;'> MLOps project - MedMNIST dataset Image Classifier</h1>")
+    #app message/information )
+    gr.Markdown("This is a Gradio web application for MLOps course project. Given images are stored in our dataset. " \
+    "By uploading images you agrree that they will be stored by us and insures that they can be stored by us. " \
+    "If you somewhat passed the login and are not connected to the project, please do not upload any images. " )
+    #app spine layout
+    with gr.Column():
+        #title of load segment
+        gr.Markdown("<h2 style='text-align: center;'> Upload Images</h2>")
+        #images loading component
+        images_input = gr.File(file_count="multiple", file_types=["image"], label="Upload Images")
+        #buttons row for app functionality
+        with gr.Row():
+            submit_btn = gr.Button("Classify")
+            reset_btn = gr.Button("Reset")
+        #title of results segment
+        gr.Markdown("<h2 style='text-align: center;'> Results</h2>")
+        #classification results output component
+        output = gr.HTML(label="Results")
+    #getting callable reset function
+    def reset():
+        return None, ""
+    #linking buttons to functions
+    submit_btn.click(classify_images, inputs=images_input, outputs=output)
+    reset_btn.click(reset, outputs=[images_input, output])
+
+
+    #just launch
+    server_name = os.getenv("GRADIO_SERVER_NAME", "127.0.0.1")
+    demo.launch(
+        server_name=server_name,
+        auth=[("user", AUTH_TOKEN)] if AUTH_TOKEN else None
+        )

model.py

@@ -0,0 +1,205 @@
+import torch
+from torch import nn
+
+
+class BasicBlock(nn.Module):
+    """Basic building block for ResNet-18/34"""
+    expansion = 1
+
+    def __init__(self, in_channels: int, out_channels: int, stride: int = 1, downsample: nn.Module = None):
+        super().__init__()
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(out_channels)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(out_channels)
+        self.downsample = downsample
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    """Bottleneck building block for ResNet-50/101/152"""
+    expansion = 4
+
+    def __init__(self, in_channels: int, out_channels: int, stride: int = 1, downsample: nn.Module = None):
+        super().__init__()
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(out_channels)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(out_channels)
+        self.conv3 = nn.Conv2d(out_channels, out_channels * self.expansion, kernel_size=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(out_channels * self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet(nn.Module):
+    """ResNet model for image classification
+
+    Supports ResNet-18, ResNet-34, ResNet-50, ResNet-101, ResNet-152
+    Adapted for small images like MedMNIST (28x28)
+    """
+
+    def __init__(
+        self,
+        block: type[BasicBlock | Bottleneck],
+        layers: list[int],
+        num_classes: int = 11,
+        in_channels: int = 1,
+    ):
+        super().__init__()
+        self.in_channels = 64
+
+        # Initial convolution layer (adapted for small 28x28 images)
+        self.conv1 = nn.Conv2d(in_channels, 64, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.relu = nn.ReLU(inplace=True)
+        # Removed maxpool for small images
+
+        # ResNet layers
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
+
+        # Global average pooling and classifier
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        self.fc = nn.Linear(512 * block.expansion, num_classes)
+
+        # Initialize weights
+        self._initialize_weights()
+
+    def _make_layer(self, block: type[BasicBlock | Bottleneck], out_channels: int, blocks: int, stride: int = 1) -> nn.Sequential:
+        downsample = None
+        if stride != 1 or self.in_channels != out_channels * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.in_channels, out_channels * block.expansion, kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(out_channels * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.in_channels, out_channels, stride, downsample))
+        self.in_channels = out_channels * block.expansion
+        for _ in range(1, blocks):
+            layers.append(block(self.in_channels, out_channels))
+
+        return nn.Sequential(*layers)
+
+    def _initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+        x = self.fc(x)
+
+        return x
+
+
+def resnet18(num_classes: int = 11, in_channels: int = 1) -> ResNet:
+    """ResNet-18 model
+
+    Args:
+        num_classes: Number of output classes (default: 11 for organamnist)
+        in_channels: Number of input channels (default: 1 for grayscale)
+
+    Returns:
+        ResNet-18 model
+    """
+    return ResNet(BasicBlock, [2, 2, 2, 2], num_classes=num_classes, in_channels=in_channels)
+
+
+def resnet50(num_classes: int = 11, in_channels: int = 1) -> ResNet:
+    """ResNet-50 model
+
+    Args:
+        num_classes: Number of output classes (default: 11 for organamnist)
+        in_channels: Number of input channels (default: 1 for grayscale)
+
+    Returns:
+        ResNet-50 model
+    """
+    return ResNet(Bottleneck, [3, 4, 6, 3], num_classes=num_classes, in_channels=in_channels)
+
+
+# Keep the old Model class for backward compatibility
+class Model(nn.Module):
+    """Just a dummy model to show how to structure your code"""
+    def __init__(self):
+        super().__init__()
+        self.layer = nn.Linear(1, 1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.layer(x)
+
+
+if __name__ == "__main__":
+    # Test ResNet-18
+    model18 = resnet18(num_classes=11, in_channels=1)
+    x = torch.rand(4, 1, 28, 28)  # Batch of 4 grayscale 28x28 images
+    output = model18(x)
+    print(f"ResNet-18 output shape: {output.shape}")  # Should be [4, 11]
+
+    # Test ResNet-50
+    model50 = resnet50(num_classes=11, in_channels=1)
+    output50 = model50(x)
+    print(f"ResNet-50 output shape: {output50.shape}")  # Should be [4, 11]
+
+    # Count parameters
+    params18 = sum(p.numel() for p in model18.parameters())
+    params50 = sum(p.numel() for p in model50.parameters())
+    print(f"ResNet-18 parameters: {params18:,}")
+    print(f"ResNet-50 parameters: {params50:,}")