File size: 2,927 Bytes
bfabfea |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
import cv2
from ultralytics import YOLO
import os
from dotenv import load_dotenv
from pathlib import Path
env_path = Path('.') / '.env'
load_dotenv(dotenv_path=env_path)
from functions import cropBlackBackground, enhanceImage, generateMask
path = {
'SEG_MODEL_PATH': str(os.getenv('SEG_MODEL_PATH')),
'DET_MODEL_PATH': str(os.getenv('DET_MODEL_PATH')),
'IMG_DIR_PATH': str(os.getenv('IMG_DIR_PATH')),
'INFERENCE_FOLDER': str(os.getenv('INFERENCE_FOLDER')),
}
#import models
seg_model = YOLO(path['SEG_MODEL_PATH'])
det_model = YOLO(path['DET_MODEL_PATH'])
CONF = 0.7
# do inference for detection and store croped images in folder
for img in os.listdir(path['IMG_DIR_PATH']):
img_file = cv2.imread(os.path.join(path['IMG_DIR_PATH'],img),0)
cv2.imwrite(os.path.join(path['IMG_DIR_PATH'],img),img_file)
det_model(os.path.join(path['IMG_DIR_PATH'],img),conf = CONF, save=True , save_crop=True , name=path['INFERENCE_FOLDER'],exist_ok = True)
#do inference for image segmentation and store image in folder
for img in os.listdir(path['IMG_DIR_PATH']):
img_file = cv2.imread(os.path.join(path['IMG_DIR_PATH'],img),0)
cv2.imwrite(os.path.join(path['IMG_DIR_PATH'],img),img_file)
result = seg_model(os.path.join(path['IMG_DIR_PATH'],img),save = True,name = path['INFERENCE_FOLDER'],exist_ok = True)
original_img = cv2.imread(os.path.join(path['IMG_DIR_PATH'],img))
for res in result:
crop_img, mask = generateMask(res, original_img)
image = None
if crop_img is not None:
# Convert to gray scale image
crop_img = cv2.cvtColor(crop_img, cv2.COLOR_BGR2GRAY)
crop_img = cropBlackBackground(crop_img)
image = enhanceImage(crop_img)
# save to file
if image is not None:
try:
Path('runs').mkdir(parents=True, exist_ok=True)
Path(os.path.join('runs', 'segment')).mkdir(parents=True, exist_ok=True)
Path(os.path.join('runs', 'segment', 'inference')).mkdir(parents=True, exist_ok=True)
Path(os.path.join('runs', 'segment', 'inference', 'crops_seg')).mkdir(parents=True, exist_ok=True)
Path(os.path.join('runs', 'segment', 'inference', 'enhanced')).mkdir(parents=True, exist_ok=True)
Path(os.path.join('runs', 'segment', 'inference', 'masks')).mkdir(parents=True, exist_ok=True)
except OSError as error:
print(error)
pass
cv2.imwrite(os.path.join('runs', 'segment', 'inference', 'masks', img), mask)
cv2.imwrite(os.path.join('runs', 'segment', 'inference', 'crops_seg', img), crop_img )
cv2.imwrite(os.path.join('runs', 'segment', 'inference', 'enhanced', img), image )
|