Upload app.py with huggingface_hub

app.py

@@ -1,217 +1,217 @@
-# import torch
-# import gradio as gr
-# import torchvision.transforms as transforms
-# from PIL import Image
-# from huggingface_hub import hf_hub_download
-# from C2D.models.resnet import SupCEResNet
-
-# # Define class labels
-# class_labels = [
-#     "T-shirt", "Shirt", "Knitwear", "Chiffon", "Sweater",
-#     "Hoodie", "Windbreaker", "Jacket", "Down Coat", "Suit",
-#     "Shawl", "Dress", "Vest", "Underwear"
-# ]
-
-# # Load model from Hugging Face Hub
-# def load_model_from_huggingface(repo_id="tfarhan10/Clothing1M-Pretrained-ResNet50", filename="model.pth"):
-#     try:
-#         print("Downloading model from Hugging Face...")
-#         checkpoint_path = hf_hub_download(repo_id=repo_id, filename=filename)
-
-#         # Load checkpoint
-#         checkpoint = torch.load(checkpoint_path, map_location=torch.device('cpu'),weights_only=False)
-
-#         # Extract state_dict if stored in a dictionary
-#         if isinstance(checkpoint, dict) and "model" in checkpoint:
-#             state_dict = checkpoint["model"]
-#         else:
-#             state_dict = checkpoint
-
-#         # Fix "module." prefix issue
-#         new_state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
-
-#         # Initialize model
-#         model = SupCEResNet(name='resnet50', num_classes=14, pool=True)
-
-#         # Load weights
-#         model.load_state_dict(new_state_dict, strict=False)  # `strict=False` allows minor mismatches
-#         model.eval()  # Set model to evaluation mode
-
-#         print("✅ Model loaded successfully from Hugging Face!")
-#         return model
-
-#     except Exception as e:
-#         print(f"❌ Error loading model: {e}")
-#         return None
-
-
-# # Load the model
-# model = load_model_from_huggingface()
-
-# def classify_image(image):
-#     """Process and classify an uploaded PIL image accurately."""
-    
-#     # Convert image to RGB to avoid grayscale or RGBA issues
-#     if image.mode != "RGB":
-#         image = image.convert("RGB")
-
-#     # Define the same preprocessing pipeline as training
-#     transform_test = transforms.Compose([
-#         transforms.Resize(256),  # Resize the shorter side to 256
-#         transforms.CenterCrop(224),  # Center crop to 224x224 (expected input size)
-#         transforms.ToTensor(),  # Convert to Tensor
-#         transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),  # Normalize
-#     ])
-
-#     # Apply transformations
-#     image_tensor = transform_test(image).unsqueeze(0)  # Add batch dimension
-
-#     # Ensure tensor is on the same device as model
-#     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-#     model.to(device)
-#     image_tensor = image_tensor.to(device)
-
-#     # Run inference
-#     with torch.no_grad():
-#         output = model(image_tensor)
-#         _, pred = torch.max(output, 1)  # Get predicted class index
-
-#     # Map predicted class index to label
-#     predicted_label = class_labels[pred.item()]
-#     print(pred.item())  
-#     return f"Predicted Category: {predicted_label}"
-
-
-# # Create Gradio Interface
-# example = "https://huggingface.co/tfarhan10/Clothing1M-Pretrained-ResNet50/blob/main/content/drive/MyDrive/CS5930/download.jpeg"
-# interface = gr.Interface(
-#     fn=classify_image,
-#     inputs=gr.Image(type="pil"),  # Accept image input
-#     outputs="text",
-#     title="Clothing Image Classifier",
-#     description="Upload an image and the model will classify it into one of 14 clothing categories.",
-#     allow_flagging="never",  # Disable flagging feature
-#     examples = [[example]]
-# )
-
-# # Launch the app
-# if __name__ == "__main__":
-#     interface.launch()
-
-
+"""ResNet in PyTorch.
+ImageNet-Style ResNet
+[1] Kaiming He, Xiangyu Zhang, Shaoqing Ren, Jian Sun
+    Deep Residual Learning for Image Recognition. arXiv:1512.03385
+Adapted from: https://github.com/bearpaw/pytorch-classification
+"""
 import torch
