Update modInference.py

modInference.py

@@ -3,80 +3,26 @@ import cv2
 import torch
 import glob as glob
 import os
+import time
+
 import matplotlib.pyplot as plt
 
-from models.create_fasterrcnn_model import create_model
 from utils.annotations import CNNpostAnnotations
 #from utils.annotations import inference_annotations
-from utils.general import set_infer_dir
 from utils.transforms import infer_transforms
 
-import numpy as np
-from skimage import transform
-import os
-from keras.models import Model
-from keras.optimizers import Adam
-from keras.applications.vgg16 import VGG16, preprocess_input
-from keras.layers import Dense, Dropout, Flatten
-import numpy as np
-
-
-conv_base = VGG16(include_top=False,
-                  weights='imagenet', 
-                  input_shape=(200,200,3))
-
-
-if 2 > 0:
-    for layer in conv_base.layers[:-2]:
-        layer.trainable = False
-else:
-    for layer in conv_base.layers:
-        layer.trainable = False
-
-
-top_model = conv_base.output
-top_model = Flatten(name="flatten")(top_model)
-top_model = Dense(4096, activation='relu')(top_model)
-top_model = Dense(1048, activation='relu')(top_model)
-top_model = Dense(256, activation='relu')(top_model)
-top_model = Dense(128, activation='relu')(top_model)
-top_model = Dense(64, activation='relu')(top_model)
-top_model = Dropout(0.2)(top_model)
-output_layer = Dense(5, activation='softmax')(top_model)
-
-CNN = Model(inputs=conv_base.input, outputs=output_layer)
-CNN.load_weights("CNN.hdf5")
-
-
-
-def main(weightUrl, input):
+def main(CNN, model, input):
     np.random.seed(42)
 
-    NUM_CLASSES = 2
-    CLASSES = ['__background__', 'Cell']
-
-    DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
-    OUT_DIR = set_infer_dir()
-
-    checkpoint = torch.load(weightUrl, map_location=DEVICE)
-    data_configs = True
-    NUM_CLASSES = checkpoint['config']['NC']
-    CLASSES = checkpoint['config']['CLASSES']
-    build_model = create_model[checkpoint['model_name']]
-
-    model = build_model(num_classes=NUM_CLASSES, coco_model=False)
-    model.load_state_dict(checkpoint['model_state_dict'])
-    model.to(DEVICE).eval()
-
-    COLORS = np.random.uniform(0, 255, size=(len(CLASSES), 3))
-
-
     image = input
     orig_image = image.copy()
     image = cv2.cvtColor(orig_image, cv2.COLOR_BGR2RGB)
     image = infer_transforms(image)
     image = torch.unsqueeze(image, 0)
 
+    DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+    CLASSES = ['__background__', 'Cell']
+
     outputs = model(image.to(DEVICE))
     
     # Load all detection to CPU for further operations.