Spaces:
Configuration error
Configuration error
Upload 4 files
Browse files
app.py
ADDED
@@ -0,0 +1,69 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import gradio as gr
|
2 |
+
import torch
|
3 |
+
import torchvision.transforms as transforms
|
4 |
+
from PIL import Image
|
5 |
+
from resnet import SupCEResNet # Ensure the correct import path
|
6 |
+
|
7 |
+
# ✅ Define class labels (from Clothing1M)
|
8 |
+
class_labels = [
|
9 |
+
"T-Shirt", "Shirt", "Knitwear", "Chiffon", "Sweater", "Hoodie",
|
10 |
+
"Windbreaker", "Jacket", "Downcoat", "Suit", "Shawl", "Dress",
|
11 |
+
"Vest", "Underwear"
|
12 |
+
]
|
13 |
+
|
14 |
+
# ✅ Function to load the model
|
15 |
+
def create_model_selfsup(net='resnet50', num_class=14, checkpoint_path='/content/ckpt_clothing_resnet50.pth'):
|
16 |
+
"""Loads a self-supervised pretrained model for Clothing1M classification"""
|
17 |
+
print(f"🔄 Loading model from: {checkpoint_path}")
|
18 |
+
|
19 |
+
# Load the checkpoint safely
|
20 |
+
checkpoint = torch.load(checkpoint_path, map_location="cuda" if torch.cuda.is_available() else "cpu", weights_only=False)
|
21 |
+
|
22 |
+
# Remove 'module.' prefix if using DataParallel
|
23 |
+
state_dict = {k.replace('module.', ''): v for k, v in checkpoint['model'].items()}
|
24 |
+
|
25 |
+
# Initialize and load model
|
26 |
+
model = SupCEResNet(net, num_classes=num_class, pool=True)
|
27 |
+
model.load_state_dict(state_dict, strict=False)
|
28 |
+
|
29 |
+
# Move model to GPU if available
|
30 |
+
model = model.to("cuda" if torch.cuda.is_available() else "cpu")
|
31 |
+
model.eval() # Set model to evaluation mode
|
32 |
+
|
33 |
+
print("✅ Model loaded successfully!")
|
34 |
+
return model
|
35 |
+
|
36 |
+
# ✅ Load the model once
|
37 |
+
model = create_model_selfsup()
|
38 |
+
|
39 |
+
# ✅ Define image preprocessing function
|
40 |
+
def preprocess_image(image):
|
41 |
+
"""Transforms input image for the model"""
|
42 |
+
transform = transforms.Compose([
|
43 |
+
transforms.Resize(256),
|
44 |
+
transforms.CenterCrop(224),
|
45 |
+
transforms.ToTensor(),
|
46 |
+
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
|
47 |
+
])
|
48 |
+
return transform(image).unsqueeze(0).to("cuda" if torch.cuda.is_available() else "cpu")
|
49 |
+
|
50 |
+
# ✅ Define inference function
|
51 |
+
def predict_clothing(image):
|
52 |
+
"""Runs inference on an uploaded image"""
|
53 |
+
image = Image.fromarray(image) # Convert numpy array to PIL Image
|
54 |
+
image = preprocess_image(image) # Preprocess image
|
55 |
+
|
56 |
+
with torch.no_grad():
|
57 |
+
output = model(image)
|
58 |
+
predicted_class = torch.argmax(output, dim=1).item() # Get class index
|
59 |
+
|
60 |
+
return class_labels[predicted_class] # Return class name
|
61 |
+
|
62 |
+
# ✅ Create Gradio Interface
|
63 |
+
gr.Interface(
|
64 |
+
fn=predict_clothing,
|
65 |
+
inputs=gr.Image(type="numpy"),
|
66 |
+
outputs=gr.Textbox(label="Predicted Clothing Type"),
|
67 |
+
title="Clothing1M Classification",
|
68 |
+
description="Upload an image to classify clothing into one of 14 categories."
|
69 |
+
).launch()
|
model.pth
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:a4d2f883bffa3447651ffcd7561245014edfa0c7b979ae22dee7960df9a94c11
|
3 |
+
size 224033335
|
requirments.txt
ADDED
@@ -0,0 +1,6 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
torch==2.2.0
|
2 |
+
torchvision==0.17.0
|
3 |
+
gradio==4.21.0
|
4 |
+
Pillow==10.2.0
|
5 |
+
numpy==1.26.4
|
6 |
+
scikit-learn==1.4.1.post1
|
resnet.py
ADDED
@@ -0,0 +1,217 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
"""ResNet in PyTorch.
|
2 |
+
ImageNet-Style ResNet
|
3 |
+
[1] Kaiming He, Xiangyu Zhang, Shaoqing Ren, Jian Sun
|
4 |
+
Deep Residual Learning for Image Recognition. arXiv:1512.03385
|
5 |
+
Adapted from: https://github.com/bearpaw/pytorch-classification
|
6 |
+
"""
|
7 |
+
import torch
|
8 |
+
import torch.nn as nn
|
9 |
+
import torch.nn.functional as F
|
10 |
+
|
11 |
+
|
12 |
+
class BasicBlock(nn.Module):
|
13 |
+
expansion = 1
|
14 |
+
|
15 |
+
def __init__(self, in_planes, planes, stride=1, is_last=False):
|
16 |
+
super(BasicBlock, self).__init__()
|
17 |
+
self.is_last = is_last
|
18 |
+
self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
|
19 |
+
self.bn1 = nn.BatchNorm2d(planes)
|
20 |
+
self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
|
21 |
+
self.bn2 = nn.BatchNorm2d(planes)
|
22 |
+
|
23 |
+
self.shortcut = nn.Sequential()
|
24 |
+
if stride != 1 or in_planes != self.expansion * planes:
|
25 |
+
self.shortcut = nn.Sequential(
|
26 |
+
nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False),
|
27 |
+
nn.BatchNorm2d(self.expansion * planes)
|
28 |
+
)
|
29 |
+
|
30 |
+
def forward(self, x):
|
31 |
+
out = F.relu(self.bn1(self.conv1(x)))
|
32 |
+
out = self.bn2(self.conv2(out))
|
33 |
+
out += self.shortcut(x)
|
34 |
+
preact = out
|
35 |
+
out = F.relu(out)
|
36 |
+
if self.is_last:
|
37 |
+
return out, preact
|
38 |
+
else:
|
39 |
+
return out
|
40 |
+
|
41 |
+
|
42 |
+
class Bottleneck(nn.Module):
|
43 |
+
expansion = 4
|
44 |
+
|
45 |
+
def __init__(self, in_planes, planes, stride=1, is_last=False):
|
46 |
+
super(Bottleneck, self).__init__()
|
47 |
+
self.is_last = is_last
|
48 |
+
self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=1, bias=False)
|
49 |
+
self.bn1 = nn.BatchNorm2d(planes)
|
50 |
+
self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
|
51 |
+
self.bn2 = nn.BatchNorm2d(planes)
|
52 |
+
self.conv3 = nn.Conv2d(planes, self.expansion * planes, kernel_size=1, bias=False)
|
53 |
+
self.bn3 = nn.BatchNorm2d(self.expansion * planes)
|
54 |
+
|
55 |
+
self.shortcut = nn.Sequential()
|
56 |
+
if stride != 1 or in_planes != self.expansion * planes:
|
57 |
+
self.shortcut = nn.Sequential(
|
58 |
+
nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False),
|
59 |
+
nn.BatchNorm2d(self.expansion * planes)
|
60 |
+
)
|
61 |
+
|
62 |
+
def forward(self, x):
|
63 |
+
out = F.relu(self.bn1(self.conv1(x)))
|
64 |
+
out = F.relu(self.bn2(self.conv2(out)))
|
65 |
+
out = self.bn3(self.conv3(out))
|
66 |
+
out += self.shortcut(x)
|
67 |
+
preact = out
|
68 |
+
out = F.relu(out)
|
69 |
+
if self.is_last:
|
70 |
+
return out, preact
|
71 |
+
else:
|
72 |
+
return out
|
73 |
+
|
74 |
+
|
75 |
+
class ResNet(nn.Module):
|
76 |
+
def __init__(self, block, num_blocks, in_channel=3, zero_init_residual=False, pool=False):
|
77 |
+
super(ResNet, self).__init__()
|
78 |
+
self.in_planes = 64
|
79 |
+
|
80 |
+
if pool:
|
81 |
+
self.conv1 = nn.Conv2d(in_channel, 64, kernel_size=7, stride=2, padding=3, bias=False)
|
82 |
+
else:
|
83 |
+
self.conv1 = nn.Conv2d(in_channel, 64, kernel_size=3, stride=1, padding=1, bias=False)
|
84 |
+
self.bn1 = nn.BatchNorm2d(64)
|
85 |
+
|
86 |
+
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) if pool else nn.Identity()
|
87 |
+
self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1)
|
88 |
+
self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2)
|
89 |
+
self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2)
|
90 |
+
self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2)
|
91 |
+
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
|
92 |
+
|
93 |
+
for m in self.modules():
|
94 |
+
if isinstance(m, nn.Conv2d):
|
95 |
+
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
|
96 |
+
elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
|
97 |
+
nn.init.constant_(m.weight, 1)
|
98 |
+
nn.init.constant_(m.bias, 0)
|
99 |
+
|
100 |
+
# Zero-initialize the last BN in each residual branch,
|
101 |
+
# so that the residual branch starts with zeros, and each residual block behaves
|
102 |
+
# like an identity. This improves the model by 0.2~0.3% according to:
|
103 |
+
# https://arxiv.org/abs/1706.02677
|
104 |
+
if zero_init_residual:
|
105 |
+
for m in self.modules():
|
106 |
+
if isinstance(m, Bottleneck):
|
107 |
+
nn.init.constant_(m.bn3.weight, 0)
|
108 |
+
elif isinstance(m, BasicBlock):
|
109 |
+
nn.init.constant_(m.bn2.weight, 0)
|
110 |
+
|
111 |
+
def _make_layer(self, block, planes, num_blocks, stride):
|
112 |
+
strides = [stride] + [1] * (num_blocks - 1)
|
113 |
+
layers = []
|
114 |
+
for i in range(num_blocks):
|
115 |
+
stride = strides[i]
|
116 |
+
layers.append(block(self.in_planes, planes, stride))
|
117 |
+
self.in_planes = planes * block.expansion
|
118 |
+
return nn.Sequential(*layers)
|
119 |
+
|
120 |
+
def forward(self, x, layer=100):
|
121 |
+
out = self.maxpool(F.relu(self.bn1(self.conv1(x))))
|
122 |
+
out = self.layer1(out)
|
123 |
+
out = self.layer2(out)
|
124 |
+
out = self.layer3(out)
|
125 |
+
out = self.layer4(out)
|
126 |
+
out = self.avgpool(out)
|
127 |
+
out = torch.flatten(out, 1)
|
128 |
+
return out
|
129 |
+
|
130 |
+
|
131 |
+
def resnet18(**kwargs):
|
132 |
+
return ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)
|
133 |
+
|
134 |
+
|
135 |
+
def resnet34(**kwargs):
|
136 |
+
return ResNet(BasicBlock, [3, 4, 6, 3], **kwargs)
|
137 |
+
|
138 |
+
|
139 |
+
def resnet50(**kwargs):
|
140 |
+
return ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
|
141 |
+
|
142 |
+
|
143 |
+
def resnet101(**kwargs):
|
144 |
+
return ResNet(Bottleneck, [3, 4, 23, 3], **kwargs)
|
145 |
+
|
146 |
+
|
147 |
+
model_dict = {
|
148 |
+
'resnet18': [resnet18, 512],
|
149 |
+
'resnet34': [resnet34, 512],
|
150 |
+
'resnet50': [resnet50, 2048],
|
151 |
+
'resnet101': [resnet101, 2048],
|
152 |
+
}
|
153 |
+
|
154 |
+
|
155 |
+
class LinearBatchNorm(nn.Module):
|
156 |
+
"""Implements BatchNorm1d by BatchNorm2d, for SyncBN purpose"""
|
157 |
+
|
158 |
+
def __init__(self, dim, affine=True):
|
159 |
+
super(LinearBatchNorm, self).__init__()
|
160 |
+
self.dim = dim
|
161 |
+
self.bn = nn.BatchNorm2d(dim, affine=affine)
|
162 |
+
|
163 |
+
def forward(self, x):
|
164 |
+
x = x.view(-1, self.dim, 1, 1)
|
165 |
+
x = self.bn(x)
|
166 |
+
x = x.view(-1, self.dim)
|
167 |
+
return x
|
168 |
+
|
169 |
+
|
170 |
+
class SupConResNet(nn.Module):
|
171 |
+
"""backbone + projection head"""
|
172 |
+
|
173 |
+
def __init__(self, name='resnet50', head='mlp', feat_dim=128, pool=False):
|
174 |
+
super(SupConResNet, self).__init__()
|
175 |
+
model_fun, dim_in = model_dict[name]
|
176 |
+
self.encoder = model_fun(pool=pool)
|
177 |
+
if head == 'linear':
|
178 |
+
self.head = nn.Linear(dim_in, feat_dim)
|
179 |
+
elif head == 'mlp':
|
180 |
+
self.head = nn.Sequential(
|
181 |
+
nn.Linear(dim_in, dim_in),
|
182 |
+
nn.ReLU(inplace=True),
|
183 |
+
nn.Linear(dim_in, feat_dim)
|
184 |
+
)
|
185 |
+
else:
|
186 |
+
raise NotImplementedError(
|
187 |
+
'head not supported: {}'.format(head))
|
188 |
+
|
189 |
+
def forward(self, x):
|
190 |
+
feat = self.encoder(x)
|
191 |
+
feat = F.normalize(self.head(feat), dim=1)
|
192 |
+
return feat
|
193 |
+
|
194 |
+
|
195 |
+
class SupCEResNet(nn.Module):
|
196 |
+
"""encoder + classifier"""
|
197 |
+
|
198 |
+
def __init__(self, name='resnet50', num_classes=10, pool=False):
|
199 |
+
super(SupCEResNet, self).__init__()
|
200 |
+
model_fun, dim_in = model_dict[name]
|
201 |
+
self.encoder = model_fun(pool=pool)
|
202 |
+
self.fc = nn.Linear(dim_in, num_classes)
|
203 |
+
|
204 |
+
def forward(self, x):
|
205 |
+
return self.fc(self.encoder(x))
|
206 |
+
|
207 |
+
|
208 |
+
class LinearClassifier(nn.Module):
|
209 |
+
"""Linear classifier"""
|
210 |
+
|
211 |
+
def __init__(self, name='resnet50', num_classes=10):
|
212 |
+
super(LinearClassifier, self).__init__()
|
213 |
+
_, feat_dim = model_dict[name]
|
214 |
+
self.fc = nn.Linear(feat_dim, num_classes)
|
215 |
+
|
216 |
+
def forward(self, features):
|
217 |
+
return self.fc(features)
|