project_DF / stamp_detection.py
Ahmedhisham's picture
Upload 6 files
2ae875f verified
from ultralytics import YOLO
import os
import csv
model=YOLO('mainpipeline/models/stamp_detection_model.pt')
img_dir='cache/GB'
#output_csv='cache/extracted_data.csv'
#set a detection flag = false
output_csv='cache/output/appended_data.csv'
revision_status='unrevised'
for img_name in os.listdir(img_dir):
img_path=os.path.join(img_dir,img_name)
if os.path.isfile(img_path):
results=model.predict(img_path,conf=0.25,save=False)
if len(results)>0 :
revision_status="Revised"
break
with open(output_csv,mode='a',newline='') as file:
writer=csv.writer(file)
#writer.writerow(['Revision status',revision_status])
file.seek(0,os.SEEK_END)
if file.tell()==0:
writer.writerow(["Revision_status"])
writer.writerow(["revision_status",revision_status])