-import gradio as gr
-import torchvision.transforms as transforms
-from PIL import Image
-from huggingface_hub import hf_hub_download
-import requests
-from io import BytesIO
-from resnet import SupCEResNet
-
-# Define class labels
-class_labels = [
-    "T-shirt", "Shirt", "Knitwear", "Chiffon", "Sweater",
-    "Hoodie", "Windbreaker", "Jacket", "Down Coat", "Suit",
-    "Shawl", "Dress", "Vest", "Underwear"
-]
-
-# Load model from Hugging Face Hub
-def load_model_from_huggingface(repo_id="tfarhan10/Clothing1M-Pretrained-ResNet50", filename="model.pth"):
-    try:
-        print("Downloading model from Hugging Face...")
-        checkpoint_path = hf_hub_download(repo_id=repo_id, filename=filename)
-
-        # Load checkpoint
-        checkpoint = torch.load(checkpoint_path, map_location=torch.device('cpu'),weights_only=False)
-
-        # Extract state_dict if stored in a dictionary
-        if isinstance(checkpoint, dict) and "model" in checkpoint:
-            state_dict = checkpoint["model"]
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, in_planes, planes, stride=1, is_last=False):
+        super(BasicBlock, self).__init__()
+        self.is_last = is_last
+        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+
+        self.shortcut = nn.Sequential()
+        if stride != 1 or in_planes != self.expansion * planes:
+            self.shortcut = nn.Sequential(
+                nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(self.expansion * planes)
+            )
+
+    def forward(self, x):
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = self.bn2(self.conv2(out))
+        out += self.shortcut(x)
+        preact = out
+        out = F.relu(out)
+        if self.is_last:
+            return out, preact
         else:
-            state_dict = checkpoint
-
-        # Fix "module." prefix issue
-        new_state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
-
-        # Initialize model
-        model = SupCEResNet(name='resnet50', num_classes=14, pool=True)
-
-        # Load weights
-        model.load_state_dict(new_state_dict, strict=False)  # `strict=False` allows minor mismatches
-        model.eval()  # Set model to evaluation mode
+            return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, in_planes, planes, stride=1, is_last=False):
+        super(Bottleneck, self).__init__()
+        self.is_last = is_last
+        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv3 = nn.Conv2d(planes, self.expansion * planes, kernel_size=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(self.expansion * planes)
+
+        self.shortcut = nn.Sequential()
+        if stride != 1 or in_planes != self.expansion * planes:
+            self.shortcut = nn.Sequential(
+                nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(self.expansion * planes)
+            )
+
+    def forward(self, x):
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = F.relu(self.bn2(self.conv2(out)))
+        out = self.bn3(self.conv3(out))
+        out += self.shortcut(x)
+        preact = out
+        out = F.relu(out)
+        if self.is_last:
+            return out, preact
+        else:
+            return out
 
-        print("✅ Model loaded successfully from Hugging Face!")
-        return model
 
-    except Exception as e:
-        print(f"❌ Error loading model: {e}")
-        return None
+class ResNet(nn.Module):
+    def __init__(self, block, num_blocks, in_channel=3, zero_init_residual=False, pool=False):
+        super(ResNet, self).__init__()
+        self.in_planes = 64
 
-# Load the model
-model = load_model_from_huggingface()
+        if pool:
+            self.conv1 = nn.Conv2d(in_channel, 64, kernel_size=7, stride=2, padding=3, bias=False)
+        else:
+            self.conv1 = nn.Conv2d(in_channel, 64, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(64)
+
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) if pool else nn.Identity()
+        self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1)
+        self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2)
+        self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2)
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+
+        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.GroupNorm)):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        # Zero-initialize the last BN in each residual branch,
+        # so that the residual branch starts with zeros, and each residual block behaves
+        # like an identity. This improves the model by 0.2~0.3% according to:
+        # https://arxiv.org/abs/1706.02677
+        if zero_init_residual:
+            for m in self.modules():
+                if isinstance(m, Bottleneck):
+                    nn.init.constant_(m.bn3.weight, 0)
+                elif isinstance(m, BasicBlock):
+                    nn.init.constant_(m.bn2.weight, 0)
+
+    def _make_layer(self, block, planes, num_blocks, stride):
+        strides = [stride] + [1] * (num_blocks - 1)
+        layers = []
+        for i in range(num_blocks):
+            stride = strides[i]
+            layers.append(block(self.in_planes, planes, stride))
+            self.in_planes = planes * block.expansion
+        return nn.Sequential(*layers)
+
+    def forward(self, x, layer=100):
+        out = self.maxpool(F.relu(self.bn1(self.conv1(x))))
+        out = self.layer1(out)
+        out = self.layer2(out)
+        out = self.layer3(out)
+        out = self.layer4(out)
+        out = self.avgpool(out)
+        out = torch.flatten(out, 1)
+        return out
+
+
+def resnet18(**kwargs):
+    return ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)
+
+
+def resnet34(**kwargs):
+    return ResNet(BasicBlock, [3, 4, 6, 3], **kwargs)
+
+
+def resnet50(**kwargs):
+    return ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
+
+
+def resnet101(**kwargs):
+    return ResNet(Bottleneck, [3, 4, 23, 3], **kwargs)
+
+
+model_dict = {
+    'resnet18': [resnet18, 512],
+    'resnet34': [resnet34, 512],
+    'resnet50': [resnet50, 2048],
+    'resnet101': [resnet101, 2048],
+}
+
+
+class LinearBatchNorm(nn.Module):
+    """Implements BatchNorm1d by BatchNorm2d, for SyncBN purpose"""
+
+    def __init__(self, dim, affine=True):
+        super(LinearBatchNorm, self).__init__()
+        self.dim = dim
+        self.bn = nn.BatchNorm2d(dim, affine=affine)
+
+    def forward(self, x):
+        x = x.view(-1, self.dim, 1, 1)
+        x = self.bn(x)
+        x = x.view(-1, self.dim)
+        return x
+
+
+class SupConResNet(nn.Module):
+    """backbone + projection head"""
+
+    def __init__(self, name='resnet50', head='mlp', feat_dim=128, pool=False):
+        super(SupConResNet, self).__init__()
+        model_fun, dim_in = model_dict[name]
+        self.encoder = model_fun(pool=pool)
+        if head == 'linear':
+            self.head = nn.Linear(dim_in, feat_dim)
+        elif head == 'mlp':
+            self.head = nn.Sequential(
+                nn.Linear(dim_in, dim_in),
+                nn.ReLU(inplace=True),
+                nn.Linear(dim_in, feat_dim)
+            )
+        else:
+            raise NotImplementedError(
+                'head not supported: {}'.format(head))
 
