Upload app.py with huggingface_hub

app.py

@@ -1,217 +1,217 @@
-"""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
+# 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()
+
+
 import torch
-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:
-            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
+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"]
         else:
-            return out
+            state_dict = checkpoint
 
+        # Fix "module." prefix issue
+        new_state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
 
-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
+        # Initialize model
+        model = SupCEResNet(name='resnet50', num_classes=14, pool=True)
 
-        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))
+        # Load weights
+        model.load_state_dict(new_state_dict, strict=False)  # `strict=False` allows minor mismatches
+        model.eval()  # Set model to evaluation mode
 
-    def forward(self, x):
-        feat = self.encoder(x)
-        feat = F.normalize(self.head(feat), dim=1)
-        return feat
+        print("✅ Model loaded successfully from Hugging Face!")
+        return model
 
+    except Exception as e:
+        print(f"❌ Error loading model: {e}")
+        return None
 
-class SupCEResNet(nn.Module):
-    """encoder + classifier"""
+# Load the model
+model = load_model_from_huggingface()
 
-    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)
+def classify_image(image):
+    """Process and classify an uploaded PIL image accurately."""
 
-    def forward(self, x):
-        return self.fc(self.encoder(x))
+    # 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 LinearClassifier(nn.Module):
-    """Linear classifier"""
+    # Apply transformations
+    image_tensor = transform_test(image).unsqueeze(0)  # Add batch dimension
 
-    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)
+    # 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, features):
-        return self.fc(features)
+    # 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}"
+
+# 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 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()