Add application file

Nets.py

@@ -0,0 +1,141 @@
+import torch
+import torch.nn as nn
+import math
+from collections import OrderedDict
+
+################## AlexNet ##################
+def bn_relu(inplanes):
+    return nn.Sequential(nn.BatchNorm2d(inplanes), nn.ReLU(inplace=True))
+
+def bn_relu_pool(inplanes, kernel_size=3, stride=2):
+    return nn.Sequential(nn.BatchNorm2d(inplanes), nn.ReLU(inplace=True), nn.MaxPool2d(kernel_size=kernel_size, stride=stride))
+
+class AlexNet(nn.Module):
+    def __init__(self, num_classes=1):
+        super(AlexNet, self).__init__()
+        self.conv1 = nn.Conv2d(3, 96, kernel_size=11, stride=4, bias=False)
+        self.relu_pool1 = bn_relu_pool(inplanes=96)
+        self.conv2 = nn.Conv2d(96, 192, kernel_size=5, padding=2, groups=2, bias=False)
+        self.relu_pool2 = bn_relu_pool(inplanes=192)
+        self.conv3 = nn.Conv2d(192, 384, kernel_size=3, padding=1, groups=2, bias=False)
+        self.relu3 = bn_relu(inplanes=384)
+        self.conv4 = nn.Conv2d(384, 384, kernel_size=3, padding=1, groups=2, bias=False)
+        self.relu4 = bn_relu(inplanes=384)
+        self.conv5 = nn.Conv2d(384, 256, kernel_size=3, padding=1, groups=2, bias=False)
+        self.relu_pool5 = bn_relu_pool(inplanes=256)
+        # classifier
+        self.conv6 = nn.Conv2d(256, 256, kernel_size=5, groups=2, bias=False)
+        self.relu6 = bn_relu(inplanes=256)
+        self.conv7 = nn.Conv2d(256, num_classes, kernel_size=1, bias=False)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.relu_pool1(x)
+        x = self.conv2(x)
+        x = self.relu_pool2(x)
+        x = self.conv3(x)
+        x = self.relu3(x)
+        x = self.conv4(x)
+        x = self.relu4(x)
+        x = self.conv5(x)
+        x = self.relu_pool5(x)
+        x = self.conv6(x)
+        x = self.relu6(x)
+        x = self.conv7(x)
+        x = x.view(x.size(0), -1)
+        return x
+
+
+
+################## ResNet ##################
+def conv3x3(in_planes, out_planes, stride=1):
+    # 3x3 convolution with padding
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False)
+
+class BasicBlock(nn.Module):
+    expansion = 1
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(BasicBlock, self).__init__()
+        m = OrderedDict()
+        m['conv1'] = conv3x3(inplanes, planes, stride)
+        m['bn1'] = nn.BatchNorm2d(planes)
+        m['relu1'] = nn.ReLU(inplace=True)
+        m['conv2'] = conv3x3(planes, planes)
+        m['bn2'] = nn.BatchNorm2d(planes)
+        self.group1 = nn.Sequential(m)
+        self.relu = nn.Sequential(nn.ReLU(inplace=True))
+        self.downsample = downsample
+
+    def forward(self, x):
+        if self.downsample is not None:
+            residual = self.downsample(x)
+        else:
+            residual = x
+        out = self.group1(x) + residual
+        out = self.relu(out)
+
+        return out
+
+class ResNet(nn.Module):
+    def __init__(self, block, layers, num_classes=1000):
+        self.inplanes = 64
+        super(ResNet, self).__init__()
+
+        m = OrderedDict()
+        m['conv1'] = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
+        m['bn1'] = nn.BatchNorm2d(64)
+        m['relu1'] = nn.ReLU(inplace=True)
+        m['maxpool'] = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.group1= nn.Sequential(m)
+
+        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)
+
+        self.avgpool = nn.Sequential(nn.AvgPool2d(7))
+        self.group2 = nn.Sequential(
+            OrderedDict([
+                ('fullyconnected', nn.Linear(512 * block.expansion, num_classes))
+            ])
+        )
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, math.sqrt(2. / n))
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.Linear):
+                torch.nn.init.xavier_uniform_(m.weight.data)
+                torch.nn.init.constant_(m.bias.data, 0)
+
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes * block.expansion, kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.group1(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        x = self.avgpool(x)
+        x = x.view(x.size(0), -1)
+        x = self.group2(x)
+
+        return x

app.py

@@ -1,5 +1,6 @@
 import streamlit as st
 from PIL import Image
+import core
 
 
 # 加载模型和其他必要的库
@@ -17,7 +18,7 @@ def main():
 
         # 使用模型进行颜值打分
         # score = model.predict(image)
-        score = 100
+        score = core.fun(uploaded_image)
 
         # 在界面上显示颜值分数
         st.write("颜值打分：", score)

core.py

@@ -0,0 +1,58 @@
+import torchvision.transforms as transforms
+from PIL import Image
+import Nets
+
+
+def load_model(pretrained_dict, new):
+    model_dict = new.state_dict()
+    # 1. filter out unnecessary keys
+    pretrained_dict = {k: v for k, v in pretrained_dict['state_dict'].items() if k in model_dict}
+    # 2. overwrite entries in the existing state dict
+    model_dict.update(pretrained_dict)
+    new.load_state_dict(model_dict)
+
+
+import torch
+
+# 检查GPU可用性
+if not torch.cuda.is_available():
+    device = torch.device("cuda")
+else:
+    device = torch.device("cpu")
+    print("use cpu")
+
+# model_ckpt_path = "./models/resnet18.pth"
+model_ckpt_path = "https://huggingface.co/M4869/beauty_prediction_fpb5k/blob/main/resnet18.pth"
+
+net = Nets.ResNet(block=Nets.BasicBlock, layers=[2, 2, 2, 2], num_classes=1).to(device)
+load_model(torch.load(model_ckpt_path, encoding='latin1'), net)
+net.eval()
+
+transform = transforms.Compose([
+    transforms.Resize(256),
+    transforms.CenterCrop(224),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ])
+
+
+def fun(img_path):
+    img = Image.open(img_path).convert('RGB')
+    img = transform(img)
+    with torch.no_grad():
+        img = img.unsqueeze(0).to(device)
+        output = net(img).squeeze(1).cpu().numpy()[0]
+        return output
+
+# def main():
+#     for i in range(6, 7):
+#         img = Image.open("./data2/%d.jpg" % i).convert('RGB')
+#         img = transform(img)
+#
+#         with torch.no_grad():
+#             img = img.unsqueeze(0).to(device)
+#             output = net(img).squeeze(1).cpu().numpy()[0]
+#             print(i, output * 20)
+
+
+# if __name__ == '__main__':
+# main()