Spaces:
No application file
No application file
| 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]) | |