Upload New File

annotate_with_yolo.py

+#predict
+from ultralytics import YOLO
+from PIL import Image
+import pandas as pd
+from glob import glob
+import os
+
+model = YOLO('path_to_the_finetuned_yolo8.pt')
+
+parentpath = 'path_to_the_images_annotated_with_owl'
+i = 0
+ext = '.jpg'
+imgPaths = glob(parentpath +  "/*" + ext)
+
+for fname in imgPaths:
+	try:
+		raw_image = Image.open(fname)
+	except:
+		continue
+
+	txtfile = fname.replace("jpg", "txt")
+	txtfile = txtfile.replace("images", "labels")
+	with open(txtfile, "r") as f:
+		lines = f.readlines()
+	current_classes = []
+	for line in lines:
+		current_classes.append(int(line[0]))
+
+	results=model.predict(source=raw_image, save=True)
+
+	a=results[0].boxes.data
+	a = a.detach().cpu().numpy()  
+	px=pd.DataFrame(a).astype("float")
+
+	height, width = results[0].orig_shape
+
+	lines=[]
+             
+	for index,row in px.iterrows():
+		x1 = row[0]
+		y1 = row[1]
+		x2 = row[2]
+		y2 = row[3]
+		d  = int(row[5])
+		xcen = float((x1 + x2)) / 2 / width
+		ycen = float((y1 + y2)) / 2 / height
+		w = float((x2 - x1)) / width
+		h = float((y2 - y1)) / height
+		if d not in current_classes:
+			lines.append([d, xcen, ycen, w, h])
+	if os.path.exists(txtfile):
+		with open(txtfile, "a") as f:
+        		for line in lines:
+	             		f.write("%d %.06f %.06f %.06f %.06f\n" % (line[0], line[1], line[2], line[3], line[4]))
+	
+	i += 1
+	print(i)