orin.py

DIMENSION = 512
DATABASE_NAME = "./OrinEdgeAI.db"
COLLECTION_NAME = "OrinEdgeAI"
PATH = "/home/jetson/unstructureddata/images/"
BLIP_MODEL = "Salesforce/blip-image-captioning-large"

extractor = FeatureExtractor("resnet34")

milvus_client = MilvusClient(DATABASE_NAME)
fields = [
    FieldSchema(name='id', dtype=DataType.INT64, is_primary=True, auto_id=True),
    FieldSchema(name='caption', dtype=DataType.VARCHAR, max_length=512),
    FieldSchema(name='filename', dtype=DataType.VARCHAR, max_length=512),
    FieldSchema(name='currenttime', dtype=DataType.VARCHAR, max_length=512),
    FieldSchema(name='vector', dtype=DataType.FLOAT_VECTOR, dim=DIMENSION)
]
schema = CollectionSchema(fields=fields)
milvus_client.create_collection(COLLECTION_NAME, DIMENSION, schema=schema, metric_type="COSINE", auto_id=True)
index_params = milvus_client.prepare_index_params()
index_params.add_index(field_name = "vector", metric_type="COSINE")
milvus_client.create_index(COLLECTION_NAME, index_params)

cam = cv2.VideoCapture(0)
result, image = cam.read()
strfilename = PATH + 'orin{0}.jpg'.format(uuid.uuid4())

if result:
    cv2.imwrite(strfilename, image)
else:
    print("No image")

currenttimeofsave = datetime.datetime.now().strftime('%m/%d/%Y %H:%M:%S')
hostname = os.uname()[1]

processor = BlipProcessor.from_pretrained(BLIP_MODEL)
model = BlipForConditionalGeneration.from_pretrained(BLIP_MODEL)

inputs = processor(image, return_tensors="pt")
out = model.generate(**inputs)
caption = (processor.decode(out[0], skip_special_tokens=True))

try:
    imageembedding = extractor(strfilename)
    milvus_client.insert( COLLECTION_NAME, {"vector": imageembedding, "currenttime":currenttimeofsave,
            "filename": strfilename, "caption": str(caption)})
    milvus_client.close()
except Exception as e:
    print("An error:", e)