NoteCrawling / process.py
nyonyong
Directory fix
e285fc9
raw
history blame
4.05 kB
import os
import shutil
import subprocess
from PIL import Image
import directories as Dir
def clearDir():
#์ƒ์„ฑ๋˜๋Š” ํŒŒ์ผ๋“ค์„ ์‚ญ์ œํ•จ์œผ๋กœ์จ ์ค‘๋ณต ์ถœ๋ ฅ ๋ฐฉ์ง€
# 1. User input
# "/TextDetection/cookie/user_input.jpg" ์‚ญ์ œ
# 2. Cropped images
# TextDetection/runs/detect/ ํ•˜์œ„ ํŒŒ์ผ ์ „๋ถ€ ์‚ญ์ œ
# 3. Recognition Result
# /TextRecognition/log_demo_result.txt ์‚ญ์ œ
user_input = "/home/user/app"+ Dir.yolo_dir +"/cookie/user_input.jpg"
crops = "/home/user/app" + Dir.yolo_dir + Dir.cropped_img_path
result = "/home/user/app" + Dir.txt_file_path
if os.path.isfile(user_input):
os.remove(user_input)
return print('Message: user input clear')
else:
print('Message: no user input')
#cropped_img_path = "/runs/detect/" + cropped_img_folder_name
if os.path.isdir(crops):
shutil.rmtree(crops) #'/runs/detect/user_output'
return print('Message: crop folder clear')
else:
print('Message: no crop folder')
#txt_file_path = "/log_demo_result.txt"
if os.path.isfile(result):
os.remove(result)
return print('Message: result clear')
else:
print('Message: no result')
def textDetection(im):
#change dir to yolo folder
#yolo_dir = "/HCR/TextDetection/"
#subprocess.call('pwd', shell=True)
#subprocess.call('ls', shell=True)
## subprocess.call('sudo cd '+ Dir.yolo_dir, shell=True)
#transfrom ndarray type to PIL type
im = Image.fromarray(im)
# save input image to cookie folder
## subprocess.call('cd cookie', shell=True)
im.save("/home/user/app"+Dir.yolo_dir+"/cookie/user_input.jpg", 'JPEG')
#yolo_dir = "/HCR/TextDetection/"
## subprocess.call('cd '+ Dir.yolo_dir, shell=True)
# (Shell) run detect.py to get cropped word images
subprocess.call(['python', "/home/user/app"+Dir.yolo_dir+'/detect.py',
#User Input Data : /text_detection/cookie
'--source',"/home/user/app"+Dir.yolo_dir+'/cookie',
#Text Detection Model : /runs/wordDetection/weights/best.pt
'--weights', "/home/user/app"+Dir.detect_model_dir,
'--conf','0.25',
#Output Images Save Directory /runs/detect/user_output
'--name', Dir.cropped_img_folder_name,
'--save-crop',
'--save-conf'])
#g = (size / max(im.size)) # gain
#im = im.resize((int(x * g) for x in im.size), Image.ANTIALIAS) # resize
#results = model(im) # inference
#results.render() # updates results.imgs with boxes and labels
#return Image.fromarray(results.imgs[0])
def textRearrange():
## subprocess.call('cd ' + Dir.DBSCAN_dir, shell=True)
subprocess.call(['python', "/home/user/app"+Dir.DBSCAN_dir+'/DBSCAN.py'])
def textRecognition():
#%cd /content/drive/MyDrive/KITA/Text/lmdb/deep-text-recognition-benchmark
## subprocess.call('cd '+Dir.recog_dir, shell=True)
#!CUDA_VISIBLE_DEVICES=0 python3 demo.py --Transformation TPS --FeatureExtraction ResNet --SequenceModeling BiLSTM --Prediction Attn --image_folder /content/drive/MyDrive/KITA/Text/YOLO/runs/detect/youtube_data2/crops/word --saved_model /content/drive/MyDrive/KITA/Text/best_accuracy_s/best_accuracy_s.pth
subprocess.call('CUDA_VISIBLE_DEVICES="" python3 '+ "/home/user/app"+ Dir.recog_dir +'/demo.py --Transformation TPS --FeatureExtraction ResNet --SequenceModeling BiLSTM --Prediction Attn --image_folder /home/user/app' + Dir.yolo_dir + Dir.cropped_img_path + '/crops/word --saved_model /home/user/app'+Dir.recog_model_dir, shell=True)
def getHcrResult(file_path):#*#
texts = ""
with open(file_path, 'r') as file:
lines = file.readlines()
for line in lines[3:]:
line = line.replace("\t","*",1)
line = line.replace(" ","*",1)
parts = line.replace(" ","")
parts2 = parts.split("*",2)
#print(len(parts2))
texts = texts +" "+ str(parts2[1:2])[2:-2]
return texts