1
# 观察到数据往往是来自同一个视频
2
# 首先进行数据清洗，手动根据训练集的标题，归到分别的文件夹中
3

4
# train_raw
5
# ----类别1
6
#      --------01文件夹
7
#                 ---------npy
8
#      --------02文件夹
9
#                 ---------npy
10
# ----类别2
11
#      --------01文件夹
12
#                 ---------npy
13
#      --------02文件夹
14
#                 ---------npy
15

16
# 手动删去train中非影像、文字装饰太多的图片（共7张）
17
# 包括：
18
# ./data/train/Cyst/04/05 (4)到10 (4).npy （非影像演示图）
19
# ./data/train/Vascular/01/05.npy （文字装饰太多）
20

21
# 这个代码用来查看一个类中的所有图像，便于手动删除
22

23
# Anomalies       01-11 220 -> 220
24
# Cyst            01-11 180 -> 174
25
# Inflammation    01-11 180 -> 180
26
# Tumor           01-13 250 -> 250
27
# Vascular        01-11 159 -> 158
28

29
import numpy as np
30
import mediapy as media
31
import glob, os
32

33
for i in range(1,12):
34
    image_paths = sorted(glob.glob(f"./data/train/Anomalies/{i:02}/*.npy"))
35
    images = []
36
    print(i)
37
    for image_path in image_paths:
38
        images.append(np.load(os.path.join(image_path)))
39
    media.show_images(images, columns=10)
40

41
# 针对训练集每一个类进行内部聚类，获取类内的聚类点特征
42
# 减轻类内的数据不平衡、一些图片基本一致的影响
43

44
train_class_and_folder_num = {
45
    "Anomalies": 11,
46
    "Cyst": 11,
47
    "Inflammation": 11,
48
    "Tumor": 13,
49
    "Vascular": 11
50
}
51

52
# 抽取图像特征
53
# 获取resnet50特征
54

55
import torch
56
import torch.nn as nn
57
import torchvision.models as models
58
import torchvision.transforms as transforms
59
from PIL import Image
60

61
class DemoNet(nn.Module):
62
    def __init__(self):
63
        super(DemoNet, self).__init__()
64

65
        model = models.resnet50(pretrained=True)
66
        model.fc = torch.nn.Identity()
67
        self.model = model
68

69
    def forward(self, img):
70
        out = self.model(img)
71
        return out
72

73
extract = DemoNet().cuda().eval()
74

75
transform = transforms.Compose([
76
    transforms.Resize((256, 256)),
77
    transforms.ToTensor(),
78
])
79

80
# 抽取图像特征
81
# 遍历每个类
82

83
with torch.no_grad():
84
    for idx, class_name in enumerate(train_class_and_folder_num):
85
        features = []
86
        folder_num = train_class_and_folder_num[class_name]
87
        for i in range(1, folder_num+1):
88
            image_paths = sorted(glob.glob(f"./data/train/{class_name}/{i:02}/*.npy"))
89
            for image_path in image_paths:
90
                image = np.load(os.path.join(image_path))
91
                image = Image.fromarray(image)
92
                image = transform(image).unsqueeze(0).cuda()
93
                feature = extract(image)
94
                features.append(feature)
95
        features = torch.cat(features, dim=0)
96
        torch.save(f"./data/train_class_{idx}.pt", features)