-
Notifications
You must be signed in to change notification settings - Fork 0
/
loading.py
63 lines (48 loc) · 1.93 KB
/
loading.py
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
from pymilvus import Collection, CollectionSchema, FieldSchema, DataType, connections
from tensorflow.keras.applications.resnet50 import ResNet50, preprocess_input
from tensorflow.keras.preprocessing import image
import numpy as np
IMG_PATHS = [
'Schloss_Sigmaringen_2022.jpg',
'Canary_Wharf_from_Limehouse_London_June_2016_HDR.jpg',
'Daslook._Allium_ursinum,_zwellende_bloemknop._18-04-2022_(actm.)_04.jpg',
]
connections.connect(host='localhost', port='19530')
# Define the fields for the collection
fields = [
FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=2048)
]
# Define the schema for the collection
schema = CollectionSchema(fields, description="Image collection")
# Specify the index parameters
index_params = {
"index_type": "IVF_FLAT", # Choose the index type suitable for your use case
"params": {"nlist": 1024}, # Index parameters (vary based on index type)
"metric_type": "L2" # Distance metric (e.g., L2, IP)
}
# Create the collection in Milvus
collection_name = "image_collection"
collection = Collection(name=collection_name, schema=schema)
collection.create_index(field_name="embedding", index_params=index_params)
collection.load()
# Load the ResNet50 model
model = ResNet50(weights='imagenet', include_top=False, pooling='avg')
for img_path in IMG_PATHS:
# Load an image
img = image.load_img(img_path, target_size=(224, 224))
# Preprocess the image
x = image.img_to_array(img)
x = np.expand_dims(x, axis=0)
x = preprocess_input(x)
# Extract features
features = model.predict(x)
# Prepare the data for insertion
data = [
{"embedding": features.flatten().tolist()}
]
# Insert the data into the collection
mr = collection.insert(data)
# Remember to save the ids returned by Milvus for future reference
ids = mr.primary_keys
print(ids, img_path)