image.py 4.3 KB
Newer Older
lipengfei's avatar
lipengfei committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
import sys
import re
import paddle
import subprocess

def check_version():
    try:
        version = paddle.__version__
        version_pattern = r"(\d+)\.(\d+)\.(\d+)"
        match = re.match(version_pattern, version)
        major, minor, patch_version = map(int, match.groups())
        return major, minor
    except ImportError:
        print("PaddlePaddle is not installed.")
        sys.exit(1)

major, minor = check_version()

# PaddlePaddle 2.4及以上版本
if major >= 3 or (major == 2 and minor >= 4):
    from ppdiffusers import StableDiffusionPipeline

    pipe = StableDiffusionPipeline.from_pretrained("./stable-diffusion-2")

    paddle.seed(5232132133)
    prompt = "a portrait of shiba inu with a red cap growing on its head. intricate. lifelike. soft light. sony a 7 r iv 5 5 mm. cinematic post - processing "
    image = pipe(prompt, guidance_scale=7.5, height=768, width=768).images[0]

    image.save("shiba_dog_with_a_red_cap.png")

# PaddlePaddle 2.0到2.3版本
elif 2 <= major < 3 and 0 <= minor < 4:
    import paddle
    import paddle.vision.transforms as T
    from PIL import Image
    import numpy as np
    import json

    model = paddle.vision.models.resnet50(pretrained=True)
    model.eval()

    def preprocess_image(image_path):
        transform = T.Compose([
            T.Resize(256),
            T.CenterCrop(224),
            T.ToTensor(),
            T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])
        image = Image.open(image_path).convert('RGB')
        image = transform(image)
        image = paddle.unsqueeze(image, axis=0)
        return image

    def predict(model, image_tensor):
        with paddle.no_grad():
            logits = model(image_tensor)
            probs = paddle.nn.functional.softmax(logits, axis=1).numpy()
        return probs

    def load_labels(filepath):
        with open(filepath, "r") as f:
            labels = json.load(f)
        return labels

    def display_top_predictions(probs, labels):
        top5_idx = np.argsort(probs[0])[-5:][::-1]
        for idx in top5_idx:
            print(f"Class: {labels[idx]}, Probability: {probs[0][idx]:.4f}")

    if __name__ == "__main__":
        image_path = "human.jpg"
        labels_path = "imagenet-simple-labels.json"
        image_tensor = preprocess_image(image_path)
        labels = load_labels(labels_path)
        probs = predict(model, image_tensor)
        display_top_predictions(probs, labels)

# PaddlePaddle 2.0以下版本
else:
    import paddle.fluid as fluid
    import numpy as np
    from PIL import Image

    class SimpleCNN(fluid.dygraph.Layer):
        def __init__(self, num_classes=10):
            super(SimpleCNN, self).__init__()
            self.conv1 = fluid.dygraph.Conv2D(num_channels=3, num_filters=32, filter_size=3, stride=1, padding=1)
            self.pool1 = fluid.dygraph.Pool2D(pool_size=2, pool_stride=2, pool_type='max')
            self.conv2 = fluid.dygraph.Conv2D(num_channels=32, num_filters=64, filter_size=3, stride=1, padding=1)
            self.pool2 = fluid.dygraph.Pool2D(pool_size=2, pool_stride=2, pool_type='max')
            self.fc = fluid.dygraph.Linear(input_dim=64*8*8, output_dim=num_classes)

        def forward(self, x):
            x = self.conv1(x)
            x = fluid.layers.relu(x)
            x = self.pool1(x)
            x = self.conv2(x)
            x = fluid.layers.relu(x)
            x = self.pool2(x)
            x = fluid.layers.flatten(x, 1)
            x = self.fc(x)
            return x

    def load_image(image_path):
        img = Image.open(image_path)
        img = img.resize((32, 32), Image.ANTIALIAS)
        img = np.array(img).astype('float32') / 255.0
        img = img.transpose((2, 0, 1))
        img = np.expand_dims(img, axis=0)
        return img

    def infer(image_path, model):
        with fluid.dygraph.guard():
            img = load_image(image_path)
            img = fluid.dygraph.to_variable(img)
            prediction = model(img)
            predicted_class = np.argmax(prediction.numpy())
            return predicted_class

    if __name__ == "__main__":
        image_path = "human.jpg"
        with fluid.dygraph.guard():
            model = SimpleCNN(num_classes=10)
            model.eval()
            predicted_label = infer(image_path, model)
            print(f"Predicted label: {predicted_label}")
print("finish")