-
Notifications
You must be signed in to change notification settings - Fork 0
/
image_processing.py
80 lines (62 loc) · 2.55 KB
/
image_processing.py
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
67
68
69
70
71
72
73
74
75
76
77
78
79
80
import cv2
import numpy as np
# For anomally detect
import sys
import os
sys.path.append("./svdd_anubis/")
from PIL import Image, ImageDraw
from deepSVDD import DeepSVDD
from utils.config import Config
from datasets.preprocessing import global_contrast_normalization
import torchvision.transforms as transforms
import torch
def processImage_main(image_in):
"""!@brief Function is used to process the input image and for
its return to the processing queue.
@param[image_in] Input image
"""
# Create output image
image_out = image_in
#image_out = edge_detect(image_in)
image_out = anomally_detect(image_in)
return image_out
def edge_detect(image):
# Blur the image
image_blur = cv2.GaussianBlur(image[0], (3,3), 0)
# Canny Edge Detection
image[0] = cv2.Canny(image=image_blur, threshold1=100, threshold2=200)
return image
def anomally_detect(image):
model_dir_path="./svdd_anubis/model"
load_config=os.path.join(model_dir_path,"config.json")
load_model=os.path.join(model_dir_path,"model.tar")
useCUDA = torch.cuda.is_available()
if useCUDA:
dev = torch.device('cuda')
else:
dev = torch.device('cpu')
# Get configuration
cfg = Config(locals().copy())
cfg.load_config(import_json=load_config)
# Data transform information
min_max = [(-1.3942, 6.5378),
(-1.9069, 18.5352),
(-2.6035, 29.8313),
(-2.3845, 12.4581)]
trans=transforms.Compose([transforms.Resize((480,480)),
transforms.ToTensor(),
transforms.Lambda(lambda x: global_contrast_normalization(x, scale='l1')),
transforms.Normalize([min_max[cfg.settings['normal_class']][0]] * 3,
[min_max[cfg.settings['normal_class']][1] - min_max[cfg.settings['normal_class']][0]] * 3)])
deep_SVDD = DeepSVDD(cfg.settings['objective'], cfg.settings['nu'])
deep_SVDD.set_network(cfg.settings['net_name'])
deep_SVDD.load_model(model_path=load_model, load_ae=False, device = dev)
img = np.repeat(image[0],3, axis=-1)
im = Image.fromarray(img,'RGB')
img = trans(im) #np.squeeze(image[0], axis=-1)
scores = deep_SVDD.test_image(img, dev)
#print(scores)
draw = ImageDraw.Draw(im)
draw.rectangle([(0,0), im.size], outline='#ffffff' if scores[0]> 1000 else '#000000', width=10) #'#ff0000' if scores[0]> 1 else '#00ff00'
image[0] = np.array(im.convert('L'))
return image