-def classify_image(image):
-    """Process and classify an uploaded PIL image accurately."""
+    def forward(self, x):
+        feat = self.encoder(x)
+        feat = F.normalize(self.head(feat), dim=1)
+        return feat
 
-    # Ensure image is in RGB format
-    if image.mode != "RGB":
-        image = image.convert("RGB")
 
-    # Define preprocessing transformations (same as training)
-    transform_test = transforms.Compose([
-        transforms.Resize(256),  # Resize the shorter side to 256
-        transforms.CenterCrop(224),  # Center crop to 224x224 (expected input size)
-        transforms.ToTensor(),  # Convert to Tensor
-        transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),  # Normalize
-    ])
+class SupCEResNet(nn.Module):
+    """encoder + classifier"""
 
-    # Apply transformations
-    image_tensor = transform_test(image).unsqueeze(0)  # Add batch dimension
+    def __init__(self, name='resnet50', num_classes=10, pool=False):
+        super(SupCEResNet, self).__init__()
+        model_fun, dim_in = model_dict[name]
+        self.encoder = model_fun(pool=pool)
+        self.fc = nn.Linear(dim_in, num_classes)
 
-    # Ensure tensor is on the same device as model
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model.to(device)
-    image_tensor = image_tensor.to(device)
+    def forward(self, x):
+        return self.fc(self.encoder(x))
 
-    # Run inference
-    with torch.no_grad():
-        output = model(image_tensor)
-        _, pred = torch.max(output, 1)  # Get predicted class index
 
-    # Map predicted class index to label
-    predicted_label = class_labels[pred.item()]
-    return f"Predicted Category: {predicted_label}"
+class LinearClassifier(nn.Module):
+    """Linear classifier"""
 
-# Load example image from Hugging Face repository
-example_url = "https://huggingface.co/tfarhan10/Clothing1M-Pretrained-ResNet50/resolve/main/content/drive/MyDrive/CS5930/download.jpeg"
+    def __init__(self, name='resnet50', num_classes=10):
+        super(LinearClassifier, self).__init__()
+        _, feat_dim = model_dict[name]
+        self.fc = nn.Linear(feat_dim, num_classes)
 
-def load_example_image():
-    """Download and return an example image from Hugging Face"""
-    try:
-        response = requests.get(example_url)
-        if response.status_code == 200:
-            return Image.open(BytesIO(response.content)).convert("RGB")
-        else:
-            print("⚠️ Failed to fetch example image.")
-            return None
-    except Exception as e:
-        print(f"⚠️ Error loading example image: {e}")
-        return None
-
-# Example image
-example_image = load_example_image()
-
-# Create Gradio Interface
-interface = gr.Interface(
-    fn=classify_image,
-    inputs=gr.Image(type="pil"),  # Accept image input
-    outputs="text",
-    title="Clothing Image Classifier",
-    description="Upload an image or use the example below. The model will classify it into one of 14 clothing categories.",
-    allow_flagging="never",  # Disable flagging feature
-    examples=[[example_image]] if example_image else None  # Use example image if available
-)
-
-# Launch the app
-if __name__ == "__main__":
-    interface.launch()
+    def forward(self, features):
+        return self.fc(features)